第5回 Java WebアプリでTensorFlow（実装編） （2/4）

2) 学習済みモデルの書き出し

学習済みモデルをTensorFlow Servingに載せるには、SavedModel形式で書き出す必要があります。
第3回で作成したPythonコードに対して必要な変更は以下の2点です：

• Estimator作成時に書き出し用の情報を付加するようにする
• 実際に書き出すコードを追加する

(1) Estimator作成部分の変更

Estimator作成のための関数のうち最後に値を返す部分は、学習を行う場合と学習はせず利用のみする場合で条件分岐していました。
今回行いたい「書き出し」操作にも学習は必要ないので、学習はせず利用のみする場合の分岐を利用し、書き出し用の情報を付加するような変更を加えます。

<             return tf.estimator.EstimatorSpec(mode=mode, predictions=probabilities)
---
>             output = tf.estimator.export.ClassificationOutput(scores=probabilities)
>             return tf.estimator.EstimatorSpec(mode=mode, predictions=probabilities,
>                                               export_outputs={'result': output})

上記の通り、EstimatorSpecオブジェクトの作成時にClassificationOutputを渡すように変更しています。これにより、このモデルがデータの分類(Classification)タスクであり、出力としてprobabilities変数の指すもの、つまり犬猫それぞれの確率を返すということを定義しています。
この定義はTensorFlow Servingで学習済みモデルを動作させる際に利用されます。

上記の変更を適用すると、mydnn_tf.pyは以下のようになります。

import tensorflow as tf
 
def MyDNN(input_shape=(32, 32, 1), output_size=10, learning_rate=0.001,
          keep_prob=0.5, model_dir='tfmodel'):
    def mydnn_fn(features, labels, mode):
        input_layer = tf.reshape(features["img"], [-1] + list(input_shape))
 
        layer = tf.layers.conv2d(filters=20, kernel_size=5, strides=2,
                                 activation=tf.nn.relu,
                                 inputs=input_layer)
        layer = tf.layers.max_pooling2d(pool_size=3, strides=2,
                                        inputs=layer)
        layer = tf.layers.batch_normalization(inputs=layer)
          
        layer = tf.layers.conv2d(filters=50, kernel_size=5, strides=2,
                                 activation=tf.nn.relu,
                                 inputs=layer)
        layer = tf.layers.max_pooling2d(pool_size=3, strides=2,
                                        inputs=layer)
        layer = tf.layers.batch_normalization(inputs=layer)
      
        layer = tf.contrib.layers.flatten(layer)
        layer = tf.layers.dense(units=100, activation=tf.nn.relu,
                                inputs=layer)
        layer = tf.layers.dropout(rate=(1 - keep_prob),
                                  training=(mode == tf.estimator.ModeKeys.TRAIN),
                                  inputs=layer)
        output_layer = tf.layers.dense(units=output_size, 
                                       inputs=layer)
 
        # for prediction
        if mode == tf.estimator.ModeKeys.PREDICT:
            probabilities = tf.nn.softmax(output_layer)
            output = tf.estimator.export.ClassificationOutput(scores=probabilities)
            return tf.estimator.EstimatorSpec(mode=mode, predictions=probabilities,
                                              export_outputs={'result': output})
 
        # for training
        labels = tf.reshape(labels, [-1, output_size])
        loss = tf.losses.softmax_cross_entropy(onehot_labels=labels,            
                                               logits=output_layer)
        optimizer = tf.train.AdamOptimizer(learning_rate, epsilon=1e-1)       
        train_op = optimizer.minimize(loss=loss, global_step=tf.train.get_global_step())
 
        # for evaluation
        classes = tf.argmax(input=output_layer, axis=1)
        eval_metric_ops = {
          "accuracy": tf.metrics.accuracy(labels=tf.argmax(labels, axis=1),
                                          predictions=classes)
        }
 
        return tf.estimator.EstimatorSpec(mode=mode, loss=loss,
                                          train_op=train_op,
                                          eval_metric_ops=eval_metric_ops)
 
 
    estimator = tf.estimator.Estimator(model_fn=mydnn_fn, model_dir=model_dir)
 
    return estimator

(2) 実際に書き出すコードの追加

モデルを書き出すには、Estimatorオブジェクトのexport_savedmodelメソッドを呼び出します。
第1引数は保存先ディレクトリ、第2引数は入力データの受け取り方を定義する関数です。

    def input_receiver_fn():
        # ここでtf.estimator.export.ServingInputReceiverを作成して返す
 
    model.export_savedmodel(export_dir, input_receiver_fn)

このコードでは第2引数にわたす関数としてinput_receiver_fnを定義しています。
最終的にはtf.estimator.export.ServingInputReceiverオブジェクトを生成して返すのがこの関数の役目です。
ServingInputReceiverは名前の通りTensorFlow Servingに送られてきた入力データの処理の仕方を定義するためのものです。

では、input_receiver_fnの中身を見てみましょう。

    def input_receiver_fn():
        input_example = tf.placeholder(dtype=tf.string)
        input_size_xyc = input_size_x * input_size_y * input_size_c
        feature_spec = {'img' : tf.FixedLenFeature(dtype=tf.float32,
                                                   shape=(input_size_xyc))}
        features = tf.parse_example(input_example, feature_spec)
        receiver_tensors = {'examples': input_example}
        return tf.estimator.export.ServingInputReceiver(features, receiver_tensors)

ServingInputReceiverのイニシャライザの第1引数featuresは、クライアントから送られてくるデータを受け止めてEstimatorで利用可能な形に加工したものです
（※ 正確には、そのような加工処理を行うTensorFlowの計算工程を示すものです）。
クライアントからは文字列の形にシリアライズされたデータが送られてきますので、それをうまく構造化データにもどしてやる必要があるわけです。

