PyTorch 分类模型转换并在 OpenCV C++ 中启动

上一个教程： 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

操作

在此部分中，我们将介绍以下要点

1. 创建一个分类模型转换管道
2. 提供推理，处理预测结果

模型转换管道

本子章节中的代码位于samples/dnn/dnn_model_runner模块中，可以使用以下代码行执行

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

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

1. 初始化 PyTorch 模型
2. 将 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

福克斯松鼠，东部福克斯松鼠，黑松鼠

分类模型输入图像

对于获得预测的标签解码，我们还需要 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 的要点

1. 使用 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 斑块
blobFromImage(frame, blob, scale, Size(inpWidth, inpHeight), mean, swapRB, crop);
 
// 检查 std 值。
if (std.val[0] != 0.0 && std.val[1] != 0.0 && std.val[2] != 0.0)
{
    // 将斑块除以 std。
divide(blob, std, blob);
}

在此步骤中，我们使用 cv::dnn::blobFromImage 函数准备模型输入。我们设置 Size(rszWidth, rszHeight)，--initial_width=256 --initial_height=256 以便初步调整图像大小，如 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 - 狐松鼠，东狐松鼠，黑松鼠

ResNet50 OpenCV C++ 推断输出