Model FIT using Keras works well only on CPU, on GPU loss NaN

Viewed 346

I have prepared this code, for image colorization:

train = dataGenerator.flow_from_directory("simple_images/data/", target_size=(256, 256), batch_size=340, class_mode=None)
for img in train[0]:
  try:
      lab = rgb2lab(img)
      X.append(lab[:,:,0]) 
      Y.append(lab[:,:,1:] / 128)
  except:
     print('error')
L = np.array(L)
ab = np.array(ab)
L = L.reshape(L.shape+(1,))

And then this model:

#Encoder
model = Sequential()
model.add(Conv2D(32, (3, 3), activation='relu', padding='same', input_shape=(256, 256, 1)))
model.add(Conv2D(64, (3, 3), activation='relu', padding='same', strides=2))
model.add(Conv2D(64, (3, 3), activation='relu', padding='same'))
model.add(Conv2D(128, (3, 3), activation='relu', padding='same', strides=2))
model.add(Conv2D(128, (3,3), activation='relu', padding='same'))
model.add(Conv2D(256, (3, 3), activation='relu', padding='same', strides=2))
model.add(Conv2D(256, (3,3), activation='relu', padding='same'))
#Decoder
model.add(Conv2D(128, (3,3), activation='relu', padding='same'))
model.add(UpSampling2D((2, 2)))
model.add(Conv2D(64, (3,3), activation='relu', padding='same'))
model.add(UpSampling2D((2, 2)))
model.add(Conv2D(32, (3,3), activation='relu', padding='same'))
model.add(Conv2D(16, (3,3), activation='relu', padding='same'))
model.add(Conv2D(2, (3, 3), activation='tanh', padding='same'))
model.add(UpSampling2D((2, 2)))
model.compile(optimizer='adam', loss='mse' , metrics=['accuracy'])
model.summary()
history = model.fit(L,ab,validation_split=0.1, epochs=300, batch_size=64)

And when i run it, on CPU it works perfect, but it will took so much time, to train. So instead of CPU, i want to use my GPU RTX 3070. But when i run same code, on GPU i have got this:

Epoch 1/2
 2/20 [==>...........................] - ETA: 1s - loss: nan - accuracy: 0.3925   WARNING:tensorflow:Callbacks method `on_train_batch_end` is slow compared to the batch time (batch time: 0.0089s vs `on_train_batch_end` time: 0.0678s). Check your callbacks.
20/20 [==============================] - 2s 78ms/step - loss: nan - accuracy: 0.3167 - val_loss: nan - val_accuracy: 0.2474
Epoch 2/2
20/20 [==============================] - 1s 69ms/step - loss: nan - accuracy: 0.3106 - val_loss: nan - val_accuracy: 0.2474
INFO:tensorflow:Assets written to: colorize_autoencoder.model\assets

For demonstration i used only two epochs. How to fix it?

1 Answers

Providing solution here (answer section) for the benefit of community. Thank you @Martin Poláček for the update.

Please refer working code for CPU and GPU as shown below

#Encoder
model = Sequential()
model.add(Conv2D(32, (3, 3), activation='relu', padding='same', input_shape=(256, 256, 1)))
model.add(Conv2D(64, (3, 3), activation='relu', padding='same', strides=2))
model.add(Conv2D(64, (3, 3), activation='relu', padding='same'))
model.add(Conv2D(128, (3, 3), activation='relu', padding='same', strides=2))
model.add(Conv2D(128, (3,3), activation='relu', padding='same'))
model.add(Conv2D(256, (3, 3), activation='relu', padding='same', strides=2))
model.add(Conv2D(256, (3,3), activation='relu', padding='same'))
#Decoder
model.add(Conv2D(128, (3,3), activation='relu', padding='same'))
model.add(UpSampling2D((2, 2)))
model.add(Conv2D(64, (3,3), activation='relu', padding='same'))
model.add(UpSampling2D((2, 2)))
model.add(Conv2D(32, (3,3), activation='relu', padding='same'))
model.add(Conv2D(16, (3,3), activation='relu', padding='same'))
model.add(Conv2D(2, (3, 3), activation='tanh', padding='same'))
model.add(UpSampling2D((2, 2)))
model.compile(optimizer='adamax', loss='mse' , metrics=['accuracy'])
model.summary()
history = model.fit(L,ab,validation_split=0.1, epochs=300, batch_size=64)
Related