coremltools 4.0でPyTorchモデルを変換する

Core ML 4（iOS 14）に合わせて、coremltools 4.0がリリースされた。coremltools 4.0の目玉となるアップデートが、**PyTorchモデルがONNXを経由することなく直接Core MLモデルに変換できるようになった**点。

> New converter path to directly convert PyTorch models without going through ONNX. 

（coremltools 4.0b1のWhat's Newより）

ドキュメントも一新され、Examplesのページに"Converting a PyTorch Segmentation Model to Core ML"というセグメンテーションモデルを変換するチュートリアルがあったので、そちらで変換を試してみた。

https://coremltools.readme.io/docs/pytorch-conversion-examples#converting-a-pytorch-segmentation-model-to-core-ml

このモデルを利用すると、iOSやmacOSアプリでこんな感じでセグメンテーションを行える。

![picture_pc_207c6f68a223eb5847191761c9e79acd.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3180/342716ff-7e8e-cf81-a010-f1f05ccc7105.png)

ちなみにDeepLabV3のCore MLモデルは公式配布のものが既にあるが、PyTorchを用いて独自データセットで学習させたセグメンテーションモデルをiOS上で動かしたい場合、本記事の手順でCore MLモデルに変換する必要が出てくる。


## Unified Conversion API

以前のcoremltoolsでは、Kerasモデルを変換する場合は`convertes.keras.convert`、TensorFlow 1.xモデルを変換する場合は`tfcoreml`といった感じで変換を担うコンバーターがバラバラだった。

これが**coremltools 4.0では一つのコンバーターに統合され、これひとつでTensorFlowモデル（1.x, 2.x, tf.keras）とPyTorchモデルを変換できるように**なった。

> Convert TensorFlow or Pytorch models to Core ML model format. Whether a parameter is required may differ between frameworks (see below). Note that this function is aliased as ct.convert in the tutorials.

こちらは**Unified Conversion API**と呼ばれている。

https://coremltools.readme.io/docs/unified-conversion-api


このUnified Conversion APIがサポートするフォーマットは以下：

![Screen Shot 2020-12-22 at 5.00.41.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3180/2751f94c-cb14-9a2d-6a98-86ef6cd43781.png)



## coremltools.converters.convert

新しいUnifiedなコンバータのリファレンスを見ると、`convert`メソッドの定義は次のようになっていて、

```python
coremltools.converters.convert(model, source='auto', inputs=None, outputs=None, classifier_config=None, minimum_deployment_target=None, **kwargs)
```

第1引数`model`の説明によると、Pytorch向けには、以下の2種類をモデルを示す値として受け取るようだ。

- A TorchScript object
- Path to a .pt file


## TorchScriptとtorch.jit.trace

TorchScriptというのは、「PyTorchの中間表現コードとその変換機構」を表すもので、次のような利点があるらしい：


> - TorchScriptコードは独自のインタプリタで実行・解釈され、Pythonインタプリタのグローバルインタプリタロック (GIL) とも無縁なので、マルチスレッドで並列計算できる
- コードとパラメータをまるごと保存でき、Python以外の実行環境でもロードできる
- 中間コードへ落とし込むコンパイラにて最適化しやすくなる

[PyTorchのモデルをTorchScriptへ変換する - け日記](https://ohke.hateblo.jp/entry/2019/12/21/230000)

で、変換には`torch.jit.trace`を使う。

https://pytorch.org/docs/master/generated/torch.jit.trace.html

モデルの`forward`メソッドを実行し、その処理を記録 (トレース) することで変換する、ということらしい。

## coremltools 4.0へのアップデート

```zsh
$ pip install -U coremltools
```


バージョン確認

```python
import coremltools as ct

ct.__version__    # 4.0
```


## PyTorchモデルをCore MLモデルに変換するための最小実装

[こちらの手順](https://coremltools.readme.io/docs/pytorch-conversion-examples#converting-a-pytorch-segmentation-model-to-core-ml
)のうち、実はなくても大丈夫な手順をカットした最小実装を提示する。


### 1. モデルのロード

DeepLabV3モデルのロード自体は`torch.hub.load`を用いて次のように行えるのだが、

```python
import torch
import torchvision

model = torch.hub.load('pytorch/vision:v0.6.0', 'deeplabv3_resnet101', pretrained=True).eval()
```


これを`torch.jit.trace`に渡すとエラーになってしまう。出力でdictionaryを返しているのが問題らしい。


> When running this code, the tracer outputs an error that Only tensors or tuples of tensors can be output from traced functions, but our model returns a dictionary.

ので、チュートリアルでは次のようにラッパークラスを定義して、出力時にdictionaryからtensorを抽出して返すようにするワークアラウンドが提示されている。

> Get around this limitation by wrapping the model in a module that extracts the output we want from the dictionary:

```python
import torch.nn as nn

class WrappedDeeplabv3Resnet101(nn.Module):

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

   def forward(self, x):
       res = self.model(x)
       # Extract the tensor we want from the output dictionary
       x = res["out"]
       return x
```

```python
traceable_model = WrappedDeeplabv3Resnet101().eval()
```

### 2. traceに渡す入力テンソルを生成

公式チュートリアルでは1枚の入力画像を読み込んでそれを加工したものをtraceに渡しているが、

```python
input_image = Image.open("dog_and_cat.jpg")
```

```python
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)
```

実は**テンソルのshapeさえ合っていればランダムで良い**らしい。

> Note: A random input of the same shape would also work.


ので、次のようにrandで1x3x448x448のテンソルを生成するコードで置き換えた。

```python
example_input = torch.rand(1, 3, 448, 448)
```

### 3. TorchScriptオブジェクトを生成

converterにわたすためのTorchScriptオブジェクトを`torch.jit.trace`を用いて生成する。

```python
trace = torch.jit.trace(traceable_model, example_input)
```

### 4. Core MLモデルに変換

```python
import coremltools as ct

mlmodel = ct.convert(
   trace,
   inputs=[ct.TensorType(name="input", shape=example_input.shape)],
)
```

第1引数に`TorchScript`オブジェクト、第2引数で入力テンソルの名前やshapeを渡している。

### 5. mlmodelファイルとして保存

```python
mlmodel.save('./PyTorchDeepLabV3.mlmodel')
```

### 6. 動作確認

#### 画像読み込み

```python
from PIL import Image

input_image = Image.open("dog_and_cat.jpg")
display(input_image)
```

#### 前処理

```python
from torchvision import transforms

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)
```

#### 推論処理の実行

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

#### 結果のビジュアライズ

```python
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)

display_segmentation(input_image, torch_predictions)
```

#### 生成したモデルをXcodeでプレビュー

![picture_pc_ce0aecfd71c391945a7d27a22785f124.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3180/72b98a2f-62a6-60a8-25c8-aec262cfc28a.png)

（coremltoolsのバージョンが4.0b2となっているのはb2の頃に試した際のスクショを使用しているため）