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 中。

推理管道

現在我們可以使用 models/resnet50.onnx 來進行使用 OpenCV C/C++ API 的推理管道。實現的管道可以在 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 的關鍵點

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 blob
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)
{
    // 將 blob 除以 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 - 狐松鼠，東部狐松鼠，Sciurus niger

ResNet50 OpenCV C++ 推理輸出