Error in layer of Discriminator Model while making a GAN model

Question

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)

Answer 1

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))