import os, time
import numpy as np
import tensorflow as tf
import matplotlib.pyplot as plt
from tqdm import tqdm

2024-02-06 17:27:29.903751: E tensorflow/compiler/xla/stream_executor/cuda/cuda_dnn.cc:9342] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered

2024-02-06 17:27:29.903790: E tensorflow/compiler/xla/stream_executor/cuda/cuda_fft.cc:609] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered

2024-02-06 17:27:29.903822: E tensorflow/compiler/xla/stream_executor/cuda/cuda_blas.cc:1518] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered

2024-02-06 17:27:29.910330: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.

To enable the following instructions: AVX2 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.

2024-02-06 17:27:30.610234: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT

(train_images, _), (_, _) = tf.keras.datasets.fashion_mnist.load_data()

print('Original shape of dataset:', train_images.shape)

Original shape of dataset: (60000, 28, 28)

train_images = (train_images - 127.5) / 127.5

train_images = train_images.reshape(train_images.shape[0], 28, 28, 1).astype('float32')

print('New shape of the dataset:', train_images.shape)

New shape of the dataset: (60000, 28, 28, 1)

print('Shape of 1 image:', train_images[6969].shape)
plt.imshow(train_images[6969].squeeze(), cmap='gray')
plt.axis('off')

Shape of 1 image: (28, 28, 1)

(-0.5, 27.5, 27.5, -0.5)

buffer_size = train_images.shape[0]
batch_size = 16

train_dataset = tf.data.Dataset.from_tensor_slices(train_images).shuffle(buffer_size).batch(batch_size)

2024-02-06 17:27:35.986173: E tensorflow/compiler/xla/stream_executor/cuda/cuda_driver.cc:268] failed call to cuInit: CUDA_ERROR_NO_DEVICE: no CUDA-capable device is detected

2024-02-06 17:27:35.987796: W tensorflow/tsl/framework/cpu_allocator_impl.cc:83] Allocation of 188160000 exceeds 10% of free system memory.

def generator_model():
    model = tf.keras.Sequential()
    model.add(tf.keras.layers.Dense(64, input_dim=100))
    model.add(tf.keras.layers.ReLU())

    model.add(tf.keras.layers.Dense(128))
    model.add(tf.keras.layers.ReLU())

    model.add(tf.keras.layers.Dense(256))
    model.add(tf.keras.layers.ReLU())

    model.add(tf.keras.layers.Dense(train_images.shape[1] * train_images.shape[2], activation='tanh'))
    model.add(tf.keras.layers.Reshape(train_images.shape[1:]))

    return model

generator = generator_model()

generator.summary()

Model: "sequential"

_________________________________________________________________

 Layer (type)                Output Shape              Param #   

=================================================================

 dense (Dense)               (None, 64)                6464      

                                                                 

 re_lu (ReLU)                (None, 64)                0         

                                                                 

 dense_1 (Dense)             (None, 128)               8320      

                                                                 

 re_lu_1 (ReLU)              (None, 128)               0         

                                                                 

 dense_2 (Dense)             (None, 256)               33024     

                                                                 

 re_lu_2 (ReLU)              (None, 256)               0         

                                                                 

 dense_3 (Dense)             (None, 784)               201488    

                                                                 

 reshape (Reshape)           (None, 28, 28, 1)         0         

                                                                 

=================================================================

Total params: 249296 (973.81 KB)

Trainable params: 249296 (973.81 KB)

Non-trainable params: 0 (0.00 Byte)

_________________________________________________________________

noise = tf.random.normal([1, 100])
generated_image = generator(noise, training=False)

print('Shape of generated image:', generated_image.shape)

plt.imshow(generated_image[0, :, :, 0], cmap='gray'); plt.axis('off')

Shape of generated image: (1, 28, 28, 1)

(-0.5, 27.5, 27.5, -0.5)

