tensorflowのscalar_summaryのログを展開する（tensorflowを使わない場合）

TensorflowやPyTorchによるログの出力としてTensorboardのtfeventファイルへの書き出しがある。これを可視化するにはtensorboardを使うことが多いだろう。しかしmatplotlibでのプロットなど他の用途でログを使いたい場合があると思う。Tensorflowには出力されたtfeventファイルをパースしてくれるモジュールがある。しかしPyTorchを使っている場合などはこのログを展開ためにわざわざtensorflowをインストールするのは煩わしい。

PyTorchのTensorboardモジュール実行に必要であるTensorboardだけでもログをPythonのオブジェクトに展開できる。TensorboardはEventAccumulatorを使ってログであるtfeventファイルを読み込んでいるのでこれを使えば良い。

以下のサンプルはランダムに出力されたScalar Summaryを読み込んでprintで出力している。他のデータ型も同様にできると思う。詳しくはTensorboardのソースコードを参照してほしい。

import tempfile
from pathlib import Path

import numpy as np
from tensorboard.backend.event_processing.plugin_event_accumulator import EventAccumulator
from torch.utils.tensorboard import SummaryWriter


def load_tfevent(event_filename: Path, tag: str = 'foo'):
    assert event_filename.exists()
    acc = EventAccumulator(path=str(event_filename))
    acc.Reload()

    # check `tag` exists in event
    assert tag in acc.Tags()['tensors']

    for t in acc.Tensors(tag):
        print(f'walltime = {t.wall_time}, step = {t.step}, value = {t.tensor_proto.float_val[0]}')


def parse_tflog():
    writer = None
    with tempfile.TemporaryDirectory() as dname:
        log_dir = Path(dname)
        with SummaryWriter(log_dir=log_dir) as writer:
            for i in range(100):
                writer.add_scalar('foo', np.random.random(), i)

        tfevents = list(log_dir.glob('events.out.tfevents.*'))
        assert len(tfevents) == 1
        tfevent = tfevents[0]

        load_tfevent(tfevent)


if __name__ == '__main__':
    parse_tflog()