作成日: 2022.06.27

iPhone で PyTorch で作成した深層学習モデルの推論を行う

PyTorchで作成・学習したネットワークモデルをiPhone上で推論したい場合どうすべきか。

いくつか方法はあるが、簡単なのはcoremltoolsを用いてPyTorchのモデルをCoreMLで動く形式に変換することである。

環境構築

Macでの作業を前提とする（だってiPhoneで動かすんだから、MacじゃないとXcodeうごかないでしょ)

あと、PyThonの（正確にはPyTrochなど）のバージョンは最新にするとエラーになることが多いので注意。
あくまで以下のバージョンで動作確認しているが、バージョンが変わると動かない（エラーになる）がしようがない。
PyTorchかcoremltoolsか知らんが下位互換性があんまないので。

pyenvは入っているものとする。

pyenv global 3.9.13
mkdir pytorch_ios_sample
cd pytorch_ios_sample/
python3 -m venv .venv
pyenv global system
source .venv/bin/activate
pip install pip==22.1.2

PyTorch関連のを入れる(このサンプルでは2022/06/27現在の最新で問題なかった)

pip install torch torchvision
pip install coremltools
export PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python

以降、バージョンが異なるだけで不可解なエラーが発生しやすいので、何か問題が起きたらまずはバージョンを確認し、異なっている場合は合わせた上で再度確認すること。

$ python -c "import torch;print(torch.__version__)"
1.11.0
$ python -c "import torchvision;print(torchvision.__version__)"
0.12.0
$ python -c "import coremltools;print(coremltools.__version__)"
WARNING:root:Torch version 1.11.0 has not been tested with coremltools. You may run into unexpected errors. Torch 1.10.2 is the most recent version that has been tested.
5.2.0

coremltoolsでwarningが発生しているが、このサンプルを実行する上で問題が起きなかったので、torchを1.10.2にダウングレードはしていない。

なお、macOSは12.4で、Xcodeは13.4である。

サンプルコード実行

以下の公式のサンプルコードを動かす。

PyTorchセグメンテーションモデルを変換する
https://coremltools.readme.io/docs/pytorch-conversion-examples

ページの一番下の方にコードがあるのでコピペする。
(ここではdeeplabv3_resnet101.pyと命名する)

import urllib
import warnings
warnings.simplefilter(action='ignore', category=FutureWarning)

import torch
import torch.nn as nn
import torchvision
import json

from torchvision import transforms
from PIL import Image

import coremltools as ct

# Load the model (deeplabv3)
model = torch.hub.load('pytorch/vision:v0.12.0', 'deeplabv3_resnet101', pretrained=True).eval()

# Load a sample image (cat_dog.jpg)
input_image = Image.open("cat_dog.jpg")
input_image.show()

preprocess = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225],
    ),
])

input_tensor = preprocess(input_image)
input_batch = input_tensor.unsqueeze(0)

with torch.no_grad():
    output = model(input_batch)['out'][0]
torch_predictions = output.argmax(0)

def display_segmentation(input_image, output_predictions):
    # Create a color pallette, selecting a color for each class
    palette = torch.tensor([2 ** 25 - 1, 2 ** 15 - 1, 2 ** 21 - 1])
    colors = torch.as_tensor([i for i in range(21)])[:, None] * palette
    colors = (colors % 255).numpy().astype("uint8")

    # Plot the semantic segmentation predictions of 21 classes in each color
    r = Image.fromarray(
        output_predictions.byte().cpu().numpy()
    ).resize(input_image.size)
    r.putpalette(colors)

    # Overlay the segmentation mask on the original image
    alpha_image = input_image.copy()
    alpha_image.putalpha(255)
    r = r.convert("RGBA")
    r.putalpha(128)
    seg_image = Image.alpha_composite(alpha_image, r)
    # display(seg_image) -- doesn't work
    seg_image.show()

display_segmentation(input_image, torch_predictions)

