如何使用Python编写代码调用NVIDIA 3060的CUDA核心？

如果你不想学新语言，以下几个Python工具能帮你轻松调用GPU，不用手动写CUDA代码：

1. Numba #

这是最适合新手的选择之一，完全基于Python语法，只需给函数加个装饰器，就能把它编译成CUDA代码在GPU上跑。

示例代码：

from numba import cuda
import numpy as np

# 用@cuda.jit标记这是GPU核函数
@cuda.jit
def add_arrays(a, b, result):
    # 获取当前线程的索引
    idx = cuda.grid(1)
    # 确保索引不越界
    if idx < a.size:
        result[idx] = a[idx] + b[idx]

# 创建CPU数组
a = np.random.rand(1000000)
b = np.random.rand(1000000)
result = np.zeros_like(a)

# 配置GPU的线程块和网格大小
threads_per_block = 256
blocks_per_grid = (a.size + threads_per_block - 1) // threads_per_block

# 调用GPU函数，格式是[网格大小, 线程块大小](参数)
add_arrays[blocks_per_grid, threads_per_block](a, b, result)

优势：几乎不用改动现有Python代码，自动处理CUDA底层细节，学习成本极低。

2. CuPy #

可以把它理解成GPU版的NumPy，API和NumPy完全一致，如果你平时用NumPy做数值计算，直接把import numpy as np换成import cupy as cp，大部分代码不用改就能在GPU上运行。

示例代码：

import cupy as cp

# 直接在GPU上创建数组
a = cp.random.rand(1000000)
b = cp.random.rand(1000000)

# 和NumPy一样写运算，自动在GPU执行
result = a + b

# 要是需要把结果转回CPU，用cp.asnumpy()
result_cpu = cp.asnumpy(result)

优势：零学习成本，是把现有数值计算代码迁移到GPU最简单的方式。

3. PyTorch/TensorFlow（深度学习场景） #

如果你是做深度学习，这两个框架完全用Python编写，自动利用CUDA加速，不用手动写任何CUDA代码。

PyTorch示例：

import torch

# 自动检测是否有可用的GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 创建GPU上的张量（类似NumPy数组）
a = torch.rand(1000000, device=device)
b = torch.rand(1000000, device=device)

# 直接做运算，全程在GPU上执行
result = a + b

优势：针对深度学习优化，自动管理GPU内存和计算流程，适合构建训练神经网络。

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