TensorBoardでcompute timeを見る方法

tf.RunMetadata

を使うことで各ノードの計算時間やメモリ使用状況をTensroBoardで知ることができる。

公式ドキュメントで公開されているコード

mnist_train.py

def feed_dict(train):
    """Make a TensorFlow feed_dict: maps data onto Tensor placeholders."""
    if train or FLAGS.fake_data:
      xs, ys = mnist.train.next_batch(100, fake_data=FLAGS.fake_data)
      k = FLAGS.dropout
    else:
      xs, ys = mnist.test.images, mnist.test.labels
      k = 1.0
    return {x: xs, y_: ys, keep_prob: k}

  for i in range(FLAGS.max_steps):
    if i % 10 == 0:  # 10回に1回だけtest-setでのaccuracyを保存する
      summary, acc = sess.run([merged, accuracy], feed_dict=feed_dict(False))
      test_writer.add_summary(summary, i)
      print('Accuracy at step %s: %s' % (i, acc))
    else:  # Record train set summaries, and train
      if i % 100 == 99:  # 100回に1回だけ実行時の状況を保存する(この記事の本題)
        run_options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE)
        run_metadata = tf.RunMetadata()
        summary, _ = sess.run([merged, train_step],
                              feed_dict=feed_dict(True),
                              options=run_options,
                              run_metadata=run_metadata)
        train_writer.add_run_metadata(run_metadata, 'step%d' % i)
        train_writer.add_summary(summary, i)
        print('Adding run metadata for', i)
      else:  #summaryを保存する
        summary, _ = sess.run([merged, train_step], feed_dict=feed_dict(True))
        train_writer.add_summary(summary, i)

run_options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE)
run_metadata = tf.RunMetadata()

を定義してsess.sun()に渡す。
オプションにはFULL_TRACEの他に、SOFTWARE_TRACE、HARDWARE_TRACE などがあるが、ドキュメントでも詳細な説明はない。
とりあえずFULL_TRACEを使えばよいと思われる。

train_writer.add_run_metadata(run_metadata, 'step%d' % i)

でmetadataを保存する。

TensorBordで見るには

TensorBordのSession runsから見たいstepを選択すればよい。