numpy - "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 )；对于 future 的扩展来说，它更能容忍，因为每个实例都可以为该属性具有不同的值。一个类似的、更健壮但麻烦的方法是使 xp每个函数的第一个参数。

在编写可在任何情况下使用的库时(例如，多线程代码，在单个进程中同时使用 NumPy 和 CuPy)，最好让每个函数/类为参数适本地处理命名空间。我经常用get_array_module用于该目的的实用程序。 CuPy has this function ，尽管它需要安装 CuPy。 Chainer also has it .自己写也很简单。使用此实用程序，您可以使代码与 NumPy 或 CuPy 数组一起使用，而无需全局开关。

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