I want to use the basic UNET model to reconstruct noisy 512×512 images towards a noiseless image. The TensorFlow implementation of the model is here below. The training gets stuck on the same loss and val_loss for each epoch for at least 200 epochs straight on several learning rates.

def UNET(IMG_HEIGHT, IMG_WIDTH, IMG_CHANNELS):
    #Build the model
    inputs = Input((IMG_HEIGHT, IMG_WIDTH, IMG_CHANNELS))

    #Contraction path
    x = Conv2D(64, (3, 3), activation='relu', padding='same')(inputs)
    c1 = Conv2D(64, (3, 3), activation='relu', padding='same')(x)
    p1 = MaxPooling2D((2, 2))(c1)

    x = Conv2D(128, (3, 3), activation='relu', padding='same')(p1)
    c2 = Conv2D(128, (3, 3), activation='relu', padding='same')(x)
    p2 = MaxPooling2D((2, 2))(c2)

    x = Conv2D(256, (3, 3), activation='relu', padding='same')(p2)
    c3 = Conv2D(256, (3, 3), activation='relu', padding='same')(x)
    p3 = MaxPooling2D((2, 2))(c3)

    x = Conv2D(512, (3, 3), activation='relu', padding='same')(p3)
    c4 = Conv2D(512, (3, 3), activation='relu', padding='same')(x)
    p4 = MaxPooling2D(pool_size=(2, 2))(c4)

    x = Conv2D(1024, (3, 3), activation='relu', padding='same')(p4)
    x = Conv2D(1024, (3, 3), activation='relu', padding='same')(x)

    #Expansive path
    x = Conv2DTranspose(512, (2, 2), strides=(2, 2), padding='same')(x)
    x = concatenate([x, c4])
    x = Conv2D(512, (3, 3), activation='relu', padding='same')(x)
    x = Conv2D(512, (3, 3), activation='relu', padding='same')(x)       

    x = Conv2DTranspose(256, (2, 2), strides=(2, 2), padding='same')(x)
    x = concatenate([x, c3])
    x = Conv2D(256, (3, 3), activation='relu', padding='same')(x)
    x = Conv2D(256, (3, 3), activation='relu', padding='same')(x)       

    x = Conv2DTranspose(128, (2, 2), strides=(2, 2), padding='same')(x)
    x = concatenate([x, c2])
    x = Conv2D(128, (3, 3), activation='relu', padding='same')(x)
    x = Conv2D(128, (3, 3), activation='relu', padding='same')(x)       

    x = Conv2DTranspose(64, (2, 2), strides=(2, 2), padding='same')(x)
    x = concatenate([x, c1], axis=3)
    x = Conv2D(64, (3, 3), activation='relu', padding='same')(x)
    x = Conv2D(64, (3, 3), activation='relu', padding='same')(x)        

    outputs = Conv2D(1, (1, 1), activation='relu')(x) # from sigmoid to relu

    model = Model(inputs=[inputs], outputs=[outputs])
    model.compile(optimizer='sgd', loss='mae')

    return model

then perform the following training code:

# get model, data
model = UNET(512,512,1)
trainX,trainy, testX,testy = train_test_split function from sklearn

trainY = trainX - trainy

# fit the model
history = model.fit(trainX, trainY,
            epochs=150, batch_size=1,
            validation_split=0.3, verbose=2)

My images are greyscale images that have been obtained with different resolutions (noise levels). I want to train the model to learn the noise and then subtract the model output (after training) to subtract the noise from the noisy image:

Breakdown of model use

The problem is that the model does not learn anything at each epoch. Both the loss and val_loss stay the same after each epoch, which I do not understand:
No improvement in training

But I do strongly believe that it has something to do with the training data that I put in the model because I normalized the input and ground truth images to range [0,1] first, and then I subtracted the ground truths from the input images to get the noise as targets for the training. But obviously that nets negative values in the resulting extracted noise images.
And another sign that its to do with the training data is that exactly the same happens with a different architecture with the same optimizer learning rate and training data; it gets stuck at exactly the same loss and val_loss.

An overview of this is shown in the picture below:
First: input image, Second: Corresponding GT Image, Third: Image resulting from subtraction

I believe that the negative values are ruining the learning. My supervisor said that I shouldn’t need to renormalize the resulting noise images without explanation, but if I do renormalize them, the training is not stuck anymore, but the resulting model just does not perform good.

Now, a peer student is doing the same project in PyTorch and implemented an identical UNet and got good training results with L1 (MAE) loss and Stochastic Gradient Descent optimizer with a learning_rate of 0.0005 and batch size of 1. My project supervisor said that it shouldn’t matter that I use TensorFlow because if the hyperparams are the same, it should train equally fine. He also did not renormalize the target images and he got good results, so I shouldn’t have to either, right?

I’m not very confident in choice making for CNNs since I’m still learning, which is why I listen blindly to my supervisor. But I wanted to check with the internet if I’m doing something terribly wrong from the get go in the model already.

What I tried:

1. Adding BatchNormalization() layers after each conv layer in the model
2. Tried 0.1,0.01,0.05,0.001,0.005,0.0001,0.0005 learning rates, each with learning rate scheduling also
3. As stated earlier, tried renormalizing targets, got no good results
4. Tried a different validation split
5. Tried doing it on Colab, exactly the same loss and val_loss plateau.
6. Tried using ‘linear’ and ‘tanh’ instead of ‘relu’ for last (output) activation layer but also converges to the same loss and val_loss