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/functional.py 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?