上一篇教程： PyTorch分类模型的转换和使用OpenCV Python运行


原作者	Anastasia Murzova
兼容性	OpenCV >= 4.5

目标

在本教程中，您将学习如何：

将PyTorch分类模型转换为ONNX格式
使用OpenCV C/C++ API运行转换后的PyTorch模型
进行模型推理

我们将以ResNet-50架构为例，探讨上述要点。

简介

让我们简要回顾一下PyTorch模型与OpenCV API交互流程中涉及的关键概念。将PyTorch模型转换为cv::dnn::Net 的第一步是将模型转换为ONNX 格式。ONNX旨在实现不同框架之间神经网络的互换性。PyTorch内置了用于ONNX转换的函数：torch.onnx.export。然后，将获得的.onnx模型传递到cv::dnn::readNetFromONNX 或cv::dnn::readNet。

需求

要能够使用下面的代码进行实验，您需要安装一些库。我们将为此使用带有python3.7+的虚拟环境

virtualenv -p /usr/bin/python3.7 <env_dir_path>

source <env_dir_path>/bin/activate

对于从源代码构建OpenCV-Python，请遵循OpenCV入门中的相应说明。

在开始安装库之前，您可以自定义requirements.txt，排除或包含（例如，opencv-python）一些依赖项。以下命令将启动在先前激活的虚拟环境中安装依赖项

pip install -r requirements.txt

实践

在本部分中，我们将介绍以下几点：

创建分类模型转换流程
进行推理，处理预测结果

模型转换流程

本小节中的代码位于samples/dnn/dnn_model_runner模块中，可以使用以下命令执行：

python -m dnn_model_runner.dnn_conversion.pytorch.classification.py_to_py_resnet50_onnx

以下代码包含对以下步骤的描述：

实例化PyTorch模型
将PyTorch模型转换为.onnx

# 初始化PyTorch ResNet-50模型
original_model = models.resnet50(pretrained=True)
 
# 获取转换后的ONNX PyTorch模型的路径
full_model_path = get_pytorch_onnx_model(original_model)
print("PyTorch ResNet-50模型已成功转换：", full_model_path)

get_pytorch_onnx_model(original_model)函数基于torch.onnx.export(...)调用

# 定义转换后模型保存的目录
onnx_model_path = "models"
# 定义转换后模型的名称
onnx_model_name = "resnet50.onnx"
 
# 为转换后的模型创建目录
os.makedirs(onnx_model_path, exist_ok=True)
 
# 获取转换后模型的完整路径
full_model_path = os.path.join(onnx_model_path, onnx_model_name)
 
# 生成模型输入
generated_input = Variable(
torch.randn(1, 3, 224, 224)
)
 
# 将模型导出为ONNX格式
torch.onnx.export(
original_model,
generated_input,
full_model_path,
verbose=True,
input_names=["input"],
output_names=["output"],
opset_version=11
)

成功执行上述代码后，我们将得到以下输出：

PyTorch ResNet-50模型已成功转换：models/resnet50.onnx

dnn/samples模块中提出的dnn_model_runner允许我们针对以下PyTorch分类模型重现上述转换步骤：

alexnet
vgg11
vgg13
vgg16
vgg19
resnet18
resnet34
resnet50
resnet101
resnet152
squeezenet1_0
squeezenet1_1
resnext50_32x4d
resnext101_32x8d
wide_resnet50_2
wide_resnet101_2

要获得转换后的模型，应执行以下命令：

python -m dnn_model_runner.dnn_conversion.pytorch.classification.py_to_py_cls --model_name <pytorch_cls_model_name> --evaluate False

对于ResNet-50，应运行以下命令：

python -m dnn_model_runner.dnn_conversion.pytorch.classification.py_to_py_cls --model_name resnet50 --evaluate False

转换后模型存储的默认根目录在CommonConfig模块中定义：

@dataclass
class CommonConfig
output_data_root_dir: str = "dnn_model_runner/dnn_conversion"

因此，转换后的ResNet-50将保存在dnn_model_runner/dnn_conversion/models中。

推理流程

现在，我们可以使用OpenCV C/C++ API使用models/resnet50.onnx进行推理。实现的流程可以在samples/dnn/classification.cpp中找到。构建示例后（BUILD_EXAMPLES标志值应为ON），将提供相应的example_dnn_classification可执行文件。

为了进行模型推理，我们将使用以下松鼠照片（在CC0许可下），对应于ImageNet类别ID 335

狐松鼠，东部狐松鼠，Sciurus niger

分类模型输入图像

对于获得的预测的标签解码，我们还需要imagenet_classes.txt文件，其中包含ImageNet类别的完整列表。

在本教程中，我们将从构建（samples/build）目录运行转换后的PyTorch ResNet-50模型的推理过程：

./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

让我们一步步探索classification.cpp的关键点

使用cv::dnn::readNet读取模型，初始化网络

Net net = readNet(model, config, framework);

model参数值取自--model键。在本例中，它是resnet50.onnx。

预处理输入图像

if (rszWidth != 0 && rszHeight != 0)
{
resize(frame, frame, Size(rszWidth, rszHeight));
}
 
// 从帧创建4D blob
blobFromImage(frame, blob, scale, Size(inpWidth, inpHeight), mean, swapRB, crop);
 
// 检查标准差值。
if (std.val[0] != 0.0 && std.val[1] != 0.0 && std.val[2] != 0.0)
{
    // 将blob除以标准差。
divide(blob, std, blob);
}

在此步骤中，我们使用cv::dnn::blobFromImage函数准备模型输入。我们使用--initial_width=256 --initial_height=256设置Size(rszWidth, rszHeight)进行初始图像大小调整，正如PyTorch ResNet推理管道中所述。

需要注意的是，首先在cv::dnn::blobFromImage中减去平均值，然后才将像素值乘以比例因子。因此，我们使用--mean="123.675 116.28 103.53"，这相当于将[0.485, 0.456, 0.406]乘以255.0，以重现PyTorch分类模型的原始图像预处理顺序

img /= 255.0
img -= [0.485, 0.456, 0.406]
img /= [0.229, 0.224, 0.225]

进行前向传播

net.setInput(blob);

Mat prob = net.forward();

处理预测结果

Point classIdPoint;
double confidence;
minMaxLoc(prob.reshape(1, 1), 0, &confidence, 0, &classIdPoint);
int classId = classIdPoint.x;

在这里，我们选择最可能的物体类别。在本例中，classId的结果是335——狐狸松鼠，东部狐狸松鼠，Sciurus niger

ResNet50 OpenCV C++ 推理输出