Noisy output of U-Net autoencoder - How to deal with this situations?

Hello everyone,
I have built an autoencoder based on Unet model with some modifications like changing activation functions or kernel size etc.

My model does not have any error, but when I try to train, I got completely noisy output.

I applied Canny edge-detection algorithm to make ground truth and original image is my input. I trained my model on 10 images for 10 epochs to make sure my model can overfit. But it never fits!

This is a link to colab notebook of current model.
(Sorry I didn’t pasted the code, it is too long)

Something I should point is, I use BCELoss because I have only 0/1 in my ground truth image and I used nn.Sigmoid() as the activation function in C class (final layer).

Here is what happening to loss (I use reduction='mean'):
First number = epoch, second = batch

1 , 1 loss: 0.8242199420928955
2 , 1 loss: 0.6571497917175293
3 , 1 loss: 0.5504651665687561
4 , 1 loss: 0.4780828356742859
5 , 1 loss: 0.38702648878097534
6 , 1 loss: 0.2822074592113495
7 , 1 loss: 0.2899935245513916
8 , 1 loss: 0.25027400255203247
9 , 1 loss: 0.24944481253623962
10 , 1 loss: 0.23760955035686493
Finished Training
loss: 0.461

Input (I just showing 9 of them in single image to you. In train they are separated):

Output:

I have no idea what should I do or where the problem comes from.
I am currently trying to use Visdom to visualize the process.

Another thing I should point is, I used a simple CNN network with same transform and loss function and everything looks fine and I got good results.
Here is the one:

Thanks for any advice.

1 Like