PyTorch 对 Batch 中每个样本计算损失 Loss for each sample

前言

PyTorch 的损失函数（这里我只使用与调研了 MSELoss）默认会对一个 Batch 的所有样本计算损失，并求均值。如果我需要每个样本的损失用于之后的一些计算（与优化模型参数，梯度下降无关），比如使用样本的损失做一些操作，那使用默认的损失函数做不到，搜了一下没有找到相关的资料，在 PyTorch 的论坛发现了相关的问题。

Loss for each sample in batch - PyTorch Forumshttps://discuss.pytorch.org/t/loss-for-each-sample-in-batch/36200此外，还参考了官方文档：

MSELoss — PyTorch 1.11.0 documentationhttps://pytorch.org/docs/stable/generated/torch.nn.MSELoss.html#torch.nn.MSELoss

解决方案

假设原来的损失函数实例化代码如下：

loss_fn = nn.MSELoss()

改为下面的代码：

loss_fn_each = nn.MSELoss(reduction ='none')  #MSE

假设原来计算损失阶段的代码如下：

loss = loss_fn(Ypred, Y)

则改为：

loss_each = torch.mean(loss_fn_each(Ypred, Y), 1)  # 这里算出来的损失是每个样本的
loss = torch.mean(loss_each)  # 这里算出来是整个 Batch 的平均损失

如果你需要用每个样本的损失做一些操作的话，使用 `loss_each` 即可。如果你预测的结果包括多个维度，那可能需要多次使用 mean 函数。如果需要转为 list 保存，可以直接使用 tolist() 方法。相关内容参见文档与网络。

通常情况下，如果你想要每个样本的损失，应该是在 Eval 阶段了，毕竟训练阶段每个样本损失不稳定也没有太大意义。如果在训练阶段，别忘了下面的代码。

total_loss.append(loss.item())  # 记录损失
optimizer.zero_grad()  # 训练阶段基本代码，清空梯度
loss.backward()  # 训练阶段基本代码，反向传播
optimizer.step()  # 训练阶段基本代码，优化模型参数

原理

在设置损失函数的时候，一般是下面的写法，没有任何参数。

loss_fn = nn.MSELoss()

其实参照官方文档，这个损失函数是可以传入参数 reduction 的。reduction 有三种，默认是 'mean' ，可以是 'none' | 'mean' | 'sum'。

从 none 开始说，none 是最完整的默认的损失，比如我计算出来的模型输出是 batch_size * output_dim 的 shape 的 Tensor（也就是说每个样本有 output_dim 维度的数据），那么送入 MSELoss 的两个输入（我的预测和目标）大小都是这样的， 而 MSELoss 的输出尺寸与输入尺寸是一模一样的，也是（batch_size * output_dim ）。

但是我想要的是 （batch_size * 1）的损失，也就是说对于每个样本，样本内部的维度的损失求了均值，但是样本之间不能作平均。使用 mean 的话（即默认情况），得到的是一个（ 1*1 ）的 Tensor，把 Batch 维度和样本内部维度都做了平均。

所以思路就是直接用  `nn.MSELoss(reduction ='none')` 作为损失函数计算损失，得到的 Tensor 沿 dim=1 求均值，这样每个样本就都计算出了独立的损失，而样本之间并没有求均值。