Python 基于GPU的PyTorch慢批量矩阵乘法

我在两个大小分别为

（100128128）

的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

GPU:

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

CPU在设备（GPU）上启动内核*

CPU在GPU上分配内存

CPU将输入数据复制到GPU

CPU在GPU上启动内核来处理输入数据

CPU将输出结果复制到自身

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

CPU

鉴于，0.0323秒是以下任一项所用的总时间：

CPU内通信

CPU和主存通信

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

结论

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

---

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

原来我的GPU忙于其他事情，cuda上变量和结果的初始化需要时间。我所要做的就是调用

torch.cuda.synchronize（）

来确定时间。