0
$\begingroup$

CONTEXT

I'm trying to colorize images with Variational Autoencoder. Input is 256x256 gray image. Output is 256x256x2 as I convert image to a LAB color space and then put gray channel as input and other two as outputs.

PROBLEM

My network is training, but loss is not getting smaller epoch after epoch.

Epoch 2/5 720/720 [==============================] - 34s 47ms/sample - loss: 0.1081 - val_loss: 0.0860 
Epoch 3/5 720/720 [==============================] - 35s 48ms/sample - loss: 0.1074 - val_loss: 0.0853 
Epoch 4/5 720/720 [==============================] - 34s 48ms/sample - loss: 0.1070 - val_loss: 0.0855 
Epoch 5/5 720/720 [==============================] - 34s 48ms/sample - loss: 0.1068 - val_loss: 0.0854

After training I can put gray image into the network to get color channels, but I always get the same sort of orange dot in the middle of the screen. No matter how long I train or how many training images I use.

Result:

vae_curie

Tech stuff

Network I use:

            dropout = 0.25

            input_data = tensorflow.keras.layers.Input(shape=(IMG_HEIGHT, IMG_WIDTH, 1))

            encoder = tensorflow.keras.layers.Conv2D(64, (3, 3), activation='relu', padding='same')(input_data)
            # encoder = tensorflow.keras.layers.MaxPooling2D((2, 2))(encoder)
            encoder = tensorflow.keras.layers.Dropout(dropout)(encoder)

            encoder = tensorflow.keras.layers.Conv2D(128, (3, 3), activation='relu', padding='same')(encoder)
            encoder = tensorflow.keras.layers.MaxPooling2D((2, 2))(encoder)
            encoder = tensorflow.keras.layers.Dropout(dropout)(encoder)

            encoder = tensorflow.keras.layers.Conv2D(128, (3, 3), activation='relu', padding='same')(encoder)
            # encoder = tensorflow.keras.layers.MaxPooling2D((2, 2))(encoder)
            encoder = tensorflow.keras.layers.Dropout(dropout)(encoder)

            encoder = tensorflow.keras.layers.Conv2D(256, (3, 3), activation='relu', padding='same')(encoder)
            encoder = tensorflow.keras.layers.MaxPooling2D((2, 2))(encoder)
            encoder = tensorflow.keras.layers.Dropout(dropout)(encoder)

            encoder = tensorflow.keras.layers.Conv2D(256, (3, 3), activation='relu', padding='same')(encoder)
            # encoder = tensorflow.keras.layers.MaxPooling2D((2, 2))(encoder)
            encoder = tensorflow.keras.layers.Dropout(dropout)(encoder)

            encoder = tensorflow.keras.layers.Conv2D(512, (3, 3), activation='relu', padding='same')(encoder)
            encoder = tensorflow.keras.layers.MaxPooling2D((2, 2))(encoder)
            encoder = tensorflow.keras.layers.Dropout(dropout)(encoder)

            encoder = tensorflow.keras.layers.Flatten()(encoder)
            encoder = tensorflow.keras.layers.Dense(16)(encoder)


            def sample_latent_features(distribution):
                distribution_mean, distribution_variance = distribution
                batch_size = tensorflow.shape(distribution_variance)[0]
                random = tensorflow.keras.backend.random_normal(
                    shape=(batch_size, tensorflow.shape(distribution_variance)[1]))
                return distribution_mean + tensorflow.exp(0.5 * distribution_variance) * random


            distribution_mean = tensorflow.keras.layers.Dense(2, name='mean')(encoder)
            distribution_variance = tensorflow.keras.layers.Dense(2, name='log_variance')(encoder)
            latent_encoding = tensorflow.keras.layers.Lambda(sample_latent_features)(
                [distribution_mean, distribution_variance])

            encoder_model = tensorflow.keras.Model(input_data, latent_encoding)
            encoder_model.summary()

            decoder_input = tensorflow.keras.layers.Input(shape=(2))
            decoder = tensorflow.keras.layers.Dense(512)(decoder_input)
            decoder = tensorflow.keras.layers.Reshape((1, 1, 512))(decoder)
            decoder = tensorflow.keras.layers.Conv2D(512, (3, 3), activation='relu', padding='same')(decoder)
            decoder = tensorflow.keras.layers.UpSampling2D((2, 2))(decoder)

            decoder = tensorflow.keras.layers.Conv2D(256, (3, 3), activation='relu', padding='same')(decoder)
            decoder = tensorflow.keras.layers.UpSampling2D((2, 2))(decoder)

            decoder = tensorflow.keras.layers.Conv2D(128, (3, 3), activation='relu', padding='same')(decoder)
            decoder = tensorflow.keras.layers.UpSampling2D((2, 2))(decoder)

            decoder = tensorflow.keras.layers.Conv2D(64, (3, 3), activation='relu', padding='same')(decoder)
            decoder = tensorflow.keras.layers.UpSampling2D((2, 2))(decoder)
            
            decoder = tensorflow.keras.layers.Conv2D(32, (3, 3), activation='relu', padding='same')(decoder)
            decoder = tensorflow.keras.layers.UpSampling2D((2, 2))(decoder)
            
            decoder = tensorflow.keras.layers.Conv2D(32, (3, 3), activation='relu', padding='same')(decoder)
            decoder = tensorflow.keras.layers.UpSampling2D((2, 2))(decoder)  #
            decoder = tensorflow.keras.layers.Conv2D(32, (3, 3), activation='relu', padding='same')(decoder)
            decoder = tensorflow.keras.layers.UpSampling2D((2, 2))(decoder)  #
            decoder = tensorflow.keras.layers.Conv2D(16, (3, 3), activation='relu', padding='same')(decoder)
            decoder = tensorflow.keras.layers.UpSampling2D((2, 2))(decoder)

            decoder_output = tensorflow.keras.layers.Conv2D(2, (3, 3), activation='tanh', padding='same')(
                decoder)

            decoder_model = tensorflow.keras.Model(decoder_input, decoder_output)
            decoder_model.summary()

            encoded = encoder_model(input_data)
            decoded = decoder_model(encoded)

            vae = tensorflow.keras.models.Model(input_data, decoded)

            def get_loss(encoder_mu, encoder_log_variance):
                def vae_r_loss(y_true, y_predict):
                    reconstruction_loss_factor = 10
                    reconstruction_loss = tensorflow.keras.backend.mean(
                        tensorflow.keras.backend.square(y_true - y_predict), axis=[1, 2, 3])
                    return reconstruction_loss_factor * reconstruction_loss

                def vae_kl_loss(y_true, y_pred):
                    kl_loss = -0.5 * tensorflow.keras.backend.sum(
                        1.0 + encoder_log_variance - tensorflow.keras.backend.square(
                            encoder_mu) - tensorflow.keras.backend.exp(encoder_log_variance), axis=1)
                    return kl_loss

                def vae_loss(y_true, y_pred):
                    r_loss = vae_r_loss(y_true, y_pred)
                    kl_loss = vae_kl_loss(y_true, y_pred)
                    return r_loss + kl_loss

                return vae_loss


            vae.compile(loss=get_loss(distribution_mean, distribution_variance), optimizer='adam',
                        experimental_run_tf_function=False)
            vae.summary()

            history = vae.fit(X, y, epochs=5, batch_size=8, validation_split=0.2)

