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.