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

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


我担心以后会遇到问题。这是好习惯吗?

最佳答案

当脚本很小并且可以在启动时固定要使用的名称空间时,我经常使用名为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.onescupy.random.*)。

10-07 15:35