nb.oの日記: Edge TPUで独自モデルを作る（その１）

Edge TPUで独自のモデルを作った（MNIST Edge TPU）ときにいろいろわかった点、つまずいた点をまとめてみたい。ここに記載のコードはMNIST Edge TPUのものです。
（これらは自分の理解です。間違っていたらすみません。）

まずは、学習モデルの構築〜FreezeGraphまでを説明。
その２はこちら。

Kerasではだめなの？

現時点:2019/5/31でKerasはQuantization-aware trainingをサポートしていない #27880。
Issueでは2019の2Q Endに対応するとあるので、気長に待つしかない。

Edge TPU Modelを作るまでの流れ

学習モデルの構築からEdge TPU Modelを生成する流れを図にすると以下になる。

モデルの構築の流れ

それぞれ生成されるファイルの内容は

CheckPoints:
学習した結果の重みが保存されてるファイル。
GraphDef:
モデルの構造のみが記録されたファイル。
FreezeGraph:
CheckPoints＋GraphDef。
モデルの構造と学習した結果の重みを保存する。
ProtocolBuffers形式である。
TF-Lite Model:
量子化したTF-Liteモデル。TensorFlow Lite Converterで生成。
Edge TPU Model
Edge TPU で動作するモデル。Edge TPU Compilerで生成。

Quantization-aware training

Edge TPUモデルを生成する場合は"量子化学習"が必要になる。
詳細は、Quantization-aware trainingを参照。

学習モデルと推論モデルの構築時、モデル（Graph）を書き換える（量子化を意識したノードを追加）ためにAPIを呼び出す必要がある。
モデルを書き換える前に呼び出しても意味がない。

学習モデルの構築時

学習モデルの構築は通常と何ら代わりはない。量子化学習を行う場合は、tf.contrib.quantize.create_training_graphでモデルの書き換えを行う。
これにより、量子化学習に必要なノードがモデルに追加される。

loss関数を指定した直後（tf.losses.sparse_softmax_cross_entropyなど)に指定を行う。
（勘違いかもしれないが、training Opsの指定後だとうまく変換できなかった？）
引数のquant_delayは量子化学習を開始するstep数を指定する。
通常、学習が十分学習が進んだころ（曲線が寝たぐらい）から開始するのがよいとある。または、すでに通常の学習モデルがある場合は、同じデータセットで最初から（quant_delay=0）である程度の学習を行ってもよいと思われる（Fine-tuningのようなもの？DeepLabはquant_delay=0でPASCAL-VOCの再学習を行っている）。

    with tf.Graph().as_default():
        # Generate placeholders for the images and labels.
        images_placeholder, labels_placeholder = placeholder_inputs(batch_size)

        # Build a Graph that computes predictions from the inference model.
        logits = mnist_inference(images_placeholder, hidden1, hidden2)

        # Add to the Graph the Ops for loss calculation.
        loss = mnist_loss(logits, labels_placeholder)
    
        # Quantization-aware training
        if quantize_delay >= 0:
            print('tf.contrib.quantize.create_training_graph')
            tf.contrib.quantize.create_training_graph(quant_delay=quantize_delay)
            
        # Add to the Graph the Ops that calculate and apply gradients.
        train_op = mnist_training(loss, learning_rate)

推論モデルの構築

推論モデルも同様にモデル構築した直後に指定を行う。

tf.contrib.quantize.create_eval_graphでモデルの書き換えを行う。

with tf.Graph().as_default() as graph:
    # Generate placeholders for the images
    placeholder = tf.placeholder(name='input_tensor', dtype=tf.float32, shape=(1, IMAGE_PIXELS))
    
    # Build a Graph that computes predictions from the inference model.
    logits = mnist_inference(placeholder, hidden1, hidden2)

    tf.nn.softmax(logits, name="softmax_tensor")

    # Rewrites an eval input_graph in place for simulated quantization.
    if quantize_delay >= 0:
        tf.contrib.quantize.create_eval_graph()

推論モデルを構築するときに注意すること

あまりサンプルもなくちょっと戸惑ったところ（Checkpointがあれば推論はできるので問題ないから？）。

基本的には学習モデルを構築したコードを流用すればよい。以下の点に注意。

tf.constantでモデルを構築しない。
学習モデルと同じようにtf.Variableで構築する。
tf.constantで構築すると、create_eval_graphで量子化モデルに書き換わらない。
学習モデルでは、活性化関数とloss関数がセットになったAPIを指定している
例）
TensorFlowのMNISTの場合、tf.losses.sparse_softmax_cross_entropyを使っている。この場合、活性化関数はsoftmaxなので、推論モデルではtf.nn.softmaxを指定する。
プレースホルダー、活性化関数（出力）には名前をつける。
FreezeGraph, TF-Liteへの変換の際、必要になるので予め指定しておいたほうがよい。

FreezeGraph

freeze_graphコマンドでCheckPointsとGraphDefをマージする。

基本的に量子化学習に関係することはなく、ほとんど通常と同じである。
freeze_graphコマンドはPython APIも用意されている。
DeepLabのdeeplab/export_model.pyはfreeze_graph_with_def_protosを使ってFreeze Graphを行っている（ただ、APIリファレンスはない？）。

（その2に続く）

nb.oの日記

2019年6月1日土曜日

Edge TPUで独自モデルを作る（その１）