如何从Flux.train!中直接获取并打印损失以避免重复计算？

架构老林

2026-8-19

在Flux中直接获取训练批次损失（避免重复计算）

核心结论

Flux.train!默认不会返回每个批次的损失值，但你可以通过两种方式直接捕获训练时的批次损失，无需额外调用损失函数：

方案1：手动实现训练循环（最灵活可控）

直接拆解Flux.train!的内部逻辑，自己编写训练循环，这样每一步计算的损失可以直接用于反向传播和打印/记录，完全没有重复开销：

# 遍历训练数据加载器的每个批次
for (train_data, train_targets) in train_data_loader
    # 计算梯度的同时获取损失
    grads = Flux.gradient(Flux.params(model)) do
        pred = model(train_data)
        batch_loss = logitcrossentropy(pred, train_targets)
        # 直接打印当前批次的训练损失
        println("Batch loss: ", batch_loss)
        return batch_loss
    end
    # 更新模型参数
    Flux.update!(opt, Flux.params(model), grads)
end

这个方式里，batch_loss只计算一次，既用于梯度计算，又能直接输出，没有额外的损失计算开销。

方案2：包装损失函数实现实时记录

如果你不想手动写循环，可以把损失函数包装一下，让它在每次计算（也就是Flux.train!处理每个批次时）自动记录或打印损失：

# 定义带打印逻辑的损失函数
function train_loss(x, y)
    pred = model(x)
    batch_loss = logitcrossentropy(pred, y)
    # 打印当前批次损失
    println("Batch loss: ", batch_loss)
    return batch_loss
end

# 用包装后的损失函数训练
Flux.train!(train_loss, Flux.params(model), train_data_loader, opt)

Flux.train!在每个批次都会调用一次这个损失函数来计算梯度，所以这里打印的就是该批次的训练损失，没有额外计算。