Updated Pytorch DCGAN Tutorial with New Video

DCGAN (Deep Convolutional Generative Adversarial Networks) is a type of generative model that can generate high-quality images. In this tutorial, we will learn how to implement DCGAN using PyTorch.

Before we start, please make sure you have PyTorch installed on your system. You can install it using pip:

pip install torch torchvision

Now, let’s start by importing the necessary libraries:

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

Next, we will define the generator and discriminator models. The generator takes a random noise vector as input and generates an image, while the discriminator takes an image as input and predicts whether it is real or fake.

class Generator(nn.Module):
    def __init__(self, nz, ngf, nc):
        super(Generator, self).__init__()
        self.main = nn.Sequential(
            nn.ConvTranspose2d(nz, ngf*4, 4, 1, 0, bias=False),
            nn.BatchNorm2d(ngf*4),
            nn.ReLU(True),
            nn.ConvTranspose2d(ngf*4, ngf*2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf*2),
            nn.ReLU(True),
            nn.ConvTranspose2d(ngf*2, ngf, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf),
            nn.ReLU(True),
            nn.ConvTranspose2d(ngf, nc, 4, 2, 1, bias=False),
            nn.Tanh()
        )

    def forward(self, input):
        return self.main(input)

class Discriminator(nn.Module):
    def __init__(self, nc, ndf):
        super(Discriminator, self).__init__()
        self.main = nn.Sequential(
            nn.Conv2d(nc, ndf, 4, 2, 1, bias=False),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(ndf, ndf*2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ndf*2),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(ndf*2, ndf*4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ndf*4),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(ndf*4, 1, 4, 1, 0, bias=False),
            nn.Sigmoid()
        )

    def forward(self, input):
        return self.main(input)

Now, we will define the hyperparameters and create the models:

# Hyperparameters
nz = 100
ngf = 64
ndf = 64
nc = 1

# Create the models
generator = Generator(nz, ngf, nc)
discriminator = Discriminator(nc, ndf)

Next, we will define the loss function and optimization algorithm:

criterion = nn.BCELoss()
optimizer_G = optim.Adam(generator.parameters(), lr=0.0002, betas=(0.5, 0.999))
optimizer_D = optim.Adam(discriminator.parameters(), lr=0.0002, betas=(0.5, 0.999))

Now, we will load the dataset and prepare the data loaders:

transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

dataset = datasets.MNIST(root='./data', train=True, transform=transform, download=True)
dataloader = DataLoader(dataset, batch_size=128, shuffle=True)

Finally, we will train the models:

# Training loop
num_epochs = 5
for epoch in range(num_epochs):
    for i, data in enumerate(dataloader):
        # Update discriminator
        discriminator.zero_grad()
        real_images = data[0].to(device)
        real_labels = torch.ones(real_images.size(0), 1).to(device)
        output = discriminator(real_images)
        loss_real = criterion(output, real_labels)

        fake_images = generator(torch.randn(batch_size, nz, 1, 1).to(device))
        fake_labels = torch.zeros(fake_images.size(0), 1).to(device)
        output = discriminator(fake_images.detach())
        loss_fake = criterion(output, fake_labels)

        loss_D = loss_real + loss_fake
        loss_D.backward()
        optimizer_D.step()

        # Update generator
        generator.zero_grad()
        output = discriminator(fake_images)
        loss_G = criterion(output, real_labels)
        loss_G.backward()
        optimizer_G.step()

        if i % 100 == 0:
            print(f'Epoch [{epoch}/{num_epochs}], Step [{i}/{len(dataloader)}], Loss D: {loss_D.item()}, Loss G: {loss_G.item()}')

# Save the models
torch.save(generator.state_dict(), 'generator.pt')
torch.save(discriminator.state_dict(), 'discriminator.pt')

That’s it! You have successfully implemented a DCGAN using PyTorch. Feel free to experiment with different hyperparameters and datasets to generate high-quality images.