Rename pages/23_Gan.py to pages/23_GANs.py

pages/23_GANs.py

@@ -0,0 +1,114 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import torchvision.transforms as transforms
+import torchvision.datasets as datasets
+from torch.utils.data import DataLoader
+import numpy as np
+import matplotlib.pyplot as plt
+import streamlit as st
+
+# Define the Generator
+class Generator(nn.Module):
+    def __init__(self, input_dim, output_dim):
+        super(Generator, self).__init__()
+        self.model = nn.Sequential(
+            nn.Linear(input_dim, 128),
+            nn.ReLU(),
+            nn.Linear(128, 256),
+            nn.ReLU(),
+            nn.Linear(256, output_dim),
+            nn.Tanh()
+        )
+
+    def forward(self, x):
+        return self.model(x)
+
+# Define the Discriminator
+class Discriminator(nn.Module):
+    def __init__(self, input_dim):
+        super(Discriminator, self).__init__()
+        self.model = nn.Sequential(
+            nn.Linear(input_dim, 256),
+            nn.LeakyReLU(0.2),
+            nn.Linear(256, 128),
+            nn.LeakyReLU(0.2),
+            nn.Linear(128, 1),
+            nn.Sigmoid()
+        )
+
+    def forward(self, x):
+        return self.model(x)
+
+# Hyperparameters
+latent_dim = 100
+image_dim = 28 * 28  # MNIST images are 28x28 pixels
+lr = 0.0002
+batch_size = 64
+epochs = 50
+
+# Prepare the data
+transform = transforms.Compose([
+    transforms.ToTensor(),
+    transforms.Normalize([0.5], [0.5])
+])
+
+dataset = datasets.MNIST(root='mnist_data', train=True, transform=transform, download=True)
+dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
+
+# Initialize the models
+generator = Generator(latent_dim, image_dim)
+discriminator = Discriminator(image_dim)
+
+# Optimizers
+optimizer_G = optim.Adam(generator.parameters(), lr=lr)
+optimizer_D = optim.Adam(discriminator.parameters(), lr=lr)
+
+# Loss function
+criterion = nn.BCELoss()
+
+# Streamlit interface
+st.title("GAN with PyTorch and Hugging Face")
+st.write("Training a GAN to generate MNIST digits")
+
+train_gan = st.button("Train GAN")
+
+if train_gan:
+    # Training loop
+    for epoch in range(epochs):
+        for i, (imgs, _) in enumerate(dataloader):
+            # Prepare real and fake data
+            real_imgs = imgs.view(imgs.size(0), -1)
+            real_labels = torch.ones(imgs.size(0), 1)
+            fake_labels = torch.zeros(imgs.size(0), 1)
+            z = torch.randn(imgs.size(0), latent_dim)
+            fake_imgs = generator(z)
+            
+            # Train Discriminator
+            optimizer_D.zero_grad()
+            real_loss = criterion(discriminator(real_imgs), real_labels)
+            fake_loss = criterion(discriminator(fake_imgs.detach()), fake_labels)
+            d_loss = real_loss + fake_loss
+            d_loss.backward()
+            optimizer_D.step()
+            
+            # Train Generator
+            optimizer_G.zero_grad()
+            g_loss = criterion(discriminator(fake_imgs), real_labels)
+            g_loss.backward()
+            optimizer_G.step()
+            
+        st.write(f"Epoch [{epoch+1}/{epochs}] | D Loss: {d_loss.item():.4f} | G Loss: {g_loss.item():.4f}")
+
+    st.write("Training completed")
+
+    # Generate and display images
+    z = torch.randn(16, latent_dim)
+    generated_imgs = generator(z).view(-1, 1, 28, 28).data
+    grid = np.transpose(np.array([generated_imgs[i].numpy() for i in range(16)]), (1, 2, 0))
+    
+    fig, ax = plt.subplots(figsize=(8, 8))
+    ax.imshow(np.squeeze(grid), cmap="gray")
+    st.pyplot(fig)
+else:
+    st.write("Click the button to start training the GAN")

pages/23_Gan.py

@@ -1,84 +0,0 @@
-import torch
-import torch.nn as nn
-import torch.optim as optim
-import torchvision.transforms as transforms
-import torchvision.utils as vutils
-import streamlit as st
-
-# Define the Generator
-class Generator(nn.Module):
-    def __init__(self):
-        super(Generator, self).__init__()
-        self.main = nn.Sequential(
-            nn.ConvTranspose2d(100, 256, 4, 1, 0, bias=False),
-            nn.BatchNorm2d(256),
-            nn.ReLU(True),
-            nn.ConvTranspose2d(256, 128, 4, 2, 1, bias=False),
-            nn.BatchNorm2d(128),
-            nn.ReLU(True),
-            nn.ConvTranspose2d(128, 64, 4, 2, 1, bias=False),
-            nn.BatchNorm2d(64),
-            nn.ReLU(True),
-            nn.ConvTranspose2d(64, 1, 4, 2, 1, bias=False),
-            nn.Tanh()
-        )
-
-    def forward(self, input):
-        return self.main(input)
-
-# Define the Discriminator
-class Discriminator(nn.Module):
-    def __init__(self):
-        super(Discriminator, self).__init__()
-        self.main = nn.Sequential(
-            nn.Conv2d(1, 64, 4, 2, 1, bias=False),
-            nn.LeakyReLU(0.2, inplace=True),
-            nn.Conv2d(64, 128, 4, 2, 1, bias=False),
-            nn.BatchNorm2d(128),
-            nn.LeakyReLU(0.2, inplace=True),
-            nn.Conv2d(128, 256, 4, 2, 1, bias=False),
-            nn.BatchNorm2d(256),
-            nn.LeakyReLU(0.2, inplace=True),
-            nn.Conv2d(256, 1, 4, 1, 0, bias=False),
-            nn.Sigmoid()
-        )
-
-    def forward(self, input):
-        return self.main(input)
-
-# Initialize the models
-netG = Generator()
-netD = Discriminator()
-
-# Loss function
-criterion = nn.BCELoss()
-
-# Optimizers
-optimizerD = optim.Adam(netD.parameters(), lr=0.0002, betas=(0.5, 0.999))
-optimizerG = optim.Adam(netG.parameters(), lr=0.0002, betas=(0.5, 0.999))
-
-# Device
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-netG.to(device)
-netD.to(device)
-criterion.to(device)
-
-# Function to generate and save images
-def generate_images(num_images, noise_dim):
-    netG.eval()
-    noise = torch.randn(num_images, noise_dim, 1, 1, device=device)
-    fake_images = netG(noise)
-    return fake_images
-
-# Streamlit interface
-st.title("Simple GAN with Streamlit")
-st.write("Generate images using a simple GAN")
-
-num_images = st.slider("Number of images to generate", min_value=1, max_value=64, value=8)
-noise_dim = 100
-
-if st.button("Generate Images"):
-    with st.spinner("Generating images..."):
-        fake_images = generate_images(num_images, noise_dim)
-        grid = vutils.make_grid(fake_images.cpu(), padding=2, normalize=True)
-        st.image(grid.permute(1, 2, 0).numpy(), caption="Generated Images")