推理评估脚本

[GuarDream]

你好，注意到TIGER论文里有提到CPU和GPU的推理时间和显存占用，请问这部分评估代码方便更新吗，这样方便后续论文与TIGER论文的指标进行对比。

[JusperLee]

# Params and MACs
with torch.cuda.device(0):
    a = torch.randn(1, 1, 16000).cuda()
    total_macs = 0
    total_params = 0
    model = audiomodel.cuda()
    with torch.no_grad():
        macs, params = get_model_complexity_info(
            model, (16000,), as_strings=False, print_per_layer_stat=True, verbose=False
        )
        # print(model(a).shape)
    total_macs += macs
    total_params += params
    print("MACs: ", total_macs / 10.0 ** 9)
    print("Params: ", total_params / 10.0 ** 6)

torch.cuda.empty_cache()

# Inference time CPU (4s)
data = torch.randn(4, 1, 16000)
audiomodel = audiomodel.cpu()
audiomodel.eval()
start_time = time.time()
with torch.no_grad():
    for i in tqdm(range(100)):
        audiomodel(data)
print("Inference CPU Time: ", (time.time() - start_time) / 100)

torch.cuda.empty_cache()

# Inference time GPU (4s)
data = data.cuda()
audiomodel = audiomodel.cuda()
audiomodel.eval()
# 设置用于测量时间的 cuda Event, 这是PyTorch 官方推荐的接口,理论上应该最靠谱
starter, ender = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
# 初始化一个时间容器
timings = np.zeros((1000, 1))
with torch.no_grad():
    for i in tqdm(range(1000)):
        starter.record()
        audiomodel(data)
        ender.record()
        torch.cuda.synchronize() # 等待GPU任务完成
        curr_time = starter.elapsed_time(ender) # 从 starter 到 ender 之间用时,单位为毫秒
        timings[i] = curr_time
print("Inference GPU Time: ", timings.sum()/1000)

torch.cuda.empty_cache()

# Inference GPU Memory
data = torch.randn(4, 1, 16000).cuda()
audiomodel = audiomodel.cuda()
# 清空缓存并同步
torch.cuda.empty_cache()
torch.cuda.synchronize()
# 记录初始GPU内存使用情况
start_memory_allocated = torch.cuda.memory_allocated()
start_memory_reserved = torch.cuda.memory_reserved()
with torch.no_grad():
    audiomodel(data)
# 记录推理期间最大GPU内存使用情况
max_memory_allocated = torch.cuda.max_memory_allocated()
max_memory_reserved = torch.cuda.max_memory_reserved()

# 计算总内存消耗
total_allocated_memory = max_memory_allocated - start_memory_allocated
print(f"Test GPU memory: {total_allocated_memory / (1024 ** 2):.2f} MB")

torch.cuda.empty_cache()

# Backward time GPU (4s)
data = torch.randn(1, 1, 16000).cuda()
target = torch.randn(1, 2, 16000).cuda()
audiomodel = audiomodel.cuda()
audiomodel.train()
opt = torch.optim.Adam(audiomodel.parameters(), lr=0.001)
loss = hydra.utils.instantiate(cfg.loss)
# 设置用于测量时间的 cuda Event, 这是PyTorch 官方推荐的接口,理论上应该最靠谱
starter, ender = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
# 初始化一个时间容器
timings = np.zeros((1000, 1))
for i in tqdm(range(1000)):
    starter.record()
    opt.zero_grad()
    out = audiomodel(data)
    l = loss(out.view(1, 2, 16000), target)
    l.backward()
    opt.step()
    ender.record()
    torch.cuda.synchronize() # 等待GPU任务完成
    curr_time = starter.elapsed_time(ender) # 从 starter 到 ender 之间用时,单位为毫秒
    timings[i] = curr_time
print("Training Time: ", timings.sum()/1000) 

torch.cuda.empty_cache()

# Backward GPU Memory
data = torch.randn(1, 1, 16000).cuda()
target = torch.randn(1, 2, 16000).cuda()
audiomodel = audiomodel.cuda()
audiomodel.train()
opt = torch.optim.Adam(audiomodel.parameters(), lr=0.001)
loss = hydra.utils.instantiate(cfg.loss)
# 清空缓存并同步
torch.cuda.empty_cache()
torch.cuda.synchronize()
# 记录初始GPU内存使用情况
start_memory_allocated = torch.cuda.memory_allocated()
start_memory_reserved = torch.cuda.memory_reserved()
opt.zero_grad()
out = audiomodel(data)
l = loss(out.view(1, 2, 16000), target)
l.backward()
opt.step()
# 记录推理期间最大GPU内存使用情况
max_memory_allocated = torch.cuda.max_memory_allocated()
max_memory_reserved = torch.cuda.max_memory_reserved()

# 计算总内存消耗
total_allocated_memory = max_memory_allocated - start_memory_allocated
print(f"Train GPU Memory: {total_allocated_memory / (1024 ** 2):.2f} MB")