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 2 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):

image.jpg768×768 192 KB

Output:

image.jpg768×768 232 KB

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:

out.png768×768 354 KB

Thanks for any advice.