如何在Julia Flux中为小批量训练传入样本权重

安全阿岳

2026-8-9

问题描述

我目前使用Julia Flux实现小批量单步训练的代码片段如下：

loss, grad = Flux.withgradient(modelDQN.evalParameters) do
    qEval = modelDQN.evalModel(evalInput)
    Flux.mse(qEval, qTarget)
end
Flux.update!(modelDQN.optimizer, modelDQN.evalParameters, grad)

我需要添加一个与qEval、qTarget样本长度相等的权重向量ISWeights，实现加权损失计算（调整每个样本对总误差及梯度的影响），效果等同于TensorFlow Keras的fit函数中sample_weight参数的作用，请问如何用Flux实现？

解决方案

Flux没有内置带样本权重的MSE损失，但可以手动实现加权逻辑，核心是对每个样本的误差赋予对应权重后再计算总损失。

修改后的代码

loss, grad = Flux.withgradient(modelDQN.evalParameters) do
    qEval = modelDQN.evalModel(evalInput)
    # 计算每个样本的平方误差
    sq_error = (qEval .- qTarget).^2
    # 调整权重维度以匹配误差张量（根据你的数据维度调整reshape的方式）
    # 假设qEval是(N, batch_size)，ISWeights是(batch_size,)，则reshape为(1, batch_size)
    weighted_sq_error = sq_error .* reshape(ISWeights, 1, :)
    # 加权平均损失，和Keras sample_weight逻辑一致
    weighted_mse = mean(weighted_sq_error)
    weighted_mse
end
Flux.update!(modelDQN.optimizer, modelDQN.evalParameters, grad)

关键细节说明

维度匹配：根据你的张量形状调整reshape的参数。如果qEval是(batch_size, feature_num)（批量维度在前），则把ISWeights转为(batch_size, 1)：
```
weighted_sq_error = sq_error .* reshape(ISWeights, :, 1)
```
确保权重能和每个样本的所有特征误差对应相乘。
损失逻辑选择：
- 用mean(weighted_sq_error)等价于Keras的加权平均损失；
- 若需要加权求和（不除以样本数），替换为sum(weighted_sq_error)即可。
权重归一化：如果ISWeights的总和不为1，加权平均会放大/缩小损失量级。若要保持和原MSE的量级一致，可先归一化权重：
```
normalized_weights = ISWeights ./ sum(ISWeights)
weighted_sq_error = sq_error .* reshape(normalized_weights, 1, :)
```