受け取った文字列はまず文字列形式のplaceholderであるinput_exampleで受け止めることとします。
そしてその文字列をtf.parse_example関数に渡してパースします（ちなみにexampleという名前は、TensorFlow Servingのサーバ／クライアント間でデータの受け渡しに使われるクラスtf.train.Exampleに由来しています）。
parse_exampleの第2引数として与えているfeature_specは、受け取るデータは固定長のfloatデータであり、Estimatorに渡すためにimgというキーを添えることを指定しています。
parse_exampleの返り値としてfeaturesが得られます。

ただし、このままでは、input_exampleは数多くあるplaceholderのひとつに過ぎませんので、クライアントから受け取ったデータをinput_exampleで受け止めたいということを明示的に伝える必要があります。
それを行っているのがServingInputReceiverのイニシャライザの第2引数receiver_tensorsです。

全体として、cat_dog_dnn_tf.pyのコードは以下のようになります。
既存の「学習モード」「判定モード」に加えて「書き出しモード」を選択できるようにし、上記の書き出し部分のコードを追加しています。

import argparse
import tensorflow as tf
import os, sys
import numpy as np
from datetime import datetime
from mydnn_tf_layers import MyDNN
 
ImageDataGenerator = tf.keras.preprocessing.image.ImageDataGenerator
load_img = tf.keras.preprocessing.image.load_img
img_to_array = tf.keras.preprocessing.image.img_to_array
 
input_size_x = 224
input_size_y = 224
batch_size = 20
input_size_c = 3
output_size = 2
model_dir = 'ckpt/tf/'
 
export_dir = 'tf_savedmodels/'
 
parser = argparse.ArgumentParser()
parser.add_argument("--export", action="store_true", default=False,
                    help="学習済みモデルを本番向けに書き出す")
parser.add_argument("--infer", action="store", nargs="?", type=str,
                    help="学習は行わず、このオプションで指定した画像ファイルの判定処理を行う")
parser.add_argument("--epochs", action="store", nargs="?", default=10, type=int,
                    help="学習データ全体を何周するか")
args = parser.parse_args()
 
epochs = args.epochs
 
model = MyDNN(input_shape=(input_size_x, input_size_y, input_size_c),
              output_size=output_size, model_dir=model_dir)
 
train = not (args.infer or args.export)
if train:
    print("学習モード")
    # 学習データの読み込み
    keras_idg = ImageDataGenerator(rescale=1.0 / 255)
    train_generator = keras_idg.flow_from_directory('data/train',
                          target_size=(input_size_x, input_size_y),
                          batch_size=1,
                          class_mode='categorical',
                          shuffle=True)
    valid_generator = keras_idg.flow_from_directory('data/valid',
                          target_size=(input_size_x, input_size_y),
                          batch_size=1,
                          class_mode='categorical')
 
    # 学習の実行
    num_data_train_dog = len(os.listdir('data/train/dog'))
    num_data_train_cat = len(os.listdir('data/train/cat'))
    num_data_train = num_data_train_dog + num_data_train_cat
 
    num_data_valid_dog = len(os.listdir('data/valid/dog'))
    num_data_valid_cat = len(os.listdir('data/valid/cat'))
    num_data_valid = num_data_valid_dog + num_data_valid_cat
 
    steps_per_epoch = num_data_train / batch_size
    validation_steps = num_data_valid / batch_size
 
    def my_input_fn(generator):
        gen_fn = lambda: generator 
        dataset = tf.data.Dataset.from_generator(gen_fn, (tf.float32, tf.float32))
        dataset = dataset.map(lambda f, l: ({"img": f}, l)).batch(batch_size)
        return dataset.make_one_shot_iterator().get_next()
 
    for epoch in range(epochs):
        print("epoch ",  epoch)
        print("training...")
        model.train(input_fn=lambda: my_input_fn(train_generator),
                    steps=steps_per_epoch)
        print("evaluation:")
        eval_results = model.evaluate(input_fn=lambda: my_input_fn(valid_generator),
                                      steps=validation_steps)
        print(eval_results)
 
elif args.infer:
    print("判定モード")
    # 判定する画像の読み込み
    image_infer = load_img(args.infer, target_size=(input_size_x, input_size_y))
    data_infer = img_to_array(image_infer)
    data_infer = np.expand_dims(data_infer, axis=0)
    data_infer = data_infer / 255.0
    print(data_infer.shape)
  
    # 判定処理の実行
    predict_input_fn = tf.estimator.inputs.numpy_input_fn(x={"img": data_infer},
                                                          batch_size=1,
                                                          shuffle=False)
    result_generator = model.predict(predict_input_fn) 
    result = next(result_generator) * 100
 
    # 判定結果の出力
    if result[0] > result[1]:
        print('Cat (%.1f%%)' % result[0])
    else:
        print('Dog (%.1f%%)' % result[1])
 
else:
    print("書き出しモード")
    def input_receiver_fn():
        input_example = tf.placeholder(dtype=tf.string)
        input_size_xyc = input_size_x * input_size_y * input_size_c
        feature_spec = {'img' : tf.FixedLenFeature(dtype=tf.float32,
                                                   shape=(input_size_xyc))}
        features = tf.parse_example(input_example, feature_spec)
        receiver_tensors = {'examples': input_example}
        return tf.estimator.export.ServingInputReceiver(features, receiver_tensors)
 
    model.export_savedmodel(export_dir, input_receiver_fn)