I am trying to train a model that colorize grayscale pictures but i had an error that i dont understand

General Discussion
, ,
#1

This is my code and i have a train and test sets with a subdirectories as ‘gray’ and ‘color’ as i show in the picture;
Screenshot_1

import tensorflow as tf
import numpy
import matplotlib

# Define the input and output layers
inputs = tf.keras.layers.Input(shape=(256, 256, 1))
outputs = tf.keras.layers.Input(shape=(256, 256, 3))

# Define the encoder part of the network
conv1 = tf.keras.layers.Conv2D(64, 3, activation='relu', padding='same', kernel_initializer='he_normal')(inputs)
conv1 = tf.keras.layers.Conv2D(64, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv1)
pool1 = tf.keras.layers.MaxPooling2D(pool_size=(2, 2))(conv1)

conv2 = tf.keras.layers.Conv2D(128, 3, activation='relu', padding='same', kernel_initializer='he_normal')(pool1)
conv2 = tf.keras.layers.Conv2D(128, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv2)
pool2 = tf.keras.layers.MaxPooling2D(pool_size=(2, 2))(conv2)

conv3 = tf.keras.layers.Conv2D(256, 3, activation='relu', padding='same', kernel_initializer='he_normal')(pool2)
conv3 = tf.keras.layers.Conv2D(256, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv3)
pool3 = tf.keras.layers.MaxPooling2D(pool_size=(2, 2))(conv3)

conv4 = tf.keras.layers.Conv2D(512, 3, activation='relu', padding='same', kernel_initializer='he_normal')(pool3)
conv4 = tf.keras.layers.Conv2D(512, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv4)
pool4 = tf.keras.layers.MaxPooling2D(pool_size=(2, 2))(conv4)

# Define the decoder part of the network
up5 = tf.keras.layers.Conv2D(256, 2, activation='relu', padding='same', kernel_initializer='he_normal')(tf.keras.layers.UpSampling2D(size = (2,2))(pool4))
up6 = tf.keras.layers.Conv2D(256, 2, activation='relu', padding='same', kernel_initializer='he_normal')(tf.keras.layers.UpSampling2D(size = (2,2))(up5))


