はじめに

TensorBoardで損失関数や正答率などの平均を計算する際に、メモリ容量の都合でミニバッチに分けて計算することがあると思います。その場合のバッチをまたいだ平均の取り方の一例を示します。同時に訓練データとテストデータに分けてTensorBoardに表示できるようにします。

参考

バッチをまたいだ平均の取り方
tf.metricsを使ってみる
tf.metrics 公式ドキュメント
訓練データとテストデータを分けてTensorBoardに表示する方法
TensorBoardでTrainとTest, Validationを分けて学習状況を把握する

tf.metricsの使い方

通常、TensorFlowの計算グラフはfeed_dictで与えたバッチデータに対してのみ計算を行いますが、tf.metricsを使うと過去に与えたすべてのデータに対しての平均などの計算ができるようになります。本記事では単純な加算平均を取るtf.metrics.meanを使ったコードを示します。他にも正答率を計算するtf.metrics.accuracy等がありますが、使い方は同じです。

tf.metrics.meanはmeanとupdate_opの二つの返り値を持ちます。
meanが平均値、update_opは平均値を計算するために必要となる値の総和と総数を更新する操作です。
バッチごとにsess.run(update_op, feed_dict=mini_batch)を実行すると総和と総数が更新され、その後、sees.run(mean)を実行すると平均値が取得できます。
次章に示したサンプルコードのevaluationという関数の最後の部分に相当します。

    # バッチをまたいだ平均の計算
    for i in range(0, data_num, batch_size):
        sess.run(op, feed_dict={x: val_x[i:i+batch_size], y: val_y[i:i+batch_size]})
    # 平均値取得
    return sess.run([val_loss, summary])

ただし、総和や総数をリセットしないと学習が始まってからすべてのバッチに対しての平均を取る、訓練データとテストデータに分けて平均を計算できない、などの問題があります。そこで、平均値をとる前に総和と総数を0にリセットします。そのためにはリセットする操作を取得しておいて、平均値の計算前に実行しなければいけません。
サンプルコードでは名前空間を使ってリセット操作を取得しています。

reset_vars = [i for i in tf.local_variables() if i.name.split('/')[0] in ['train_summary', 'test_summary']]
reset_op = [tf.variables_initializer(reset_vars)]

これでsess.run(reset_op)を実行すると平均を出すための総和と総数がリセットされるようになります。

その他の注意点として、tf.local_variables_initializer()を実行する必要があります。

訓練データとテストデータに分けてsummaryを生成

tf.metrics.meanを呼ぶ際にnameを指定しておくと同じlossでも複数の計算結果を保持できるので、これを利用して訓練データ用とテストデータ用に分けてsummaryを作っておきます。こうすることでTensorBoardに個別にグラフを表示できるようになります。

サンプルコード

ここでは、$y=2x+(雑音)$というデータを生成し、$y=wx$で多項式近似する（wを学習する）サンプルを載せています。

import tensorflow as tf
import numpy as np

# データ生成用の関数
def data(num):
    w = 2.0
    x = np.random.random(size=(num, 1))
    n = np.random.random(size=(num, 1))
    y = w * x + (n-0.5)
    return x, y

# データ生成
train_num = 100
test_num = 100
batch_size = 10
train_x, train_y = data(train_num)
test_x, test_y = data(test_num)

# TensorFlowのグラフ構成
x = tf.placeholder(shape=(None, 1), dtype=tf.float64)
y = tf.placeholder(shape=(None, 1), dtype=tf.float64)
w = tf.Variable(tf.random_uniform([1], dtype=tf.float64))
loss = tf.reduce_mean((w*x-y)**2)
train = tf.train.GradientDescentOptimizer(0.5).minimize(loss)

# TensorBoard用のsummary
# tf.metrics.meanを使ってlossに対して、異なる入力データからの平均値を取得できるようにする
with tf.name_scope("train_summary"):
    train_loss, train_loss_op = tf.metrics.mean(loss, name="train_loss")
    train_summary = tf.summary.scalar("train_loss", train_loss)
with tf.name_scope("test_summary"):
    test_loss, test_loss_op = tf.metrics.mean(loss, name="test_loss")
    test_summary = tf.summary.scalar("test_loss", test_loss)
# metrics.meanをリセットできるようにする
reset_vars = [i for i in tf.local_variables() if i.name.split('/')[0] in ['train_summary', 'test_summary']]
reset_op = [tf.variables_initializer(reset_vars)]

# 変数初期化
sess = tf.Session()
sess.run(tf.global_variables_initializer())
sess.run(tf.local_variables_initializer()) # tf.metrics.meanの計算に必要
writer = tf.summary.FileWriter("logs", sess.graph)

