Dataset を使った TensorFlow のモデルを SavedModel を使用して保存・復元する
はじめに
TensorFlow の tf.data.Dataset を使用したモデルを SavedModel 保存・復元する方法について説明したいと思います。
保存したモデルを C++ で復元・利用する方法についてはこちらで説明しています。
本記事で説明しているソースコードは こちら に載せてあります。
実行環境
- Windows
- Python 3.6
- tensorflow 1.10.0
モデルの学習から保存まで
はじめに学習用データとその正解ラベル、検証用のデータを作成します。
今回は y = 2 * x を学習してみることにします。
# 学習用データ
train = np.array(np.random.sample((100, 1)) * 10, dtype='float32')
train_label = 2 * train
# 検証用データ
test_data = np.array([[1.0], [2.0], [3.0]], dtype='float32')
test_label = test_data
実行時に指定されるデータのための placeholder を作成します。
後で利用しやすいように、各々に name も付けておきます。
batch_size = tf.placeholder(tf.int64, name='batch_size')
x = tf.placeholder(tf.float32, shape=[None, 1], name='input')
y = tf.placeholder(tf.float32, shape=[None, 1], name='target')
上記の placeholder を用いて Dataset と Iterator を生成します。
Iterator が返すデータ自体は改めて指定するので、ここでは tf.data.Iterator.from_structure を使用して形だけを作成しておきます。
dataset = tf.data.Dataset.from_tensor_slices((x, y)).batch(batch_size).repeat()
iter = tf.data.Iterator.from_structure(dataset.output_types, dataset.output_shapes)
dataset_init_op = iter.make_initializer(dataset, name='dataset_init')
上記の Iterator からモデルへの入力となるデータ(tensor)を取得します。
x_, y_ = iter.get_next()
上記のデータを使用してモデルを作成します。
# make a simple model
prediction = tf.identity(tf.layers.dense(x_, 1), name='output')
# Optimize loss
loss = tf.reduce_mean(tf.square(prediction - y_), name='loss')
optimizer = tf.train.GradientDescentOptimizer(learning_rate=0.01)
train_op = optimizer.minimize(loss, name='train')
ここから学習を開始していきます。初めに dataset_init_op を実行して、学習に使用するデータを指定します。
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
sess.run(dataset_init_op, feed_dict={ x: train, y: train_label, batch_size: BATCH_SIZE })
次に、実際の学習を進めていきます。学習量データは入力済みのため、ここでは学習の実行(train_op)と loss の取得(loss)のみ run の引数で指定しています。
for i in range(EPOCHS):
_, loss_value = sess.run([train_op, loss])
print("Iter: {}, Loss: {:.4f}".format(i, loss_value))
学習が完了したら dataset_init_op を再度実行して、検証用のデータを指定します。
更に prediction を実行して検証結果を取得・表示します。
sess.run(dataset_init_op, feed_dict={ x: test_data, y: test_label, batch_size: test_data.shape[0] })
results = sess.run(prediction)
for i in range(len(results)):
print("{} {}".format(test_data[i], results[i]))
最後に tf.saved_model.simple_save を使用して model というフォルダにモデルを保存します。
inputs_dict = {
"x": x,
"y": y,
"batch_size": batch_size
}
outputs_dict = {
"output": prediction
}
tf.saved_model.simple_save(sess, './model', inputs_dict, outputs_dict)
上記を実行すると以下のような結果が得られます。
Iter: 0, Loss: 81.2936
Iter: 1, Loss: 11.0675
Iter: 2, Loss: 0.0215
Iter: 3, Loss: 0.0091
Iter: 4, Loss: 0.0231
Iter: 5, Loss: 0.0119
Iter: 6, Loss: 0.0125
Iter: 7, Loss: 0.0121
Iter: 8, Loss: 0.0144
Iter: 9, Loss: 0.0160
[1.] [2.1919134]
[2.] [4.1599956]
[3.] [6.1280775]
モデルの復元と推論
学習時と同様に入力データを生成します。
test_data = np.array([[2.0], [3.0], [4.0]], dtype='float32')
test_label = test_data
model フォルダからモデルのデータを復元します。
graph = tf.Graph()
with graph.as_default():
with tf.Session(graph=graph) as sess:
tf.saved_model.loader.load(
sess,
[tag_constants.SERVING],
'./model'
)
次に、iterator を初期化する dataset_init をグラフから取得します。
dataset_init_op = graph.get_operation_by_name('dataset_init')
この dataset_init と上記の入力データを使用して iterator の初期化を行います。
sess.run(dataset_init_op, feed_dict={
'input:0': test_data,
'target:0': test_label,
'batch_size:0': test_data.shape[0]
}
)
最後に推論を行い、結果を取得します。
results = sess.run('output:0')
for i in range(len(results)):
print("{} {}".format(test_data[i], results[i]))
実行すると以下のような結果が得られます。
[2.] [4.1599956]
[3.] [6.1280775]
[4.] [8.09616]
ハマりどころ1
上記の結論に達するまでに、下記の例外が発生する問題がなかなか解決できずに苦労しました。
FailedPreconditionError (see above for traceback): GetNext() failed because the iterator has not been initialized
以下のようなコードを書いていたところ
test_data = np.array([[1.0], [2.0], [3.0]], dtype='float32')
test_label = test_data
dataset = tf.data.Dataset.from_tensor_slices((test_data, test_label))
iter = dataset.make_initializable_iterator()
x_, y_ = iterator.get_next()
graph = tf.Graph()
with graph.as_default():
with tf.Session(graph=graph) as sess:
tf.saved_model.loader.load(
sess,
[tag_constants.SERVING],
'./model'
)
sess.run(iter.initializer, feed_dict={
'input:0': test_data,
'target:0': test_label,
'batch_size:0': test_data.shape[0]
}
)
results = sess.run('output:0')
sess.run('output:0') の実行時に下記の例外が発生してしまいました。
FailedPreconditionError (see above for traceback): GetNext() failed because the iterator has not been initialized
これは、上記のコード内で iter.initializer を実行した場合でも、モデルの中に保存されている Iterator を初期化することにはならない為のようです。
ハマりどころ2
初めは SavedModel ではなく、freeze した .pb 形式のモデルを保存・復元しようとしていたのですがこのやり方だと Iterator の保存・復元が上手くできませんでした。（やり方はあるかも知れませんが、できるという情報もできないという情報も見つけられなかったため、途中で諦めてしまいました）