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