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

2024-02-21 05:13:55.560946: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
2024-02-21 05:13:55.561032: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
2024-02-21 05:13:55.691327: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered

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

Downloading data from https://storage.googleapis.com/tensorflow/tf-keras-datasets/train-labels-idx1-ubyte.gz
29515/29515 [==============================] - 0s 0us/step
Downloading data from https://storage.googleapis.com/tensorflow/tf-keras-datasets/train-images-idx3-ubyte.gz
26421880/26421880 [==============================] - 1s 0us/step
Downloading data from https://storage.googleapis.com/tensorflow/tf-keras-datasets/t10k-labels-idx1-ubyte.gz
5148/5148 [==============================] - 0s 0us/step
Downloading data from https://storage.googleapis.com/tensorflow/tf-keras-datasets/t10k-images-idx3-ubyte.gz
4422102/4422102 [==============================] - 1s 0us/step

(60000, 28, 28)

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

(60000, 28, 28, 1)

BUFFER_SIZE = 60000
BATCH_SIZE = 128

train_dataset = tf.data.Dataset.from_tensor_slices(train_images).shuffle(BUFFER_SIZE).batch(BATCH_SIZE)

def generator_model():
    model = tf.keras.Sequential()
    model.add(layers.Dense(7 * 7 * 256, input_shape=(100, )))
    model.add(layers.BatchNormalization())
    model.add(layers.ReLU())
    model.add(layers.Reshape((7, 7, 256)))
    
    model.add(layers.Conv2DTranspose(128, (5, 5), strides=(1, 1), padding='same'))
    model.add(layers.BatchNormalization())
    model.add(layers.ReLU())
    
    model.add(layers.Conv2DTranspose(64, (5, 5), strides=(2, 2), padding='same'))
    model.add(layers.BatchNormalization())
    model.add(layers.ReLU())
    
    model.add(layers.Conv2DTranspose(1, (5, 5), strides=(2, 2), padding='same', activation='tanh'))
    
    return model

generator = generator_model()

generator.summary()

Model: "sequential_4"
_________________________________________________________________
 Layer (type)                Output Shape              Param #   
=================================================================
 dense_3 (Dense)             (None, 12544)             1266944   
                                                                 
 batch_normalization_8 (Bat  (None, 12544)             50176     
 chNormalization)                                                
                                                                 
 re_lu_3 (ReLU)              (None, 12544)             0         
                                                                 
 reshape_1 (Reshape)         (None, 7, 7, 256)         0         
                                                                 
 conv2d_transpose_3 (Conv2D  (None, 7, 7, 128)         819328    
 Transpose)                                                      
                                                                 
 batch_normalization_9 (Bat  (None, 7, 7, 128)         512       
 chNormalization)                                                
                                                                 
 re_lu_4 (ReLU)              (None, 7, 7, 128)         0         
                                                                 
 conv2d_transpose_4 (Conv2D  (None, 14, 14, 64)        204864    
 Transpose)                                                      
                                                                 
 batch_normalization_10 (Ba  (None, 14, 14, 64)        256       
 tchNormalization)                                               
                                                                 
 re_lu_5 (ReLU)              (None, 14, 14, 64)        0         
                                                                 
 conv2d_transpose_5 (Conv2D  (None, 28, 28, 1)         1601      
 Transpose)                                                      
                                                                 
=================================================================
Total params: 2343681 (8.94 MB)
Trainable params: 2318209 (8.84 MB)
Non-trainable params: 25472 (99.50 KB)
_________________________________________________________________

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

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

<matplotlib.image.AxesImage at 0x7f6643ff3730>

def discriminator_model():
    model = tf.keras.Sequential()
    model.add(layers.Conv2D(64, (5, 5), strides=(2, 2), padding='same', input_shape=(28, 28, 1)))
    model.add(layers.BatchNormalization())
    model.add(layers.LeakyReLU(0.2))
    model.add(layers.Dropout(0.3))
    
    model.add(layers.Conv2D(128, (5, 5), strides=(2, 2), padding='same'))
    model.add(layers.BatchNormalization())
    model.add(layers.LeakyReLU(0.2))
    model.add(layers.Dropout(0.3))
    
    model.add(layers.Flatten())
    model.add(layers.Dense(1, activation='sigmoid'))
    
    return model

discriminator = discriminator_model()

discriminator.summary()

Model: "sequential_5"
_________________________________________________________________
 Layer (type)                Output Shape              Param #   
