如何进一步加速Numpy大规模数组运算？已尝试Numba仍需优化

我在处理形状为(1000,3,3000)的大规模数组aa时，需要完成以下逻辑：对每个(i,j)位置的三个值，选择最大值对应的位置标记为1，其余为0；若存在平局，优先级为索引2 > 索引1 > 索引0（即多个值同为最大值时，仅保留最靠后的索引为1）。

最初用Numpy实现，但运算速度太慢；改用Numba加速后有提升，但由于该计算是外层循环的一部分，仍需进一步优化，求各位提供优化建议。

原Numpy实现代码 #

import numpy as np

def get_prob(aa):
    allmax = aa.max(axis=1)[:, None]
    findmax = aa - allmax
    mask = ((findmax[:,1,:]==0)&(findmax[:,2,:]==0))
    findmax[:, 1, :][mask] = -1

    mask = ((findmax[:, 0, :] == 0) & (findmax[:, 1, :] == 0))
    findmax[:, 0, :][mask] = -1

    mask = ((findmax[:, 0, :] == 0) & (findmax[:, 1, :] == 0) & (findmax[:, 2, :] == 0))
    findmax[:, 0, :][mask] = -1
    findmax[:, 1, :][mask] = -1

    p = np.where(findmax < 0, 0.0, 1.0).transpose(0,2,1)
    return p

原Numba实现代码 #

import numba as nb
import numpy as np
import time

@nb.jit(nopython=True)
def get_prob_nb(aa,num_params,num_action):
    p=np.zeros_like(aa)

    for i in range(num_params):
        for j in range(num_action):
            a1 = aa[i, 0, j]
            a2 = aa[i, 1, j]
            a3 = aa[i, 2, j]
            if a1>a2 and a1>a3:
                p[i, 0, j] = 1.
            elif a2>=a1 and a2>a3:
                p[i, 1, j] = 1.
            elif a3>=a2 and a3>=a1:
                p[i, 2, j] = 1.

    p = p.transpose(0, 2, 1)
    return p

# 测试代码
rng = np.random.default_rng()
aa=rng.uniform(0.0, 1.0, 9000000)
aa=aa.reshape(1000,3,3000)
start = time.time()
get_prob_nb(aa, 1000, 3000)
print("elapse", time.time()-start)

优化建议与实现 #

1. 优化Numba版本（重点方向） #

Numba的核心优化点是利用并行和优化内存访问：

• 开启多核并行：用prange替代range处理外层循环，因为每个i的计算完全独立
• 去掉冗余参数：直接从数组形状获取维度，避免手动传参出错
• 简化条件判断：最后一个elif可改为else，逻辑等价且减少判断次数
• 优化内存布局：确保数组是连续内存（如np.ascontiguousarray），提升访问速度

优化后的Numba代码：

@nb.jit(nopython=True, parallel=True)
def get_prob_nb_opt(aa):
    num_params, _, num_action = aa.shape
    p = np.zeros_like(aa)
    # 并行处理外层循环
    for i in nb.prange(num_params):
        for j in range(num_action):
            a1 = aa[i, 0, j]
            a2 = aa[i, 1, j]
            a3 = aa[i, 2, j]
            if a1 > a2 and a1 > a3:
                p[i, 0, j] = 1.0
            elif a2 >= a1 and a2 > a3:
                p[i, 1, j] = 1.0
            else:
                p[i, 2, j] = 1.0
    return p.transpose(0, 2, 1)

额外优化：如果精度允许，将数组转为float32，可大幅减少内存占用和带宽压力：

aa = aa.astype(np.float32)

2. 简化Numpy版本逻辑 #

原Numpy代码通过修改findmax数组处理平局，可改为直接用布尔数组标记，逻辑更清晰且减少内存操作：

def get_prob_opt(aa):
    # 获取每个(i,j)对应的最大值
    allmax = aa.max(axis=1, keepdims=True)
    # 标记所有等于最大值的位置
    is_max = aa == allmax
    
    # 处理平局：优先保留索引2，再索引1，最后索引0
    # 情况1：索引0和1同为最大值，取消索引0的标记
    mask = is_max[:,0,:] & is_max[:,1,:]
    is_max[:,0,:] &= ~mask
    # 情况2：索引1和2同为最大值，取消索引1的标记
    mask = is_max[:,1,:] & is_max[:,2,:]
    is_max[:,1,:] &= ~mask
    
    # 转换为目标形状并返回
    return is_max.astype(np.float64).transpose(0,2,1)

3. 其他通用优化 #

• 提前编译Numba函数：如果外层循环多次调用该函数，可提前触发编译（比如第一次调用后，后续调用会直接用编译好的机器码）
• 避免重复transpose：如果外层循环能接受p的原始形状，可去掉transpose操作，减少内存拷贝

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