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), # img 目標大小
mean=mean,
swapRB=True, # BGR -> RGB
crop=True # 中心裁剪
)
# 3. 除以 std
input_blob[0] /= np.asarray(std, dtype=np.float32).reshape(3, 1, 1)

在此步驟中，我們讀取影像並使用 cv.dnn.blobFromImage 函式準備模型輸入，該函式返回 4 維 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
===== 使用以下引數執行模型的評估
* val 資料位置：./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 中表示。 classification.py 可以使用 --input 中提供的轉換後的模型和填充的 cv.dnn.blobFromImage 引數自主執行。

要從頭開始重現“模型轉換管道”中描述的 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 推理輸出