TensorFlow2:TensorBoardのグラフがうまく表示されず困った件

前回は、kerasのsummaryなどでグラフの可視化をしましたが、今回はTensorboardのグラフを使ってみます。

使用するモデルは、TensorFlow2:MNISTやってみる（keras model class）で作成したものです。

1.xのときと違って、Summaryを保存するだけではグラフは出力されないようです。 ということで、変更を加えて出力してみたけど、うまく動かなかったり、いまいちきれいに出なかったという話。

Reference: Examining the TensorFlow Graph

グラフ出力のための変更

紆余曲折が長いので、まずは最終的な形を記載します。

今回はsubclassedなモデルで、model.fit()は使用していないので、Graphs of tf.functions セクションで記載のある方法を試す。

計算開始時にtf.summary.trace_on(graph=True)を呼び、終了時にtf.summary.trace_export()を呼ぶ。

注意点としては、trace_export()は、writer.as_default()のブロック内で呼び出す必要がある。 必然的に、SummaryWriterのインスタンスが必要になる（Writerが書き出すファイル内にグラフが入るので、当たり前かもしれないが一応）。

以下の「# !!! ADD」コメント部分が追加した箇所。

def training(model, optimizer, epochs, ds_train, ds_test, logdir=None):
    # ...skip

    writer = tf.summary.create_noop_writer()
    if logdir:
        logdir.mkdir(parents=True, exist_ok=True)
        writer = tf.summary.create_file_writer(str(logdir))

    global_step = 1
    tf.summary.trace_on(graph=True)  # !!! ADD
    for epoch in range(1, epochs+1):
        # do training.
        ts_train = datetime.datetime.now()
        for i, (images, labels) in enumerate(ds_train, 1):
            loss_step = step_train(model, optimizer, loss_obj, images, labels, metr_train_loss, metr_train_acc)  # type: ignore # noqa
            if logdir and global_step == 1:  # !!! ADD this block
                with writer.as_default():
                    tf.summary.trace_export("graph", step=1)
            if i % 200 == 0:
                print("step {:8d}, loss: {:.3f}".format(i, loss_step))
                with writer.as_default():
                    tf.summary.scalar("train_loss_step", loss_step, global_step)
            global_step += 1
        ts_train = datetime.datetime.now() - ts_train
        # do test.
        # ...skip

    writer.close()

TensorBoardの表示。※この形にするには、ちょっと手を加える必要があります（後述）

変更１：はじまり

素直にtrace_on()とtrace_export()の呼び出してみた。 最初は、training全体の最初と最後に追加した。

def training(model, optimizer, epochs, ds_train, ds_test, logdir=None):
    # ...skip

    writer = tf.summary.create_noop_writer()
    if logdir:
        logdir.mkdir(parents=True, exist_ok=True)
        writer = tf.summary.create_file_writer(str(logdir))
        tf.summary.trace_on(graph=True)  # !!! ADD

    for epoch in range(1, epochs+1):
        # ...skip

    if logdir:
        with writer.as_default(): # !!! ADD this block
            tf.summary.trace_export("graph", step=0)

    writer.close()

しかし、グラフは表示されない…ッ！

TensorBoard上には、「GRAPHS」タブが表示されるようになった。

しかし、TensorBoardの「GRAPHS」ページにはエラーメッセージが表示されるのみ。

Graph visualization failed.

Error: Malformed GraphDef. This can sometimes be caused by a bad 
network connection or difficulty reconciling multiple GraphDefs; for 
the latter case, please refer to https://github.com/tensorflow/tensorboard/issues/1929.

TensorBoardのコンソールには以下のエラーが表示される。

ValueError: Cannot combine GraphDefs because nodes share a name but contents are different: assignaddvariableop_resource

変更２：エラー回避を試みる

エラーメッセージや、issueの記載を見る感じ、@tf.functionの複数回実行がトリガとなっているようだ。

今回のケースでは、学習を行うstep_train()と、テストを行うstep_testの両関数が、@tf.function付きになっている。 内部では、同一のモデルのインスタンスを利用しているので、これがnodes share a name but contents are differentになっていると思われる。

ということで、@tf.functionのアノテーションを、step_trainのみにしてみる。

お、表示された。

モデルはどこに行った？

すみっこにいた。

とりあえず、いるのは分かったので、「右クリック」→「Add to main graph」で左側に動かすことができる。

変更３：やっぱりパフォーマンスも気になりだす

traceのon/offはtrainの最初と最後がきれいかなと思っていたが、step_testはグラフ化が行われないのでの処理時間が落ちることが気になりだす。

昔はsession.run(train_op)ってやってたしと思い直し、trace_on()する区間を、step_train()のみに限定する。

※コードは冒頭の最終系を参照してください。

グラフの出力は１回のみでよいので、global_stepを用いた判定を入れた。 ループ内に分岐が入ってしまったが、仕方ないか。

この形になると、当初外したstep_testの@tf.functionも、外しても動くようになる。 結局issueの対策と同じになった。

変更４：ちらかった部分をまとめる

モデルがmain graphに移動したのはよいが、いまいち流れが見えない。

ごちゃごちゃしている計算を見るに、tf.keras.metricsあたりに見える。昔ながらのtf.name_scope()でまとめてみる。

@tf.function
def step_train(model, optimizer, loss_obj, images, labels, metr_loss, metr_acc):
    with tf.GradientTape() as tape:
        predictions = model(images, training=True)
        loss = loss_obj(labels, predictions)
        gradients = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    with tf.name_scope("metrics"):  # !!! ADD name scope
        metr_loss(loss)
        metr_acc(labels, predictions)
    return loss

あとは、Optimizerや入力のImageなど主要部分をmain graphに移動して…

うん。満足。

おわりに

eagerの挙動が難しい。