Load TensorBoard logging files to Pandas.

tensorboard2pandas.py

#!/usr/bin/env python3
"""Simple script to extract metrics calculated by HuggingFace tansformers from
TensorBoard logs.
"""

import pandas as pd
import tensorboard as tb
import tensorboard.data_compat
import tensorflow as tf

from argparse import ArgumentParser, Namespace
from pathlib import Path
from typing import List, Optional


def convert(srcs: List[Path], dst: Optional[Path] = None) -> pd.DataFrame:
    """Function convert reads TensorBoard (written by HuggingFace transformers
    library), filters it contents, and write it to CSV file.
    """
    records = []
    for i, el in enumerate(tf.data.TFRecordDataset(srcs)):
        event = tb.compat.proto.event_pb2.Event.FromString(el.numpy())
        event = tb.data_compat.migrate_event(event)
        if event.summary is None:
            continue
        for value in event.summary.value:
            prefix, _ = value.tag.split('/', 1)
            if prefix not in ('eval', 'train'):
                continue
            records.append((event.step, value.tag, value.tensor.float_val[0]))

    df = pd.DataFrame(data=records, columns=('step', 'tag', 'value'))
    df = df.set_index(['tag', 'step'])
    df = df.sort_index()
    if dst:
        df.to_csv(dst)
    return df


def main(args: Namespace):
    convert(args.filenames, args.output)


parser = ArgumentParser()
parser.add_argument('-o', '--output', type=Path, help='Output CSV-file.')
parser.add_argument('filenames', type=Path, nargs='+', help='Log files.')


if __name__ == '__main__':
    main(parser.parse_args())