# Wrap the Model to Allow Tracing*
class WrappedDeeplabv3Resnet101(nn.Module):

    def __init__(self):
        super(WrappedDeeplabv3Resnet101, self).__init__()
        self.model = torch.hub.load('pytorch/vision:v0.12.0', 'deeplabv3_resnet101', pretrained=True).eval()

    def forward(self, x):
        res = self.model(x)
        x = res["out"]
        return x

# Trace the Wrapped Model
traceable_model = WrappedDeeplabv3Resnet101().eval()
trace = torch.jit.trace(traceable_model, input_batch)

# Convert the model
mlmodel = ct.convert(
    trace,
    inputs=[ct.TensorType(name="input", shape=input_batch.shape)],
)

# Save the model without new metadata
mlmodel.save("SegmentationModel_no_metadata.mlmodel")

# Load the saved model
mlmodel = ct.models.MLModel("SegmentationModel_no_metadata.mlmodel")

# Add new metadata for preview in Xcode
labels_json = {"labels": ["background", "aeroplane", "bicycle", "bird", "board", "bottle", "bus", "car", "cat", "chair", "cow", "diningTable", "dog", "horse", "motorbike", "person", "pottedPlant", "sheep", "sofa", "train", "tvOrMonitor"]}

mlmodel.user_defined_metadata["com.apple.coreml.model.preview.type"] = "imageSegmenter"
mlmodel.user_defined_metadata['com.apple.coreml.model.preview.params'] = json.dumps(labels_json)

mlmodel.save("SegmentationModel_with_metadata.mlmodel")

実行するには cat_dog.jpg が必要なのでダウンロードしておく。

wget -O cat_dog.jpg https://files.readme.io/690356d-cat_dog.jpg

実行。結構待たされる。

$ python deeplabv3_resnet101.py
WARNING:root:Torch version 1.11.0 has not been tested with coremltools. You may run into unexpected errors. Torch 1.10.2 is the most recent version that has been tested.
Converting Frontend ==> MIL Ops: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▊| 844/845 [00:00<00:00, 911.47 ops/s]
Running MIL Common passes: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 34/34 [00:01<00:00, 20.21 passes/s]
Running MIL Clean up passes: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 9/9 [00:00<00:00, 34.23 passes/s]
Translating MIL ==> NeuralNetwork Ops: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 944/944 [01:14<00:00, 12.67 ops/s]

で、SegmentationModel_no_metadata.mlmodelとSegmentationModel_with_metadata.mlmodelの２つのモデルができあがる。
SegmentationModel_no_metadata.mlmodelで問題ないが、これだとXcode上でプレビューができない。
(メタデータ付きモデルだとXcode上でプレビューができるようになる)
別にプレビューいらんがな、という人はSegmentationModel_no_metadata.mlmodelでなんら問題ない。

なぜクラスをラッピングしているのか？

WrappedDeeplabv3Resnet101(nn.Module): というクラスでラッピングしているがなぜか？
理由はPyTorchのTorchScriptがDictonaryに対応していないから。
よってわざわざラッピングして、outキーの値であるTensorだけを返却するようにしている。
いつか治るかも知らんが。

MIL形式

ほとんど意識しなくて良いが、Torch ScriptモデルからCoreMLモデルに変換する際、coremltools内ではMIL形式に変換している。
勝手にやっているので気にする必要がないが、変換においてサポートされていない演算子エラーが出た場合、このMILプログラムを書いて演算子を追加するなどする。
現状、MaskRCNNなどでエラーが生じるので、その辺りをMILで書き換えればうまく動かしたりできたりもする。
(MILを書き換えInstanceSegmentationを実現位しているサンプルコードは別途紹介するhttps://github.com/xta0/CoreML-MaskRCNN）

torch.jit.traceとtorch.jit.script

どちらもPyTorchモデルをTorchScriptモデルに変換するツール。

torch.jit.traceが昔からあるけど、最近はtorch.jit.scriptを使った方が良いらしい。

ので torch.jit.script を使ったらエラーになった。

    assert str(node.output().type()) == "Tensor"
AssertionError

理不尽。

iPhoneで動かす

セグメンテーションモデルを動かすiOSサンプルアプリはAppleから提供されていないので、ネットから適当に探した以下を使う。

---

---

---

---