PyTorch在GPU上进行慢速批量矩阵乘法

Question

我在每个尺寸为(100 , 128 , 128 )的2个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将输出结果复制到自身

* Kernel是一个串行代码，它只是原始代码的一小部分。

中央处理器

然而，0.0323秒是两者之间的总时间：

1. CPU内部通信
2. CPU和主存储器通信

它们非常快，并且还必须执行更少的命令才能完成工作。

结论

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

如果有更多的计算要做，那么GPU加速将压倒CPU-GPU通信所需的额外时间。