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)
