如何借助NumPy、Numba及GPU提升十进制转21进制Python脚本性能

原脚本通过递归实现十进制到21进制（对应elements列表长度）的转换，遍历数组逐个处理，面对超大数组时性能瓶颈明显。以下是利用NumPy、Numba GPU加速的具体修改方案：

核心优化方向 #

• 替换递归实现为迭代逻辑：递归无法在GPU上高效并行执行，迭代更适配硬件的并行计算模式
• 用Numba CUDA编写并行核函数：让每个GPU线程独立处理一个输入元素，最大化并行利用率
• 适配可变长度输出：不同数字转换后的结果长度不同，通过预分配内存+记录实际长度的方式统一处理

完整优化代码 #

import numpy as np
from timeit import default_timer as timer
from numba import cuda
import numba as nb

# 保留原元素映射列表
elements = [
    "n|0", "n|1", "n|2", "n|3", "n|4", "n|5", "n|6", "n|7", "n|8", "n|9", "n|10",
    "o|+", "o|*", "o|/", "om|-",
    "bl|(", "br|)",
    "e|**2", "e|**3", "e|**0.5", "e|**(1/3)",
]
elements_len = len(elements)

# 迭代实现进制转换（CPU兼容，也作为GPU实现的基础逻辑）
@nb.njit
def decimal_to_custom_iter(number):
    result = []
    if number == 0:
        result.append(0)
        return result
    while number > 0:
        remainder = number % elements_len
        result.append(remainder)
        number = number // elements_len
    # 反转结果以对齐原递归逻辑的输出顺序
    result.reverse()
    return result

# CUDA核函数：每个线程处理一个输入数字
@cuda.jit
def decimal_to_custom_gpu(input_arr, output_arr, output_lengths):
    # 获取当前线程的全局索引
    idx = cuda.grid(1)
    # 超出输入数组范围则直接返回
    if idx >= input_arr.size:
        return
    
    num = input_arr[idx]
    # 预分配局部数组存储临时结果（32位足够覆盖绝大多数大数场景）
    temp = cuda.local.array(32, dtype=nb.int64)
    count = 0
    
    if num == 0:
        temp[count] = 0
        count += 1
    else:
        n = num
        while n > 0:
            temp[count] = n % elements_len
            n = n // elements_len
            count += 1
    
    # 反转临时数组，得到正确顺序的结果
    for i in range(count // 2):
        temp[i], temp[count - 1 - i] = temp[count - 1 - i], temp[i]
    
    # 将结果写入全局输出数组
    for i in range(count):
        output_arr[idx][i] = temp[i]
    # 记录当前结果的实际长度
    output_lengths[idx] = count

# 封装GPU转换流程
def gpu_convert(decimal_numbers):
    # 转换为NumPy int64数组，确保GPU可处理
    input_arr = np.asarray(decimal_numbers, dtype=np.int64)
    # 预定义最大结果长度，可根据实际场景调整
    max_result_len = 32
    # 分配全局输出数组和长度记录数组
    output_arr = np.zeros((input_arr.size, max_result_len), dtype=np.int64)
    output_lengths = np.zeros(input_arr.size, dtype=np.int64)
    
    # 配置CUDA线程块和网格大小
    threads_per_block = 256
    blocks_per_grid = (input_arr.size + threads_per_block - 1) // threads_per_block
    
    # 启动GPU核函数
    decimal_to_custom_gpu[blocks_per_grid, threads_per_block](input_arr, output_arr, output_lengths)
    
    # 整理结果：按实际长度截取有效部分
    custom_numbers = []
    for i in range(input_arr.size):
        custom_numbers.append(output_arr[i][:output_lengths[i]].tolist())
    return custom_numbers

# 性能测试与结果验证
if __name__ == "__main__":
    # 构造超大测试数组（原输入重复10000次）
    decimal_numbers = np.array([15, 18, 28, 11, 7, 5, 41, 139, 6, 507] * 10000)
    
    # 原CPU迭代方法耗时
    start = timer()
    custom_numers_old = [decimal_to_custom_iter(num) for num in decimal_numbers]
    print(f"CPU迭代方法耗时: {timer() - start:.4f}s")
    
    # GPU方法耗时
    start = timer()
    custom_numers_gpu = gpu_convert(decimal_numbers)
    print(f"GPU加速方法耗时: {timer() - start:.4f}s")
    
    # 验证结果一致性
    assert custom_numers_old == custom_numers_gpu
    print("结果验证一致，无误差")

关键改动说明 #

1. 递归转迭代：原递归函数在GPU上无法并行，迭代逻辑通过循环取余、整除完成转换，更适合硬件执行
2. CUDA核函数设计：每个线程独立处理一个输入元素，使用cuda.local.array避免动态内存分配开销，提升执行效率
3. 可变长度处理：预分配足够大的输出数组，通过output_lengths记录每个结果的实际长度，最后截取有效数据，解决GPU对固定内存布局的要求
4. 结果对齐：迭代得到的余数是逆序，通过反转操作对齐原递归脚本的输出顺序

注意事项 #

• 需确保安装支持CUDA的Numba版本，且GPU满足CUDA运行要求（可通过numba.cuda.is_available()检查）
• 超大数组场景下GPU优势显著，小数组可能因数据传输开销抵消性能提升
• 局部数组大小32可根据实际最大输入值调整，21^32远超过常规整数范围，足够覆盖绝大多数场景

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