PyTorch分类模型的转换和使用OpenCV Python启动

上一教程： 如何运行自定义OCR模型
下一教程： PyTorch分类模型的转换和使用OpenCV C++启动

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

目标

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

• 将PyTorch分类模型转换为ONNX格式
• 使用OpenCV Python API运行转换后的PyTorch模型
• 获得PyTorch和OpenCV DNN模型的评估结果。

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

简介

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

需求

要能够使用下面的代码进行实验，您需要安装一组库。我们将为此使用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. 评估和测试分类模型

如果您只想运行评估或测试模型流程，“模型转换流程”部分可以跳过。

模型转换流程

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

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

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

1. 实例化PyTorch模型
2. 将PyTorch模型转换为.onnx
3. 使用OpenCV API读取转换后的网络
4. 准备输入数据
5. 提供推理

# 初始化PyTorch ResNet-50模型
original_model = models.resnet50(pretrained=True)
 
# 获取转换为ONNX的PyTorch模型的路径
full_model_path = get_pytorch_onnx_model(original_model)
 
# 使用OpenCV API读取转换后的.onnx模型
opencv_net = cv2.dnn.readNetFromONNX(full_model_path)
print("OpenCV模型已成功读取。层ID：\n", opencv_net.getLayerNames())
 
# 获取预处理图像
input_img = get_preprocessed_img("../data/squirrel_cls.jpg")
 
# 获取ImageNet标签
imagenet_labels = get_imagenet_labels("../data/dnn/classification_classes_ILSVRC2012.txt")
 
# 获取OpenCV DNN预测结果
get_opencv_dnn_prediction(opencv_net, input_img, imagenet_labels)
 
# 获取原始PyTorch ResNet50预测结果
get_pytorch_dnn_prediction(original_model, input_img, imagenet_labels)

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

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

分类模型输入图像

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

让我们通过预训练的PyTorch ResNet-50的例子深入研究每个步骤

• 实例化PyTorch ResNet-50模型

# 初始化PyTorch ResNet-50模型
original_model = models.resnet50(pretrained=True)

• 将PyTorch模型转换为ONNX

# 定义进一步转换模型保存的目录
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
)

成功执行上述代码后，我们将得到models/resnet50.onnx。

• 使用cv.dnn.readNetFromONNX读取转换后的网络，并将上一步中获得的ONNX模型传递给它

# 使用OpenCV API读取转换后的.onnx模型
opencv_net = cv2.dnn.readNetFromONNX(full_model_path)

• 准备输入数据

# 读取图像
input_img = cv2.imread(img_path, cv2.IMREAD_COLOR)
input_img = input_img.astype(np.float32)
 
input_img = cv2.resize(input_img, (256, 256))
 
# 定义预处理参数
mean = np.array([0.485, 0.456, 0.406]) * 255.0
scale = 1 / 255.0
std = [0.229, 0.224, 0.225]
 
# 准备输入blob以适应模型输入
# 1. 减去均值
# 2. 缩放以将像素值设置为0到1
input_blob = cv2.dnn.blobFromImage(
image=input_img,
scalefactor=scale,
size=(224, 224), # 目标图像大小
mean=mean,
swapRB=True, # BGR -> RGB
crop=True # 中心裁剪
)
# 3. 除以标准差
input_blob[0] /= np.asarray(std, dtype=np.float32).reshape(3, 1, 1)

此步骤中，我们读取图像并使用cv.dnn.blobFromImage 函数准备模型输入，该函数返回一个四维 blob。需要注意的是，首先在cv.dnn.blobFromImage 中减去平均值，然后才将像素值乘以比例因子。因此，mean 乘以 255.0 以重现原始图像预处理顺序。

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

• OpenCV cv.dnn.Net 推理

# 设置 OpenCV DNN 输入
opencv_net.setInput(preproc_img)
 
# OpenCV DNN 推理
out = opencv_net.forward()
print("OpenCV DNN 预测结果：\n")
print("* 形状：", out.shape)
 
# 获取预测的类别 ID
imagenet_class_id = np.argmax(out)
 
# 获取置信度
confidence = out[0][imagenet_class_id]
print("* 类别 ID：{}, 标签：{}".format(imagenet_class_id, imagenet_labels[imagenet_class_id]))
print("* 置信度：{:.4f}".format(confidence))

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

OpenCV DNN 预测结果
* 形状： (1, 1000)
* 类别 ID：335，标签：狐松鼠，东部狐松鼠，Sciurus niger
* 置信度：14.8308

• PyTorch ResNet-50 模型推理

original_net.eval()
preproc_img = torch.FloatTensor(preproc_img)
 
