TL;DR

PyTorchで学習したモデルをTFLiteのモデルに変換して、さらに量子化して、モバイルデバイス上で使えるようにする方法
について書きます。

前置き

PyTorchはモデルの記述のしやすさや学習の速さなどから、多くのディープラーニングモデルの学習に使われるようになってきています。私も何かディープでPoCしようと思ったらPyTorchが第一選択です。また、少し前まではデプロイにおいての優位性という点でTensorFlowに分があったりしましたが、最近は onnxruntime といったONNX形式のモデルをservingできる選択肢も現れ始め、プロダクションレベルでもPyTorchを学習からデプロイまで使い倒せるようになってきました。

そんな中、TensorFlowが明らかにPyTorchに対して優位な点として、以下のようなものが挙げられます。

TPUを使って学習できる
TFLite や TensorFlow.js など、フロントエンドでのデプロイを容易にするAPIが充実している

特に、PyTorchは後者が弱いです。 Caffe2 のモデルにONNX形式を介して変換すれば一応モバイル環境でも使うことができますが、公式のドキュメントが充実しておらず、モデルの量子化も簡単には行えません。

しかし、やはりPyTorchで学習したモデルをそのままモバイルデバイスで使いたい時ってあると思います。そういう時に、PyTorchで学習したモデルをTFLiteのモデルに変換する方法があると嬉しいですよね。というわけで、今回はそれについて現時点（2019/7/7）でわかっている一つの方法を紹介します。

環境やバージョン

Python 3.6.8
使用するPythonライブラリ:
- Keras==2.2.4
- onnx==1.5.0
- onnx2keras==0.0.3
- tensorflow==1.14.0
- torch==1.1.0

やり方

大まかな流れを説明すると、以下のような感じになります。（ここでは、画像処理系のモデルを想定しています。現時点では、モバイルでは画像系のモデルを使うことがほとんどなんじゃないかと思います。）

PyTorchで、いい感じのモデルを定義して学習する
学習済みモデルをONNX形式に変換して保存する
ONNX形式のモデルをKerasのモデルに変換する
Kerasのモデルを、channel_first形式からchannel_last形式に変換する
Kerasのモデルをsaved_model.pbとして保存する
tflite_convertでその保存されたモデルをTFLiteモデルに変換し、さらに量子化して保存する
変換前と後のモデルの出力を比べる

なぜこんな回り道をしないといけないかと言うと、主に下記のような理由によってです。

PyTorchのモデルはNCHWというchannel_first形式でモデルが定義されるが、TFLiteはNHWCというchannel_last形式のモデルじゃないとそもそも変換できない（NCHWに対応するオペレーションが定義されてない）
ONNXからTensorFlowモデルへの変換はonnx-tfを使うと楽だが、TensorFlowのモデル（つまり計算グラフ）をNHWCのモデルに変換するのはとてつもなく骨が折れる
KerasのモデルでNHWCに変換するのは比較的楽であり、onnx2kerasを使えば変換できるが、onnx-tfと比べるとまだサポートされていないPyTorchのオペレーションが結構あるため、いい感じにモデルを最初に定義しないと変換できない

とてつもなく遠回りなので、もっと楽な方法を見つけた方がいらっしゃったら教えてください🥺

MobileNetV3Largeの学習済みモデルをTFLiteに変換する

では実際に具体的な方法について書いていきます。折角なので、VGGみたいな実際には使い物にならないモデルではなく、もっと現実的なケースとしてモバイルでの利用を見据えて、今回はMobileNetV3Largeというモデルの変換を行っていきます。下記レポジトリの作者さんが親切にもモデルの実装と学習済みモデルの重みを置いておいてくれています。これを使っていきます。

また、今回の実装はすべてこのレポジトリにまとめています。python main.pyでとりあえず動きます。

1. PyTorchでいい感じのモデルを定義して学習する

今回は学習済みモデルを使うので学習部分は割愛しますが、このMobileNetV3Largeは、そのままではKerasのモデルに変換できません。詰まりポイントは以下の2点です。

AdaptiveAvgPool2dが各所に使われている
SELayer内のtensor.view() によるreshapeがなぜかKerasでのオペレーションに変換できない

これ、この2つともonnx-tfを使えば普通に大丈夫なんですが、onnx2kerasでは引っかかって変換できません。それぞれ、以下のようなワークアラウンドを入れてやる必要があります。

