GPU 上的 PyTorch Slow Batch 矩阵乘法

PyTorch Slow Batch matrix multiplication on GPU

我在每个大小为 (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

GPU:

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

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

*kernel是串口代码，是原代码的一小部分

CPU

然而，0.0323 秒是以下两者所花费的总时间：

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

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

结论

因此，这是命令的复杂性（CPU 命令更复杂）和由于并行性和命令的简单性质（GPU 可以触发更多并行线程，它们是本质上要简单得多）。

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