Python中利用GPU对大型数组快速计算tan()与arctan()的最优方案

需对超大型数组执行tan()和arctan()运算：数据共8000万行、300万列，可分批处理（每批约20行，即规模为(20, 3000000)的数组）。此前尝试用Numba的@njit结合np.tan()/np.arctan()，但代码仅运行在CPU；math库仅支持标量运算，无法处理数组。需明确：

• Numba能否实现GPU加速的数组三角函数运算？
• PyTorch是否适用，是否存在过多开销？
• 针对这种分批的大型数组运算，最快方案是什么？选CPU还是GPU？用Numba还是PyTorch？用@njit还是@vectorize？

一、Numba实现GPU加速 #

Numba的CUDA后端可实现数组的GPU并行三角函数运算，有两种常用方式：

1. 手写CUDA核函数 #

手动管理线程索引，直接操作数组元素：

import numpy as np
import numba as nb

@nb.cuda.jit
def tan_arctan_kernel(input_arr, tan_out, arctan_out):
    idx = nb.cuda.grid(1)
    if idx < input_arr.size:
        val = input_arr.flat[idx]
        tan_out.flat[idx] = np.tan(val)
        arctan_out.flat[idx] = np.arctan(val)

def numba_gpu_process(arr):
    # 将数据拷贝至GPU内存
    d_arr = nb.cuda.to_device(arr)
    d_tan = nb.cuda.device_array_like(arr)
    d_arctan = nb.cuda.device_array_like(arr)
    
    # 配置线程块（通常每块256/512线程）
    threads_per_block = 512
    blocks_per_grid = (arr.size + threads_per_block - 1) // threads_per_block
    
    # 启动核函数
    tan_arctan_kernel[blocks_per_grid, threads_per_block](d_arr, d_tan, d_arctan)
    
    # 将结果拷贝回CPU并释放GPU内存
    tan_result = d_tan.copy_to_host()
    arctan_result = d_arctan.copy_to_host()
    d_arr.deallocate()
    d_tan.deallocate()
    d_arctan.deallocate()
    
    return tan_result, arctan_result

2. 使用@vectorize生成向量化GPU函数 #

无需手动处理线程，代码更简洁：

import numpy as np
import numba as nb

@nb.vectorize(['float32(float32)', 'float64(float64)'], target='cuda')
def gpu_tan(x):
    return np.tan(x)

@nb.vectorize(['float32(float32)', 'float64(float64)'], target='cuda')
def gpu_arctan(x):
    return np.arctan(x)

def numba_vectorize_process(arr):
    d_arr = nb.cuda.to_device(arr)
    d_tan = gpu_tan(d_arr)
    d_arctan = gpu_arctan(d_arr)
    tan_result = d_tan.copy_to_host()
    arctan_result = d_arctan.copy_to_host()
    return tan_result, arctan_result

二、PyTorch实现GPU加速 #

PyTorch对GPU张量的三角函数运算支持非常成熟，API简洁且开销极低：

import torch

def pytorch_gpu_process(arr):
    # 将NumPy数组转为CUDA张量
    tensor = torch.tensor(arr, device='cuda', dtype=torch.float32)
    # 全程在GPU执行运算
    tan_result = torch.tan(tensor)
    arctan_result = torch.atan(tensor)
    # 转换为NumPy数组返回（按需选择）
    return tan_result.cpu().numpy(), arctan_result.cpu().numpy()

PyTorch会自动管理GPU内存，内置函数经过高度优化，对于你每批6亿元素的规模，数据拷贝的开销完全会被GPU的计算优势覆盖。

三、CPU vs GPU选择 #

• 优先选GPU：你的每批数据规模（6亿元素）属于计算密集型场景，GPU的海量并行计算核心能带来数量级的速度提升。
• 若GPU显存不足：可将每批再拆分为更小的子块（比如拆成4个(5, 3000000)的子数组），避免显存溢出。

四、工具选择建议 #

• 首选PyTorch：代码最简洁，无需手动管理CUDA线程与内存，内置函数优化程度高，上手成本低。
• 若依赖Numba生态：选择@vectorize(target='cuda')方案，比手写核函数更简洁，性能接近。
• 仅当无GPU可用时考虑CPU：可使用Numba的@njit(parallel=True)结合NumPy向量化运算，或直接用原生NumPy，但速度远低于GPU。

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