为什么在 pytorch 中批量计算比 for 循环花费更多时间？

Question

我的代码有问题吗？ 当我一起运行以下代码时，所花费的时间有很大差异： 批处理时间：1.3649392127990723 时间：0.7864551544189453 当我单独运行它们时，它们花费的时间几乎相同

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()



    torch.cuda.synchronize()
    start_time = time.time()
    out = a(b)
    torch.cuda.synchronize()
    end_time = time.time()
    print("batch time :", end_time - start_time)



    torch.cuda.synchronize()
    start_time = time.time()
    for i in range(20):
        out = a(c[i])
    torch.cuda.synchronize()
    end_time = time.time()
    print("for time:", end_time - start_time)

Answer 1

时间差异是由于代码示例具有不同的内存分配。

你的第一个例子：

# ex1
torch.cuda.synchronize()
start_time = time.time()
out = a(b)
torch.cuda.synchronize()
end_time = time.time()
print("batch time :", end_time - start_time)

out

张量的大小为

(20, 5000, 5000)

。

在你的第二个例子中：

# ex2
torch.cuda.synchronize()
start_time = time.time()
for i in range(20):
    out = a(c[i])
torch.cuda.synchronize()
end_time = time.time()
print("for time:", end_time - start_time)

out

张量的大小为

(5000, 5000)

。

ex1

显得较慢，因为它必须为更大的输出张量分配内存。进程第一次分配内存时，此开销也会更高（请参阅下面的时间结果）。

考虑第三个示例，其中我们使用

ex2

中的代码，但具有与

ex1

相同的输出大小：

# ex3
torch.cuda.synchronize()
start_time = time.time()
out = torch.empty(20, 5000, 5000)
for i in range(20):
    out[i] = a(c[i])
torch.cuda.synchronize()
end_time = time.time()
print("for time:", end_time - start_time)

在我的 GPU 上，我得到以下结果：

• ex1

  首次运行：0.75秒（进程首次内存分配的额外开销）

• ex1

  第二次运行：0.26 秒

• ex2

  ：0.28秒

• ex3

  ：1.13秒