Previous tutorial: Convert a PyTorch classification model and publish with OpenCV Python
Original author | Anastasia Murzova |
---|---|
Compatibility | OpenCV >= 4.5 |
Goals
In this tutorial you will learn how to
- Convert PyTorch classification model to ONNX format
- Run the converted PyTorch model using the OpenCV C/C++ API
- Provide model inference
We will discuss the above points using the ResNet-50 architecture as an example.
Introduction
Let’s briefly review the key concepts involved in the process of converting a PyTorch model using the OpenCV API. The first step in converting a PyTorch model to cv::dnn::Net is to convert the model to ONNX format. ONNX aims to enable the interchange of neural networks between different frameworks. There is a built-in function in PyTorch for converting ONNX: torch.onnx.export
. The resulting .onnx
model is then passed into cv::dnn::readNetFromONNX or cv::dnn::readNet.
Requirements
To experiment with the code below, you need to install a set of libraries. For this we will use a virtual environment with python3.7+:
virtualenv -p /usr/bin/python3.7 <env_dir_path> source <env_dir_path>/bin/activate
To build OpenCV-Python from source, follow the appropriate instructions in **Introduction to OpenCV**.
Before you start installing the library, you can customize requirements.txt to exclude or include (for example, opencv-python
) certain dependencies. The following line installs the startup requirements into the previously activated virtual environment:
pip install -r requirements.txt
Practice
In this part we will cover the following points:
- Create a classification model transformation pipeline
- Provide inference and process prediction results
Model conversion pipeline
The code in this section is located in the samples/dnn/dnn_model_runner
module and can be executed from the following command line:
python -m dnn_model_runner.dnn_conversion.pytorch.classification.py_to_py_resnet50_onnx
The following code contains instructions for the following steps:
- Instantiate a PyTorch model
- Convert PyTorch model to
.onnx
# Initialize PyTorch ResNet-50 model original_model = models.resnet50(pretrained=True) # Get the path converted to ONNX PyTorch model full_model_path = get_pytorch_onnx_model(original_model) print("PyTorch ResNet-50 model conversion successful:", full_model_path)
The get_pytorch_onnx_model(original_model)
function is called based on torch.onnx.export(...)
:
# Define the saving directory for further converted models onnx_model_path = "models" # Define names for further converted models # Define names for further converted models onnx_model_name = "resnet50.onnx" # Create a directory for further converted models # Create directories for further converted models os.makedirs(onnx_model_path, exist_ok=True) # Get the full path of the converted model full_model_path = os.path.join(onnx_model_path, onnx_model_name) # Generate model input generated_input = Variable( torch.randn(1, 3, 224, 224) ) # Export the model to ONNX format torch.onnx.export( original_model, generated_input, full_model_path, verbose=True, input_names=["input"]、 output_names=["output"]、 opset_version=11 )
After successfully executing the above code, we will get the following output:
PyTorch ResNet-50 model has been successfully converted: models/resnet50.onnx
With the dnn/samples
module dnn_model_runner
we can reproduce the above transformation steps for the following PyTorch classification models:
- alexnet
- vgg11
- vgg13
- vgg16
- vgg19
- resnet18
- resnet34
- resnet50
- resnet101
- resnet152
- squeezenet1_0
- squeezenet1_1
- resnext50_32x4d
- resnext101_32x8d
- wide_resnet50_2
- wide_resnet101_2
To get the converted model, the following line should be executed:
python -m dnn_model_runner.dnn_conversion.pytorch.classification.py_to_py_cls --model_name <pytorch_cls_model_name> --evaluate False
For the ResNet-50 case, the following line should be run:
python -m dnn_model_runner.dnn_conversion.pytorch.classification.py_to_py_cls --model_name resnet50 --evaluate False
The default root directory for converted model storage is defined in the CommonConfig
module:
@dataclass class CommonConfig: output_data_root_dir: str = "dnn_model_runner/dnn_conversion" (output data root directory)
Therefore, the converted ResNet-50 will be saved in dnn_model_runner/dnn_conversion/models
.
Inference Pipeline
We can now use models/resnet50.onnx
for inference via the OpenCV C/C++ API. The implemented pipeline can be found in samples/dnn/classification.cpp. After building the sample (the BUILD_EXAMPLES
flag value should be ON
), the corresponding example_dnn_classification
executable is provided.
To provide model inference, we will use the following photo of a squirrel corresponding to ImageNet class ID 335 (licensed under CC0 license):
Fox squirrel, Eastern fox squirrel, Sciurus niger
Classification model input image
In order to label decode the predictions, we also need the imagenet_classes.txt file, which contains the complete list of ImageNet classes.
In this tutorial, we will run the inference process of the converted PyTorch ResNet-50 model in the build (samples/build
) directory:
./dnn/example_dnn_classification --model=../dnn/models/resnet50.onnx --input=../data/squirrel_cls.jpg --width=224 --height=224 --rgb=true - -scale="0.003921569" --mean="123. 675 116.28 103.53" --std="0.229 0.224 0.225" --crop=true --initial_width=256 --initial_height=256 --classes=../data/ dnn/classification_classes_ILSVRC2012.txt
Let’s explore the key points of classification.cpp
step by step:
- Use cv::dnn::readNet to read the model and initialize the network:
net = readNet(model,config,framework);
Model parameter values are taken from the --model
key. In our case it is resnet50.onnx
.
- Preprocess the input image: if (rszWidth != 0 & amp; & rszHeight != 0) {<!-- --> resize(frame, frame, Size(rszWidth, rszHeight)); } //Create 4D blob from frame blobFromImage(frame, blob, scale, Size(inpWidth, inpHeight), mean, swapRB, crop); // Check std value. if (std.val[0] != 0.0 & amp; & amp; std.val[1] != 0.0 & amp; & amp; std.val[2] != 0.0) {<!-- --> // Divide blob by std. divide(blob, std, blob); }
In this step, we prepare the model input using the cv::dnn::blobFromImage function. We set Size(rszWidth,rszHeight)
, where --initial_width=256 --initial_height=256
is used to adjust the initial image size, as in thePyTorch ResNet inference pipeline as described in .
Note that the average value in cv::dnn::blobFromImage is subtracted first, and only then the pixel value is multiplied by the scale. Therefore, we use --mean="123.675 116.28 103.53"
, which is equivalent to [0.485, 0.456, 0.406]
multiplied by 255.0
to re- The original image preprocessing sequence of the PyTorch classification model:
img /= 255.0 img -= [0.485, 0.456, 0.406] img /= [0.229, 0.224, 0.225]
- Make a prequel:
net.setInput(blob); Mat prob = net.forward();
- Process prediction results:
Point classIdPoint; double confidence; minMaxLoc(prob.reshape(1, 1), 0, & amp;confidence, 0, & amp;classIdPoint); int classId = classIdPoint.x;
- Here we select the most likely object class. The
classId
result for this example is 335 – Fox Squirrel, Eastern Fox Squirrel, Sciurus niger:
ResNet50 OpenCV C++ inference output