【教程】PyTorch Timer计时器

转载请注明出处：小锋学长生活大爆炸[xfxuezhang.cn]

OpenCV的Timer计时器可以看这篇：Python Timer和TimerFPS计时工具类

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' # 保存每次的elapsedclass Timer(object):def __init__(self, device: torch.device):super(Timer, self).__init__()self._device = deviceself._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] = 0if 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 Nonedef store(self, name):'''计算指定name的累计耗时'''if not self._record_sxf.get(name):print(f'>> 不存在此计时器[{name}]，请先start')return Noneif not self._record_sxf[name].get(TimerKeys.ELAPSED):print(f'>> 请先stop')return Noneself._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_sxfdef peak_item(self, name, key=None):if key:return self._record_sxf[name].get(key) if self._record_sxf.get(name) else Nonereturn 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')# xxxxxxtime_used = timer.stop('a')timer.show_store()