Spaces:
Running
Running
File size: 3,491 Bytes
1b6b756 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 |
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")
|