pyt_timer.py

import torch, time, gc, contextlib

# Timing utilities
start_time = None

@contextlib.contextmanager
def add_timer(local_msg):
    try:
        start_timer()
        yield
    finally:
        end_timer(local_msg)


def start_timer():
    global start_time
    gc.collect()
    torch.cuda.empty_cache()
    torch.cuda.reset_max_memory_allocated()
    torch.cuda.synchronize()
    start_time = time.perf_counter()

def end_timer(local_msg=""):
    torch.cuda.synchronize()
    end_time = time.perf_counter()
    print("\n" + local_msg)
    print("Total execution time = {:.3f} ms".format((end_time - start_time)*1000))
    print("Max memory used by tensors = {:.3f} GB".format(torch.cuda.max_memory_allocated()/(1024*1024*1024)))

Usage:

from pyt_timer import add_timer

with add_timer():
    <code_block>