# ミニバッチを使ったときの平均の計算
def evaluation(isTrain):
    # 引数で訓練データとテストデータを切り替え
    if isTrain:
        data_num = train_num
        op = train_loss_op
        val_x = train_x
        val_y = train_y
        val_loss = train_loss
        summary = train_summary
    else:
        data_num = test_num
        op = test_loss_op
        val_x = test_x
        val_y = test_y
        val_loss = test_loss
        summary = test_summary
    # 平均計算をリセット
    sess.run(reset_op)
    # バッチをまたいだ平均の計算
    for i in range(0, data_num, batch_size):
        sess.run(op, feed_dict={x: val_x[i:i+batch_size], y: val_y[i:i+batch_size]})
    # 平均値取得（損失関数の平均値とそのsummaryを返す）
    return sess.run([val_loss, summary])

# 訓練と評価
t=0 # 勾配更新回数
for epoch in range(5):
    batch_idx = np.random.permutation(train_num)
    for i in range(0, train_num, batch_size):
        X = train_x[batch_idx[i:i+batch_size]]
        Y = train_y[batch_idx[i:i+batch_size]]
        sess.run(train, feed_dict={x: X, y: Y})

        #　評価
        w_ = sess.run(w) # 学習したwの値
        tr_l, tr_s = evaluation(True) # 訓練データの損失関数平均値
        te_l, te_s = evaluation(False) # テストデータの損失関数平均値
        writer.add_summary(tr_s, t)
        writer.add_summary(te_s, t)
        print(w_, tr_l, te_l, t)
        t += 1

# jupyter notebookではflush()しないとtensorboardが更新されないことがある
writer.flush()

結果

    [0.81023639] 0.50949466 0.54334414 0
    [1.11114108] 0.3181283 0.34896314 1
    [1.31760125] 0.21976742 0.24668348 2
    [1.45633257] 0.16873363 0.19217044 3
    [1.5928319] 0.1303354 0.14968534 4
    [1.73803411] 0.10235226 0.116632186 5
    [1.77873357] 0.09688808 0.109613255 6
    [1.93128664] 0.08567798 0.09205474 7
    [1.97550464] 0.08516492 0.08954786 8
    [1.99003648] 0.08526479 0.0889774 9
    [1.98071531] 0.08518545 0.08932889 10
    [1.93413139] 0.08560797 0.09185852 11
    [2.00279574] 0.085461974 0.08857987 12
    [1.99528732] 0.08533353 0.088802114 13
    [1.99148938] 0.08528207 0.08892726 14
    [1.98968659] 0.08526082 0.08898966 15
    [2.01771783] 0.085822485 0.08823757 16
    [1.9590682] 0.08521204 0.090344176 17
    [1.96545566] 0.08517354 0.09001566 18
    [1.95934922] 0.0852098 0.0903292 19
    [1.94439742] 0.085397616 0.09119038 20
    [1.96005929] 0.08520438 0.090291604 21
    [1.94928473] 0.085320756 0.0908943 22
    [1.93590982] 0.085566774 0.091738306 23
    [1.88190214] 0.08770468 0.09622655 24
    [1.877801] 0.087941974 0.0966381 25
    [1.91255662] 0.0862661 0.093466565 26
    [1.97842536] 0.08517433 0.089423135 27
    [2.04861815] 0.0870143 0.08794899 28
    [2.04226268] 0.086720124 0.087962046 29
    [1.9276418] 0.08577515 0.09231316 30
    [2.03801546] 0.08653769 0.08798413 31
    [1.98790605] 0.08524184 0.089053184 32
    [2.01068972] 0.085635245 0.088382326 33
    [2.03210527] 0.08630269 0.08803268 34
    [2.07455445] 0.088478245 0.08814432 35
    [1.94044388] 0.08547077 0.09144029 36
    [1.93696516] 0.085543275 0.09166785 37
    [1.92977176] 0.08571736 0.09216119 38
    [1.95681684] 0.08523172 0.09046573 39
    [1.96895851] 0.085163325 0.08984579 40
    [1.97275074] 0.08516096 0.08967011 41
    [1.96138375] 0.08519511 0.090222284 42
    [1.91241975] 0.0862712 0.09347764 43
    [1.88666088] 0.0874426 0.095761515 44
    [1.9198232] 0.08601174 0.09289411 45
    [2.02702417] 0.086118236 0.08809101 46
    [2.01134852] 0.08565147 0.08836751 47
    [2.00511529] 0.085508816 0.088517986 48
    [1.99727061] 0.085364014 0.088740155 49

TensorBoardには次のような感じで表示されます。名前空間ごとにtrain_lossとtest_lossがそれぞれ表示されています。