search

jeya-maria-jose / KiU-Net-pytorch

Official Pytorch Code of KiU-Net for Image/3D Segmentation - MICCAI 2020 (Oral), IEEE TMI
https://sites.google.com/view/kiunet/kiu-net
MIT License
357 stars 80 forks source link

Issue with loss function. #19

Closed astorfi closed 3 years ago

astorfi commented 3 years ago

Hi,

I think there is an issue with the loss function implementation:

  1. In line https://github.com/jeya-maria-jose/KiU-Net-pytorch/blob/master/arch/ae.py#L242, a softmax operation is in place.
  2. According to the loss function implementation, you used cross-entropy loss which requires raw outputs (no log or softmax).

I think this loss function implementation is not correct. You may get the results as with logsoftmax of the outputs, still, the loss function can do a descent job.

Please investigate.

jeya-maria-jose commented 3 years ago

Hi,

Thank you for pointing it out. In this implementation of the code, an extra log(softmax) has been used. The CE loss already has log(softmax). So effectively, the code actually does log(softmax(log(softmax))).

I ran some experiments removing the extra terms and the results were mostly the same (margin of difference being 0.2 dice ). Technically, using an extra log(softmax) does not change the scale of the values; so we observe that the training is as stable as using the normal CE. However, I agree that it is not needed. I will correct the code in the next commit.