tensorflowで綺麗なログ出力をする（tqdmを使おう） #Python

はじめに

tensorflowで学習している時のログに、何の変数をどのタイミングで出力させるか考えるのがめんどくさいと思ったことはないですか？
kerasはfit関数などを使うとプログレスバーを出力してくれてログが綺麗にまとまってくれますが、tensorflowを使って学習する時のログ出力はみなさん我流で書いていることが多いと思います。
今回はtensorflowで学習する時はtqdmというライブラリを使うことをオススメする記事です。（すでに使っている方も多いと思いますが）

インストールコマンド

$ pip install tqdm

実装環境

Python 3.6.4
tensorflow 1.9.0
tqdm

扱うタスクと使用するモデル

MNISTのクラス分類モデルを扱います。モデルはシンプルなCNNです。

    def _model_fn(self):
        x = tf.keras.layers.Conv2D(8, (3, 3), activation='relu', padding='same')(self.x)
        x = tf.keras.layers.BatchNormalization()(x)
        x = tf.keras.layers.Conv2D(8, (3, 3), activation='relu', padding='same')(x)
        x = tf.keras.layers.BatchNormalization()(x)
        x = tf.keras.layers.MaxPooling2D((2, 2))(x)

        x = tf.keras.layers.Conv2D(16, (3, 3), activation='relu', padding='same')(x)
        x = tf.keras.layers.BatchNormalization()(x)
        x = tf.keras.layers.Conv2D(16, (3, 3), activation='relu', padding='same')(x)
        x = tf.keras.layers.BatchNormalization()(x)
        x = tf.keras.layers.MaxPooling2D((2, 2))(x)

        x = tf.keras.layers.Flatten()(x)

        x = tf.keras.layers.Dense(64, activation='relu')(x)
        x = tf.keras.layers.Dense(10, activation='softmax')(x)

        self.cross_entropy_loss = -tf.reduce_sum(self.y * tf.log(x))
        self.train_step = tf.train.AdamOptimizer(1e-4).minimize(self.cross_entropy_loss)

        correct_prediction = tf.equal(tf.argmax(x, 1), tf.argmax(self.y, 1))
        self.accuracy = tf.reduce_mean(tf.cast(correct_prediction, "float"))

        tf.summary.scalar('loss', self.cross_entropy_loss)
        tf.summary.scalar('acc', self.accuracy)
        self.summary_op = tf.summary.merge_all()

ログ出力部分のコード

tqdmはfor文の配列部分を囲うだけで綺麗なプログレスバーを出力してくれます。また、set_postfix関数を使うと好きな値をプログレスバーの中に一緒に出力することができます。kerasと同じようにlossとaccuracyは各epoch中の平均の値を出力するようにしています。

def train(sess, model, x_train, y_train, batch_size):
    # Use tqdm for progress bar
    t = tqdm(range(len(x_train) // batch_size))
    acces, losses = [], []
    for i in t:
        batch_x = x_train[i * batch_size:(i + 1) * batch_size]
        batch_y = y_train[i * batch_size:(i + 1) * batch_size]
        feed_dict = {model.x: batch_x, model.y: batch_y}
        _, loss, acc = sess.run([model.train_step, model.cross_entropy_loss, model.accuracy], feed_dict)
        # print(loss, acc)
        losses.append(loss)
        acces.append(acc)
        t.set_postfix(train_loss=sum(losses) / len(losses), train_acc=sum(acces) / len(acces))
    return sum(losses) / len(losses), sum(acces) / len(acces)

実行時の出力は以下のようになります。

Epoch 1/10
100%|██████████| 1875/1875 [01:03<00:00, 29.73it/s, train_acc=0.827, train_loss=18.1]
train/acc: 0.8272, train/loss: 18.1091
valid/acc: 0.9120, valid/loss: 8.7919
Epoch 2/10
100%|██████████| 1875/1875 [01:03<00:00, 29.62it/s, train_acc=0.948, train_loss=5.5]
train/acc: 0.9478, train/loss: 5.4978
valid/acc: 0.9525, valid/loss: 4.7187
.
.
.
Epoch 10/10
100%|██████████| 1875/1875 [01:09<00:00, 26.89it/s, train_acc=0.988, train_loss=1.25]
train/acc: 0.9883, train/loss: 1.2526
valid/acc: 0.9829, valid/loss: 1.6424

終わり

tqdmを使うと実装を少し変えるだけで綺麗に出力できると思います。毎ステップごとに一行ずつ表示していた時にログがずらずらと長くなってしまって嫌な気持ちになることもないです。今回のコードはgithubに載せているので参考になればと思います。