Hello Joe!
No, this is not right. Your target (“ground truth”) should be (a
LongTensor) of shape (batch, H, W). For any given sample
within the batch, the values of the “pixels” in your “H x W” “image”
should be the integer class labels in [0, 5).
You don’t use multiple “Layers” to specify which class a given
pixel is in – you just use an integer class label as the value of
that pixel.
Note, the input to CrossEntropyLoss is not of the same shape
as the target; the input is of shape
(batch, nClasses, H, W). For any given sample in the batch,
each pixel in the H x W input “image” has nClasses values. These
are the “logits” (“scores” in -infinity to +infinity) that tell you how likely
the model thinks a given pixel is to be in each of the nClasses
classes.
This error message looks reasonable in light of what I said above
(except that it looks like you are actually using batchsize = 4,
not batchsize = 5).
Assuming batchsize = 4, nClasses = 5, H = 224, and
W = 224, CrossEntropyLoss will be expecting the input
(prediction) you give it to be a FloatTensor of shape
(4, 5, 244, 244), and the target (ground truth) to be a
LongTensor of shape (4, 244, 244).
Good luck.
K. Frank