Extract events from tensorboards tfrecords events file

extract_tensorboard_tensorflow_rfrecords.py

import tensorflow as tf
from collections import defaultdict

def process_tensorboard_events_file(tensorboard_events_filename, tf_version=2, tags=None):
    assert isinstance(tags, list)
    for v in tags:
        assert isinstance(v, str)    
    it = _get_events_iterator(tensorboard_events_filename, tf_version)
    ans = defaultdict(list)
    for event in it:
        for v in event.summary.value:            
            if (tags is None) or v.tag in tags:
                if hasattr(v, 'simple_value'):
                    data = ans[v.tag]
                    data.append(v.simple_value)
    return ans

def _get_events_iterator(tensorboard_events_filename, tf_version):
    if tf_version==1:
        return _get_events_iterator_tf1(tensorboard_events_filename)
    elif tf_version==2:
        return _get_events_iterator_tf2(tensorboard_events_filename)
    else:
        raise ValueError('tf_version must be 1 or 2')  

def _get_events_iterator_tf1(tensorboard_events_filename):   
    for e in tf.compat.v1.train.summary_iterator(tensorboard_events_filename):
        yield e
        # print(e)
        # for v in e.summary.value:
        #    tag = v.tag
        #    val = tf.make_ndarray(v.tensor)
        #    print('tag:',tag)
        #    print('val:',val)

def _get_events_iterator_tf2(tensorboard_events_filename):    
    from tensorflow.core.util import event_pb2
    
    serialized_examples = tf.data.TFRecordDataset(tensorboard_events_filename)
    for serialized_example in serialized_examples:
        event = event_pb2.Event.FromString(serialized_example.numpy())
        yield event
        # for value in event.summary.value:
        #     t = tf.make_ndarray(value.tensor)
        #     print(value.tag, event.step, t, type(t))    

if __name__=='__main__':
  fn = '/path/to/some/filename/events.out.tfevents.1652171232.whi-7.441.0'
  ans = process_tensorboard_events_file(fn, tags=[
        'val_loss',
        'val_rmse',
        'val_pearson',
        'val_spearman',
  ])
  print(ans)