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

テクノロジー開発部の村上です。現在はアーキテクチャ周りを担当しています。
弊社で社内向けにお酒判定iOSアプリを作成したので、そこで使った技術を3回に渡って紹介したいと思います。

Skorch (Pytorchを使ったライブラリ) でCNNモデル作成
Pytorchモデルを -> ONNX -> CoreMLモデルとiOSで使える形に変換
CoreMLモデルをSwiftのPlaygroundで検証してみる

お酒画像は公開されていないので、今回はその代わりに犬猫判定器を作ろうと思います。
お気づきの方もいるかと思いますが、これは私がfast.aiで勉強して、その教材として犬猫判定が使われているからです。
fast.aiは理論面も実際の現場での応用面も押さえて説明しているので、これから勉強する人には非常にオススメです。

使用技術の選定理由

興味がない人は飛ばしてしまって大丈夫ですが、何か参考になればと思います。
まず、今回の開発の要件として、以下の3つがありました。

社内端末がiPhoneなのでiOSアプリで
オフラインで全て動くようにしてほしい
- 現場からの要望
開発期間は一人で約一ヶ月

そこで、以下のように対処しました。

CoreMLを使用
CoreML、Realm、Amazon Mobile Hubを使用して、オフラインで動きオンラインだとデータ収集が出来るように
以下で対処
- iOSアプリの開発は初めてなので、社内の以前のプロジェクトで使ったライブラリ(RxSwift、Eurekaなど)を活用
- 使用経験のあるPytorchでCNNモデル作成

Amazon Mobile Hub、RxSwift、Eurekaなどはそのうち紹介する機会があるかもしれません。

1. Skorch (Pytorchを使ったライブラリ) でCNNモデル作成

巷ではGoogleのTensorflowが有名ですが、今回はFacebookが中心に開発しているPytorchを使いました。(バージョンは1.0です)
Pytorchの利点としては、

普通のPythonプログラムのように書けるので、学習コストが低い
- Define by Run というタイプの深層学習フレームワークだからです
  - Tensforflowも途中から対応しています
世界的に(特に学術界で)よく使われていて、情報ソースが多い
ONNXというフレームワーク横断のモデル定義形式に対応しているので、iOSモデルにも変換可能
- 逆に言うとまだプロダクションではあまりPytorchは使われていません

しかし利点1の弊害として、学習ループを自分でforループで書かないといけません。
そこで今回はPytorchを使いやすいインターフェースでラップしたskorchを使用します。(バージョンは0.3です)
skorchはインターフェースをscikit-learnと揃えているので、scikit-learnを使ってきた人は特に楽だと思います。

モデル作成の仕組み

転送学習 (transfer learning)という手法を使います。自前のデータだけでは画像が足りないので、既存の画像で学習したモデルが公開されていてそれを使うことで画像の不足を補います。
ただし解きたい問題が異なるので、そのモデルの一部を再学習して自分の問題に合わせるのが転送学習です。
平たくいうと、人間が未知の物体でも色や輪郭を掴めるように、画像から物体の特徴を理解するような普遍的な部分は問題が違っても流用出来るということです。

データのロード

http://files.fast.ai/data/dogscats.zip からダウンロードしてください。
dogscats/trainを見ると、dogディレクトリとcatディレクトリが含まれています。
Pytorchのdataloaderは、dog/1.jpg、dog/2.jpg、cat/1.jpg、cat/2.jpgのようにラベルごとにディレクトリが分かれていると、自動的にアルファベット順に、catは0、dogは1とラベリングして、それぞれに対応付けて画像を読み込んでくれます。
また、画像ロード時に画像を加工することも出来ます。後述するモデルではImageNetという画像データセットで事前に学習されたモデルを使うので、そのモデルの作成時の作法に合わせて以下の加工が必要です。

224×224のサイズに縮小・中央クロッピングする
画像を決まった方法で正規化する

trainの方だけRandomというprefixがついた加工をしていますが、これはData augmentationと呼ばれる、データを人工的に増やす手法です。
微妙に角度を変えたり、左右を反転させたりして、データを増やしています。ポイントとしては、実際にありえそうな変換に抑えることです。例えば、天井に張り付いた猫の写真はめったにないと思うので、上下反転はかえって精度が落ちます。

from pathlib import Path

import torch
from torch import optim
from torchvision import datasets, transforms
import torchvision.models as models 

path = "dogscats"

