(ML) How to Convert PyTorch MobileNetV2 to CoreML and Perform Inference on iPhone

Photo by Kevin Bhagat on Unsplash

環境：Macbook Air M1 = 13.2 + Pytorch = 2.1 + Coremltools = 5.0b5

目的：學習將 pytorch 中預訓練模型 MobileNetV2 轉換為 CoreML ，並在 iphone 實機上進行推理

MobileNetV2

(ML) MobileNetV2: Inverted Residuals and Linear Bottlenecks

Generate a TorchScript Version

import torch
import torchvision

# Load a pre-trained version of MobileNetV2
torch_model = torchvision.models.mobilenet_v2(pretrained=True)

# Set the model in evaluation mode.
torch_model.eval()

# Trace the model with random data.
example_input = torch.rand(1, 3, 224, 224) 
traced_model = torch.jit.trace(torch_model, example_input)
out = traced_model(example_input)

Convert to Core ML

import coremltools as ct

# Using image_input in the inputs parameter:
# Convert to Core ML neural network using the Unified Conversion API.
model = ct.convert(
    traced_model,
    inputs=[ct.TensorType(shape=example_input.shape)]
 )

model.save("mobile.mlmodel")

以上轉換流程是參考 coremltools 教學

Converting from PyTorch

在下方專案中使用的是 Apple Developer 所提供的模型，我們會將專案中的模型改為自己轉換後的模型

(SwiftUI) 使用 CoreML 進行圖像辨識

下載專案後，將 mobile.mlmodel 模型拖移到 Xcode 中

跟專案中的 Resnet50Int8LUT 做個比較

在 mobile 中少了 Preview 與 Class Labels 欄位

嘗試將專案中的 Resnet50Int8LUT 模型改為 mobile 進行 inference

import SwiftUI
import CoreML
import UIKit
import UniformTypeIdentifiers

struct ContentView: View {
    ...

    private var model: mobile?

    init() {
        // 初始化模型
        do {
            // 創建 Resnet50Int8LUT 模型實例，並使用空的 MLModelConfiguration 進行初始化
            model = try mobile(configuration: MLModelConfiguration())
        } 
        ...
    }

Error message

使用 Netron 將官方所提供的 MobileNetV2 與 mobile 進行可視化比較

Netron

可視化 MobileNetV2.mlmodel

可視化 mobile.mlmodel

自行轉換的模型，出現以下問題

1. 在 Xcode 沒有出現 Preview 與 ClassLabels 欄位
2. 轉換後的模型輸出沒有出現 classLabel 、classLabelProbs、 softmax
3. 執行 Xcode 會出現 error

重新轉換模型

從 torchvision 下載模型

import torch
import torchvision

# Load a pre-trained version of MobileNetV2
torch_model = torchvision.models.mobilenet_v2(pretrained=True)

創建一個 MobileNetV2ClassificationModel，並新增 softmax 用於歸一化

import torch.nn as nn 

class MobileNetV2ClassificationModel(nn.Module):
    def __init__(self):
        super(MobileNetV2ClassificationModel, self).__init__()
        self.layers = nn.Sequential(
            torch_model,
            nn.Softmax(dim=1)
        )
        
    def forward(self, x):
        return self.layers(x)

model = MobileNetV2ClassificationModel().eval()

若是將 MobileNetV2ClassificationModel().train() 進行模型轉換，轉換後的模型在可視化會出現 Batch Normalize 層，但是在轉換過程也會跳出 Model is not in eval mode. Consider calling ‘.eval()’ on your model prior to conversion，並且最後在 inference 時，會無法預測出類別

下載 ImageNet 類標籤，將其製成列表，然後從列表中創建一個 CoreML 分類器配置

import urllib
import coremltools as ct

label_url = 'https://storage.googleapis.com/download.tensorflow.org/data/ImageNetLabels.txt'
class_labels = urllib.request.urlopen(label_url).read().splitlines()
class_labels = class_labels[1:] # remove the first class which is background assert len(class_labels) == 1000


for i, label in enumerate(class_labels):
    if isinstance(label, bytes):
        class_labels[i] = label.decode("utf8")

        
classifier_config = ct.ClassifierConfig(class_labels)

轉換為 CoreML

ex_input=torch.randn((1,3,224,224))

traced = torch.jit.trace(model,ex_input)

image_input = ct.ImageType(name="image",
                           shape=ex_input.shape,                      
                           bias=[-0.485/0.229,-0.456/0.224,-0.406/0.225],
                           scale=1.0/255.0/0.226
                          )

output = ct.TensorType(name="classLabelProbs")


mlmodel = ct.convert(traced,
                     inputs=[image_input], 
                     outputs=[output],                     
                     classifier_config=classifier_config,
                     )


mlmodel.save("mobile_coreML.mlmodel")

Netron 可視化

不管是在輸入還是輸出，轉換後的模型已經相當接近官方所提供的 MobileNetV2，但是還是有一個小問題，也就是在 pytorch 中的 softmax 轉換後，會轉換成 softmaxND

修復轉換模型後產生的 softmaxND

(ML) Fix softmax after converting CoreML model becomes softmaxND

Xcode 模型測試

將 mobile_coreML.mlmodel 加入到 Xcode 中，並使用 Preivew 進行測試

Inference on iphone

將專案中的 Resnet50Int8LUT 模型改為 mobile_coreML 進行 inference

import SwiftUI
import CoreML
import UIKit
import UniformTypeIdentifiers

struct ContentView: View {
    ...
    private var model: mobile_coreML?
    init() {
        // 初始化模型
        do {
            // 創建 Resnet50Int8LUT 模型實例，並使用空的 MLModelConfiguration 進行初始化
            model = try mobile_coreML(configuration: MLModelConfiguration())
        } 
        ...
    }

執行結果