PyTorch中批量计算为何比for循环耗时更长？代码是否有问题？

代码本身无功能问题，耗时差异源于测试顺序的干扰 #

你遇到的批量计算比for循环慢的情况，本质是CUDA初始化耗时的分配差异，而非批量计算本身效率更低：

• 连续运行测试时，第一个执行的批量计算会触发一系列CUDA初始化操作：包括Linear层权重向GPU的迁移、CUDA kernel首次编译加载、GPU内存分配优化等，这些额外的初始化耗时全部被计入了批量计算的总时间。
• 当轮到for循环测试时，GPU已经完成初始化（处于"预热"状态），此时仅统计纯推理耗时，因此看起来比批量计算更快。

而单独运行任意一个测试时，每个测试都会独自承担初始化的额外耗时，所以两者总耗时几乎一致。

优化测试方法，获取真实性能对比 #

理论上GPU的批量并行计算效率远高于循环单样本计算，你可以通过以下方式修正测试：

1. 增加预热步骤，提前完成CUDA初始化，避免干扰正式计时；
2. 加入torch.no_grad()关闭自动求导，减少计算图构建的额外开销。

优化后的测试代码如下：

import time
import torch
import torch.nn as nn
if __name__ == '__main__':
    a= nn.Linear(5000, 5000).cuda()
    b = torch.randn(20, 5000, 5000).cuda()
    c = torch.randn(20, 5000, 5000).cuda()

    # 预热：提前触发CUDA初始化和kernel编译
    with torch.no_grad():
        a(b[:1])
        torch.cuda.synchronize()

    # 测试批量计算
    torch.cuda.synchronize()
    start_time = time.time()
    with torch.no_grad():
        out = a(b)
    torch.cuda.synchronize()
    end_time = time.time()
    print("batch time :", end_time - start_time)

    # 测试for循环计算
    torch.cuda.synchronize()
    start_time = time.time()
    with torch.no_grad():
        for i in range(20):
            out = a(c[i])
    torch.cuda.synchronize()
    end_time = time.time()
    print("for time:", end_time - start_time)

调整后再测试，你会发现批量计算的耗时会明显低于for循环，符合GPU并行计算的预期。

内容的提问来源于stack exchange，提问作者lvguima