merge5 = tf.keras.layers.concatenate([conv3,up6], axis = 3)
conv5 = tf.keras.layers.Conv2D(256, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(merge5)
conv5 = tf.keras.layers.Conv2D(256, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv5)

up7 = tf.keras.layers.Conv2D(128, 2, activation='relu', padding='same', kernel_initializer='he_normal')(tf.keras.layers.UpSampling2D(size = (2,2))(conv5))
merge6 = tf.keras.layers.concatenate([conv2,up7], axis = 3)
conv6 = tf.keras.layers.Conv2D(128, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(merge6)
conv6 = tf.keras.layers.Conv2D(128, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv6)

up8= tf.keras.layers.Conv2D(64, 2, activation='relu', padding='same', kernel_initializer='he_normal')(tf.keras.layers.UpSampling2D(size = (2,2))(conv6))
merge7 = tf.keras.layers.concatenate([conv1,up8], axis = 3)
conv7 = tf.keras.layers.Conv2D(64, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(merge7)
conv7 = tf.keras.layers.Conv2D(64, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv7)
conv7 = tf.keras.layers.Conv2D(3, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv7)
print(conv7.shape)
model = tf.keras.Model(inputs = inputs, outputs = conv7)

model.compile(optimizer = 'adam', loss = 'mae', metrics = ['accuracy'])


# Load the train and test datasets
train_ds = tf.keras.preprocessing.image_dataset_from_directory(
    'C:/Users/Admin/Desktop/pythonProjects/colorizedPictures-CNN/train/',
    labels='inferred',
    label_mode='binary',
    class_names=['gray', 'color'],
    color_mode='grayscale',
    batch_size=32,
    image_size=(256, 256))


test_ds = tf.keras.preprocessing.image_dataset_from_directory(
    'C:/Users/Admin/Desktop/pythonProjects/colorizedPictures-CNN/test/',
    labels='inferred',
    label_mode='binary',
    class_names=['gray', 'color'],
    color_mode='grayscale',
    batch_size=32,
    image_size=(256, 256))

# Fit the model on the training dataset
model.fit(train_ds, epochs=5, validation_data=test_ds)

But when i run this code i get this error;

Traceback (most recent call last):
  File "C:\Users\Admin\Desktop\pythonProjects\colorizedPictures-CNN\colorizedPictures.py", line 72, in <module>
    model.fit(train_ds, epochs=5, validation_data=test_ds)
  File "C:\Users\Admin\Desktop\pythonProjects\colorizedPictures-CNN\lib\site-packages\keras\utils\traceback_utils.py", line 70, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "C:\Users\Admin\Desktop\pythonProjects\colorizedPictures-CNN\lib\site-packages\tensorflow\python\eager\execute.py", line 52, in quick_execute
    tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
tensorflow.python.framework.errors_impl.InvalidArgumentError: Graph execution error:

Detected at node 'gradient_tape/mean_absolute_error/sub/BroadcastGradientArgs' defined at (most recent call last):
    File "C:\Users\Admin\Desktop\pythonProjects\colorizedPictures-CNN\colorizedPictures.py", line 72, in <module>
      model.fit(train_ds, epochs=5, validation_data=test_ds)
    File "C:\Users\Admin\Desktop\pythonProjects\colorizedPictures-CNN\lib\site-packages\keras\utils\traceback_utils.py", line 65, in error_handler
      return fn(*args, **kwargs)
    File "C:\Users\Admin\Desktop\pythonProjects\colorizedPictures-CNN\lib\site-packages\keras\engine\training.py", line 1650, in fit
      tmp_logs = self.train_function(iterator)
    File "C:\Users\Admin\Desktop\pythonProjects\colorizedPictures-CNN\lib\site-packages\keras\engine\training.py", line 1249, in train_function
      return step_function(self, iterator)
    File "C:\Users\Admin\Desktop\pythonProjects\colorizedPictures-CNN\lib\site-packages\keras\engine\training.py", line 1233, in step_function
      outputs = model.distribute_strategy.run(run_step, args=(data,))
    File "C:\Users\Admin\Desktop\pythonProjects\colorizedPictures-CNN\lib\site-packages\keras\engine\training.py", line 1222, in run_step
      outputs = model.train_step(data)
    File "C:\Users\Admin\Desktop\pythonProjects\colorizedPictures-CNN\lib\site-packages\keras\engine\training.py", line 1027, in train_step
      self.optimizer.minimize(loss, self.trainable_variables, tape=tape)
    File "C:\Users\Admin\Desktop\pythonProjects\colorizedPictures-CNN\lib\site-packages\keras\optimizers\optimizer_experimental\optimizer.py", line 526, in minimize
      grads_and_vars = self.compute_gradients(loss, var_list, tape)
    File "C:\Users\Admin\Desktop\pythonProjects\colorizedPictures-CNN\lib\site-packages\keras\optimizers\optimizer_experimental\optimizer.py", line 259, in compute_gradients
      grads = tape.gradient(loss, var_list)
Node: 'gradient_tape/mean_absolute_error/sub/BroadcastGradientArgs'
Incompatible shapes: [32,256,256,3] vs. [32,1]
	 [[{{node gradient_tape/mean_absolute_error/sub/BroadcastGradientArgs}}]] [Op:__inference_train_function_4769]

Process finished with exit code 1

How can i fix it? Since im new at this area im literally stuck here and have no idea what to do.

#3

Hi @Emre_Akbas I think the error might be occurring due to your output being of a different shape than the one expected by the model, your output image is of shape (32,1) whereas the model expects (32,256,256,3) .

You could look into the summary of the model to see how the shape of your input image changes with each layer and change the layer parameters accordingly in order to get the expected output.

1 Like