从 python 列表创建一个 cupy 数组很慢

creatiing a cupy array from python list is slow

我正在使用 cupy 执行以下操作,这非常快:

import cupy as cp

shape = (256, 170, 256)

deformation = cp.meshgrid(cp.arange(shape[0]),
                          cp.arange(shape[1]),
                          cp.arange(shape[2]),
                          indexing='ij')

但是,如果我将其转换为数组:

deformation = cp.array(cp.meshgrid(cp.arange(shape[0]),
                                   cp.arange(shape[1]),
                                   cp.arange(shape[2]),
                                   indexing='ij'))

这似乎很慢或者只是挂起(我在 5 分钟后放弃了)。我不确定我在这里做错了什么。

我也尝试将 copy=False 传递给 cp.array 调用,但这并没有改变任何东西。

我认为不支持将 cupy 数组列表转换为 cupy 数组。如果我让你的 shape 小得多,例如(8,8,8) 我收到 python 错误。

如果我们为 cupy.meshgrid 研究 the documentation,我们会发现它 returns:

Returns: list of cupy.ndarray

cupy documentation 具体说:

Currently, cupy.array() or cupy.asarray() cannot create an array from Python object containing CuPy array (e.g., a list of CuPy arrays). Use cupy.stack() instead.

使用那里的建议,这对我来说似乎工作得相对较快:

$ cat t6.py
import cupy as cp

shape = (256, 170, 256)

deformation = cp.stack(cp.meshgrid(cp.arange(shape[0]),
                          cp.arange(shape[1]),
                          cp.arange(shape[2]),
                          indexing='ij'))
$ time python t6.py

real    0m1.281s
user    0m0.608s
sys     0m0.492s
$