"import numpy as cp"处理非 GPU 情况的好做法吗？

Question

我正在使用CuPy编写理想地在GPU上运行以提高速度的代码。但是，我希望代码能够使用 numpy 实现运行(尽管速度更慢(。 目前我正在执行以下操作：

import numpy as np
if gpu_present:
import cupy as cp
else:
import numpy as cp

我担心以后可能会遇到问题。这是好的做法吗？

Answer 1

当脚本很小并且要使用的命名空间可以在启动时修复时，我经常使用名为xp的全局变量(与您的解决方案相同(。我有时也使用的类似模式是使其成为类的实例属性(再次命名为xp(;对于将来的扩展，它更容易容忍，因为每个实例都可以为该属性具有不同的值。一种类似、更健壮但繁琐的方法是将xp作为每个函数的第一个参数。

在编写可以在任何情况下使用的库时(例如，多线程代码，在单个进程中同时使用 NumPy 和 CuPy(，最好让每个函数/类为参数适当地处理命名空间。我经常为此目的使用get_array_module实用程序。CuPy具有此功能，尽管它需要安装CuPy。链器也有。自己写也很简单。使用此实用程序，可以使代码可用于 NumPy 或 CuPy 数组，而无需全局开关。

另请注意，NumPy>=1.17 可以将 CuPy 数组调度到适当的 CuPy 例程，因此在大多数情况下，您可以将 CuPy 数组直接传递给numpy.*函数。如果你的代码只对给定的数组进行计算，你甚至根本不需要使用cupy命名空间(你仍然需要使用它来创建一个新数组，而不给另一个数组，如cupy.ones和cupy.random.*(。