def discriminator_model():
    model = tf.keras.Sequential()
    model.add(tf.keras.layers.Input(shape=train_images.shape[1:]))
    model.add(tf.keras.layers.Flatten())

    model.add(tf.keras.layers.Dense(256))
    model.add(tf.keras.layers.LeakyReLU(0.2))
    model.add(tf.keras.layers.Dropout(0.5))

    model.add(tf.keras.layers.Dense(128))
    model.add(tf.keras.layers.LeakyReLU(0.2))
    model.add(tf.keras.layers.Dropout(0.3))

    model.add(tf.keras.layers.Dense(64))
    model.add(tf.keras.layers.LeakyReLU(0.2))
    model.add(tf.keras.layers.Dropout(0.2))

    model.add(tf.keras.layers.Dense(1, activation='sigmoid'))

    return model

discriminator = discriminator_model()

discriminator.summary()

Model: "sequential_1"

_________________________________________________________________

 Layer (type)                Output Shape              Param #   

=================================================================

 flatten (Flatten)           (None, 784)               0         

                                                                 

 dense_4 (Dense)             (None, 256)               200960    

                                                                 

 leaky_re_lu (LeakyReLU)     (None, 256)               0         

                                                                 

 dropout (Dropout)           (None, 256)               0         

                                                                 

 dense_5 (Dense)             (None, 128)               32896     

                                                                 

 leaky_re_lu_1 (LeakyReLU)   (None, 128)               0         

                                                                 

 dropout_1 (Dropout)         (None, 128)               0         

                                                                 

 dense_6 (Dense)             (None, 64)                8256      

                                                                 

 leaky_re_lu_2 (LeakyReLU)   (None, 64)                0         

                                                                 

 dropout_2 (Dropout)         (None, 64)                0         

                                                                 

 dense_7 (Dense)             (None, 1)                 65        

                                                                 

=================================================================

Total params: 242177 (946.00 KB)

Trainable params: 242177 (946.00 KB)

Non-trainable params: 0 (0.00 Byte)

_________________________________________________________________

output = discriminator(generated_image)
print('Probability of image to be real:', output)

Probability of image to be real: tf.Tensor([[0.5209325]], shape=(1, 1), dtype=float32)

bce = tf.keras.losses.BinaryCrossentropy()

def discriminator_loss(real_output, fake_output):
    real_loss = bce(tf.ones_like(real_output), real_output)
    fake_loss = bce(tf.zeros_like(fake_output), fake_output)
    total_loss = real_loss + fake_loss
    
    return total_loss

def generator_loss(fake_output):
    gen_loss = bce(tf.ones_like(fake_output), fake_output)
    return gen_loss

generator_optimizer = tf.keras.optimizers.Adam(learning_rate=1e-3)
discriminator_optimizer = tf.keras.optimizers.Adam(learning_rate=1e-3)

checkpoint_dir = './training_checkpoints'
checkpoint_prefix = os.path.join(checkpoint_dir, 'ckpt')
checkpoint = tf.train.Checkpoint(generator_optimizer=generator_optimizer,
                                 discriminator_optimizer=discriminator_optimizer,
                                 generator=generator,
                                 discriminator=discriminator)

epochs = 50
noise_dim = 100
num_examples_to_generate = 16

seed = tf.random.normal([num_examples_to_generate, noise_dim])

@tf.function
def train_step(images):
    noise = tf.random.normal([batch_size, noise_dim])

    with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape:
        generated_images = generator(noise, training=True)

        real_output = discriminator(images, training=True)
        fake_output = discriminator(generated_images, training=True)

        disc_loss = discriminator_loss(real_output, fake_output)
        gen_loss = generator_loss(fake_output)

    gradients_of_generator = gen_tape.gradient(gen_loss, generator.trainable_variables)
    gradients_of_discriminator = disc_tape.gradient(disc_loss, discriminator.trainable_variables)

    generator_optimizer.apply_gradients(zip(gradients_of_generator, generator.trainable_variables))
    discriminator_optimizer.apply_gradients(zip(gradients_of_discriminator, discriminator.trainable_variables))

    return gen_loss, disc_loss, tf.reduce_mean(real_output), tf.reduce_mean(fake_output)

