背景

すでに chainer で組まれた機械学習アプリがある
モバイルで推論を動かしたい. tensorflow-lite ではモバイル対応がされている
- Caffe2 などもあるが, 2019 年 4 月時点で, CPU 実行でよいで production で安定して使えそうなのは tflite くらい?
- 以下にあるように ONNX 経由だと遅そうだし protobuf は好きくない.
一旦 ONNX に変換を試そうとしたが, モデルデータが 200 MB を超えるようであると, ONNX では読み込みや変換に数分かかったりで, デバッグが検証がやりずらい
- いずれにせよ, モバイルで動かすときにはモデルデータの読み込みは高速にしなければならない
- TensorFlow-Lite では flatbuffers でモデルデータをシリアライズしており, flatbuffers では読み込み(C/C++)が高速であることはすでに確認している
Chainer の全部の function/link をコンバートする必要はない

Chainer から tflite(Flatbuffers)に変換するコンバーターを書くのが手っ取り早いという判断になりました.

コンバーター自体は python で書きます.
(いずれ Chainer C/C++(?)が出てきたら, コンバーターの C++ 化も検討したい)

個々の op を動かすくらいまではできました.

Chainer

Chainer では, ネットワークの定義を python で, ネットワーク構造と重みのシリアライズを h5/npz で行います.
(ネットワークの定義をシリアライズする標準機能はあるのかな? pickle を使えばいいのかしら?)

onnx-chainer と, 優秀な機械学習若人さまが作成になられた chainer -> 未公開の機械学習ライブラリのコンバータを参考にして作成したいと思います.

重みは, h5 ではなく npz フォーマットで保存するものとします.
参考までに, npz は, 各 function/op のパラメータの重みをディレクトリとして作って npy 形式で保存し, zip でまとめていまず. npy 自体はたとえば CNPY などで C++ でお手軽に読むことができます > https://github.com/rogersce/cnpy

tflite schema

tflite は flatbuffers でフォーマット(スキーマ)の定義がされています.

じつのところ, 定義自体はそれほど複雑ではありません.
シリアライズや各種言語で扱いやすいように, glTF 2.0 のようにインデックスでのアクセスをベースにしていますので, 優秀な glTF 2.0 若人さまにおかれましては簡単にわかってしまいますね.
(欲を言えば, glTF みたいに, スキーマ定義だけでなくドキュメントもほしいですね)

以下で, tflite での主要なコンポーネントです.
naming は tflite schema 定義のものでは無く, python/C++ 生成コードのほうに合わせています.

Model

ルートとなる, モデルの定義です.

Subgraphs : グラフの構造があります.
OperatorCodes : Op(function/link. e.g. CONV_2D など) のリストを格納しています
Buffers : テンソルのデータ(Buffer)の配列.
- 0 番目の Buffer は空になります(Placeholder など, 実データの無いテンソルを定義したいときなど用に予約)

拡張データ格納用に Metadata が用意されています.

Operator codes

BuiltinCode : tflite で定義されている Op コード名になります(BuiltinOperator)

カスタム定義の場合は CustomCode に定義(string)があります.

SubGraph

Tensors : テンソル(重み)の配列があります
Inputs : 入力の op のインデックス
Outputs : 出力の op のインデックス
Operators : op のリスト

Operators

Chainer でいう function/link ですかね.

OpcodeIndex : Model にある OperatorCodes 配列へのインデックスになります.
Inputs : 入力の op のインデックスのリスト. SubGraph にある Tensor 配列へのインデックス. optional な input(たとえば FullyConnected の bias 項)の場合は, -1 を設定します.
Outputs 出力の op のインデックスのリスト. SubGraph にある Tensor 配列へのインデックス.
MutatingVariableInputs : T.B.W.
CustomOptions : custom
BuiltinOptionsType : op の option(パラメータ). たとえば Conv2DOptions.
BuiltinOptions : option(パラメータ)の定義.