# 推理
out = original_net(preproc_img)
print("\nPyTorch 模型预测结果：\n")
print("* 形状：", out.shape)
 
# 获取预测的类别 ID
imagenet_class_id = torch.argmax(out, axis=1).item()
print("* 类别 ID：{}, 标签：{}".format(imagenet_class_id, imagenet_labels[imagenet_class_id]))
 
# 获取置信度
confidence = out[0][imagenet_class_id]
print("* 置信度：{:.4f}".format(confidence.item()))

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

PyTorch 模型预测结果
* 形状：torch.Size([1, 1000])
* 类别 ID：335，标签：狐松鼠，东部狐松鼠，Sciurus niger
* 置信度：14.8308

原始 ResNet-50 模型和cv.dnn.Net 的推理结果相同。为了对模型进行扩展评估，我们可以使用 dnn_model_runner 模块的 py_to_py_cls。该模块部分将在下一小节中描述。

模型评估

samples/dnn 中提出的 dnn_model_runner 模块允许在 ImageNet 数据集上运行完整的评估流程，并测试以下 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>

从列表中选择的分类模型将被读入 OpenCV cv.dnn.Net 对象。PyTorch 和 OpenCV 模型的评估结果（准确率、推理时间、L1）将写入日志文件。推理时间值也将以图表的形式显示，以便概括获得的模型信息。

必要的评估配置定义在test_config.py 中，可以根据实际数据位置的路径进行修改。

@dataclass
class TestClsConfig
batch_size: int = 50
frame_size: int = 224
img_root_dir: str = "./ILSVRC2012_img_val"
    # 图像类别匹配的位置
img_cls_file: str = "./val.txt"
bgr_to_rgb: bool = True

要启动 PyTorch ResNet-50 的评估，请运行以下命令：

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

脚本启动后，将在 dnn_model_runner/dnn_conversion/logs 中生成包含评估数据的日志文件。

PyTorch resnet50 模型已成功获取并转换为 OpenCV DNN resnet50。
===== 使用以下参数运行模型评估
* 验证数据位置：./ILSVRC2012_img_val
* 日志文件位置：dnn_model_runner/dnn_conversion/logs/PyTorch_resnet50_log.txt

测试模式

以下行表示在测试模式下运行模块，即它提供了模型推理的步骤：

python -m dnn_model_runner.dnn_conversion.pytorch.classification.py_to_py_cls --model_name <pytorch_cls_model_name> --test True --default_img_preprocess <True/False> --evaluate False

这里的 default_img_preprocess 键定义了是否要使用某些特定值或默认值来参数化模型测试过程，例如 scale、mean 或 std。

测试配置在test_config.py 的 TestClsModuleConfig 类中表示。

@dataclass
class TestClsModuleConfig
cls_test_data_dir: str = "../data"
test_module_name: str = "classification"
test_module_path: str = "classification.py"
input_img: str = os.path.join(cls_test_data_dir, "squirrel_cls.jpg")
model: str = ""
 
frame_height: str = str(TestClsConfig.frame_size)
frame_width: str = str(TestClsConfig.frame_size)
scale: str = "1.0"
mean: List[str] = field(default_factory=lambda: ["0.0", "0.0", "0.0"])
std: List[str] = field(default_factory=list)
crop: str = "False"
rgb: str = "True"
rsz_height: str = ""
rsz_width: str = ""
classes: str = os.path.join(cls_test_data_dir, "dnn", "classification_classes_ILSVRC2012.txt")

默认图像预处理选项定义在default_preprocess_config.py 中。例如：

BASE_IMG_SCALE_FACTOR = 1 / 255.0
PYTORCH_RSZ_HEIGHT = 256
PYTORCH_RSZ_WIDTH = 256
 
pytorch_resize_input_blob = {
    "mean": ["123.675", "116.28", "103.53"],
    "scale": str(BASE_IMG_SCALE_FACTOR),
    "std": ["0.229", "0.224", "0.225"],
    "crop": "True",
    "rgb": "True",
    "rsz_height": str(PYTORCH_RSZ_HEIGHT),
    "rsz_width": str(PYTORCH_RSZ_WIDTH)
}

模型测试的基础代码位于 samples/dnn/classification.py。 提供转换后的模型路径给--input参数，并填写cv.dnn.blobFromImage 的相关参数后，即可独立运行classification.py。

要根据“模型转换流程”中描述的步骤，使用dnn_model_runner从头开始复现OpenCV步骤，请执行以下命令：

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

网络预测结果将显示在输出窗口的左上角。

ResNet50 OpenCV 推理输出