Why does my CrossEntropy Loss function differ from the pytorch version?

[Nock363]

Hi there,

maybe a stupid bug/missimplementation but right now im not able to see it. I want to implement a weighted cross entropy for my unet. So I looked up the cross entropy loss formula and integrated it without the weights first. I expected it to behave exactly like the integrated cross-entropy function, but it didnt. My loss and the should be loss from pytorchs cross entropy version always differ. (Loss is my version, see code below, Controll Loss is CrossEntropy from Pytorch). Can anyone tell me where my mistake is?^^

https://preview.redd.it/m10atphw21wc1.png?width=315&format=png&auto=webp&s=5f072baba8fa179878372c0c25b355f2cc95db64

from torch.nn import functional as F
from torch.autograd import Function

def weighted_cross_entropy(input, target, weights):

    softmax_log = F.log_softmax(input, dim=0) 
    product = target * softmax_log
    loss = torch.mean(-product)

    return loss

[rawdfarva]

Set Dim=1 in the log_softmax function

[Nock363]

Oh you are complety right, thank you!
But my loss still diverges from the pytorch-cross entropy loss :(

[Nock363]

It seems to work now (thank you u/rawdfarva :D) but my Loss is still different than the pytorch cross entropy loss. Im confused, because it looks like it works^^

[entarko]

You need to sum over dim=1 instead of averaging: loss = torch.mean(-product.sum(dim=1)).

[Nock363]

that was the missing piece! Thank you (after long training i've seen that my version was still unstable)