Target out of bounds for UNET

I’m trying to implement a UNET for my dataset of electron microscopy images. My segmentation masks are the same size of my images. In the mask, I’ve set the value to 0 and three types of defects as 1, 2, and 3. So, there are 4 possible classes - 0, 1, 2, and 3 , where 0 refers to the background, and the rest refer to different type of defects.

I’m using a basic UNET learner defined as follows:
learn = unet_learner(data, models.resnet18, metrics=dice)
where dice is defined as:
def dice(pred, targs):
pred = (pred>0).float()
return 2. * (pred*targs).sum() / (pred+targs).sum()

On CPU, I get the following error:

/usr/local/lib/python3.6/dist-packages/torch/nn/ in nll_loss(input, target, weight, size_average, ignore_index, reduce, reduction)
2262 .format(input.size(0), target.size(0)))
2263 if dim == 2:
-> 2264 ret = torch._C._nn.nll_loss(input, target, weight, _Reduction.get_enum(reduction), ignore_index)
2265 elif dim == 4:
2266 ret = torch._C._nn.nll_loss2d(input, target, weight, _Reduction.get_enum(reduction), ignore_index)

IndexError: Target 3 is out of bounds.

Does this mean that the numbers in my mask must be modified?

Since Pytorch expects 0-n when calculating the loss (and in turn fastai), you can either adjust the values in your mask or set your loss function to include an ignore_index of 0, which the latter should work for you.

IE: loss_func = CrossEntropyLossFlat(axis=1, ignore_index=0)

Hey, thanks for your reply. The problem seems to persist even after using ignore_index = 0. Also, I made a small mistake in my post. The classes in the mask are 0 to 3, where 0 is the background. How exactly would I have to change my mask values to solve this problem?

There’s actually a helpful thread if you look hard enough :wink: