How can I improve the performance on unseen data for semantic segmentation using an auto-encoder?

I am using simple autoencoders for the task of semantic segmentation on the VOC2012 dataset. I am currently using a simple autoencoder based model. It is trained on adam optimizer with cross-entropy loss on 21 classes 0 - 20. You can find the code here: https://github.com/parthv21/VOC-Semantic-Segmentation

My Architecture:

   self.encoder = nn.Sequential(
nn.LeakyReLU(),
nn.LeakyReLU(),
nn.LeakyReLU(),
)

self.decorder = nn.Sequential(
nn.LeakyReLU(),
nn.LeakyReLU(),
nn.LeakyReLU(),
)


After 200 iterations I am getting the following output

Validation Data

Is a more complex architecture the only way I can fix this problem? Or can I fix this with a different loss function like dice or more regularization? The same issue happened after training for 100 iterations. So the model is not generalizing for some reason.

Edit

I also tried adding weights to CrossEntropy such that w_label = 1 - frequency(label). The idea was that 0 label for the background which was more common would contribute less to the loss, and other labels which were rare, would contribute more to the loss. But that did not help:

Another thing I tried was ignoring label 0 for background in the loss. But that created horrible results even for training data:

• Why are you using a simple auto-encoder? Have you tried to use more sophisticated models for the task? – nbro Nov 12 '20 at 10:50
• I wanted to see what mileage I get from a simple auto-encoder. Is it absolutely impossible to get good validation performance using auto-encorder? I was thinking of testing skip connections as a more advanced architecture. – Parth Tamane Nov 12 '20 at 15:36