CODEWITHSUNDEEP
pytorchone-hot-encodinglosscross-entropy
Adar Cohen
Adar Cohen

Reputation: 306

Proper way to use Cross entropy loss with one hot vector in Pytorch

I've encountered many suggestions by searching online, but I don't understand the proper way to do it. Lets say my output of the model is 4 neurons, and the target labels are 1000 0100 0010 0001.

In tensorflow, I added a softmax layer at the end, and model.fit with CategoricalCrossEntropy.

What is the proper way to do it in pytorch?

Right now, what I did (and its working) is that I insert my model output, and torch.argmax() of the target one hot vector. for example model output is (0.7,0.1,0.1,0.1) and the torch.argmax of the target vector is 0. on that I'm activating nn.CrossEntropyLoss().

is this the proper way to do it?

Upvotes: 0

Views: 1635

Answers (1)

DerekG
DerekG

Reputation: 3958

Another way to do this would be to use BCELoss(), which is the same as cross-entropy loss except that a target vector in the range [0,1] is expected as well as an output vector. This just saves you having to do the torch.argmax() step.

Upvotes: 1

Related Questions