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

Question

我正在使用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 调用，但这并没有改变任何东西。

Answer 1

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

如果我们为cupy.meshgrid 研究the documentation，我们会看到它返回：

返回：cupy.ndarray 列表

丘比特documentation具体说：

目前，cupy.array() 或 cupy.asarray() 无法从包含 CuPy 数组（例如，CuPy 数组列表）的 Python 对象创建数组。请改用cupy.stack()。

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

$ 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
$