Hello Everyone,
I am training an Autoencoder based on Resnet-Unet Architecture. Here the loss remains constant through out training.  I tried varying the learning rate, Used learning rate scheduler, played arround with different optimizers and loss functions(SSE, BCE etc). Used normalized and unnormalized data .I followed the suggestions provided by in the pytorch forum. But was unable to fix the problem. It would be great if someone can point out where i am going wrong.
Thank you
Code
batch_size=3
def convrelu(in_channels, out_channels, kernel, padding):
    return nn.Sequential(
        nn.Conv2d(in_channels, out_channels, kernel, padding=padding),
        nn.ReLU(inplace=True),
    )


class ResNetUNet(nn.Module):
    def __init__(self, n_class):
        super().__init__()

        self.base_model = models.resnet18(pretrained=True)
        self.base_layers = list(self.base_model.children())
        self.layer0 = nn.Sequential(*self.base_layers[:3]) # size=(N, 64, x.H/2, x.W/2)
        self.layer0_1x1 = convrelu(64, 64, 1, 0)
        self.layer1 = nn.Sequential(*self.base_layers[3:5]) # size=(N, 64, x.H/4, x.W/4)
        self.layer1_1x1 = convrelu(64, 64, 1, 0)
        self.layer2 = self.base_layers[5]  # size=(N, 128, x.H/8, x.W/8)
        self.layer2_1x1 = convrelu(128, 128, 1, 0)
        self.layer3 = self.base_layers[6]  # size=(N, 256, x.H/16, x.W/16)
        self.layer3_1x1 = convrelu(256, 256, 1, 0)
        self.layer4 = self.base_layers[7]  # size=(N, 512, x.H/32, x.W/32)
        self.layer4_1x1 = convrelu(512, 512, 1, 0)

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

        self.conv_up3 = convrelu(256 + 512, 512, 3, 1)
        self.conv_up2 = convrelu(128 + 512, 256, 3, 1)
        self.conv_up1 = convrelu(64 + 256, 256, 3, 1)
        self.conv_up0 = convrelu(64 + 256, 128, 3, 1)

        self.conv_original_size0 = convrelu(3, 64, 3, 1)
        self.conv_original_size1 = convrelu(64, 64, 3, 1)
        self.conv_original_size2 = convrelu(64 + 128, 64, 3, 1)

        self.conv_last = nn.Conv2d(64, n_class, 1)

    def forward(self, input):
        
        x_original = self.conv_original_size0(input)
        
        x_original = self.conv_original_size1(x_original)
        
        
        layer0 = self.layer0(input)
        
        #save_image(layer0,"input_layer0.png")
        
        layer1 = self.layer1(layer0)
        layer2 = self.layer2(layer1)
        layer3 = self.layer3(layer2)
        layer4 = self.layer4(layer3)
        
        #
        
        layer4 = self.layer4_1x1(layer4)
        x = self.upsample(layer4)
        print(x.shape)
        save_image(x[1,1,:,:],"upsample_layer.png")
        
        save_image(layer3[1,0,:,:],"layer_3.png")
        save_image(x[1,0,:,:],"layer_3_x.png")
        layer3 = self.layer3_1x1(layer3)
        x = torch.cat([x, layer3], dim=1)
        x = self.conv_up3(x)

        x = self.upsample(x)
        save_image(layer2[1,0,:,:],"layer_2.png")
        save_image(x[1,0,:,:],"layer_2_x.png")
        layer2 = self.layer2_1x1(layer2)
        x = torch.cat([x, layer2], dim=1)
        x = self.conv_up2(x)

        x = self.upsample(x)
        save_image(layer1[1,0,:,:],"layer_1.png")
        save_image(x[1,0,:,:],"layer_1_x.png")
        layer1 = self.layer1_1x1(layer1)
        x = torch.cat([x, layer1], dim=1)
        x = self.conv_up1(x)

        x = self.upsample(x)
        save_image(layer0[1,0,:,:],"layer_0.png")
        save_image(x[1,0,:,:],"layer_0_x.png")
        layer0 = self.layer0_1x1(layer0)
        x = torch.cat([x, layer0], dim=1)
        x = self.conv_up0(x)
        

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

        out = self.conv_last(x)
        
        save_image(out[1,0,:,:],"Last_layer_ch1.png")
        save_image(out[1,1,:,:],"Last_layer_ch2.png")
        save_image(out[1,2,:,:],"Last_layer_ch3.png")
        return out
    
#**********************************************************************************************************************************************************
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = ResNetUNet(n_class=3)
summary(model, input_size=(3, 224, 224))
model = model.to(device)


transform = transforms.Compose([ transforms.Resize((224,224), interpolation=2),transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ])
trainset=torchvision.datasets.ImageFolder("../data/train", transform=transform, target_transform=None)
trainloader = torch.utils.data.DataLoader(trainset, shuffle = True , batch_size = batch_size , num_workers = 2, drop_last=True)


testset=torchvision.datasets.ImageFolder("../data/test", transform=transform, target_transform=None)
testloader = torch.utils.data.DataLoader(testset, shuffle = True , batch_size = batch_size , num_workers = 2, drop_last=True)

#autoencoder_criterion = nn.MSELoss()
optimizer_ft = optim.Adam(model.parameters(), lr = 1e-2)
#optimizer_ft = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=20, gamma=0.1)

#**********************************************************************************************************************************************************

def calc_loss(pred, target):
    autoencoder_criterion = nn.MSELoss()
    loss = autoencoder_criterion(pred, target)
    #m = nn.Sigmoid()
    #autoencoder_criterion = nn.BCELoss()
    #loss = autoencoder_criterion(m(pred), target)
    return loss

def train_model(model, optimizer, scheduler, num_epochs=25):

    #model.train()  # Set model to training mode
    for epoch in range(num_epochs):
        print('Epoch {}/{}'.format(epoch, num_epochs - 1))
        print('-' * 10)
        since = time.time()
        
        
        #model.train()  # Set model to training mode

        run_loss = 0
        
        for data  in trainloader:
            
            optimizer.zero_grad()
            inputs,_ = data
            inputs = inputs.to(device)
            outputs = model(inputs)
            exit(0)
            loss = calc_loss(outputs, inputs)
            loss.backward(retain_graph=True)
            optimizer.step()
            #exit(0)
            time_elapsed = time.time() - since
            #print("loss=",loss.item())
            run_loss +=  loss.item()*inputs.size(0)
        
        run_loss = run_loss / len(trainset)
        print("scheduler=",scheduler.get_last_lr())
        scheduler.step()
        print('Epoch: {} \tTraining Loss: {:.6f}'.format(epoch, run_loss))
        print('{:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))


    for idx in np.arange(3):
        save_image(outputs[idx],"Train_pred_output_image%d.png"%idx)
        save_image(inputs[idx],"Train_pred_input_image%d.png"%idx)


#**********************************************************************************************************************************************************
    
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(device)

num_class = 3
model = ResNetUNet(num_class).to(device)
train_model(model, optimizer_ft, exp_lr_scheduler, num_epochs=50)

#**********************************************************************************************************************************************************
model.eval()   # Set model to the evaluation mode

inputs_test, _ = next(iter(testloader))
inputs_test = inputs_test.to(device)
pred = model(inputs_test)
#pred = F.sigmoid(pred)
for idx in np.arange(3):
    save_image(inputs_test[idx],"result_input_image%d.png"%idx)
    save_image(pred[idx],"result_output_image%d.png"%idx)