BuiltinOptions は flatbuffers の Table 形式です.
ここから実際の options のクラスを生成したい場合は, 以下のようにします.

    if (op.BuiltinOptionsType() == tflite.BuiltinOptions.BuiltinOptions.Conv2DOptions):                                 
        assert(type(op.BuiltinOptions()) == flatbuffers.table.Table)                                                    
                                                                                                                        
        # Init from Table                                                                                               
        opt = tflite.Conv2DOptions.Conv2DOptions()                                                                      
        opt.Init(op.BuiltinOptions().Bytes, op.BuiltinOptions().Pos)                                                    
        print("padding ", opt.Padding())

Tensor

Name : 名前. ユニークであるのが望ましい.
Shape : shape(int array)
IsVariable : T.B.W.
Quantization : T.B.W.
Buffer : 実際のテンソルのデータがある Model の Buffers へのインデックス
Type : TesnsorType の id. e.g. 0 = float32

.tflite を, tensorflow-lite 提供の python の interpreter で読み込む場合は, name が設定されていないと Interpreter.get_input_details() あたりでエラーがでます. Tensor には, (ユニークな)名前を設定してシリアライズしましょう.

コンバート

基本的には Chainer のモデルをトラーバスし, それぞれ Chainer の function/link の定義と重みデータを tflite のスキーマに合わせてコンバートしていくのをひたらすらやっていけばよさそうなのがわかります.

Tensor を重みデータなどで埋める(最終的な input と output の部分は除く)
op(function/link)のパラメータを ***Optionsで埋める
op と Tensor との接続を埋める
SubGraph として書き出し

でいけそうです.

TFlite 側参考情報

tensorflow/tensorflow/python/ops/nn_ops.py で TensorFlow レイヤでの op 変換があります. たとえば dropout はここでいくつかの op に分解されます.

