Hi kliu!



 kliu:

I am trying to predict some binary image.


If I understand your use case, you should start with
BCEWithLogitsLoss as your loss function (and only change
to something else if you have a good reason and testing shows
that the change is for the better).


a simple two layer neural net with ReLU activations after each layer


Just to be clear, for BCEWithLogitsLoss, the last layer of your
network should be a Linear, and it should not be followed by
an activation function (e.g., neither ReLU nor Sigmoid).


outputs some 2x2 matrix [[0.01, 0.9], [0.1, 0.2]].


Note, these values – between 0.0 and 1.0 – are not values you
would typically get from a final Linear layer.  They could occur,
but in general, you would expect values running from -inf to
inf (that would be interpreted as raw-score logits).


I am currently using BCE loss with logits over the prediction and the ground truth, however I was wondering if I should be using the plain BCE loss instead (no logits)? If neither of these are suitable, should I be using CrossEntropyLoss? Not quite sure what the difference between these is in this context.


Yes, you should be using BCEWithLogitsLoss.
Sigmoid followed by BCELoss is mathematically equivalent to
BCEWithLogitsLoss, but numerically less stable.
CrossEntropyLoss (which would better be called
“CategoricalCrossEntropyWithLogitsLoss”) is essentially the same as
BCEWithLogitsLoss, but requires making some small modifications
to your network and your ground-truth labels that add a small amount
of unnecessary redundancy to your network.
Best.
K. Frank