python - Pytorch 速度比较——GPU 比 CPU 慢

Question

我试图找出 GPU 张量运算是否真的比 CPU 张量运算更快。所以，我在下面编写了这个特定的代码来依次实现 CPU 张量和 GPU cuda 张量的简单 2D 添加，以查看速度差异:

import torch
import time

###CPU
start_time = time.time()
a = torch.ones(4,4)
for _ in range(1000000):
    a += a
elapsed_time = time.time() - start_time

print('CPU time = ',elapsed_time)

###GPU
start_time = time.time()
b = torch.ones(4,4).cuda()
for _ in range(1000000):
    b += b
elapsed_time = time.time() - start_time

print('GPU time = ',elapsed_time)

令我惊讶的是，CPU 时间为 0.93 秒，GPU 时间高达 63 秒。我是否正确地执行了 cuda 张量操作，或者 cuda 张量的概念是否仅在非常复杂的操作(如神经网络)中工作得更快？

注意:我的 GPU 是 NVIDIA 940MX，torch.cuda.is_available() 调用返回 True。

Answer 1

GPU 加速通过计算的大量并行化工作。在 GPU 上，你有大量的内核，每个内核都不是很强大，但这里的大量内核很重要。

像 PyTorch 这样的框架尽其所能使并行计算成为可能。一般来说，矩阵运算非常适合并行化，但并不总是可以并行化计算!

在你的例子中你有一个循环:

b = torch.ones(4,4).cuda()
for _ in range(1000000):
    b += b

您有 1000000 操作，但由于代码的结构，不可能并行化这些计算中的大部分。如果您考虑一下，要计算下一个 b，您需要知道上一个(或当前) b。

因此您有 1000000 个操作，但每个操作都必须一个接一个地计算。可能的并行化仅限于张量的大小。这个尺寸虽然在你的例子中不是很大:

torch.ones(4,4)

因此每次迭代只能并行化 16 个操作(加法)。由于 CPU 有几个，但有更多强大的核心，对于给定的例子来说它要快得多!

但如果您更改张量的大小，情况就会发生变化，然后 PyTorch 能够并行化更多的整体计算。我将迭代次数更改为 1000(因为我不想等那么久:)，但您可以输入任何您喜欢的值，CPU 和 GPU 之间的关系应该保持不变。

以下是不同张量大小的结果:

#torch.ones(4,4)       - the size you used
CPU time =  0.00926661491394043
GPU time =  0.0431208610534668

#torch.ones(40,40)     - CPU gets slower, but still faster than GPU
CPU time =  0.014729976654052734
GPU time =  0.04474186897277832

#torch.ones(400,400)   - CPU now much slower than GPU
CPU time =  0.9702610969543457
GPU time =  0.04415607452392578

#torch.ones(4000,4000) - GPU much faster then CPU 
CPU time =  38.088677167892456
GPU time =  0.044649362564086914

因此，正如您所见，在可以并行化东西的地方(这里添加了张量元素)，GPU 变得非常强大。
对于给定的计算，GPU 时间根本没有改变，GPU 可以处理更多!
(只要它不会耗尽内存:)