tensorflow/tensorflow/lite/kernel/** で Builtin op の実装があります.
tensorflow/tensorflow/lite/toco/import_tensorflow.cc で TensorFlow op の tflite コンバートの実装があります.

NCHW or NCHW?

tflite は現状 NHWC のみのサポートなります.
Chainer は NCHW になります. 最近 NHWC フォーマットも対応したようではあります.
必要に応じて transpose などで NCHW <-> NHWC 変換が必要になります.

padding

Chainer では pooling, conv function/link で padding 幅を設定することができますが, TensorFlow では, pooling, conv op では padding 幅をサポートしていません(TensorFlow での padding は, 'SAME' or 'VALID' で畳み込みの振る舞いを設定するだけ).

たとえば, Chainer で padding=1 で zero padding している場合は, 一旦 Pad op を挿入して padding させる必要があります. また, tflite 側での pooling, conv での padding パラメータ設定は VALID に設定します(ところで, tf の SAME/VALID という用語はわかりにくいですね).

疑問点として, tflite には, PAD と PADV2 と二つの pad op がありますが, 動作の違いが不明です.
(実装を見る感じでは PADV2 がメインで PAD は deprecated のように見える)

とりあえずは PADV2 を使っておけば問題が無いようには見えます.

Sigmoid

Chainer sigmoid は, tflite では LOGISTIC op を使います(動作は同じ).

Conv2D

tflite の CONV_2D は, filter の shape が [outC, kh, kw, inC] となっています(NHWC と合わせるため?). これは tf.nn.conv2d の filter の shape 定義とは異なりますので注意ください.

bias は optional ですが, ランタイム側で bias tensor を inputs に指定しないと現状エラーになるので,

work around として, zero valued tensor で bias を埋めておきます.

本来は tensor id -1 を指定して optional 扱いでも OK のはずですが,
現状ランタム側でエラーになります. これは https://github.com/tensorflow/tensorflow/issues/28306 で issue にあげています.

stride = 1. kernel width = 3 のときは tflite 側の CONV_2D で padding を SAME にすることで Pad op を挿入せずにすみます.
(stride>1 のときも SAME にできるかな?)

Deconvolution(Conv2DTranspose)

tflite では CONV_2D_TRANSPOSE は廃止?されて TRANSPOSE_CONV になっています.
ドキュメントでは CONV_2D_TRANAPOSE が残ったままになっているので注意しましょう.

Average pooling

pooling では, たとえば 2x2 filter size の場合でも, tf SAME としても結果が異なりますので, zero padding して対処します.

より具体的には, filter size 2x2 だと, Chainer では bottom と right の 1 pixel 分が増えます.
(e.g. [1, 1, 5, 5] だと, [1, 1, 6, 6] の shape のテンソルが生成される)

Max pooling

Max pooling では, -inf を padding する値に設定しないといけません(そうしないと負の値に対する max をあつかえない).
現状(r1.13.1)の Pad では, ドキュメントでは, padding の constant value を指定できないとありますが, 一応任意の constant value に対応しています(CPU 実行で確認).

padding の constant value は, 適当な tensor(1D でよい)を作って inputs に追加します(ドキュメントには未記載).

GPU 実行などで, ゼロ以外の定数の padding がうまく動かないときは, CONCATENATION を縦と横で二回呼ぶことで代用できそうです.

pooling op の output shape

tflite では, 実行時 pooling op の output shape の width と height は, input shape と op パラメータから算出されます.

tflite にシリアライズされた output shape は, 実行時に上書きされますので注意ください. tflite シリアライズ時に間違った output shape 情報(ComputePaddingHeightWidth で求まる値と異なる width, height)を書き込んでしまうと, 実行時に後続の op で pooling op を参照するところで shape が合わないエラーになります.
(netron などで shape 情報は期待する(output shape)のが出るのに, 実行時に shape エラーでなぜ? とこの原因を見つけるのに 5 営業日くらい消費してしまいました)

hstack, vstack

Pack と Concatenate で実現できそうです.

tflite での Pack のドキュメントが不明瞭ですが, ソースコードを見ると,

複数の inputs
1 個の output
inputs の個数は values_count parameter と同じにする
pack する軸は axis パラメータで指定

となります.

ResizeImages

Chainer の ResizeImages は bilinear 補間のみになります.

tflite で resize_bilinear + align_corner=True に変換することで, Chainer と tflite の結果が一致します.

ところで, 詳細は上にありますが, Chainer での bilinear の振る舞いには問題があります.
本来ですと tf 2.0(or pytorch 0.4+)のような補間をするのを新規で Chainer 側に追加するのがよいでしょう.

Unpooling2D

Chainer の Unpooling2D では, nearest neighbor で upsample で, 補間しての upsample には対応していないようです(補間する場合は deconvolution?)

RESIZE_NEAREST_NEIGHBOR にマッピングします.

LocalResponseNormalization

LOCAL_RESPONSE_NORMALIZAION にマッピングします.

TensorFlow では depth = C channel(NHWC) なので transpose しておきます(Chainer は NCHW).
また, Chainer では Normalization window width ですが, TensorFlow では radius なので半分に割ります.

DepthwiseConv2D

Chainer では, depwise_conv2d のビルトインはなくて, Weight tensor に変換をかけて, convolution2d で実現しています.
tflite では DEPTHWISE_CONV_2D op があります.

Chainer での DepthwiseConvolution2D を直接 tflite に変換できるとよいのですが, 現状では conv2d のパラメータなどを推定して, depthwise に変換する処理が必要になります.

tflite では, filter(weight) は, tf.nn.depthwise_conv2d の [y, x, in_channels, channel_multiplier]と異なり,

[1, y, x, out_channels] の shape を持ちます. このあたりドキュメントには記載されていませんので注意しましょう.
ここで, out_channels は in_channels * depth_multiplier です(channel_multiplier == depth_multiplier?).

toco のテストソースコードには [in_channels, y, x, out_channels] とあるのですが,
これだと input と convolve するときに shape が合いませんので, [1, y, x, out_channels] が正しいでしょう(test コードでは in_channels = 1 なのでテスト自体は通る).

kernel のコードを見ても filter[0] の値を取得しているところはありませんでした.

depth_multiplier は, DepthwiseConv2DOption のパラメータにも書き出しします. そうしないと tflite のランタイムチェックでエラーになります.

BatchNormalization, FixedBatchNormalization

tflite では BatchNorm は Builtin op にはありません.

また, 推論においてはこれは基本的に a * (b / sqrt(v)) の計算だけする実装が標準的のようです.

ここで, なんらかの方法で, パラメータは Chainer から与えられるものとします.
(一番よいのはネットワークの構築の時点で train のときだけ BatchNormalization を定義するようにして, 推論では使わない)

FusedBatchNorm

Chainer では BatchNorm で代用?
いずれにせよ Chainer, tflite 両方に定義はありません.

Dropout

tflite では乱数テンソル(RandomUniform) builtin op は無いため, 乱数は事前計算して定数テンソルにします. これにより deterministic な動作になります.

精度がどれくらい変わるかはモデルを動かしてみないとわからないかもですね. CPU で実行の場合は実際に推論時に乱数を生成してもよいかもしれません.

tensorflow では, dropout はいくつかの op に分解されます.

r1.13 までは, 本来の dropout の実装とはちょっと違い, floor を使い比較演算子を使わない方法でした.
https://github.com/tensorflow/tensorflow/blob/r1.13/tensorflow/python/ops/nn_ops.py#L3044

推論用ですでに rate が固定と仮定してよい場合は, 0, 1 のマスクテンソルの部分まで事前計算することができます.

r1.14 からは, 比較演算子を使う方法に戻って?います.
https://github.com/tensorflow/tensorflow/blob/r1.14/tensorflow/python/ops/nn_ops.py#L4221

推論の呼び出し回数のカウンターを元に, 事前計算した何個かの乱数テンソルを切り替えて, ある程度ランダムネスを持たせるという手もあるかもしれません.

RNN 系

RNN

RNN builtin op が Tensorflow Lite にあります.
仕様などドキュメントはありませんが, さほど難しくはありません(RNN 論文参照?)
tflite では, control flow を含む graph を RNN で実現したりしているようです.

Chainer には L.NstepRNNRelu などがあります. これもいくつかの function に分解されて実装されているようなので, いろいろうまく tflite に変換できればいけそうかもです.

tflite には BIDIRECTIONAL_SEQUENCE_RNN もあります.

GRU, StagelessGRU

Chainer には L.GRU, L.StatelessGRU, L.StatefullGRU があります.
Tensorflow, tflite では GRU builtin op はありません.

Chainer では sigmod, tanh などのいくつかの function に分解されますので, これを tflite にコンバートできればよさそうです.

LSTM

Chainer には L.LSTM があります.
TensorFlow では LSTM は直接の builtin op は無くて, いくつかの op で subgraph で実現されています.
(tf.rnn.contrib.BasicLSTMCell 参照)
tflite では LSTM builtin op が定義されています(ドキュメントは無い)
toco では, IdentifyLstmCell で, LSTM の subgraph を transform して tflite LSTM に変換しています.

したがって Chainer LSTM の仕様を tflite LSTM に変換できると完了かな?

また, tflite には UNIDIRECTIONAL_SEQUENCE_LSTM, BIDIRECTIONAL_SEQUENCE_LSTM もあります.

Where

T.B.W.

Flatbuffer

pip で入る flatbuffer 1.11 (2019 年 04 月 25 日時点. master branch も同様）では, python で file_identifier の書き出しに対応していません.

TensorFlow-Lite の実装側では, file_identifier(TFL3)をチェックするため, 上記 issue にあるような work around を行う必要があります.

Placeholder(input)

Shape の情報はあるが, Buffer は empty の Tensor を作ります.

C++ inteprter で,

tflite::PrintInterpreterState(interpreter.get());

で, Tensor の情報を dump できます.

Tensor  12 input                kTfLiteFloat32  kTfLiteArenaRw       3136 bytes

Placeholder(input) は kTfLiteArenaRw になっていることを確認しましょう.

データサイズが 0 の Buffer を作ると, これは empty と解釈されず, ゼロサイズのテンソルと見なされるようで, kTfLiteMmapRo とリードオンリーの Tensor が生成されるので注意ください. kTfLiteMmapRo のテンソルに値をセットしようとすると, seg fault します.

バイナリデータのシリアライズ

flatbuffers 自体には, 直接にはバイナリデータを処理する API はなさそうです.
Tensor のデータなどは, 以下のようにしてシリアライズします.

Faster Python Serialization
https://github.com/google/flatbuffers/issues/4668

            # Serialize tensor data: [uint8]                                                                                                                                   
            buffer_start = tflite.Buffer.BufferStartDataVector(                                                         
                self.builder, data_len)                                                                                 
                                                                                                                        
            # We need to seek the header to correct place before writing into                                           
            # Bytes array                                                                                               
            self.builder.head = self.builder.head - data_len                                                            
            self.builder.Bytes[self.builder.head: (                                                                     
                self.builder.head + data_len)] = data                                                                   
                                                                                                                        
            tf_data = self.builder.EndVector(data_len)

データを 1 byte　ずつ追加していくのに比べれば早いはず?

配列のシリアライズ

Flatbuffers では, データを prepend して追加(シリアライズ)していくため, 配列データは逆順にして追加するようにします.

また, 配列のメンバ変数を追加する場合は, まず配列のオブジェクトを作り, その後オフセット値を追加します.

配列データでは, primitive(e.g. int)型の場合は PrependInt32 などのメソッドを使い, table 型の場合は PrependUOffsetTRelative メソッドを使います.

以下に例を示します.

        # Serialize shape: [int32]                                                                                      
        tflite.Tensor.TensorStartShapeVector(self.builder, len(shape))                                                  
                                                                                                                        
        for i in reversed(range(len(shape))):                                                                           
            self.builder.PrependInt32(shape[i])                                                                         
        tf_shape = self.builder.EndVector(len(shape))  

        tflite.Tensor.TensorStart(self.builder)                                                                         
        tflite.Tensor.TensorAddShape(self.builder, tf_shape)   
        ...
        tflite.Tensor.TensorEnd(self.builder)

エンディアンについて

flatbuffers では, エンディアンを考慮しているとあります.

SPARC や POWER などのビックエンディアンアーキテクチャで推論を実行する場合, アプリ側でエンディアン変換を明示的に行わなくても動くかもですね.

シリアライズしたデータの確認

flatbuffers は, データの中身はさほどチェックしません. たとえば間違ったオフセットやインデックスがあるとおかしくな結果になります(実行時に out-of-bounds アクセスしてしまう, など).

一応 Verify 関数があるようですが, データの中身を詳しくまでは検証しないです.

tensoflow lite の python バインディングで, .tflite モデルの読み込みと推論実行を確認したり, インデックスが正しいか確認したり, C++ で推論する場合は各種 sanitizer(asan. ubsan, etc) などでの実行時チェックを強固にしましょう.

tflite で気になった点

MatMul op が無い. FullyConnected で代用(もしくは TRANSPOSE + MUL の組み合わせ?)

Python でコンバーターを書いた所感

機械学習のネットワーク(グラフ構造)は, シェーダ言語やコンパイラの AST ほど量は多くないのですが, 以外と自前でやることがあり(id の割り当てとか, グラフの最適化とか), めんどくさいことがわかりました.

Chainer の構造を一旦 TensorFlow protobuf text or freezed graph に変換し, toco で tflite に変換したほうが, テキスト処理で済むので楽かも? とはいえ, toco でうまく変換できないものが出てくるかもしれませんので, 直に書き出せる方法があったほうが良い時もありそうです. 併用するといいかもしれませんね.

おまけ: tensorflow-lite micro

TensorFlow-lite を, より実装を軽くして embedded デバイス向けにした micro が experimental であります.

softmax
fully-connected
depthwise convolution

くらいの対応になりますが, chainer で学習したモデルを IoT で動かせるかもですね.

TODO

LinearFunction + ReLU を FullyConnected(+ activation function を ReLU)と一つにまとめる
分岐やコントロールフローのあるモデルの書き出しに対応する(簡単なものであれば Select や Where op で行けるか?. TensorFLow 側での待ちのものものある https://www.tensorflow.org/lite/guide/roadmap)
tflite のスキーマを間借りして, 独自の推論実装に食わせたい
Chainer -> tflite -> TensorFlow で, TensorFlow で学習させる