Hello Yoonho -
Let me try to answer part of your question.
(I haven’t looked at your dataloader or model code, so I won’t comment
on it.)
Let me try to give you some general context here.



 Yolkandwhite:

I want to get predicted brain cancer mask out of brain MRIs and Brain cancer masks image file.
brain MRI has shape of 1x3x256x256(RGB)


The shape of your input doesn’t really matter to your loss function – it
just has to match what your model is expecting for its input.
But I interpret this as meaning you have an input image of 256x256
pixels with three (RGB color) channels.  I assume that the “1”, means
that you have a batch size of nBatch = 1.  (There is nothing necessarily
wrong with using a batch size of 1, but you might get better throughput
and / or training with a larger batch size.)
So a batch of input samples will be a tensor of shape
(nBatch, 3, 256, 256).
Next, I assume that you are performing binary image segmentation.
That is, you want to label each pixel in your image with a “0” (this
pixel is not part of a cancerous tumor) or “1” (this pixel is cancerous).


and mask has shape of 1x1x256x256(Black and white).


Similarly I assume that your “black-and-white” image masks – the
“targets” that you use for your training – are also per-pixel "0"s and
"1"s.  So a batch of targets will have shape (nBatch, 256, 256)
(and each value of such a tensor will be 0 or 1).
Lastly, for a single input sample, the output of your network should
be a 256x256-pixel “image,” so that when applied to a batch of
input samples, your network will produce output tensors of shape
(nBatch, 256, 256).  Each “pixel” of your output tensor (your
model’s prediction) will be a “raw-score” logit that runs from -inf
to +inf.  As you can see your model’s predictions that get fed into
your BCEWithLogits and your targets that also get fed into your
loss function are indeed of the same shape.


At first I tried to use nn.CrossEntropyLoss(), It failed maybe because of my poor understanding of dimensions.


Based on my understanding of your use case (black-and-white image
masks), BCEWithLogitsLoss (rather than CrossEntropyLoss) will be
the appropriate loss function.


but couldn’t really understand the meaning of targets getting size of [batch_size, H, W] containing the class indices. Does this mean actually the target shape isn’t [batch_size, H, W] and instead of it,  Its something like [batch_size, indices, H, W] ?


No.  Neither for BCEWithLogitsLoss nor CrossEntropyLoss will
your target have shape (nBatch, nClass, 256, 256).  It means
that each pixel (in your batch of shape (nBatch, 256, 256)) will
have a value equal to the “class index.”  For BCEWithLogitsLoss
this value will be 0 or 1 (and in [0, nClass - 1], inclusive, for
CrossEntropyLoss).  (If you were to use CrossEntropyLoss,
however, your prediction would have a class dimension, that is,
would have shape (nBatch, nClass, 256, 256).)


anyways, I just found another loss function, which is nn.BCEWithLogitsLoss()
It looked easier than using nn.CrossEntropyLoss


Not only might it be easier, it is the right loss function for your use case.


because it just needed the Inputs and Targets with same shape.


Yes, as discussed above, you will want both your predictions (inputs to
the loss function) and your targets to have shape (nBatch, 256, 256).


but using this made another problem.
Here’s my full code below …


I can’t tell you whether your code is producing the right shapes and
values (or if i might have other issues), but check, at least, that your
predictions and targets are consistent with what I outlined above.
Best.
K. Frank