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][1]][1]
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.ReLU(),
nn.MaxPool2d(2, stride=2),
nn.Conv2d(32, 64, 3, padding=1), # Input: 32 channels, Output: 64 channels
nn.ReLU(),
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.ReLU(),
nn.ConvTranspose2d(32, 1, 3, stride=2, padding=1, output_padding=1), # Input: 32 channels, Output: 1 channel
nn.Sigmoid()
)
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.