UNet Segmentation Results not good

username: Amasu
cooked: <p>Hi, I am trying to work on Segmentation with Unet on CamVid Dataset. Data has 32 classes, masks available are of [n,3,h,w]. I converted the masks in [n,32,h,w]. While during training using CrossEntropyloss Ii game target to criterion as target=target.argmax(1) which makes target of shape [n,1,h,w] while my model predictions are of shape [n,32,h,w]. But when test on test data result is not good.<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/d/d/dd8a2ad5d4c099f10c3a4d1e81ec05057243f6c2.png" data-download-href="https://discuss.pytorch.org/uploads/default/dd8a2ad5d4c099f10c3a4d1e81ec05057243f6c2" title="download (1)"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/d/d/dd8a2ad5d4c099f10c3a4d1e81ec05057243f6c2_2_690x183.png" alt="download (1)" data-base62-sha1="vBPxFFzn6UG5BL9nW5aXMmCTuIW" width="690" height="183" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/d/d/dd8a2ad5d4c099f10c3a4d1e81ec05057243f6c2_2_690x183.png, https://discuss.pytorch.org/uploads/default/original/3X/d/d/dd8a2ad5d4c099f10c3a4d1e81ec05057243f6c2.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/d/d/dd8a2ad5d4c099f10c3a4d1e81ec05057243f6c2.png 2x" data-dominant-color="6A6C6C"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">download (1)</span><span class="informations">851×226 285 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div><br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/1/3/1373c7817412b4766ad7469b840fbb230f50f4d1.png" data-download-href="https://discuss.pytorch.org/uploads/default/1373c7817412b4766ad7469b840fbb230f50f4d1" title="download (2)"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/1/3/1373c7817412b4766ad7469b840fbb230f50f4d1_2_690x183.png" alt="download (2)" data-base62-sha1="2M57VYxVe6mv8mDrZAT26UrsRvH" width="690" height="183" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/1/3/1373c7817412b4766ad7469b840fbb230f50f4d1_2_690x183.png, https://discuss.pytorch.org/uploads/default/original/3X/1/3/1373c7817412b4766ad7469b840fbb230f50f4d1.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/1/3/1373c7817412b4766ad7469b840fbb230f50f4d1.png 2x" data-dominant-color="6D5451"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">download (2)</span><span class="informations">851×226 55 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div><br>
My code is Given:</p>
<pre><code>def Double_Conv(input_channel, output_channel):
return nn.Sequential(
    nn.Conv2d(input_channel, output_channel, kernel_size=3, padding=1),
    nn.BatchNorm2d(output_channel),
    nn.ReLU(inplace=True),
    nn.Conv2d(output_channel, output_channel, kernel_size=3, padding=1),
    nn.BatchNorm2d(output_channel),
    nn.ReLU(inplace=True)
)

class UNet(nn.Module):
def __init__(self, no_of_cls):
  super().__init__()

  self.down1 = Double_Conv(3, 64)
  self.down2 = Double_Conv(64, 128)
  self.down3 = Double_Conv(128, 256)
  self.down4 = Double_Conv(256, 512)
  self.down5 = Double_Conv(512, 1024)

  self.max_pool = nn.MaxPool2d(2)
  self.upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)

  self.up4 = Double_Conv(512 + 1024, 512)
  self.up3 = Double_Conv(256 + 512, 256)
  self.up2 = Double_Conv(128 + 256, 128)
  self.up1 = Double_Conv(64 + 128, 64)

  self.lastConv = nn.Conv2d(64, no_of_cls, 1)

def forward(self, x):
  conv1 = self.down1(x)
  x = self.max_pool(conv1)

  conv2 = self.down2(x)
  x = self.max_pool(conv2)

  conv3 = self.down3(x)
  x = self.max_pool(conv3)

  conv4 = self.down4(x)
  x = self.max_pool(conv4)

  x = self.down5(x)

  x = self.upsample(x)
  x = torch.cat([x, conv4], dim=1)

  x = self.up4(x)
  x = self.upsample(x)
  x = torch.cat([x, conv3], dim=1)

  x = self.up3(x)
  x = self.upsample(x)
  x = torch.cat([x, conv2], dim=1)

  x = self.up2(x)
  x = self.upsample(x)
  x = torch.cat([x, conv1], dim=1)

  x = self.up1(x)
  out = self.lastConv(x)
  return out
</code></pre>
<p>Training loop:</p>
<pre><code>  model.train()
  stats = []
  print('Training Started.....')
for epoch in range(epochs):
   train_loss, valid_loss = 0,0
   metrics = defaultdict(float)
   for i, data in enumerate(trainloader):
     inputs, mask, _ = data
     inputs, mask = inputs.to(device), mask.to(device)

     optimizer.zero_grad()
     output = model(inputs.float())
     target = mask.argmax(1)
     loss = criterion(output, target.long())
     loss.backward()
     optimizer.step()

    train_loss += loss.item() * inputs.size(0)

  train_loss = train_loss / len(trainloader.dataset)
  with torch.no_grad():
    model.eval()
    for i, val_data in enumerate(validloader):
        inp, masks, _ = val_data
        inp, masks = inp.to(device), masks.to(device)

        out = model(inp.float())
        val_target = masks.argmax(1)
        val_loss = criterion(out, val_target.long())
        valid_loss += val_loss.item() * inp.size(0)

    valid_loss = valid_loss / len(validloader.dataset)
    model.optimizer = optimizer
    stats.append([train_loss, valid_loss])
print('Epoch',epoch,':',f' Training Loss: {train_loss:.4f},', f' Validation Loss: {valid_loss:.4f}')
stat = pd.DataFrame(stats, columns=['train_loss', 'valid_loss'])
</code></pre>