计算一下pytorch中Resnet34模型前传一次所需要的时间

来源:互联网 发布:n9软件下载 编辑:程序博客网 时间:2024/05/16 01:57
import pdbimport timeimport torchimport torchvision.models as modelsfrom torch.autograd import Variableclass Timer(object):    """A simple timer."""    def __init__(self):        self.total_time = 0.        self.calls = 0        self.start_time = 0.        self.diff = 0.        self.average_time = 0.    def tic(self):        # using time.time instead of time.clock because time time.clock        # does not normalize for multithreading        self.start_time = time.time()    def toc(self, average=True):        self.diff = time.time() - self.start_time        self.total_time += self.diff        self.calls += 1        self.average_time = self.total_time / self.calls        if average:            return self.average_time        else:            return self.diffGPUID = 1resnet34 = models.resnet34(pretrained=True)resnet34.cuda(GPUID)x = torch.rand(1,3,400,400)x = Variable(x.cuda(GPUID))# preheaty = resnet34(x)timer = Timer()timer.tic()for i in xrange(100):  y = resnet34(x)timer.toc()print ('Do once forward need {:.3f}ms ').format(timer.total_time*1000/100.0)
阅读全文
0 0
原创粉丝点击