Setting Up Unconditional Image Generation on Ubuntu using PyTorch

This guide provides detailed instructions on how to set up an unconditional image generation model on Ubuntu using PyTorch. We’ll use a Generative Adversarial Network (GAN) as an example for unconditional image generation.

1. Install System Prerequisites

Update your Ubuntu system and install the necessary dependencies. Open a terminal and run the following commands:

sudo apt update
sudo apt upgrade
sudo apt install python3 python3-pip git
    

2. Install PyTorch and Additional Libraries

Install PyTorch and torchvision, which are required for working with deep learning models in PyTorch:

pip install torch torchvision
    

3. Clone or Create a GAN Model

You can create a simple GAN model or use a pre-existing GAN model. For this example, we’ll use a simple GAN structure.

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision.utils import save_image

# Define the generator network
class Generator(nn.Module):
    def __init__(self, nz):
        super(Generator, self).__init__()
        self.main = nn.Sequential(
            nn.Linear(nz, 256),
            nn.ReLU(True),
            nn.Linear(256, 512),
            nn.ReLU(True),
            nn.Linear(512, 1024),
            nn.ReLU(True),
            nn.Linear(1024, 784),
            nn.Tanh()
        )
    
    def forward(self, input):
        return self.main(input).view(-1, 1, 28, 28)  # Reshape for image output

# Define the discriminator network
class Discriminator(nn.Module):
    def __init__(self):
        super(Discriminator, self).__init__()
        self.main = nn.Sequential(
            nn.Linear(784, 512),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Linear(512, 256),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Linear(256, 1),
            nn.Sigmoid()
        )
    
    def forward(self, input):
        return self.main(input.view(-1, 784))
    

This code defines a simple GAN structure with a generator and discriminator model.

4. Train the GAN Model

Once the model is defined, we can train the GAN on random noise to generate images. The following script trains the GAN:

# Parameters
batch_size = 64
learning_rate = 0.0002
nz = 100  # Size of latent vector (input noise)

# Initialize models and optimizer
netG = Generator(nz)
netD = Discriminator()
criterion = nn.BCELoss()
optimizerD = optim.Adam(netD.parameters(), lr=learning_rate)
optimizerG = optim.Adam(netG.parameters(), lr=learning_rate)

# Training loop
num_epochs = 50
for epoch in range(num_epochs):
    for i in range(600):  # Assume 600 batches per epoch
        # Train Discriminator
        netD.zero_grad()
        real_data = torch.randn(batch_size, 1, 28, 28)  # Replace with real dataset if available
        label_real = torch.ones(batch_size, 1)
        output_real = netD(real_data).view(-1)
        loss_real = criterion(output_real, label_real)
        
        noise = torch.randn(batch_size, nz)
        fake_data = netG(noise)
        label_fake = torch.zeros(batch_size, 1)
        output_fake = netD(fake_data.detach()).view(-1)
        loss_fake = criterion(output_fake, label_fake)

        lossD = loss_real + loss_fake
        lossD.backward()
        optimizerD.step()
        
        # Train Generator
        netG.zero_grad()
        label_g = torch.ones(batch_size, 1)  # Generator wants to make discriminator classify fake as real
        output_g = netD(fake_data).view(-1)
        lossG = criterion(output_g, label_g)
        lossG.backward()
        optimizerG.step()
    
    print(f"Epoch [{epoch+1}/{num_epochs}], Loss D: {lossD.item()}, Loss G: {lossG.item()}")
    
    # Save generated images every few epochs
    if (epoch + 1) % 10 == 0:
        with torch.no_grad():
            generated_images = netG(torch.randn(64, nz)).cpu()
            save_image(generated_images, f"generated_images_epoch_{epoch+1}.png", normalize=True)
    

This training loop trains the GAN for 50 epochs and saves generated images every 10 epochs.

5. Run the Script

Save the code into a file called gan_training.py and run it in the terminal:

python3 gan_training.py
    

This will start training the GAN, and generated images will be saved in the current directory every 10 epochs.

6. Inspect Generated Images

After training, check the saved images (e.g., generated_images_epoch_10.png) in your project directory. These images are generated from random noise, demonstrating unconditional image generation.

7. Troubleshooting

If you encounter any issues, consider the following:

• Ensure all libraries (torch, torchvision) are installed correctly.
• If training is slow, consider using a GPU if available.
• Adjust learning rates and the latent vector size for better results.

8. Conclusion

You have successfully set up an unconditional image generation system on Ubuntu using PyTorch. You can experiment with different GAN architectures or train on larger datasets for more complex image generation.