如何优雅编写优先用CuPy、自动回退到NumPy/SciPy的模块？

我正在编写一个模块，计划使用CuPy（基于CUDA实现NumPy和SciPy API的库）来加速工作中的数据分析，但希望该模块也能在无法使用CuPy的电脑上运行。

原本考虑的写法：

try:
    import cupy as np
    import cupyx.scipy as scipy
except ModuleNotFoundError:
    import numpy as np
    import scipy

但这种写法存在两个问题：

• 仅覆盖了CuPy未安装的情况，未考虑CuPy已安装但CUDA无法运行或GPU无法检测的场景；
• CuPy并非NumPy的完美替代品，例如将CuPy数组传入仅支持NumPy数组的非CuPy函数时，需调用.get()方法，否则会报错。

后来考虑给每个函数添加use_gpu参数，通过条件判断执行对应CuPy或NumPy的代码，示例如下：

import numpy # 有时即使CuPy可用也需要使用NumPy函数
try:
    import cupy as np
except ModuleNotFoundError:
    import numpy as np
    print("CuPy not available, only use_gpu=False will work")
import pandas as pd

def some_function(arg1, arg2, use_gpu=True):
    """Just some example function doing something arbitrary."""
    x = np.func1(arg1)
    y = np.func2(arg2)
    res = np.hstack(x, y)
    if use_gpu:
        return pd.DataFrame(res.get())
    else:
        return pd.DataFrame(res)

请问是否存在更优雅的实现方式？

1. 完善CuPy可用性检测 #

不要只检查是否安装，还要验证CUDA环境能否正常运行，把这部分逻辑封装成工具代码：

import numpy as np

# 全局标记CuPy是否可用（安装且CUDA能正常运行）
CUPY_AVAILABLE = False
try:
    import cupy as cp
    # 执行简单操作验证CUDA环境
    cp.array([1,2,3])
    CUPY_AVAILABLE = True
except (ImportError, RuntimeError):
    # 捕获导入失败或CUDA运行错误
    pass

# 统一获取数组处理库的工具函数
def get_array_lib(use_gpu=True):
    if use_gpu and CUPY_AVAILABLE:
        return cp
    else:
        return np

这样就能覆盖“CuPy已安装但CUDA不可用”的场景，避免后续代码执行时触发意外错误。

2. 封装数组转换工具 #

针对CuPy数组转NumPy的需求，写一个统一的转换函数，不用在每个业务函数里重复写条件判断：

def to_numpy(arr):
    """将CuPy数组转为NumPy数组，NumPy数组直接返回"""
    if CUPY_AVAILABLE and isinstance(arr, cp.ndarray):
        return arr.get()
    return arr

之后在需要传入非CuPy兼容函数（比如Pandas）时，直接调用to_numpy(res)即可：

def some_function(arg1, arg2, use_gpu=True):
    np_lib = get_array_lib(use_gpu)
    x = np_lib.func1(arg1)
    y = np_lib.func2(arg2)
    res = np_lib.hstack([x, y])
    return pd.DataFrame(to_numpy(res))

3. 用装饰器统一处理GPU/CPU切换 #

如果有大量函数需要支持use_gpu参数，可以写一个装饰器，自动处理数组库的选择逻辑：

from functools import wraps

def gpu_supported(func):
    @wraps(func)
    def wrapper(*args, use_gpu=True, **kwargs):
        np_lib = get_array_lib(use_gpu)
        # 把数组处理库作为参数传给业务函数
        return func(*args, np_lib=np_lib, **kwargs)
    return wrapper

使用装饰器后，业务函数可以简化，不用再关心库的选择：

@gpu_supported
def some_function(arg1, arg2, np_lib):
    x = np_lib.func1(arg1)
    y = np_lib.func2(arg2)
    res = np_lib.hstack([x, y])
    return pd.DataFrame(to_numpy(res))

调用时直接指定use_gpu=True/False即可，比如some_function(a, b, use_gpu=False)。

4. 全局配置开关 #

可以设置全局的GPU启用状态，避免每个函数都要传use_gpu参数：

# 全局配置：默认自动检测后启用GPU（如果可用）
USE_GPU = CUPY_AVAILABLE

def set_use_gpu(enable):
    global USE_GPU
    if enable and not CUPY_AVAILABLE:
        print("CuPy/CUDA不可用，无法启用GPU加速")
        USE_GPU = False
    else:
        USE_GPU = enable

# 修改get_array_lib，默认使用全局配置
def get_array_lib(use_gpu=None):
    if use_gpu is None:
        use_gpu = USE_GPU
    if use_gpu and CUPY_AVAILABLE:
        return cp
    else:
        return np

这样可以在模块初始化时统一设置是否启用GPU，比如set_use_gpu(False)强制用CPU；单个函数调用时也可以临时覆盖全局配置，比如some_function(a, b, use_gpu=True)。

内容的提问来源于stack exchange，提问作者Colin