Hi James!



 James_Lee:

I would like to mask out the low possibility pixels with weight=0 and keep others.


If I understand your question correctly, you can just use weights of 0.0
and 1.0 to do this.
Let’s say that the input to your model is a batch of images of shape
[nBatch, height, width], and that the output of your model (which
will be the input to BCEWithLogitsLoss) and the target you pass
to BCEWithLogitsLoss also have this shape.
Now let me assume that you somehow construct a
pixel_possibility_tensor, also of the same shape, and that this
tensor tells you the “possibility” of each pixel (whatever that means
in your use case).  If a “possibility” of, say, less than 0.35 is considered
“low possibility”, you would:
weight = (pixel_possibility_tensor > 0.35).float()
loss = (weight.numel() / weight.sum()) * torch.nn.BCEWithLogitsLoss (weight = weight) (input, target)

In the loss computation, the 0.0 values in weight will mask out
the “low possibility” pixels (and the 1.0 values will keep unchanged
the remaining pixels).  The weight.numel() / weight.sum()
pre-factor is so that loss will be the loss averaged only over the
non-masked pixels, assuming that’s what you want – without it, all
of the masked pixels would also cause a bunch of zeros to get
averaged into the loss.
Best.
K. Frank