Python PyTorch实现Timer计时器

import time
import torch
import os
from enum import Enum, unique
@unique
class TimerKeys(Enum):
    START   = 'start'   # 本次调用start函数时的时间戳
    END     = 'end'     # 本次调用end函数时的时间戳
    ELAPSED = 'elapsed' # 多次调用start和end函数的总耗时
    TOTAL   = 'total'   # 本次所耗end-start的时间
    HISTORY = 'history' # 保存每次的elapsed
class Timer(object):
    def __init__(self, device: torch.device):
        super(Timer, self).__init__()
        self._device = device
        self._record_sxf = {}
    def start(self, name, history=False):
        '''开始计时，通过name区分不同的计时器；
        在start时如果开启了history，就算下次start时没有开启history，history仍然会使用，
        除非显式调用了reset_all或reset_item，然后再调用start并不开启history，就不会记录。'''
        torch.cuda.current_stream(self._device).synchronize()
        if not self._record_sxf.get(name):
            self._record_sxf[name] = {}
            self._record_sxf[name][TimerKeys.TOTAL] = 0
            if history:
                self._record_sxf[name][TimerKeys.HISTORY] = []
        else:
            self._record_sxf[name].pop(TimerKeys.END)
            self._record_sxf[name].pop(TimerKeys.ELAPSED)
        self._record_sxf[name][TimerKeys.START] = time.time()
    def stop(self, name, store=True):
        '''计算指定name的本次耗时'''
        torch.cuda.current_stream(self._device).synchronize()
        if self._record_sxf.get(name) and self._record_sxf[name].get(TimerKeys.START):
            self._record_sxf[name][TimerKeys.END] = time.time()
            self._record_sxf[name][TimerKeys.ELAPSED] = self._record_sxf[name][TimerKeys.END] - self._record_sxf[name][TimerKeys.START]
            return self.store(name) if store else self._record_sxf[name][TimerKeys.ELAPSED]
        else:
            print(f'>> 不存在此计时器[{name}]，请先start')
        return None
    def store(self, name):
        '''计算指定name的累计耗时'''
        if not self._record_sxf.get(name):
            print(f'>> 不存在此计时器[{name}]，请先start')
            return None
        if not self._record_sxf[name].get(TimerKeys.ELAPSED):
            print(f'>> 请先stop')
            return None
        self._record_sxf[name][TimerKeys.TOTAL] += self._record_sxf[name][TimerKeys.ELAPSED]
        if self._record_sxf[name].get(TimerKeys.HISTORY) is not None:
            self._record_sxf[name][TimerKeys.HISTORY].append(self._record_sxf[name][TimerKeys.ELAPSED])
        return self._record_sxf[name][TimerKeys.TOTAL]
    def show_store(self):
        '''显示所有项目的累计耗时'''
        print(self._record_sxf)
    def pretty_show_store(self):
        print("{", end='')
        for key, value in self._record_sxf.items():
            print(f"'{key}': {{")
            for enum_key, enum_value in value.items():
                print(f"    {enum_key.name if isinstance(enum_key, Enum) else enum_key}: {enum_value},")
            print("  },")
        print("}")
    def get_store(self):
        return self._record_sxf
    def peak_item(self, name, key=None):
        if key:
            return self._record_sxf[name].get(key) if self._record_sxf.get(name) else None
        return self._record_sxf.get(name)
    def reset_item(self, name):
        self._record_sxf.pop(name)
    def reset_all(self):
        self._record_sxf = {}
if __name__ == '__main__':
    timer = Timer(torch.device('cuda:0'))
    timer.start('a')
    time.sleep(5)
    timer.stop('a')
    # {'a': {'total': 5.005435228347778, 'start': 1693419100.180317, 'end': 1693419105.1857522, 'elapsed': 5.005435228347778}}
    timer.show_store()
    print()
    timer.start('a')
    time.sleep(2)
    timer.stop('a')
    # {'a': {'total': 7.007752180099487, 'start': 1693419105.1859245, 'end': 1693419107.1882415, 'elapsed': 2.002316951751709}}
    timer.show_store()
    print()
    timer.start('b')
    time.sleep(3)
    timer.stop('b')
    # {'a': {'total': 7.007752180099487, 'start': 1693419105.1859245, 'end': 1693419107.1882415, 'elapsed': 2.002316951751709},
    #  'b': {'total': 3.0033228397369385, 'start': 1693419107.1884048, 'end': 1693419110.1917276, 'elapsed': 3.0033228397369385}}
    timer.show_store()
    timer.reset_all()
    print()
    timer.start('c')
    time.sleep(3)
    timer.start('d')
    time.sleep(3)
    timer.stop('d')
    # {'c': {'total': 0, 'start': 1693419110.1919253},
    # 'd': {'total': 3.003229856491089, 'start': 1693419113.1927958, 'end': 1693419116.1960256, 'elapsed': 3.003229856491089}}
    timer.show_store()
    timer.stop('c')
    # {'c': {'total': 6.0042500495910645, 'start': 1693419110.1919253, 'end': 1693419116.1961753, 'elapsed': 6.0042500495910645},
    #  'd': {'total': 3.003229856491089, 'start': 1693419113.1927958, 'end': 1693419116.1960256, 'elapsed': 3.003229856491089}}
    timer.show_store()
    timer.reset_all()
    print()
    timer.start('e')
    time.sleep(3)
    timer.start('f')
    time.sleep(3)
    timer.stop('e')
    # {'e': {'total': 6.004979848861694, 'start': 1693419433.8564444, 'end': 1693419439.8614242, 'elapsed': 6.004979848861694},
    #  'f': {'total': 0, 'start': 1693419436.859731}}
    timer.show_store()
    timer.stop('f')
    # {'e': {'total': 6.004979848861694, 'start': 1693419433.8564444, 'end': 1693419439.8614242, 'elapsed': 6.004979848861694},
    #  'f': {'total': 3.00180983543396, 'start': 1693419436.859731, 'end': 1693419439.8615408, 'elapsed': 3.00180983543396}}
    timer.pretty_show_store()

timer = Timer(torch.device('cuda:0'))
timer.start('a')
# xxxxxx
time_used = timer.stop('a')
timer.show_store()