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将输出结果复制到自身

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

<强> CPU

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

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

这些非常快，他们还必须执行更少的命令才能完成工作。

<强>结论

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

如果要进行更多计算，那么GPU加速将超过CPU-GPU通信所需的额外时间。