def generate_and_plot_images(model, epoch, test_input):
    predictions = model(test_input, training=False)
    
    fig = plt.figure(figsize=(8, 4))

    for i in range(predictions.shape[0]):
        plt.subplot(4, 4, i+1)
        pred = (1 + predictions[i, :, :, 0]) * 127.5
        pred = np.array(pred)
        plt.imshow(pred.astype(np.uint8), cmap='gray')
        plt.axis('off')

    plt.savefig(f'image_at_epoch_{epoch:04d}.png')
    plt.show()

def train(dataset, epochs):
    gen_loss_list = []
    disc_loss_list = []

    real_score_list = []
    fake_score_list = []

    num_batches = len(dataset)

    for epoch in tqdm(range(epochs)):
        start = time.time()
        
        print(f'Training started with epoch {epoch+1} with {num_batches} batches...')

        total_gen_loss = 0
        total_disc_loss = 0

        for batch in dataset:
            generator_loss, discriminator_loss, real_score, fake_score = train_step(batch)
            total_gen_loss += generator_loss
            total_disc_loss += discriminator_loss

        mean_gen_loss = total_gen_loss / num_batches
        mean_disc_loss = total_disc_loss / num_batches

        print(f'Epoch: {epoch+1}\tGenerator loss: {mean_gen_loss}\tDiscriminator loss: {mean_disc_loss}')

        generate_and_plot_images(generator, epoch+1, seed)

        gen_loss_list.append(mean_gen_loss)
        disc_loss_list.append(mean_disc_loss)
        real_score_list.append(real_score)
        fake_score_list.append(fake_score)

        if (epoch+1) % 10 == 0:
            checkpoint.save(file_prefix=checkpoint_prefix)
        
        print(f'Time for epoch {epoch+1} is {time.time()-start} sec')

    return gen_loss_list, disc_loss_list, real_score_list, fake_score_list

gen_loss_epochs, disc_loss_epochs, real_score_list, fake_score_list = train(train_dataset, epochs)

  0%|          | 0/50 [00:00<?, ?it/s]

Training started with epoch 1 with 3750 batches...

2024-02-06 17:27:36.491336: W tensorflow/tsl/framework/cpu_allocator_impl.cc:83] Allocation of 188160000 exceeds 10% of free system memory.

Epoch: 1	Generator loss: 2.110231637954712	Discriminator loss: 0.7386881709098816

  2%|▏         | 1/50 [00:14<11:57, 14.64s/it]

Time for epoch 1 is 14.6415855884552 sec

Training started with epoch 2 with 3750 batches...

Epoch: 2	Generator loss: 1.3505053520202637	Discriminator loss: 1.0031410455703735

  4%|▍         | 2/50 [00:27<10:46, 13.46s/it]

Time for epoch 2 is 12.636430740356445 sec

Training started with epoch 3 with 3750 batches...

Epoch: 3	Generator loss: 1.0728014707565308	Discriminator loss: 1.1759798526763916

  6%|▌         | 3/50 [00:39<10:03, 12.83s/it]

Time for epoch 3 is 12.086230278015137 sec

Training started with epoch 4 with 3750 batches...

Epoch: 4	Generator loss: 1.016257405281067	Discriminator loss: 1.2148150205612183

fig, (ax1,ax2) = plt.subplots(1, 2, figsize = (12, 8))

ax1.plot(gen_loss_epochs, label = 'Generator loss', alpha = 0.5)
ax1.plot(disc_loss_epochs, label = 'Discriminator loss', alpha = 0.5)
ax1.legend()
ax1.set_title('Training Losses')

ax2.plot(real_score_list, label = 'Real score', alpha = 0.5)
ax2.plot(fake_score_list, label = 'Fake score', alpha = 0.5)
ax2.set_title('Accuracy Scores')
ax2.legend()

<matplotlib.legend.Legend at 0x76ee3d608a00>