How to understand the parameter reduction set to 'none' in the F.cross_entropy function in pytorch?

I would like to use heat maps for pixel classification tasks. Currently, I have a classification heat map with predicted results, with dimensions (batch_num, C, H, W), as well as a ground truth heat map with dimensions (batch_num, C, H, W). At the same time, I want each pixel to have different levels of attention in different channels, so I have generated an energy matrix with dimensions (batch_num, C, H, W). In my experiment, I expect to put this energy matrix into F.cross_entropy(classification, class_heat_maps, reduction=‘none’) to get a result with dimensions (batch_num, C, H, W), which I will use for element-wise multiplication with my energy matrix.

This is my code:

classification_loss=total_class_energy*F.cross_entropy(classification,class_heat_maps,reduction=‘none’)
classification_loss=torch.sum(classification_loss,dim=1)
classification_loss=torch.mean(classification_loss)

Could you explain why the dimensions generated by my F.cross_entropy(classification, class_heat_maps, reduction=‘none’) function in the experiment are (batch_num, H, W)? Do I need to customize a cross entropy function to achieve my goal?

I would be very grateful for your help.

F.cross_entropy is used for a multi-class classification or segmentation use case where each sample or pixel is assigned to a single class only.
The output, if reduction='none' is used, will be the loss value for each sample or pixel location and thus no channel or class dimension is used. Note that this loss is not an elementwise loss since dim1 represents the class dimension in the model output where each “channel” represents the logits for this corresponding class.

1 Like

This is very helpful, thank you very much.