Semantic segmentation weighted class loss_func

I am following walk with fastai 2022 notebooks and I came across weighted class loss function.
My model can be trained without any problem with default weights, but i can’t train when i modify the loss func class weights.
Here is what i tried:

weights = torch.tensor([[0.9]*1 + [1.1]]).cuda()
loss_func = CrossEntropyLossFlat(weight=weights, axis=1)
learn = unet_learner(dls, pretrained_model, metrics=acc_camvid, self_attention=True, act_cls=Mish, opt_func=opt, cbs=callbacks, loss_func=loss_func)

classes: {‘background’: 0, ‘door’: 1}
weights: tensor([[0.9000, 1.1000]], device=‘cuda:0’)

When i run lr_finder, i get the following error:

RuntimeError: weight tensor should be defined either for all 2 classes or no classes but got weight tensor of shape: [1, 2]

Any help is appreciated!

1 Like

I think there is a pair of extra brackets around your weights tensor, which makes it of shape [1,2] rather than the expected [2]. Try:

weights = torch.tensor([0.9]*1 + [1.1]).cuda()

Hope this solves the issue.

2 Likes

How could i have not looked at that.
Thank you! That is it

1 Like