バイセル Tech Blog

バイセル Tech Blogは株式会社BuySell Technologiesのエンジニア達が知見・発見を共有する技術ブログです。

バイセル Tech Blog

Pytorch(skorch)で学習したモデルを使い、iOSで画像分類をする [2/3]

iOS DNN nmu0

テクノロジー開発部の村上です。
前回でPytorchのモデルは作成出来たので、それをCore MLのモデルに変換する部分を説明します。

Core MLはAppleが提供しているSwift・Objective-C用の機械学習用のライブラリです。
Core MLでPytorchのモデルは読み込めないので、

  1. ONNXフォーマットにモデルを変換
  2. ONNXフォーマットからCore MLモデルに変換

という手順を踏みます。
ONNXは深層学習のモデルのためのフォーマットです。
世の中には色々な深層学習のライブラリが出ていますが、
モデル自体はライブラリによらない概念なので共通フォーマットを作成することがONNXの目的です。
ONNXへの変換・ONNXからの変換をそれぞれのライブラリが用意してくれていると、今回のように変換させることが出来ます。

ONNXフォーマットにモデルを変換

ONNXに変換するのに必要なのは、モデルとダミーの入力です。
ダミーの入力の次元は batch_size x RGBの3 x 224 x 224 で、ダミーなのでbatch_sizeを1にして乱数で作成しています。

import torch
import torchvision.models as models

vgg16 = models.vgg16(pretrained=False)
num_labels = 2 # catsとdogsの2種類なので
num_features = vgg16.classifier[6].in_features
modules = list(vgg16.classifier.children())
modules.pop()
modules.append(torch.nn.Linear(num_features, num_labels))
modules.append(torch.nn.Softmax(dim=1))
new_classifier = torch.nn.Sequential(*modules)
vgg16.classifier = new_classifier
vgg16.load_state_dict(torch.load('best_model.pt'))

device = torch.device('cpu')
vgg16 = vgg16.to(device)
dummy_input = torch.randn(1, 3, 224, 224).to(device)

import torch.onnx
input_names = ['image']
output_names = ['output']
torch.onnx.export(vgg16, 
                  dummy_input, 
                  'vgg16.onnx', 
                  verbose=True, 
                  input_names=input_names, 
                  output_names=output_names)

import onnx
model = onnx.load_model('vgg16.onnx')
# 念のため、変換したモデルをチェック
onnx.checker.check_model(model)

ONNXフォーマットのモデルをCore MLモデルに変換

事前にonnx-coremlをpipでインストールしておきます。(バージョン0.3で行いました)
また、プログラム上のポイントは3つあり

  1. 入力をモデルに流すまえの事前処理を定義しておく
    • Core MLでの入力は[0, 255]なので、[0, 1]にします
    • VGGでは、画像を正規化しているのでバイアス項の値を与えます
      • Pytorchのときと符号が逆なのに注意です
  2. image_input_namesをONNXに変換するときと同じものにする
    • Ccore MLモデルの入力が配列ではなく画像に設定され、iOSでの無駄な変換が不要になります
  3. class_labelsにアルファベット順にラベルを与える
    • Core MLで推論した結果として、ここで与えたラベルが得られるようになります

バイアス項だけで良いのかと疑問を持った方もいるかと思いますが、
現時点ではスケール変換はRGBで共通のものしか出来ない仕様となっているため、
Core MLモデルに変換後にスケール変換する層を最初に加えることで対処します。

あとこれは現時点でのonnx-coremlのバグだと思うのですが(実は正しく、私の理解の方が間違っていたら申し訳ありません)、
変換後のCore MLモデルをコンパイルして設定ファイルを確認すると、バイアス値のRGBがそれぞれBGRに入れ替わっています。(キー名で判断しました)
よって、意図的にRにGのバイアス, BにRのバイアスを入れるようにしておきます。

from onnx_coreml import convert
preprocessing_args = {'is_bgr': False, 
                      'red_bias': -0.406, 
                      'green_bias': -0.456, 
                      'blue_bias': -0.485, 
                      'image_scale': 1/255}
mlmodel = convert(model, 
                  mode='classifier', 
                  preprocessing_args=preprocessing_args, 
                  image_input_names=['image'], 
                  predicted_feature_name="classLabel", 
                  class_labels=['cats', 'dogs'])
mlmodel.save('CoreMLModel.mlmodel')
# 変換したモデルの確認
spec = mlmodel.get_spec()
print(spec.description)

RGBチャンネルごとのスケール変換

前述したように、スケール変換する層を最初に加えVGGで必要な画像の正規化を実施します。
ここでもスケールの仕様がPytorchのときと異なり、1 / 0.229 のように逆数にする点に注意です。

from coremltools.models import MLModel
loaded_model = MLModel('CoreMLModel.mlmodel')
loaded_spec = loaded_model.get_spec()
nn_spec = loaded_spec.neuralNetworkClassifier

import copy
layers = nn_spec.layers
layers_copy = copy.deepcopy(layers)
del nn_spec.layers[:]

scale_layer = nn_spec.layers.add()
scale_layer.name = 'scale_layer'
scale_layer.input.append('image')
scale_layer.output.append('image_scaled')
params = scale_layer.scale
params.scale.floatValue.extend([1 / 0.229, 1 / 0.224, 1 / 0.225])
params.shapeScale.extend([3,1,1])
nn_spec.layers.extend(layers_copy)
nn_spec.layers[1].input[0] = 'image_scaled'

import coremltools
modified_mlmodel = coremltools.models.MLModel(loaded_spec)
modified_mlmodel.save('CoreMLModel.mlmodel')

こうしてCore MLモデルが作成出来ましたので、次回は実際にiOSプログラムで動かしてみます。

参考資料