如何使用NumPy向量化优化自定义加权绝对均值误差损失函数

完全可以通过NumPy实现全向量化计算，彻底替换双层Python循环，大样本场景下计算效率可提升数十到上百倍，且计算逻辑和原实现完全对齐。

核心优化思路 #

• 利用单样本真实标签仅含1个非零值的特性，直接通过非零掩码提取每个样本目标位置的预测值、真实值，跳过内层遍历所有类别的循环
• 绝对值计算、加权、求和、归一化全流程使用NumPy底层C实现的算子，跳过外层Python循环的迭代开销
• 输入自动转NumPy数组，兼容原生Python列表、NumPy数组两种输入格式

向量化实现代码 #

import numpy as np

def get_model_score_vec(preds, actuals, sw):
    # 统一转换为浮点型NumPy数组，兼容不同输入类型
    preds = np.asarray(preds, dtype=np.float64)
    actuals = np.asarray(actuals, dtype=np.float64)
    sw = np.asarray(sw, dtype=np.float64)

    # 生成真实值非零位置的掩码，每个样本仅对应1个True位置
    target_mask = actuals != 0
    # 提取所有样本目标位置的预测值、真实值，输出形状为(样本总数,)
    preds_target = preds[target_mask]
    actuals_target = actuals[target_mask]

    # 计算每个样本的加权绝对误差并求和
    total_loss = np.sum(sw * np.abs(preds_target - actuals_target))
    # 完全匹配原公式的归一化逻辑
    loss = total_loss / (np.sum(sw) * len(preds))

    return loss

逻辑一致性验证 #

用给出的示例做校验：

测试输入：
真实标签 actuals = [[0, 0, 1]]
预测值 preds = [[0.5, -0.5, 0.5]]
样本权重 sw = [1]

计算过程：

• 非零目标位置对应预测值为0.5，真实值为1，绝对误差为0.5
• 加权总误差为1 * 0.5 = 0.5
• 归一化分母为1 * 1 = 1，最终损失为0.5，和手动计算结果、原双层循环实现结果完全一致。

补充说明 #

• 该实现无任何精度损失，所有计算逻辑和原函数完全对齐
• 如果后续需要将损失函数迁移到PyTorch、TensorFlow等深度学习框架做GPU加速，仅需把NumPy API替换为对应框架的张量API即可，核心计算逻辑不需要调整

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