Pytorch 记录 每一次迭代的 数据加载、模型计算 等时长

Pytorch 记录 每一次迭代的 数据加载、模型计算 等时长

使用 time 库如何记录pytorch处理每一批次时间消耗

可结合 tqdm 工具使用 效果更佳

import time

class Process:
    """
     Pytorch 记录 每一次迭代的 数据加载、模型计算 等时长
    """
    def __init__(self) -> None:
        pass

    def record_time(self):
        self.cur_time = time.time()
        # print("self.cur_time:", self.cur_time)
        return self.cur_time

    def split_time(self):
        split_time = time.time() - self.cur_time
        self.record_time()
        return split_time

    def start(self):
        self.record_time()
        timer = dict(dataloader=0.001, model=0.001, statistics=0.001)
        # or 
        # for i,data,label in enumerate(data_lodaer):
        for i in range(3): 
            time.sleep(2)
            # do something
            # time for datalodar each batch
            timer['dataloader'] += self.split_time()

            # do something
            # time for model(data_input)
            time.sleep(3)
            timer['model'] += self.split_time()
            
            # do something else
            # ...

        print(timer)

p = Process()
p.start()

运行结果

# print(timer)
{'dataloader': 6.014460874557495, 'model': 9.015425992965698, 'statistics': 0.001}