data_transforms = {
    'train': transforms.Compose([
        transforms.RandomResizedCrop(224),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
    'valid': transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
}
image_datasets = {x: datasets.ImageFolder(Path(path).joinpath(x),
                                          data_transforms[x])
                  for x in ['train', 'valid']}
dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'valid']}
class_names = image_datasets['train'].classes

モデル作成

今回はvgg16というモデルの学習済みのものを使用します。Pytorchのvisionライブラリに色々なモデルが用意されているので、他のモデルでも大丈夫です。
https://pytorch.org/docs/stable/torchvision/models.html

vgg16は、CNN部分とFCN部分とに分かれています。モデルを出力すると、それぞれfeaturesとclassifiersと表されています。流れとしては、

features部分はそのまま使うように学習をオフにする
classifiers部分の最後の層を置き換えて、犬猫の2値分類に変える
Softmax層を一番最後に追加
- Skorchのインターフェースに合わせるためで、Pytorchなら不要です

モデルをprintして構造を確認しながら行うと良いと思います。

vgg16 = models.vgg16(pretrained=True)
vgg16

# 置き換える層以外は再学習しないようにする
for param in vgg16.parameters():
    param.requires_grad = False

num_features = vgg16.classifier[6].in_features
modules = list(vgg16.classifier.children())
modules.pop()
modules.append(torch.nn.Linear(num_features, len(class_names)))
modules.append(torch.nn.Softmax(dim=1))
new_classifier = torch.nn.Sequential(*modules)
vgg16.classifier = new_classifier

学習部分

基本的にはskorchのチュートリアル通りです。validデータを使って、一番lossが小さいモデルを得られるようにしています。
注意点は、train_split=predefined_split(image_datasets['valid'])部分です。validデータを明示的に与えています。skorchでは、labelデータと画像データを別々に与えている場合、自動的にtrainとvalidを分けてくれます。
しかし今回はPytorchのImageFolderを使っているのでその機能が使えませんし、trainとvalidを自前で持っています。よってpredefined_splitというメソッドを使って対処しているということです。

from skorch.callbacks import LRScheduler, Checkpoint
from skorch.helper import filtered_optimizer
from skorch.helper import filter_requires_grad
from skorch.helper import predefined_split
from skorch import NeuralNetClassifier

lrscheduler = LRScheduler(
    policy='StepLR', step_size=7, gamma=0.1)

checkpoint = Checkpoint(
    f_params='best_model.pt', monitor='valid_loss_best')

callbacks = [lrscheduler, checkpoint]

optimizer = filtered_optimizer(
    optim.SGD, filter_requires_grad
)

net = NeuralNetClassifier(
    vgg16,
    lr=0.003,
    batch_size=32,
    max_epochs=10,
    optimizer=optimizer,
    optimizer__momentum=0.8,
    iterator_train__shuffle=True,
    iterator_train__num_workers=4,
    iterator_valid__shuffle=False,
    iterator_valid__num_workers=4,
    train_split=predefined_split(image_datasets['valid']),
    callbacks=callbacks,
    device='cuda'
)
net.fit(image_datasets['train'], y=None)

fitを叩くと、以下のような結果が順次更新されていきます。このようにわかりやすく表示されるのがskorchを使う利点の一つです。
また、Checkpointの引数でvalid_loss_bestを指定しているので、valid_lossが最も小さいepochのモデルが保存されます。

  epoch    train_loss    valid_acc    valid_loss    cp       dur
-------  ------------  -----------  ------------  ----  --------
      1        3.9178       0.5635        2.5179     +  167.7788
      2        2.6727       0.7093        1.7409     +  167.9802
      3        2.2109       0.7472        1.3766     +  167.8943
      4        1.9680       0.7862        1.1662     +  168.0836
      5        1.8513       0.7919        1.0464     +  167.9998
      6        1.7451       0.8035        0.9523     +  168.0205
      7        1.6439       0.8279        0.8749     +  168.0122
      8        1.6038       0.8363        0.8541     +  167.9641
      9        1.5811       0.8367        0.8488     +  168.0637
     10        1.5352       0.8390        0.8420     +  168.1118

推論

最後に、実際に作成したモデルの精度をtestデータでチェックしてみます。classifierを初期化して、上で学習したデータをロードしています。
predict_probaやpredictは、人によってはscikit-learnでお馴染みのメソッドかなと思います。

net = NeuralNetClassifier(
    vgg16,
    device='cuda'
)
net.initialize()
net.load_params('best_model.pt')

path = "/path/to/test_dataset"
data_transform = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])
dataset = datasets.ImageFolder(root=path + "test", transform=data_transform)

y_test = [record[1] for record in dataset]
np.mean(predicted == y_test)