时间戳

首页 > 解决方案 > Cupy 索引很慢

python-3.x - Cupy 索引很慢

问题描述

我正在尝试对大小为 16000 的大型 cupy 数组执行操作。我发现诸如加法之类的数学运算非常快，但使用布尔掩码进行索引的速度相对较慢。例如，下面的代码：

import cupy as cp
arr = cp.random.normal(0, 1, 16000)
%timeit arr * 5
%timeit arr > 0.4
%timeit arr[arr > 0.4] = 0

给我输出：

28 µs ± 950 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
26.5 µs ± 1.61 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
104 µs ± 2.6 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

最终索引速度至少慢两倍的任何原因？我假设乘法应该比设置数组元素慢。

更新：这不适用于 numpy 索引。将cupy数组更改为numpy，我得到：

6.71 µs ± 373 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
4.42 µs ± 56.5 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
5.39 µs ± 29.6 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

标签： python-3.xperformancecupy

解决方案

在第三种情况下，cupy通过一系列操作组合结果：cupy_greater, cupy_copy, inclusive_scan_kernel, inclusive_scan_kernel, add_scan_blocked_sum_kernel, CUDA memcpy DtoH（可能提供需要设置为零的元素数量），CUDA memset（可能将数组设置为零），最后是cupy_scatter_update_mask(也许将零分散到正确的位置）。

这是一个比 复杂得多的序列，它似乎在引擎盖下arr*5运行单个。使用用户定义的内核cupy_multiply可能会做得更好：cupy

import cupy as cp
clamp_generic = cp.ElementwiseKernel(
        'T x, T c',
        'T y',
        'y = (y > x)?c:y',
        'clamp_generic')
arr = cp.random.normal(0, 1, 16000)
clamp_generic(0.4, 0, arr)

推荐阅读