File size: 887 Bytes
872b0a0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 | import os
import time
import torch
import pdb
class Profiler(object):
def __init__(self, silent=False):
self.silent = silent
torch.cuda.synchronize()
self.start = time.time()
self.cache_time = self.start
def reset(self, silent=None):
if silent is None:
silent = self.silent
self.__init__(silent=silent)
def report_process(self, process_name):
if self.silent:
return None
torch.cuda.synchronize()
now = time.time()
print('{0}\t: {1:.4f}'.format(process_name, now - self.cache_time))
self.cache_time = now
def report_all(self, whole_process_name):
if self.silent:
return None
torch.cuda.synchronize()
now = time.time()
print('{0}\t: {1:.4f}'.format(whole_process_name, now - self.start))
pdb.set_trace()
|