Network summary:

Model: "encoder"
__________________________________________________________________________________________________
Layer (type)                    Output Shape         Param #     Connected to                     
==================================================================================================
input_1 (InputLayer)            [(None, 256, 256, 1) 0                                            
__________________________________________________________________________________________________
conv2d (Conv2D)                 (None, 256, 256, 64) 640         input_1[0][0]                    
__________________________________________________________________________________________________
dropout (Dropout)               (None, 256, 256, 64) 0           conv2d[0][0]                     
__________________________________________________________________________________________________
conv2d_1 (Conv2D)               (None, 256, 256, 128 73856       dropout[0][0]                    
__________________________________________________________________________________________________
max_pooling2d (MaxPooling2D)    (None, 128, 128, 128 0           conv2d_1[0][0]                   
__________________________________________________________________________________________________
dropout_1 (Dropout)             (None, 128, 128, 128 0           max_pooling2d[0][0]              
__________________________________________________________________________________________________
conv2d_2 (Conv2D)               (None, 128, 128, 128 147584      dropout_1[0][0]                  
__________________________________________________________________________________________________
dropout_2 (Dropout)             (None, 128, 128, 128 0           conv2d_2[0][0]                   
__________________________________________________________________________________________________
conv2d_3 (Conv2D)               (None, 128, 128, 256 295168      dropout_2[0][0]                  
__________________________________________________________________________________________________
max_pooling2d_1 (MaxPooling2D)  (None, 64, 64, 256)  0           conv2d_3[0][0]                   
__________________________________________________________________________________________________
dropout_3 (Dropout)             (None, 64, 64, 256)  0           max_pooling2d_1[0][0]            
__________________________________________________________________________________________________
conv2d_4 (Conv2D)               (None, 64, 64, 256)  590080      dropout_3[0][0]                  
__________________________________________________________________________________________________
dropout_4 (Dropout)             (None, 64, 64, 256)  0           conv2d_4[0][0]                   
__________________________________________________________________________________________________
conv2d_5 (Conv2D)               (None, 64, 64, 512)  1180160     dropout_4[0][0]                  
__________________________________________________________________________________________________
max_pooling2d_2 (MaxPooling2D)  (None, 32, 32, 512)  0           conv2d_5[0][0]                   
__________________________________________________________________________________________________
dropout_5 (Dropout)             (None, 32, 32, 512)  0           max_pooling2d_2[0][0]            
__________________________________________________________________________________________________
flatten (Flatten)               (None, 524288)       0           dropout_5[0][0]                  
__________________________________________________________________________________________________
dense (Dense)                   (None, 16)           8388624     flatten[0][0]                    
__________________________________________________________________________________________________
mean (Dense)                    (None, 2)            34          dense[0][0]                      
__________________________________________________________________________________________________
log_variance (Dense)            (None, 2)            34          dense[0][0]                      
__________________________________________________________________________________________________
lambda (Lambda)                 (None, 2)            0           mean[0][0]                       
                                                                 log_variance[0][0]               
==================================================================================================
Total params: 10,676,180
Trainable params: 10,676,180
Non-trainable params: 0
__________________________________________________________________________________________________
Model: "decoder"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
input_2 (InputLayer)         [(None, 2)]               0         
_________________________________________________________________
dense_1 (Dense)              (None, 512)               1536      
_________________________________________________________________
reshape (Reshape)            (None, 1, 1, 512)         0         
_________________________________________________________________
conv2d_6 (Conv2D)            (None, 1, 1, 512)         2359808   
_________________________________________________________________
up_sampling2d (UpSampling2D) (None, 2, 2, 512)         0         
_________________________________________________________________
conv2d_7 (Conv2D)            (None, 2, 2, 256)         1179904   
_________________________________________________________________
up_sampling2d_1 (UpSampling2 (None, 4, 4, 256)         0         
_________________________________________________________________
conv2d_8 (Conv2D)            (None, 4, 4, 128)         295040    
_________________________________________________________________
up_sampling2d_2 (UpSampling2 (None, 8, 8, 128)         0         
_________________________________________________________________
conv2d_9 (Conv2D)            (None, 8, 8, 64)          73792     
_________________________________________________________________
up_sampling2d_3 (UpSampling2 (None, 16, 16, 64)        0         
_________________________________________________________________
conv2d_10 (Conv2D)           (None, 16, 16, 32)        18464     
_________________________________________________________________
up_sampling2d_4 (UpSampling2 (None, 32, 32, 32)        0         
_________________________________________________________________
conv2d_11 (Conv2D)           (None, 32, 32, 32)        9248      
_________________________________________________________________
up_sampling2d_5 (UpSampling2 (None, 64, 64, 32)        0         
_________________________________________________________________
conv2d_12 (Conv2D)           (None, 64, 64, 32)        9248      
_________________________________________________________________
up_sampling2d_6 (UpSampling2 (None, 128, 128, 32)      0         
_________________________________________________________________
conv2d_13 (Conv2D)           (None, 128, 128, 16)      4624      
_________________________________________________________________
up_sampling2d_7 (UpSampling2 (None, 256, 256, 16)      0         
_________________________________________________________________
conv2d_14 (Conv2D)           (None, 256, 256, 2)       290       
=================================================================
Total params: 3,951,954
Trainable params: 3,951,954
Non-trainable params: 0
_________________________________________________________________
Model: "vae"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
input_1 (InputLayer)         [(None, 256, 256, 1)]     0         
_________________________________________________________________
encoder (Model)              (None, 2)                 10676180  
_________________________________________________________________
decoder (Model)              (None, 256, 256, 2)       3951954   
=================================================================
Total params: 14,628,134
Trainable params: 14,628,134
Non-trainable params: 0
_________________________________________________________________
Train on 720 samples, validate on 180 samples

Questions

  1. Why my network seems to not doing any learning?
  2. Are loss functions defined properly?
$\endgroup$
2
  • $\begingroup$ How large is your training dataset? How many epochs have you trained the VAE for? Only 5 epochs? That's probably not enough. $\endgroup$
    – nbro
    Jan 14 at 13:29
  • $\begingroup$ No, I've tried training it up to 500 epochs. Result is still the same as for 2, 100 or 250. $\endgroup$ Jan 15 at 12:03

0

You must log in to answer this question.

Browse other questions tagged .