Time gpu operations in torch using `with gtimer('blockname'): ... torch code ...`

gtimer.py

# This code is in the public domain
# Author: Paul Kienzle

from contextlib import contextmanager
try:
    import torch
except ImportError:
    pass

_gtimer_block = []
@contextmanager
def gtimer(block=""):
    """
    Timer for torch cuda processing.

    Usage::

        with gtimer("block name"):
            ...

    Prints "Δt for {block name} = t ms"

    When nested, prints "... for {parent}:{block name} = ..."
    """
    _gtimer_block.append(block)
    if "torch" in globals() and torch.cuda.is_available():
        # TODO: find file/line of caller if block name is empty
        start = torch.cuda.Event(enable_timing=True)
        end = torch.cuda.Event(enable_timing=True)
        start.record()
        yield
        end.record()
        # Waits for everything to finish running
        torch.cuda.synchronize()
        elapsed = start.elapsed_time(end)
    else:
        start = time.perf_counter()
        yield
        elapsed = 1000*(time.perf_counter() - start)


    print(f"Δt for {':'.join(_gtimer_block)} = {int(elapsed+0.5)} ms")
    del _gtimer_block[-1]

@contextmanager
def notimer(block=""):
    yield