I am currently working on Semantic segmentation on VOC2012 dataset with U-net. I encounter a problem that loss is dropping a bit slow and not moving after few epoch.
Here is my training function and model
def train(args, model, device, train_loader, optimizer, epoch):
model.train()
running_loss = 0.0
total_train = 0
running_correct = 0.0
for batch_idx, (data, target) in enumerate(train_loader):
data, target = data.to(device), target.to(device)
optimizer.zero_grad()
output = model(data)
criterion = nn.CrossEntropyLoss()
loss = criterion(output, target[:, 0])
loss.backward()
optimizer.step()
running_loss += loss.item()
predicted = torch.argmax(output.data, 1)
total_train += target.nelement()
running_correct += predicted.eq(target.data).sum().item()
print(epoch, running_correct / total_train)
print(running_loss / len(train_loader.dataset))
import torch
import torch.nn as nn
import torch.nn.functional as F
class UNetEnc(nn.Module):
def __init__(self, in_channels, features, out_channels):
super().__init__()
self.up = nn.Sequential(
nn.Conv2d(in_channels, features, 3), #kernal size: 3
nn.LeakyReLU(True),
nn.Conv2d(features, features, 3),
nn.LeakyReLU(True),
nn.ConvTranspose2d(features, out_channels, 2, stride=2),
nn.LeakyReLU(True),
)
def forward(self, x):
return self.up(x)
class UNetDec(nn.Module):
def __init__(self, in_channels, out_channels, dropout=False):
super().__init__()
layers = [
nn.Conv2d(in_channels, out_channels, 3),
nn.LeakyReLU(inplace=True),
nn.Conv2d(out_channels, out_channels, 3),
nn.LeakyReLU(inplace=True),
]
if dropout:
layers += [nn.Dropout(.5)]
layers += [nn.MaxPool2d(2, stride=2, ceil_mode=True)]
self.down = nn.Sequential(*layers)
def forward(self, x):
return self.down(x)
class UNet(nn.Module):
def __init__(self, num_classes):
super().__init__()
self.dec1 = UNetDec(3, 64)
self.dec2 = UNetDec(64, 128)
self.dec3 = UNetDec(128, 256)
self.dec4 = UNetDec(256, 512, dropout=True)
self.center = nn.Sequential(
nn.Conv2d(512, 1024, 3),
nn.LeakyReLU(inplace=True),
nn.Conv2d(1024, 1024, 3),
nn.LeakyReLU(inplace=True),
nn.Dropout(),
nn.ConvTranspose2d(1024, 512, 2, stride=2),
nn.LeakyReLU(inplace=True),
)
self.enc4 = UNetEnc(1024, 512, 256)
self.enc3 = UNetEnc(512, 256, 128)
self.enc2 = UNetEnc(256, 128, 64)
self.enc1 = nn.Sequential(
nn.Conv2d(128, 64, 3),
nn.LeakyReLU(inplace=True),
nn.Conv2d(64, 64, 3),
#nn.LeakyReLU(inplace=True),
)
self.final = nn.Conv2d(64, num_classes, 1)
def forward(self, x):
dec1 = self.dec1(x)
dec2 = self.dec2(dec1)
dec3 = self.dec3(dec2)
dec4 = self.dec4(dec3)
center = self.center(dec4)
enc4 = self.enc4(torch.cat([
center, F.upsample_bilinear(dec4, center.size()[2:])], 1))
enc3 = self.enc3(torch.cat([
enc4, F.upsample_bilinear(dec3, enc4.size()[2:])], 1))
enc2 = self.enc2(torch.cat([
enc3, F.upsample_bilinear(dec2, enc3.size()[2:])], 1))
enc1 = self.enc1(torch.cat([
enc2, F.upsample_bilinear(dec1, enc2.size()[2:])], 1))
return F.upsample_bilinear(self.final(enc1), x.size()[2:])
and my optimizer
optimizer = optim.Adam(model.parameters(), lr=0.0001)
Also the output of loss
1 0.5320905093947355
2.0186073773985127
2 0.5344406446928762
1.8902393944431681
3 0.537645096729026
1.837425033883143
4 0.538817567193965
1.7639471042602735
5 0.5387050243204894
1.7177013124582658
Thank you if anyone know how can I improve the performance