I am building a denoising autoencoder to repaint lanes from a binary image. The input is a binary image that has incomplete lanes, due to vehicles getting in the way. I repaint the lanes manually so that the model can learn from the label. However the model seems to be recreating the input as shown in figure. enter image description here

The code for the model:

class Autoencoder(nn.Module):
    def __init__(self):
        super(Autoencoder, self).__init__()
        self.encoder = nn.Sequential(
            nn.Conv2d(1, 32, 3, padding=1),  # Input: 1 channel, Output: 32 channels
            nn.MaxPool2d(2, stride=2),
            nn.Conv2d(32, 64, 3, padding=1),  # Input: 32 channels, Output: 64 channels
            nn.MaxPool2d(2, stride=2)

        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(64, 32, 3, stride=2, padding=1, output_padding=1),  # Input: 64 channels, Output: 32 channels
            nn.ConvTranspose2d(32, 1, 3, stride=2, padding=1, output_padding=1),  # Input: 32 channels, Output: 1 channel

    def forward(self, x):
        x = self.encoder(x)
        x = self.decoder(x)
        return x

I'm training 5000 such images with batch size 16. I'm using BCE and adam for training.

criterion = nn.BCELoss()
optimizer = optim.Adam(model.parameters(), lr = 0.01)

The loss stagnates at 31 and doesn't seem to decrease. All images are exactly like the input and doesn't seem to learn from the labels. Is there some workaround, so that I can get the model working.

  • $\begingroup$ Try reduce the learning rate to 1e-3 or 1e-4 $\endgroup$
    – Karl
    Mar 22 at 16:05
  • $\begingroup$ Doesn't seem to work. The loss function just keeps oscillating around the same value, I tried changing the model and trying different loss functions like dice loss. Doesn't seem to work out. Can you suggest anymore changes. $\endgroup$ Mar 22 at 17:37
  • $\begingroup$ I"ve gone over 100 epochs and the loss function seems to be decreasing ever so slowly, is there any point training for more epochs? $\endgroup$ Mar 22 at 17:39
  • $\begingroup$ Why are you using 3 input channels if the image is binary? Also, have you checked that you're passing in the correct y variable? If your model recreates the input image a bug seems more likely than a model issue. $\endgroup$
    – MrMulliner
    Mar 24 at 8:33


Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.