I set up a really small toy Full Connected Network for image segmentation on Pascal VOC . The architecture looks like following ==> Conv - Conv- Conv- Transpose conv-Transpose conv-Transpose conv .I set up a per pixel cross Entropy Loss and trained it with Adam . But i am not able to lower down loss . Can anyone tell what i have fundamentally done wrong ,Is there something wrong with the way i set up the loss?.The code is below.
dataset1 = dset.VOCSegmentation(root = ‘/.data’ , download = True , transform = T.Compose([T.Resize((224,224)),T.ToTensor()])
,target_transform =  T.Compose([T.Resize((224,224)),T.ToTensor()]))
train_loader = DataLoader(dataset1, batch_size = 32 )
class Model(nn.Module):
def init(self):
super().init()
self.conv1 = nn.Conv2d(3,64,3,2,1)
self.conv2 = nn.Conv2d(64,128,3,2,1)
self.conv3 = nn.Conv2d(128,256,3,2,1)
    self.deconv1 =nn.ConvTranspose2d(256,128,4,2,1)
    self.deconv2 = nn.ConvTranspose2d(128,64,4,2,1)
    self.deconv3 = nn.ConvTranspose2d(64,64,4,2,1)
    
    self.finalconv = nn.Conv2d(64,22,3,1,1)
    
def forward(self,x):
    x = self.conv1(x).clamp(min=0)
    x = self.conv2(x).clamp(min=0)
    x = self.conv3(x).clamp(min=0)
    
    x = self.deconv1(x).clamp(min=0)
    x = self.deconv2(x).clamp(min=0)
    x = self.deconv3(x).clamp(min=0)
    
    return self.finalconv(x)

for epoch in range(num_epoch):
for image,segments in enumerate(train_loader):
    segments = segments*255
    segments = segments.int().long()
    segments[segments==255] = 21
    
    scores = model(image)
    
    scores = scores.permute(0,2,3,1).reshape(-1,22)
    
    segments = segments.reshape(segments.shape[0],segments.shape[2],segments.shape[3]).reshape(-1)
    
    loss = F.cross_entropy(scores,segments)
    
    loss.backward()
    
    optimizer.step()
    
    optimizer.zero_grad()
    
    print(loss.item())