[FIXED] Error in layer of Discriminator Model while making a GAN model

Issue

I made a GAN model for generating the images based on sample training images of animes. Where on the execution of the code I got this error.

ValueError: Input 0 of layer "discriminator" is incompatible with the layer: expected shape=(None, 64, 64, 3), found shape=(64, 64, 3)

Even changing the shape of the 1st layer of the discriminator to (None, 64, 64, 3) did not help

Code:

Preprocessing:

import numpy as np
import tensorflow as tf
from tqdm import tqdm
from tensorflow import keras
from tensorflow.keras import layers

img_h,img_w,img_c=64,64,3
batch_size=128
latent_dim=128
num_epochs=100

dir='/home/samar/Desktop/project2/anime-gan/data'

dataset = tf.keras.utils.image_dataset_from_directory(
  directory=dir,
  seed=123,
  image_size=(img_h, img_w),
  batch_size=batch_size,
  shuffle=True)

xtrain, ytrain = next(iter(dataset))
xtrain=np.array(xtrain)
xtrain=np.apply_along_axis(lambda x: x/255.0,0,xtrain)

Discriminator model:

discriminator = keras.Sequential(
    [
        keras.Input(shape=(64, 64, 3)),
        layers.Conv2D(64, kernel_size=4, strides=2, padding="same"),
        layers.LeakyReLU(alpha=0.2),
        layers.Conv2D(128, kernel_size=4, strides=2, padding="same"),
        layers.LeakyReLU(alpha=0.2),
        layers.Conv2D(128, kernel_size=4, strides=2, padding="same"),
        layers.LeakyReLU(alpha=0.2),
        layers.Flatten(),
        layers.Dropout(0.2),
        layers.Dense(1, activation="sigmoid"),
    ],
    name="discriminator",
)
discriminator.summary()

Generator Model:

generator = keras.Sequential(
    [
        keras.Input(shape=(latent_dim,)),
        layers.Dense(8 * 8 * 128),
        layers.Reshape((8, 8, 128)),
        layers.Conv2DTranspose(128, kernel_size=4, strides=2, padding="same"),
        layers.LeakyReLU(alpha=0.2),
        layers.Conv2DTranspose(256, kernel_size=4, strides=2, padding="same"),
        layers.LeakyReLU(alpha=0.2),
        layers.Conv2DTranspose(512, kernel_size=4, strides=2, padding="same"),
        layers.LeakyReLU(alpha=0.2),
        layers.Conv2D(3, kernel_size=5, padding="same", activation="sigmoid"),
    ],
    name="generator",
)
generator.summary()

Training:

opt_gen = keras.optimizers.Adam(1e-4)
opt_disc = keras.optimizers.Adam(1e-4)
loss_fn = keras.losses.BinaryCrossentropy()
for epoch in range(10):
    for idx, real in enumerate(tqdm(xtrain)):
        batch_size=real.shape[0]
        random_latent_vectors = tf.random.normal(shape=(batch_size, latent_dim))
        with tf.GradientTape() as gen_tape:
            fake = generator(random_latent_vectors)
        if idx % 100 == 0:
            img = keras.preprocessing.image.array_to_img(fake[0])
            img.save("/home/samar/Desktop/project2/anime-gan/gen_images/generated_img_%03d_%d.png" % (epoch, idx))    
        with tf.GradientTape() as disc_tape:
            loss_disc_real = loss_fn(tf.ones((batch_size,1)), discriminator(real))
            loss_disc_fake = loss_fn(tf.zeros((batch_size,1)), discriminator(fake))
            loss_disc = (loss_disc_real + loss_disc_fake) / 2
        gradients_of_discriminator = disc_tape.gradient(loss_disc, discriminator.trainable_variables)
        opt_disc.apply_gradients(zip(gradients_of_discriminator, discriminator.trainable_variables))
        with tf.GradientTape() as gen_tape:
            fake = generator(random_latent_vectors)
            output = discriminator(fake)
            loss_gen = loss_fn(tf.ones(batch_size, 1), output)

        grads = gen_tape.gradient(loss_gen, generator.trainable_weights)
        opt_gen.apply_gradients(zip(grads, generator.trainable_weights))

