Pytorch calculates the loss for each sample in the batch

Preface

PyTorch Loss function of （ Here I only use and research MSELoss） By default, a Batch Calculate the loss of all samples , And find the mean value . If I need the loss of each sample for some subsequent calculations （ And optimize model parameters , Gradient descent is irrelevant ）, For example, use the loss of samples to do some operations , You can't use the default loss function , I searched and didn't find the relevant information , stay PyTorch The forum found relevant problems .

Loss for each sample in batch - PyTorch Forumshttps://discuss.pytorch.org/t/loss-for-each-sample-in-batch/36200 Besides , Also refer to official documents ：

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

Solution

Suppose the original loss function instantiation code is as follows ：

loss_fn = nn.MSELoss()

Change it to the following code ：

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

Suppose the code of the original loss calculation stage is as follows ：

loss = loss_fn(Ypred, Y)

Is changed to ：

loss_each = torch.mean(loss_fn_each(Ypred, Y), 1)  #  The loss calculated here is for each sample 
loss = torch.mean(loss_each)  #  Here is the whole  Batch  The average loss of 

If you need to do something with the loss of each sample , Use `loss_each` that will do . If your prediction results include multiple dimensions , That may need to be used many times mean function . If necessary, switch to list preservation , You can use it directly tolist() Method . For relevant contents, please refer to documents and network .

Usually , If you want the loss of each sample , It should be Eval Stage , After all, it doesn't make much sense for each sample to lose instability in the training stage . If in training , Don't forget the following code .

total_loss.append(loss.item())  #  Record the loss 
optimizer.zero_grad()  #  Basic code of training stage , Clear the gradient 
loss.backward()  #  Basic code of training stage , Back propagation 
optimizer.step()  #  Basic code of training stage , Optimizing model parameters 

principle

When setting the loss function , It is usually written in the following way , There are no parameters .

loss_fn = nn.MSELoss()

Actually, refer to official documents , This loss function can pass in parameters  reduction  Of .reduction There are three kinds of , The default is 'mean' , It can be  'none' | 'mean' | 'sum'.

from none Start with ,none Is the most complete default loss , For example, the model output I calculated is batch_size * output_dim Of shape Of Tensor（ That is to say, each sample has output_dim Dimension data ）, Then send MSELoss Two inputs of （ My prediction and goal ） The size is like this , and MSELoss The output size of is exactly the same as the input size , It's also （batch_size * output_dim ）.

But what I want is  （batch_size * 1） The loss of , That is, for each sample , The loss of dimensions within the sample is averaged , But the samples cannot be averaged . Use mean Words （ Default ）, What you get is a （ 1*1 ） Of Tensor, hold Batch The dimensions and internal dimensions of the sample are averaged .

So the idea is to directly use   `nn.MSELoss(reduction ='none')` Calculate the loss as a loss function , Got Tensor Along the dim=1 Calculating mean , In this way, each sample calculates an independent loss , There is no mean between samples .