tensorflow keras で 学習にかかった時間を保存する

やりたいこと

tensorflow keras で各epochにかかった時間を計測して、どこかに(今回はmlflow)へ保存する。

方法

Callbackを自作して対応。

コード

import mlflow
import time
import tensorflow as tf

class SaveElapsedTime(tf.keras.callbacks.Callback):
    def __init__(self):
        self.epoch_start_time = None
        self.train_start_time = None

    def _cal_elapsed_time(self, st):
        return = time.time() - st

    def on_train_begin(self, logs=None):
        self.train_start = time.time()

    def on_train_end(self, logs=None):
        et = self._cal_elapsed_time(self.train_start_time)
        mlflow.log_param('total_train_elapsed_time', et)

    def on_epoch_begin(self, epoch, logs=None):
        self.epoch_start_time = time.time()

    def on_epoch_end(self, epoch, logs=None):
        et = self._cal_elapsed_time(self.epoch_start_time)
        mlflow.log_metric('epoch_elapsed_time', et, step=epoch+1)
        print('{} epoch fin. {:.2f} [sec]'.format(epoch+1, elapsed_time), file=sys.stderr, flush=True)