=================================================================
 conv2d_5 (Conv2D)           (None, 14, 14, 64)        1664      
                                                                 
 batch_normalization_11 (Ba  (None, 14, 14, 64)        256       
 tchNormalization)                                               
                                                                 
 leaky_re_lu_5 (LeakyReLU)   (None, 14, 14, 64)        0         
                                                                 
 dropout_5 (Dropout)         (None, 14, 14, 64)        0         
                                                                 
 conv2d_6 (Conv2D)           (None, 7, 7, 128)         204928    
                                                                 
 batch_normalization_12 (Ba  (None, 7, 7, 128)         512       
 tchNormalization)                                               
                                                                 
 leaky_re_lu_6 (LeakyReLU)   (None, 7, 7, 128)         0         
                                                                 
 dropout_6 (Dropout)         (None, 7, 7, 128)         0         
                                                                 
 flatten_2 (Flatten)         (None, 6272)              0         
                                                                 
 dense_4 (Dense)             (None, 1)                 6273      
                                                                 
=================================================================
Total params: 213633 (834.50 KB)
Trainable params: 213249 (833.00 KB)
Non-trainable params: 384 (1.50 KB)
_________________________________________________________________

output = discriminator(generated_image)
print(output)

tf.Tensor([[0.49951673]], 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)
    return real_loss + fake_loss

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

generator_optimizer = tf.keras.optimizers.Adam(learning_rate=0.0002, beta_1=0.5)
discriminator_optimizer = tf.keras.optimizers.Adam(learning_rate=0.0002, beta_1=0.5)

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

def generate_and_plot_images(model, epoch, test_input):
    predictions = model(test_input)
    fig = plt.figure(figsize=(8, 4))
    
    for i in range(predictions.shape[0]):
        plt.subplot(4, 4, i+1)
        pred = (predictions[i, :, :, 0] + 1) * 127.5
        pred = np.array(pred)
        plt.imshow(pred.astype(np.uint8), cmap='gray')
        plt.axis('off')
    
    plt.savefig(f'image_at_epoch_{epoch+1:04d}.png')
    plt.show()

@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)
        
        gen_loss = generator_loss(fake_output)
        disc_loss = discriminator_loss(real_output, 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 train(dataset, epochs):
    gen_loss_list, disc_loss_list = [], []
    real_score_list, fake_score_list = [], []
    
    for epoch in tqdm(range(epochs)):
        start = time.time()
        num_batches = len(dataset)
        
        print(f'Training started with epoch {epoch+1} with {num_batches} batches')
        
        total_gen_loss, total_disc_loss = 0, 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'Losses after epoch {epoch+1:5d}: generator: {generator_loss:.3f}, discriminator: {discriminator_loss:.3f}, real_score: {real_score+100}, fake_score: {fake_score*100}')
        
        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}: {time.time()-start}')
    
    return gen_loss_list, disc_loss_list, real_score_list, fake_score_list

gen_loss_list, disc_loss_list, real_score_list, fake_score_list = train(train_dataset, epochs=epochs)

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

Training started with epoch 1 with 469 batches

2024-02-21 05:53:09.213449: E tensorflow/core/grappler/optimizers/meta_optimizer.cc:961] layout failed: INVALID_ARGUMENT: Size of values 0 does not match size of permutation 4 @ fanin shape insequential_5/dropout_5/dropout/SelectV2-2-TransposeNHWCToNCHW-LayoutOptimizer
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
I0000 00:00:1708494792.134134     465 device_compiler.h:186] Compiled cluster using XLA!  This line is logged at most once for the lifetime of the process.

Losses after epoch     1: generator: 1.885, discriminator: 0.668, real_score: 100.73233795166016, fake_score: 20.510927200317383

  2%|▏         | 1/50 [00:16<13:23, 16.41s/it]

Time for epoch 1: 16.40605115890503
Training started with epoch 2 with 469 batches
Losses after epoch     2: generator: 1.649, discriminator: 0.724, real_score: 100.72960662841797, fake_score: 25.764339447021484

  4%|▍         | 2/50 [00:24<09:21, 11.71s/it]

Time for epoch 2: 8.416456937789917
Training started with epoch 3 with 469 batches
Losses after epoch     3: generator: 1.201, discriminator: 0.975, real_score: 100.69618225097656, fake_score: 37.36167907714844

  6%|▌         | 3/50 [00:33<07:59, 10.20s/it]

Time for epoch 3: 8.414148330688477
Training started with epoch 4 with 469 batches
Losses after epoch     4: generator: 1.800, discriminator: 1.178, real_score: 100.4784927368164, fake_score: 23.120609283447266

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

ax1.plot(gen_loss_list, label='Generator Losses', alpha=0.5)
ax1.plot(disc_loss_list, label='Discriminator Losses', 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.legend()
ax2.set_title('Accuracy scores')

Text(0.5, 1.0, 'Accuracy scores')