AdaptiveAvgPool2dを、入力tensorの(H, W)をちゃんと計算して、AvgPool2Dに置き換える
SELayerを、kernelサイズ1x1のConv2dによってDenseを置き換えることでreshapeをしなくてよくするように書き換える

SELayerはこんな感じになります。poolという新しく追加されたパラメータに、入力されるtensorの画像サイズを入れる形になります。まずは学習済みのweightをちゃんとロードするために、レイヤーにはまだ手を加えず、forward()の中身だけ変えておきます。

class SELayer(nn.Module):
    def __init__(self, channel, pool, reduction=4):
        super(SELayer, self).__init__()
        self.avg_pool = nn.AvgPool2d(pool, stride=1)
        self.fc = nn.Sequential(
                nn.Linear(channel, channel // reduction),
                nn.ReLU(inplace=True),
                nn.Linear(channel // reduction, channel),
                h_sigmoid()
        )

    def forward(self, x):
        # b, c, _, _ = x.size()
        # y = self.avg_pool(x).view(b, c)
        # y = self.fc(y).view(b, c, 1, 1)
        y = self.avg_pool(x)
        y = self.fc(y)
        return x * y

また、SELayerのDenseをConv2dに置き換えてさらに元のweightを適切に戻しておかないといけないので、下記のような後処理を入れます。

"""
MobileNetV3クラスのクラスメソッドとして実装
"""
def convert_se(self):
    for m in self.modules():
        if isinstance(m, SELayer):
            fc = m._modules['fc']
            linear_1, linear_2 = fc[0], fc[2]

            b_1, c_1 = linear_1.weight.size()
            b_2, c_2 = linear_2.weight.size()

            conv_1 = nn.Conv2d(c_1, b_1, 1, 1, 0)
            conv_1.weight.data = linear_1.weight.view(b_1, c_1, 1, 1)
            conv_1.bias.data = linear_1.bias

            conv_2 = nn.Conv2d(c_2, b_2, 1, 1, 0)
            conv_2.weight.data = linear_2.weight.view(b_2, c_2, 1, 1)
            conv_2.bias.data = linear_2.bias

            fc[0] = conv_1
            fc[2] = conv_2

# -------------------------------------------------------------------- #

model = mobilenetv3_large()

state_dict = torch.load('./data/pretrained/mobilenetv3-large-657e7b3d.pth', map_location='cpu')
model.clean_and_load_state_dict(state_dict)

model.convert_se()

上記コードのmodel.clean_and_load_state_dict(state_dict)は何をしているかというと、上の学習済みモデルの重みのstate_dictと実際のモデルにおいて、一部モジュール名が異なっていてstate_dictが読み込めないというバグに対処しています笑具体的には、こんな感じの実装になっています。

"""
MobileNetV3クラスのクラスメソッドとして実装
"""
def clean_and_load_state_dict(self, state_dict):
    state_dict['classifier.0.weight'] = state_dict['classifier.1.weight']
    state_dict['classifier.0.bias'] = state_dict['classifier.1.bias']
    state_dict['classifier.3.weight'] = state_dict['classifier.5.weight']
    state_dict['classifier.3.bias'] = state_dict['classifier.5.bias']

    del state_dict['classifier.1.weight']
    del state_dict['classifier.1.bias']
    del state_dict['classifier.5.weight']
    del state_dict['classifier.5.bias']
    self.load_state_dict(state_dict)

最終層のAdaptiveAvgPool2dも忘れずに置き換えておきましょう。これでなんとかonnx2kerasでKerasモデルへ変換できるようになりました。

2. 学習済みモデルをONNX形式に変換して保存する

これは超簡単です。これだけ。

# モデルの初期化と学習済みweightの読み込み、モデルの整形
model = mobilenetv3_large()

state_dict = torch.load('./data/pretrained/mobilenetv3-large-657e7b3d.pth', map_location='cpu')
model.clean_and_load_state_dict(state_dict)

model.convert_se()

# ONNX形式でのモデルの保存
onnx_model_path = './data/model.onnx'
dummy_input = torch.randn(1, 3, 224, 224)
input_names = ['image_array']  # ここで指定する名前が、後々saved_model.pbにする時のinput名になる
output_names = ['category']    # ここで指定する名前が、後々saved_model.pbにする時のoutput名になる

torch.onnx.export(model, dummy_input, onnx_model_path,
                  input_names=input_names, output_names=output_names)

3, 4, 5. ONNX形式のモデルをKerasのモデルに変換する　＆　Kerasのモデルを、`channel_first`形式から`channel_last`形式に変換する　＆　Kerasのモデルを`saved_model.pb`として保存する

onnx2kerasを使います。このレポジトリのドキュメントを見れば結構簡単そうですが、ちょくちょく詰まりどころがあります。具体的には、変換後、saved_model.pbにする時につまづくことが多いです。現時点で一応動くコードを下記に載せておきます。

import onnx
from onnx2keras import onnx_to_keras
import tensorflow as tf
from tensorflow.python.keras import backend as K


sess = tf.Session()
K.set_session(sess)

onnx_model = onnx.load('./data/model.onnx')

input_names = ['image_array']
# change_ordering=True で NCHW形式のモデルをNHWC形式のモデルに変換できる
k_model = onnx_to_keras(onnx_model=onnx_model, input_names=input_names,
                        change_ordering=True, verbose=False)

# 後々weightを再ロードするためにとっておく
weights = k_model.get_weights()

# saved_model.pbへと保存する
K.set_learning_phase(0)
with K.get_session() as sess:
    # FailedPreconditionErrorを回避するために必要
    init = tf.global_variables_initializer()
    sess.run(init)

    # このままだとweightの情報が消えているのでweightを再ロード
    k_model.set_weights(weights)

    tf.saved_model.simple_save(
        sess,
        str(saved_model_dir.joinpath('1')),
        inputs={'image_array': k_model.input},
        outputs={'category_id': k_model.output}
    )

上記でtf.global_variables_initializer()を使って変数を初期化しているのは、これをしないとsaved_model.pbへと保存する際にFailedPreconditionErrorという謎のエラーを吐いて落ちてしまうからです。エラーログを見てみるとどうやらconvolution層のweightが初期化されてねーぞって感じで怒られているのですが、ONNXモデルからweightはちゃんと移されています。（現に、k_model.predict(x)は動き、かつその出力はONNXモデルに同じ入力を与えた時の出力と一致する）

謎挙動ですが、内部実装を見ると、tf.keras.layers.Conv2dを初期化する際に本来は渡されることのないweightを渡していて、おそらくこれが原因でweightが初期化されてない扱いになってしまうのだと思います。（適当
これはとりあえずのワークアラウンドですが、tf.global_variables_initializer()を使っても出力は変化しないことは確かめられるので、ひとまずは大丈夫でしょう。一つ注意なのが、これで変数を初期化すると全weightが初期化された状態でsaveされてしまいます。ので、元のモデルの重みを一度取り出しておいて、初期化後に再読み込みするとかいう手間を取らないといけないのです。

6. `tflite_convert`でその保存されたモデルをTFLiteモデルに変換し、さらに量子化して保存する

TFLite はPyThon APIが提供されており、 tf.liteという形でアクセスできます。公式のドキュメントに従って、下記のように変換、量子化を行うことができます。

注意：ここで行っている量子化はtensorflow==1.14.0からしか対応していません！！

import tensorflow as tf


saved_model_dir = './data/saved_model_dir'
converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)
converter.optimizations = [tf.lite.Optimize.OPTIMIZE_FOR_SIZE]  # 量子化する時のみ
tflite_quant_model = converter.convert()

with open('../data/models/category_lens.tflite', 'wb') as f:
    f.write(tflite_quant_model)

7. 変換前と後のモデルの出力を比べる

当たり前なことですが、変換前のPyTorchのモデルの出力と変換後のTFLiteのモデルの出力を比べて、整合することをちゃんと確かめましょう。上記で量子化していない場合、何か適当な画像を入れた場合の出力はほぼ一致します。量子化した場合は、いくつか推論させて制度が大きく下がっていないかを確認してください。

まとめ

たぶん、もう少し時間が経てばこの辺はもっと簡単になると思います笑あくまで今回は、現状において、PyTorchのモデルをどうやってTFLiteのモデルに変換すればいいかの一つの方法をお伝えしました。個人的にはモデルの作成から学習までPyTorchが圧倒的に使いやすくて最高だと思っていて、TFLiteで使いたいがためだけにKerasやTensorFlowで実装・学習し直すのが癪だっただけです←

torchvisionに実装されているような大半の画像モデルはconvertできるはずなので、ぜひ皆様も試してみてください。

PyTorchで学習したモデルをTFLiteモデルに変換して使う