宏碁RTX3070笔记本Python GPU计算报错求助（Numba脚本）

本人使用搭载RTX3070的宏碁笔记本，已安装CUDA、cuDNN、TensorFlow，且TensorFlow与CUDA测试脚本均可检测到GPU并显示状态。但运行用于测试GPU计算的Numba JIT脚本时出现KeyError错误，相关脚本及错误信息如下：

测试脚本 #

from numba import jit, cuda
import numpy as np
# to measure exec time
from timeit import default_timer as timer


# normal function to run on cpu
def func(a):
    for i in range(10000000):
        a[i] += 1

    # function optimized to run on gpu


@jit(target_backend='cuda')
def func2(a):
    for i in range(10000000):
        a[i] += 1


if __name__ == "__main__":
    n = 10000000
    a = np.ones(n, dtype=np.float64)

    start = timer()
    func(a)
    print("without GPU:", timer() - start)

    start = timer()
    func2(a)
    print("with GPU:", timer() - start)

错误信息 #

"C:\Program Files\Python310\python.exe" E:\Programming\Projects\CUDA\TestingGPU.py 
without GPU: 1.1927146000000448
Traceback (most recent call last):
  File "E:\Programming\Projects\CUDA\TestingGPU.py", line 30, in <module>
    func2(a)
  File "C:\Users\Faheem\AppData\Roaming\Python\Python310\site-packages\numba\core\dispatcher.py", line 442, in _compile_for_args
    raise e
  File "C:\Users\Faheem\AppData\Roaming\Python\Python310\site-packages\numba\core\dispatcher.py", line 375, in _compile_for_args
    return_val = self.compile(tuple(argtypes))
  File "C:\Users\Faheem\AppData\Roaming\Python\Python310\site-packages\numba\core\dispatcher.py", line 905, in compile
    cres = self._compiler.compile(args, return_type)
  File "C:\Users\Faheem\AppData\Roaming\Python\Python310\site-packages\numba\core\dispatcher.py", line 80, in compile
    status, retval = self._compile_cached(args, return_type)
  File "C:\Users\Faheem\AppData\Roaming\Python\Python310\site-packages\numba\core\dispatcher.py", line 94, in _compile_cached
    retval = self._compile_core(args, return_type)
  File "C:\Users\Faheem\AppData\Roaming\Python\Python310\site-packages\numba\core\dispatcher.py", line 103, in _compile_core
    self.targetdescr.options.parse_as_flags(flags, self.targetoptions)
  File "C:\Users\Faheem\AppData\Roaming\Python\Python310\site-packages\numba\core\options.py", line 41, in parse_as_flags
    opt._apply(flags, options)
  File "C:\Users\Faheem\AppData\Roaming\Python\Python310\site-packages\numba\core\options.py", line 66, in _apply
    raise KeyError(m)
KeyError: "Unrecognized options: {'target_backend'}. Known options are dict_keys(['_dbg_extend_lifetimes', '_dbg_optnone', '_nrt', 'boundscheck', 'debug', 'error_model', 'fastmath', 'forceinline', 'forceobj', 'inline', 'looplift', 'no_cfunc_wrapper', 'no_cpython_wrapper', 'no_rewrites', 'nogil', 'nopython', 'parallel'])"

Process finished with exit code 1

问题分析与解决方案 #

核心原因 #

错误提示明确显示，target_backend不是@jit装饰器的合法参数。你混淆了Numba通用JIT与CUDA专用JIT的用法：target_backend是针对CPU JIT的配置参数，要在GPU上运行代码，必须使用Numba提供的@cuda.jit专用装饰器。

修正后的完整脚本 #

除了替换装饰器，还需要遵循GPU并行编程模型——不能直接用CPU式的串行循环，需通过线程索引分配计算任务，同时手动管理CPU与GPU之间的数据传输：

from numba import jit, cuda
import numpy as np
from timeit import default_timer as timer

# CPU版本函数
def func(a):
    for i in range(10000000):
        a[i] += 1

# GPU版本函数（采用Numba CUDA专用装饰器）
@cuda.jit
def func2(a):
    # 获取当前线程的全局索引
    i = cuda.grid(1)
    # 确保索引不越界
    if i < a.size:
        a[i] += 1

if __name__ == "__main__":
    n = 10000000
    a = np.ones(n, dtype=np.float64)
    # 将数据从CPU内存复制到GPU内存
    d_a = cuda.to_device(a)

    start = timer()
    func(a)
    print("without GPU:", timer() - start)

    start = timer()
    # 配置GPU线程块与网格大小（通用计算配置方式）
    threads_per_block = 256
    blocks_per_grid = (d_a.size + threads_per_block - 1) // threads_per_block
    # 调用GPU函数，指定网格和线程块参数
    func2[blocks_per_grid, threads_per_block](d_a)
    # 将计算结果从GPU内存复制回CPU内存
    d_a.copy_to_host(a)
    print("with GPU:", timer() - start)

关键说明 #

1. 装饰器替换：用@cuda.jit替代@jit(target_backend='cuda')，这是Numba中专门用于CUDA GPU编程的入口。
2. 并行计算模型：通过cuda.grid(1)获取线程全局索引，让每个GPU线程处理一个数组元素，充分利用GPU的并行计算能力。
3. 数据传输管理：GPU无法直接访问CPU内存，必须用cuda.to_device()和copy_to_host()手动完成数据在CPU与GPU之间的传输。

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