【Python/TensorFlow】tf.estimator.Estimatorを用いた学習のやり方(モデル作成 by Keras)

import numpy as np
import tensorflow as tf
from tensorflow.keras.models import Model
from tensorflow.keras.layers import Input, Dense, Flatten, Lambda, Conv2D
from tensorflow.keras.optimizers import RMSprop

# 保存しているSparseTensorの形状 ※ Conv2Dでは3階テンソルを要求されるのでこの形状
sparse_shape = [12, 12, 1]

# デシリアライズ用コールバック関数
def input_fn(batch_size, num_epochs, filepath_list):
    features = {"train": tf.io.SparseFeature(index_key=["rows", "cols", "channels"],
                                             value_key="values",
                                             dtype=tf.float32,
                                             size=sparse_shape),
                "label": tf.io.VarLenFeature(dtype=tf.float32)
                }

    def parse_example(example):
        train_data = tf.io.parse_single_example(example, features=features)
        x = train_data["train"]
        y = tf.sparse.to_dense(train_data["label"])
        return x, y

    ds = tf.data.TFRecordDataset(filepath_list) \
        .map(parse_example) \
        .batch(batch_size) \
        .repeat(num_epochs)

    return ds


def sample_keras_model_fn():
    # Constants
    kernel_size = (4, 4)
    num_filters = 4

    inputs = Input(batch_size=batch_size, shape=sparse_shape, sparse=True)
    X = Lambda(tf.sparse.to_dense)(inputs)
    X = Conv2D(filters=num_filters, kernel_size=kernel_size)(X)
    X = Flatten()(X)
    X = Dense(1, activation="relu")(X)
    model = tf.keras.models.Model(inputs=inputs, outputs=X)
    model.compile(loss="MSE", optimizer=RMSprop())

    return model

# メイン
if __name__ == "__main__":
    #コンソールにログを表示するために必要
    tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.INFO)

    # 各種パラメータ
    batch_size = 2
    epochs = 3

    # TFRecordsファイルへのパス & ログ保存パス
    filepath = "./sequence_conv.tfrecords"
    model_dir = "./models/"

    # Kerasモデル構築
    model = sample_keras_model_fn()

    # ログの出力先・ログ取得タイミングの設定
    config = tf.estimator.RunConfig(model_dir=model_dir, log_step_count_steps=5)

    # keras model ⇒ Estimator の変換
    estimator = tf.keras.estimator.model_to_estimator(model, model_dir=model_dir, config=config)

    # keras model ⇒ Estimator の変換
    train_spec = tf.estimator.TrainSpec(input_fn=lambda: input_fn(batch_size, epochs, [filepath]))
    eval_spec = tf.estimator.EvalSpec(input_fn=lambda: input_fn(batch_size, epochs, [filepath]))

    # 学習実行
    tf.estimator.train_and_evaluate(estimator, train_spec, eval_spec)

"""
【出力結果(一部)】
INFO:tensorflow:Saving checkpoints for 0 into ./models/model.ckpt.
INFO:tensorflow:loss = 0.28118426, step = 0
INFO:tensorflow:global_step/sec: 48.9893
INFO:tensorflow:loss = 0.27649117, step = 5 (0.102 sec)
INFO:tensorflow:global_step/sec: 559.42
INFO:tensorflow:loss = 0.21459994, step = 10 (0.010 sec)
INFO:tensorflow:Saving checkpoints for 15 into ./models/model.ckpt.
"""