Upload 5 files

README.md

@@ -1,14 +1,26 @@
----
-title: Conditional GAN MNIST
-emoji: 💻
-colorFrom: purple
-colorTo: green
-sdk: gradio
-sdk_version: 5.7.1
-app_file: app.py
-pinned: false
-license: apache-2.0
-short_description: This is a simple implementation of Conditional Generative Ad
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Generating MNIST digits using Conditional GAN
+
+This is a simple implementation of Conditional Generative Adversarial Networks (GAN) for generating MNIST digits.
+
+![cover](demos/gen_all_digits.png)
+
+I use simple BCE loss function for calculating the loss and Adam optimizer (lr=0.0001) for training.
+
+## Architecture
+
+- The **generator** is series of Linear layers with BatchNorm and ReLU activations.
+- The **discriminator** is a series of Linear layers with BatchNorm andLeakyReLU activations.
+- The Conditioning class is appended to the noise vector as a one-hot vector.
+
+## Huggingface Space
+
+You can try generating digits using this model on Huggingface Space.
+https://huggingface.co/spaces/gholamreza/Conditional-GAN-MNIST
+
+![Huggingface Space](demos/gradio_app.png)
+
+## Training History
+
+![losses_plot](demos/losses.png)
+
+Visit https://github.com/gholamrezadar/GAN-MNIST for a simpler version of this code and more details.

app.py

@@ -0,0 +1,33 @@
+import torch
+import gradio as gr
+from models import Generator
+from conditional_gan import generate_digit
+
+generator = Generator()
+
+def init():
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    
+    # Load the generator
+    generator.load_state_dict(torch.load('models/generator.pt', map_location=device))
+    generator.to(device)
+
+def generate_mnist_digit(digit):
+    return generate_digit(generator, digit)
+
+# Gradio Interface
+def gradio_generate(digit):
+    return generate_mnist_digit(digit)
+
+with gr.Blocks() as demo:
+    gr.Markdown("# MNIST Digit Generator")
+    digit = gr.Dropdown(list(range(10)), label="Select a Digit")
+    generate_button = gr.Button("Generate")
+    output_image = gr.Image(label="Generated Image", type="filepath")
+    
+    generate_button.click(gradio_generate, inputs=digit, outputs=output_image)
+
+if __name__ == '__main__':
+    init()
+    print("* Model loaded")
+    demo.launch()

conditional_gan.py

@@ -0,0 +1,47 @@
+# This files serves the neccessary functions for generating images using pretrained models
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision.utils import make_grid
+import matplotlib.pyplot as plt
+
+from models import get_noise
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+def display_image_grid(images, num_rows=5, title=""):
+    if(images.shape[-1]!=28):
+        images = images.view(-1, 1, 28, 28)
+    plt.figure(figsize=(5, 5))
+    plt.axis("off")
+    plt.title(title)
+    grid = make_grid(images.detach().cpu()[:25], nrow=num_rows).permute(1, 2, 0).numpy()
+    plt.imshow(grid)
+    plt.show()
+
+def check_generation(generator):
+    generator.eval()
+    labels = torch.tensor([0,1,2,3,4,5,6,7,8,9] * 10).to(device)
+    fake_eval_batch = generator(get_noise(100, 10, device=device), labels).view(-1, 1, 28, 28)
+    grid = make_grid(fake_eval_batch.detach().cpu(), nrow=10).permute(1, 2, 0).numpy()
+    plt.figure(figsize=(9, 9))
+    plt.title("Generated Images")
+    plt.axis('off')
+    plt.xlabel("Class")
+    plt.imshow(grid)
+    plt.show()
+
+def generate_digit(generator, digit):
+    generator.eval()
+    labels = torch.tensor([digit] * 25).to(device)
+    fake_eval_batch = generator(get_noise(25, 10, device=device), labels).view(-1, 1, 28, 28)
+    grid = make_grid(fake_eval_batch.detach().cpu(), nrow=5).permute(1, 2, 0).numpy()
+    plt.figure(figsize=(5, 5))
+    # no border
+    plt.axis('off')
+    plt.grid(False)
+    plt.xticks([])
+    plt.yticks([])
+    plt.imshow(grid)
+    plt.savefig('generated_digit.png', bbox_inches='tight', pad_inches=0)  # Save the generated image
+    return 'generated_digit.png'  # Return the image path

