GPU 上的 PyTorch 慢批矩阵乘法

Question

我在 2 个大小为 (100 , 128 , 128 ) 的 3d 张量上使用批量矩阵乘法。

import torch
a = torch.randn(100,128,128)
b = torch.randn(100,128,128)

import time

t0 = time.time()
torch.bmm(a,b)
print(time.time() - t0)

0.03233695030212402

现在如果我在 GPU 上做同样的事情需要更长的时间

a = a.cuda()
b = b.cuda()
t0 = time.time()
torch.bmm(a,b)
print(time.time() - t0)

30.574532985687256

为什么在 GPU 上求解需要这么长时间？

我有一个 GTX 1050 Ti 4GB

和处理器核心 i3-6100 3.7Ghz

Answer 1

GPU：

30.57 秒是以下步骤所花费的总时间：

1. CPU 在设备 (GPU) 上启动内核*
2. CPU 在 GPU 上分配内存
3. CPU 将输入数据复制到 GPU
4. CPU 在 GPU 上启动内核以处理输入数据
5. CPU 将输出结果复制到自身

*内核是一个序列号，是原始代码的一小部分。

CPU

而 0.0323 秒是以下任一者所花费的总时间：

1. CPU 内通信
2. CPU 和主存通信

这些速度非常快，而且它们还必须执行更少的命令来完成工作。

结论

因此，在命令的复杂性（CPU 命令更复杂）和由于命令的并行性和更简单的性质（GPU 可以触发更多并行线程并且它们在本质上更简单）导致的加速之间进行权衡。

如果要完成更多计算，那么 GPU 加速将超过 CPU-GPU 通信所花费的额外时间。