And also can you please explain me the difference between the shapes (None, 64, 64, 3) and (64, 64, 3)

Solution

The problem is you are extracting exactly one batch when running xtrain, ytrain = next(iter(train_ds)) and you are then iterating over this batch in your training loop. That is why you are missing the batch dimension (None). I am not sure what your dataset looks like, but here is a working example using tf.keras.utils.image_dataset_from_directory:

import numpy as np
import tensorflow as tf
from tqdm import tqdm
from tensorflow import keras
from tensorflow.keras import layers
import pathlib

dataset_url = "https://storage.googleapis.com/download.tensorflow.org/example_images/flower_photos.tgz"
data_dir = tf.keras.utils.get_file('flower_photos', origin=dataset_url, untar=True)
data_dir = pathlib.Path(data_dir)

img_h,img_w,img_c=64,64,3
batch_size=128
latent_dim=128
num_epochs=100

train_ds = tf.keras.utils.image_dataset_from_directory(
  data_dir,
  seed=123,
  image_size=(img_h,img_w),
  batch_size=batch_size)

normalization_layer = tf.keras.layers.Rescaling(1./255)
train_ds = train_ds.map(lambda x, y: normalization_layer(x))

discriminator = keras.Sequential(
    [
        keras.Input(shape=(64, 64, 3)),
        layers.Conv2D(64, kernel_size=4, strides=2, padding="same"),
        layers.LeakyReLU(alpha=0.2),
        layers.Conv2D(128, kernel_size=4, strides=2, padding="same"),
        layers.LeakyReLU(alpha=0.2),
        layers.Conv2D(128, kernel_size=4, strides=2, padding="same"),
        layers.LeakyReLU(alpha=0.2),
        layers.Flatten(),
        layers.Dropout(0.2),
        layers.Dense(1, activation="sigmoid"),
    ],
    name="discriminator",
)
discriminator.summary()

generator = keras.Sequential(
    [
        keras.Input(shape=(latent_dim,)),
        layers.Dense(8 * 8 * 128),
        layers.Reshape((8, 8, 128)),
        layers.Conv2DTranspose(128, kernel_size=4, strides=2, padding="same"),
        layers.LeakyReLU(alpha=0.2),
        layers.Conv2DTranspose(256, kernel_size=4, strides=2, padding="same"),
        layers.LeakyReLU(alpha=0.2),
        layers.Conv2DTranspose(512, kernel_size=4, strides=2, padding="same"),
        layers.LeakyReLU(alpha=0.2),
        layers.Conv2D(3, kernel_size=5, padding="same", activation="sigmoid"),
    ],
    name="generator",
)
generator.summary()

opt_gen = keras.optimizers.Adam(1e-4)
opt_disc = keras.optimizers.Adam(1e-4)
loss_fn = keras.losses.BinaryCrossentropy()
for epoch in range(10):
    for idx, real in enumerate(tqdm(train_ds)):
        batch_size=real.shape[0]
        random_latent_vectors = tf.random.normal(shape=(batch_size, latent_dim))
        with tf.GradientTape() as gen_tape:
            fake = generator(random_latent_vectors)
        if idx % 100 == 0:
            img = keras.preprocessing.image.array_to_img(fake[0])
        with tf.GradientTape() as disc_tape:
            loss_disc_real = loss_fn(tf.ones((batch_size,1)), discriminator(real))
            loss_disc_fake = loss_fn(tf.zeros((batch_size,1)), discriminator(fake))
            loss_disc = (loss_disc_real + loss_disc_fake) / 2
        gradients_of_discriminator = disc_tape.gradient(loss_disc, discriminator.trainable_variables)
        opt_disc.apply_gradients(zip(gradients_of_discriminator, discriminator.trainable_variables))
        with tf.GradientTape() as gen_tape:
            fake = generator(random_latent_vectors)
            output = discriminator(fake)
            loss_gen = loss_fn(tf.ones(batch_size, 1), output)

        grads = gen_tape.gradient(loss_gen, generator.trainable_weights)
        opt_gen.apply_gradients(zip(grads, generator.trainable_weights))

Answered By - AloneTogether