Modify TensorBoard event tags.

modify_tf_logfile.py

# Most of the code is taken directly from [this StackOverflow answer](https://stackoverflow.com/a/60080531/4868562).

import click
import os
import shutil
import sys
import tensorflow as tf

from pathlib import Path
from tensorflow.core.util.event_pb2 import Event
from tqdm import tqdm


def rename_events(input_path, output_path, tag_replacements):
    with tf.io.TFRecordWriter(str(output_path)) as record_writer:
        for record in tf.data.TFRecordDataset([str(input_path)]):
            event = Event()
            event.MergeFromString(record.numpy())
            if event.summary:
                for v in event.summary.value:
                    if v.tag in tag_replacements.keys():
                        v.tag = tag_replacements[v.tag]
            record_writer.write(event.SerializeToString())


def rename_events_dir(input_dir, output_dir, tags, replacements):
    assert len(tags) == len(replacements)
    assert input_dir != output_dir # this case results in errors in the output logs
    input_dir_path = Path(input_dir)
    output_dir_path = Path(output_dir) if output_dir else Path('{}n'.format(input_dir))
    output_dir_path.mkdir(parents=True, exist_ok=True)
    for ev_file in input_dir_path.glob('**/*.tfevents*'):
        out_file = Path(output_dir_path, ev_file.relative_to(input_dir_path))
        out_file.parent.mkdir(parents=True, exist_ok=True)
        rename_events(ev_file, out_file, {tag: replacement for tag, replacement in zip(tags, replacements)})
    if not output_dir: # replace
        shutil.rmtree(input_dir)
        os.rename('{}n'.format(input_dir), input_dir)


@click.group()
def cli():
    pass


@cli.command()
@click.option('-i', '--input-dir', prompt='Tensorboard logdir', help='Path to the Tensorboard logs.', type=str)
@click.option('-o', '--output-dir', type=str,
              help='Path for the modified Tensorboard logs. If not provided, the input logs will be replaced.')
@click.option('-t', '--tags', prompt='Tags to be replaced', multiple=True, type=str,
              help='For each replacement, pass an existing tag name immediately followed by the replacement tag.')
@click.option('-r', '--replacements', prompt='Replacement tags', multiple=True, type=str,
              help='Replacement tags. Must match the number and order of the --tags option.')
def rename(input_dir, output_dir, tags, replacements):
    """ Rename the event tags in all events in a Tensorboard logdir. """
    rename_events_dir(input_dir, output_dir, tags, replacements)
    click.echo('Done: {}'.format(input_dir))


@cli.command()
@click.option('-i', '--input-dir', prompt='Path to the Tensorboard logs.', type=str)
@click.option('-t', '--tags', prompt='Tags to be replaced.', multiple=True, type=str,
              help='For each replacement, pass an existing tag name immediately followed by the replacement tag.')
@click.option('-r', '--replacements', prompt='Replacement tags.', multiple=True, type=str,
              help='Replacement tags. Must match the number and order of the --tags option.')
def rename_several(input_dir, tags, replacements):
    """ Rename the event tags in all events in a directory containing several logdirs. """
    progressbar = tqdm(os.walk(input_dir), total=len(os.listdir(input_dir)))
    for path, _, _ in progressbar:
        if path != input_dir:
            progressbar.set_postfix_str(path)
            rename_events_dir(path, None, tags, replacements)


if __name__ == '__main__':
    cli()


"""
Examples:

For modifying a single logs directory:

```
rename -i "path/to/logdir" -t "original_tag_1" -t "original_tag_2" -r "replacement_for_original_tag_1" -r "replacement_for_original_tag_2"
```

For modifying multiple directories at once:

```
rename-several -i "path/containing/logdirs" -t "original_tag_1" -t "original_tag_2" -r "replacement_for_original_tag_1" -r "replacement_for_original_tag_2"
```
"""