Initial commit

Tensorflow/tutoriel35/gan.py

@@ -0,0 +1,102 @@
+import tensorflow as tf
+import glob
+import numpy as np
+import os
+from tensorflow.keras import layers, models
+import time
+import cv2
+import model
+
+batch_size=256
+epochs=500
+noise_dim=100
+tab_size=5
+num_examples_to_generate=tab_size*tab_size
+dir_images='images_gan'
+checkpoint_dir='./training_checkpoints_gan'
+checkpoint_prefix=os.path.join(checkpoint_dir, "ckpt")
+
+if not os.path.isdir(dir_images):
+    os.mkdir(dir_images)
+
+(train_images, train_labels), (test_images, test_labels)=tf.keras.datasets.mnist.load_data()
+
+train_images=train_images.reshape(-1, 28, 28, 1).astype('float32')
+train_images=(train_images-127.5)/127.5
+
+train_dataset=tf.data.Dataset.from_tensor_slices(train_images).shuffle(len(train_images)).batch(batch_size)
+
+def discriminator_loss(real_output, fake_output):
+    real_loss=cross_entropy(tf.ones_like(real_output), real_output)
+    fake_loss=cross_entropy(tf.zeros_like(fake_output), fake_output)
+    total_loss=real_loss+fake_loss
+    return total_loss
+
+def generator_loss(fake_output):
+    return cross_entropy(tf.ones_like(fake_output), fake_output)
+
+generator=model.generator_model()
+discriminator=model.discriminator_model()
+
+cross_entropy=tf.keras.losses.BinaryCrossentropy(from_logits=True)
+
+train_generator_loss=tf.keras.metrics.Mean()
+train_discriminator_loss=tf.keras.metrics.Mean()
+
+generator_optimizer=tf.keras.optimizers.Adam(1e-4)
+discriminator_optimizer=tf.keras.optimizers.Adam(1e-4)
+
+checkpoint=tf.train.Checkpoint(generator_optimizer=generator_optimizer,
+                                 discriminator_optimizer=discriminator_optimizer,
+                                 generator=generator,
+                                 discriminator=discriminator)
+
+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)
+
+      gen_loss=generator_loss(fake_output)
+      disc_loss=discriminator_loss(real_output, fake_output)
+
+      train_generator_loss(gen_loss)
+      train_discriminator_loss(disc_loss)
+      
+    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))
+
+def train(dataset, epochs):
+  for epoch in range(epochs):
+    start=time.time()
+    for image_batch in dataset:
+      train_step(image_batch)
+    generate_and_save_images(generator, epoch+1, seed)
+    if (epoch+1)%15==0:
+      checkpoint.save(file_prefix=checkpoint_prefix)
+    print ('Epoch {}: loss generator: {:.4f} loss discriminator: {:.4f}  {:.4f} sec'.format(epoch+1,
+                                                                                            train_generator_loss.result(),
+                                                                                            train_discriminator_loss.result(),
+                                                                                            time.time()-start))
+    train_generator_loss.reset_states()
+    train_discriminator_loss.reset_states()
+    
+def generate_and_save_images(model, epoch, test_input):
+  labels=tf.one_hot(tf.range(0, num_examples_to_generate, 1)%10, 10)
+  predictions=model([test_input, labels], training=False)
+  img=np.empty(shape=(tab_size*28, tab_size*28), dtype=np.float32)
+  for i in range(tab_size):
+    for j in range(tab_size):
+      img[j*28:(j+1)*28, i*28:(i+1)*28]=predictions[j*tab_size+i, :, :, 0]*127.5+127.5
+  cv2.imwrite('{}/image_{:04d}.png'.format(dir_images, epoch), img)
+  
+train(train_dataset, epochs)