models.py

@@ -0,0 +1,67 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision.utils import make_grid
+import matplotlib.pyplot as plt
+
+def get_noise(n_samples, z_dim, device='cpu'):
+    return torch.randn((n_samples, z_dim), device=device)
+
+def get_random_labels(n_samples, device='cpu'):
+    return torch.randint(0, 10, (n_samples,), device=device).type(torch.long)
+
+def get_generator_block(input_dim, output_dim):
+    return nn.Sequential(
+        nn.Linear(input_dim, output_dim),
+        nn.BatchNorm1d(output_dim),
+        nn.ReLU(inplace=True)
+    )
+
+class Generator(nn.Module):
+    def __init__(self, z_dim=10, im_dim=784, hidden_dim=128):
+        super(Generator, self).__init__()
+
+        # input is of shape (batch_size, z_dim + 10)
+        self.gen = nn.Sequential(
+            get_generator_block(z_dim + 10, hidden_dim), # 128
+            get_generator_block(hidden_dim, hidden_dim*2), # 256
+            get_generator_block(hidden_dim*2, hidden_dim*4), # 512
+            get_generator_block(hidden_dim*4, hidden_dim*8), # 1024
+            nn.Linear(hidden_dim*8, im_dim), # 784
+            nn.Sigmoid(), # output between 0 and 1
+        )
+
+    def forward(self, noise, classes):
+        '''
+        noise (batch_size, z_dim) noise vector for each image in a batch 
+        classes:long (batch_size) condition class for each image in a batch
+        '''
+        # classes = classes.type(torch.long)
+        # one-hot encode condition_class e.g. 3 -> [0,0,0,1,0,0,0,0,0,0]
+        one_hot_vec = F.one_hot(classes, num_classes=10).type(torch.float32) # (batch_size, 10)
+        conditioned_noise = torch.concat((noise, one_hot_vec), dim=1) # (batch_size, z_dim + 10)
+        return self.gen(conditioned_noise)
+
+
+def get_discriminator_block(input_dim, output_dim):
+    return nn.Sequential(
+        nn.Linear(input_dim, output_dim),
+        nn.LeakyReLU(0.2, inplace=True)
+    )
+
+class Discriminator(nn.Module):
+    def __init__(self, im_dim=784, hidden_dim=128):
+        super(Discriminator, self).__init__()
+        self.disc = nn.Sequential(
+            get_discriminator_block(im_dim + 10, hidden_dim*4), # 512
+            get_discriminator_block(hidden_dim * 4, hidden_dim * 2), # 256
+            get_discriminator_block(hidden_dim * 2, hidden_dim), # 128
+            nn.Linear(hidden_dim, 1),
+            # nn.Sigmoid(),
+            # using a sigmoid followed by BCE is less numerically stable than BCEWithLogitsLoss alone
+            # https://pytorch.org/docs/stable/generated/torch.nn.BCEWithLogitsLoss.html#torch.nn.BCEWithLogitsLoss:~:text=This%20loss%20combines%20a%20Sigmoid%20layer%20and%20the%20BCELoss%20in%20one%20single%20class.%20This%20version%20is%20more%20numerically%20stable%20than%20using%20a%20plain%20Sigmoid%20followed%20by%20a%20BCELoss%20as%2C%20by%20combining%20the%20operations%20into%20one%20layer%2C%20we%20take%20advantage%20of%20the%20log%2Dsum%2Dexp%20trick%20for%20numerical%20stability.
+        )
+
+    def forward(self, image_batch):
+        '''image_batch (batch_size, 784+10)'''
+        return self.disc(image_batch)