timing wrapper that wraps a funciton and logs its execution time to tensorboard

timing_wrapper.py

from functools import wraps
import itertools

def timing_tb(writer):
    def wrapper_outer(func):
        cont = itertools.count()
        print(f"will log under {func.__name__}")
        @wraps(func)
    
        def wrapper(*arg, **kw):        
            t1 = time.time()
            res = func(*arg, **kw)
            t2 = time.time()
            step = next(cont)
#             print("log", func.__name__, t2-t1, step)
            with writer.as_default(step=step):
                if not tf.get_static_value(tf.summary.scalar(f'metric_{func.__name__}', data=t2-t1)):
                    print("save failed")
                tf.summary.flush()
            return res
        return wrapper
    return wrapper_outer

using_it.py

file_writer = tf.summary.create_file_writer(logdir)

@timing_tb(file_writer)
def a():
  return 1