“将numpy作为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 has this function，尽管它需要安装CuPy。 Chainer also has it。自己编写也很简单。使用此实用程序，您可以使代码可用于NumPy或CuPy数组而无需全局切换。

还要注意，NumPy> = 1.17可以将CuPy数组分派到适当的CuPy例程，因此在大多数情况下，您可以将CuPy数组直接传递给numpy.*函数。如果您的代码仅在给定的数组上进行计算，则甚至根本不需要使用cupy命名空间（您仍需要使用它来创建新数组，而无需提供另一个数组，例如cupy.ones和{ {1}}。