切片300MB的CuPy阵列比NumPy慢约5倍



我的代码涉及对432x432x400个数组进行总计约1000万次的切片,以生成用于神经网络训练的批量数据。由于这些都是相当大的阵列(9200万个数据点/30MB(,我希望使用CuPy来加快速度(甚至可能通过在训练时在同一GPU上生成数据来加快训练速度(,但发现它实际上使代码慢了约5倍。

这种预期行为是由于CuPy管理费用造成的,还是我遗漏了什么?

要复制的代码:

import cupy as cp
import numpy as np
import timeit
cp_arr = cp.zeros((432, 432, 400), dtype=cp.float32)
np_arr = np.zeros((432, 432, 400), dtype=np.float32)
# numbers below are representative of my code
cp_code = 'arr2 = cp_arr[100:120, 100:120, 100:120]'
np_code = 'arr2 = np_arr[100:120, 100:120, 100:120]'
timeit.timeit(cp_code, number=8192*4, globals=globals())  # prints 0.122
timeit.timeit(np_code, number=8192*4, globals=globals())  # prints 0.027

设置:

  • GPU:NVIDIA Quadro P4000

  • CuPy版本:7.3.0

  • 操作系统:CentOS Linux 7

  • CUDA版本:10.1

  • cuDNN版本:7.6.5

NumPy和CuPy中的Slicing实际上并不是将数据复制到任何地方,而是简单地返回一个新数组,其中数据相同,但其指针偏移到新切片的第一个元素并调整形状。下面注意原始数组和切片是如何具有相同的步幅的:
In [1]: import cupy as cp
In [2]: a = cp.zeros((432, 432, 400), dtype=cp.float32)
In [3]: b = a[100:120, 100:120, 100:120]
In [4]: a.strides
Out[4]: (691200, 1600, 4)
In [5]: b.strides
Out[5]: (691200, 1600, 4)

通过用NumPy替换CuPy可以验证上述内容。

如果您想为实际的切片操作计时,最可靠的方法是为每个操作添加一个.copy(),从而强制执行内存访问/复制:

cp_code = 'arr2 = cp_arr[100:120, 100:120, 100:120].copy()'  # 0.771 seconds
np_code = 'arr2 = np_arr[100:120, 100:120, 100:120].copy()'  # 0.154 seconds

不幸的是,对于上述情况,内存模式对GPU不利,因为小块无法使内存通道饱和,因此它仍然比NumPy慢。然而,如果块能够接近内存通道饱和,CuPy可以更快,例如:

cp_code = 'arr2 = cp_arr[:, 100:120, 100:120].copy()'  # 0.786 seconds
np_code = 'arr2 = np_arr[:, 100:120, 100:120].copy()'  # 2.911 seconds

我还确认了cupy中的切片速度大约慢了5倍,而有一种更精确的方法来测量时间(例如。https://github.com/cupy/cupy/pull/2740)。

数组的大小无关紧要,因为切片操作不是复制数据,而是创建视图。以下结果类似。

cp_arr = cp.zeros((4, 4, 4), dtype=cp.float32)
cp_code = 'arr2 = cp_arr[1:3, 1:3, 1:3]'

"获取切片然后将其发送到GPU"是很自然的,因为它减少了要传输的字节。如果第一个预处理是切片,请考虑这样做。