Hey There!
I’m running a semantic segmentation model (unet) on input images of shape (3, 256, 256) and masks of shape (256,256) with pixel values of 0,1,2, and 3 (3 classes total, with pixel 0 being the background).
When retrieving the output from the model, I have:
output from model = (b_s, 3, 256, 256)
predicted mask = (b_s, 256, 256)
I use a cross-entropy loss function, but receive the following error:
1838 ret = torch._C._nn.nll_loss(input, target, weight, _Reduction.get_enum(reduction), ignore_index)
1839 elif dim == 4:
-> 1840 ret = torch._C._nn.nll_loss2d(input, target, weight, _Reduction.get_enum(reduction), ignore_index)
1841 else:
1842 # dim == 3 or dim > 4
IndexError: Target 3 is out of bounds.
This is the model I am using: source
> class DoubleConv(nn.Module):
> """(convolution => [BN] => ReLU) * 2"""
>
> def __init__(self, in_channels, out_channels, mid_channels=None):
> super().__init__()
> if not mid_channels:
> mid_channels = out_channels
> self.double_conv = nn.Sequential(
> nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1),
> nn.BatchNorm2d(mid_channels),
> nn.ReLU(inplace=True),
> nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1),
> nn.BatchNorm2d(out_channels),
> nn.ReLU(inplace=True)
> )
>
> def forward(self, x):
> return self.double_conv(x)
>
>
> class Down(nn.Module):
> """Downscaling with maxpool then double conv"""
>
> def __init__(self, in_channels, out_channels):
> super().__init__()
> self.maxpool_conv = nn.Sequential(
> nn.MaxPool2d(2),
> DoubleConv(in_channels, out_channels)
> )
>
> def forward(self, x):
> return self.maxpool_conv(x)
>
>
> class Up(nn.Module):
> """Upscaling then double conv"""
>
> def __init__(self, in_channels, out_channels, bilinear=True):
> super().__init__()
>
> # if bilinear, use the normal convolutions to reduce the number of channels
> if bilinear:
> self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
> self.conv = DoubleConv(in_channels, out_channels, in_channels // 2)
> else:
> self.up = nn.ConvTranspose2d(in_channels , in_channels // 2, kernel_size=2, stride=2)
> self.conv = DoubleConv(in_channels, out_channels)
>
>
> def forward(self, x1, x2):
> x1 = self.up(x1)
> # input is CHW
> diffY = x2.size()[2] - x1.size()[2]
> diffX = x2.size()[3] - x1.size()[3]
>
> x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2,
> diffY // 2, diffY - diffY // 2])
> # if you have padding issues, see
> # https://github.com/HaiyongJiang/U-Net-Pytorch-Unstructured-Buggy/commit/0e854509c2cea854e247a9c615f175f76fbb2e3a
> # https://github.com/xiaopeng-liao/Pytorch-UNet/commit/8ebac70e633bac59fc22bb5195e513d5832fb3bd
> x = torch.cat([x2, x1], dim=1)
> return self.conv(x)
>
>
> class OutConv(nn.Module):
> def __init__(self, in_channels, out_channels):
> super(OutConv, self).__init__()
> self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)
>
> def forward(self, x):
> return self.conv(x)
>
> class UNet(nn.Module):
> def __init__(self, n_channels, n_classes, bilinear=True):
> super(UNet, self).__init__()
> self.n_channels = n_channels
> self.n_classes = n_classes
> self.bilinear = bilinear
>
> self.inc = DoubleConv(n_channels, 64)
> self.down1 = Down(64, 128)
> self.down2 = Down(128, 256)
> self.down3 = Down(256, 512)
> factor = 2 if bilinear else 1
> self.down4 = Down(512, 1024 // factor)
> self.up1 = Up(1024, 512 // factor, bilinear)
> self.up2 = Up(512, 256 // factor, bilinear)
> self.up3 = Up(256, 128 // factor, bilinear)
> self.up4 = Up(128, 64, bilinear)
> self.outc = OutConv(64, n_classes)
>
> def forward(self, x):
> x1 = self.inc(x.float())
> x2 = self.down1(x1)
> x3 = self.down2(x2)
> x4 = self.down3(x3)
> x5 = self.down4(x4)
> x = self.up1(x5, x4)
> x = self.up2(x, x3)
> x = self.up3(x, x2)
> x = self.up4(x, x1)
> logits = self.outc(x)
> return logits
Initializing model:
model = UNet(n_channels=3, n_classes=3, bilinear=False)
optimizer = optim.RMSprop(model.parameters(), lr=0.001, weight_decay=1e-8, momentum=0.9)
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min', patience=2)
criterion = nn.CrossEntropyLoss()
The interesting thing is that when I initialize my model to n_classes = 4,5 and so on, the loop runs just fine, but the outputs are terrible. It’s when I set the n_classes = 3 when it all breaks. Any help would be greatly appreciated 