Confusing error using BCEWithLogitsLoss with weighted loss

Hi!

I’m pretty new to pytorch and neural networks in general. My project involves segmenting an image into two classes. My input is a 512x512x3 (rgb) image and my output is a 512x512x2 image where the first channel is a binary image with the positives for the first class and the second channel is a binary image with the positives for the second class.

The thing is, my dataset is pretty imbalanced. For example, in the first class there are about 11 positives for every 250 negatives. I’m using BCEWithLogitsLoss, and from browsing these forums I’ve seen that I should probably use the ‘pos_weight’ input. Here’s what the documentation for pos_weight says:

pos_weight – a weight of positive examples. Must be a vector with length equal to the number of classes.

Okay, I have two classes, so I made my vector, where the first entry is negatives/positives of class 1 and the second entry is the negatives/positives of class 2:

pw = torch.FloatTensor([21.956,0.04554])
criterion = nn.BCEWithLogitsLoss(pos_weight=pw)

But here’s the error message I get:

The size of tensor a (2) must match the size of tensor b (512) at non-singleton dimension 3

The code runs fine if I don’t try to pass anything to pos_weight, so that’s definitely what’s causing the error. I also don’t understand why it’s expecting a vector of size 512 when I only have two classes!

Any help would be greatly appreciated.

Thanks!

Hello,

I have encountered this confusing case recently. Refer ptblck’s comment from this thread.
If you are doing binary-classes segmentation, you could also use CrossEntropyLoss and pass weight argument to it.

1 Like

Thanks for the reply @MariosOreo

I’m not sure if ptblck’s suggestion was to do what I did, but this seemed to work:

I checked the shape of my label tensor, and it was (1,2,512,512)
So in order to match the pos_weight vector’s shape to the label’s shape, I did the following:

pw = torch.FloatTensor([21.956,0.04554])
pw = torch.reshape(pw,(1,2,1,1))
criterion = nn.BCEWithLogitsLoss(pos_weight=pw)

With this in mind, I think the documentation should say something like:

pos_weight – a weight of positive examples. Must be a vector with length equal to the number of classes along the same dimension as the classes.

Stuff like this might seem obvious to some people, but for newbies like me it really hurts not to have that extra bit of information.

Thanks for the help!

1 Like

Yeah, you are right. As @velikodniy explained in his post this is actually a bug.
However, I’m glad you’ve found a workaround! :slight_smile: