Post

Pytorch Representation Learning: Part Three

Posted Updated
By
12 min read
Pytorch Representation Learning: Part Three

Project && Guide

Table of Contents

Section 6. Generative Adversarial Networks (DCGAN) - PyTorch Tutorial

This post explains the code and logic for the DCGAN notebook, section by section, with annotated code blocks.

6.1. Libraries

These libraries are required for building, training, and visualizing the DCGAN model.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20

# Import time for measuring training duration
import time
# Import numpy for numerical operations
import numpy as np
from __future__ import print_function

# Import PyTorch core modules and neural network utilities
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.init as init
import torch.nn.functional as F
from torch.autograd import Variable # For automatic differentiation

# Import torchvision for datasets and image transformations
import torchvision
import torchvision.transforms

# Import matplotlib for plotting images and results
import matplotlib.pyplot as plt

6.2. Define image transformations & Initialize datasets

We define how to preprocess the MNIST images and load them into memory for training.

1
2
3
4
5
6

# Define transformation to convert images to tensors
mnist_transforms = torchvision.transforms.Compose([torchvision.transforms.ToTensor()])
# Load MNIST training dataset with the defined transform
mnist_train = torchvision.datasets.MNIST(root='./data', train=True, transform=mnist_transforms, download=True)
# Create DataLoader for efficient batch processing and shuffling
trainloader = torch.utils.data.DataLoader(mnist_train, batch_size=64, shuffle=True, num_workers=2)

6.3. Create DCGAN Generator

The Generator network takes random noise and produces images. It uses deconvolutional layers and batch normalization.

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

# Define the DCGAN Generator network
class Generator(nn.Module):
    """DCGAN Generator."""
    def __init__(self, z_dim=128, num_filters=32):
        super(Generator, self).__init__()
        self.z_dim = z_dim
        self.num_filters = num_filters
        # 1 x 1 -> 4 x 4
        self.deconv1 = nn.ConvTranspose2d(
            in_channels=z_dim, out_channels=num_filters * 4,
            kernel_size=(4, 4), bias=False
        )
        self.bn1 = nn.BatchNorm2d(num_filters * 4)
        # 4 x 4 -> 8 x 8
        self.deconv2 = nn.ConvTranspose2d(
            in_channels=num_filters * 4, out_channels=num_filters * 2,
            kernel_size=(4, 4), stride=2, padding=1, bias=False,
        )
        self.bn2 = nn.BatchNorm2d(num_filters * 2)
        # 8 x 8 -> 16 x 16
        self.deconv3 = nn.ConvTranspose2d(
            in_channels=num_filters * 2, out_channels=num_filters,
            kernel_size=(4, 4), stride=2, padding=1, bias=False,
        )
        self.bn3 = nn.BatchNorm2d(num_filters)
        # 16 x 16 -> 28 x 28
        self.deconv4 = nn.ConvTranspose2d(
            in_channels=num_filters, out_channels=1,
            kernel_size=(4, 4), stride=2, padding=3, bias=False,
        )
    def forward(self, x):
        # Pass through deconvolutional and batchnorm layers with ReLU activations
        x = F.relu(self.bn1(self.deconv1(x)))
        x = F.relu(self.bn2(self.deconv2(x)))
        x = F.relu(self.bn3(self.deconv3(x)))
        # Final layer uses tanh activation to output image
        return F.tanh(self.deconv4(x))

6.4. Create DCGAN Discriminator

The Discriminator network distinguishes real images from fake ones. It uses convolutional layers and batch normalization.

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

# Define the DCGAN Discriminator network
class Discriminator(nn.Module):
    """DCGAN Discriminator."""
    def __init__(self, num_filters=32):
        super(Discriminator, self).__init__()
        self.num_filters = num_filters
        # 28 x 28 -> 14 x 14
        self.conv1 = nn.Conv2d(
            in_channels=1, out_channels=num_filters,
            kernel_size=(4, 4), stride=2, padding=1, bias=False
        )
        self.bn1 = nn.BatchNorm2d(num_filters)
        # 14 x 14 -> 7 x 7
        self.conv2 = nn.Conv2d(
            in_channels=num_filters, out_channels=num_filters * 2,
            kernel_size=(4, 4), stride=2, padding=1, bias=False,
        )
        self.bn2 = nn.BatchNorm2d(num_filters * 2)
        # 7 x 7 -> 3 x 3
        self.conv3 = nn.Conv2d(
            in_channels=num_filters * 2, out_channels=num_filters * 4,
            kernel_size=(4, 4), stride=2, padding=1, bias=False,
        )
        self.bn3 = nn.BatchNorm2d(num_filters * 4)
        # 3 x 3 -> 1 x 1
        self.conv4 = nn.Conv2d(
            in_channels=num_filters * 4, out_channels=1,
            kernel_size=(4, 4), stride=2, padding=1, bias=False,
        )
    def forward(self, x):
        # Pass through convolutional and batchnorm layers with LeakyReLU activations
        x = F.leaky_relu(self.bn1(self.conv1(x)), 0.2)
        x = F.leaky_relu(self.bn2(self.conv2(x)), 0.2)
        x = F.leaky_relu(self.bn3(self.conv3(x)), 0.2)
        # Final layer uses sigmoid activation to output probability
        return F.sigmoid(self.conv4(x)).squeeze()

6.5. Initialize Generator & Discriminator

Instantiate the models, move them to GPU if available, and set up loss and optimizers.

1
2
3
4
5
6
7
8
9
10
11
12

# Check if CUDA (GPU) is available for faster training
cuda_available = torch.cuda.is_available()
# Instantiate generator and discriminator, move to GPU if available
generator = Generator()
discriminator = Discriminator()
if cuda_available:
    generator = generator.cuda()
    discriminator = discriminator.cuda()
# Set up loss function and optimizers for both networks
loss = nn.BCELoss()
optimizer_g = torch.optim.Adam(generator.parameters(), lr=2e-4)
optimizer_d = torch.optim.Adam(discriminator.parameters(), lr=2e-4)

6.6. Training Loop

The main loop alternately updates the discriminator and generator, tracking losses for each.

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

# Training loop for DCGAN: alternately update discriminator and generator
ctr = 0
z_dim = 128
# Track losses for monitoring
minibatch_disc_losses = []
minibatch_gen_losses = []
for epoch in range(50):
    losses = []
    # Train
    for batch_idx, (inputs, targets) in enumerate(trainloader):
        ctr += 1
        if cuda_available:
            inputs, targets = inputs.cuda(), targets.cuda()
        # Wrap inputs and targets in Variables
        inputs, targets = Variable(inputs), Variable(targets)
        # Create label tensors for real and fake samples
        zeros = Variable(torch.zeros(inputs.size(0)))
        ones = Variable(torch.ones(inputs.size(0)))
        if cuda_available:
            zeros, ones = zeros.cuda(), ones.cuda()
        ############################
        # (1) Update Discriminator 
        ############################
        # Sample z ~ N(0, 1) for generator input
        minibatch_noise = Variable(torch.from_numpy(
            np.random.randn(inputs.size(0), z_dim, 1, 1).astype(np.float32)
        ))
        if cuda_available:
            minibatch_noise = minibatch_noise.cuda()
        # Zero gradients for the discriminator
        optimizer_d.zero_grad()
        # Train with real examples
        d_real = discriminator(inputs)
        d_real_loss = loss(d_real, ones)  # Train discriminator to recognize real examples
        d_real_loss.backward()
        # Train with fake examples from the generator
        fake = generator(minibatch_noise).detach()  # Detach to prevent backpropping through the generator
        d_fake = discriminator(fake)
        d_fake_loss = loss(d_fake, zeros)  # Train discriminator to recognize generator samples
        d_fake_loss.backward()
        minibatch_disc_losses.append(d_real_loss.data[0] + d_fake_loss.data[0])
        # Update the discriminator
        optimizer_d.step()
        ############################
        # (2) Update Generator
        ############################
        # Zero gradients for the generator
        optimizer_g.zero_grad()
        # Sample z ~ N(0, 1) for generator input
        minibatch_noise = Variable(torch.from_numpy(
            np.random.randn(inputs.size(0), z_dim, 1, 1).astype(np.float32)
        ))
        if cuda_available:
            minibatch_noise = minibatch_noise.cuda()
        d_fake = discriminator(generator(minibatch_noise))
        g_loss = loss(d_fake, ones)  # Train generator to fool the discriminator into thinking these are real.
        g_loss.backward()
        # Update the generator
        optimizer_g.step()
        minibatch_gen_losses.append(g_loss.data[0])
    print('Generator loss : %.3f' % (np.mean(minibatch_gen_losses)))
    print('Discriminator loss : %.3f' % (np.mean(minibatch_disc_losses)))

6.7. Sample from the Generator

After training, sample images from the generator and visualize them in a grid.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16

# Sample images from the generator and visualize them
generator.eval() # Set generator in evaluation mode to use running means and averages for Batchnorm
# Sample z ~ N(0, 1) for generator input
minibatch_noise = Variable(torch.from_numpy(
    np.random.randn(16, z_dim, 1, 1).astype(np.float32)
))
if cuda_available:
    minibatch_noise = minibatch_noise.cuda()
fakes = generator(minibatch_noise)
# Plot generated images in a 4x4 grid
fig = plt.figure(figsize=(10, 10))
idx = 1
for ind, fake in enumerate(fakes):
    fig.add_subplot(4, 4, ind + 1)
    plt.imshow(fake.data.cpu().numpy().reshape(28, 28), cmap='gray')
    plt.axis('off')

Section 7. Creating Your Own Modules in PyTorch

This post explains how to create custom neural network modules in PyTorch, step by step, with annotated code blocks and explanations. The goal is to understand how to build custom neural network modules in PyTorch by subclassing nn.Module, and to learn how to implement, test, and update parameters for these modules.

Steps

  1. Import required libraries
  2. Implement a custom linear module mimicking nn.Linear
  3. Implement a custom linear layer from scratch
  4. Compare outputs of standard and custom linear layers
  5. Build a residual linear layer for ResNet-like skip connections
  6. Test standard and residual linear layers
  7. Build a model using Sequential and custom layers
  8. Compare CrossEntropyLoss and NLLLoss
  9. Update parameters manually and with torch.optim

7.1. Libraries

1
2
3
4
5
6

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.parameter import Parameter
import math
import numpy as np

7.2. Custom Linear module mimicking nn.Linear

This class mimics the behavior of PyTorch’s built-in nn.Linear layer, including initialization and forward pass.

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

class Linear(nn.Module):
    r"""Applies a linear transformation to the incoming data: y = Ax + b"""
    def __init__(self, in_features, out_features, bias=True):
        super(Linear, self).__init__()
        self.in_features = in_features
        self.out_features = out_features
        # Weight shape: (out_features, in_features)
        self.weight = Parameter(torch.Tensor(out_features, in_features))
        if bias:
            self.bias = Parameter(torch.Tensor(out_features))
        else:
            self.register_parameter('bias', None)
        self.reset_parameters()
    def reset_parameters(self):
        stdv = 1. / math.sqrt(self.weight.size(1))
        self.weight.data.uniform_(-stdv, stdv)
        if self.bias is not None:
            self.bias.data.uniform_(-stdv, stdv)
    def forward(self, input):
        # Linear transformation with or without bias
        if self.bias is None:
            return self._backend.Linear()(input, self.weight)
        else:
            return self._backend.Linear()(input, self.weight, self.bias)
    def __repr__(self):
        return self.__class__.__name__ + ' (' \
            + str(self.in_features) + ' -> ' \
            + str(self.out_features) + ')'

7.3. Custom implementation of a linear layer

This class implements a linear layer from scratch using matrix multiplication and optional bias.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17

class MyLinear(nn.Module):
    def __init__(self, in_features, out_features, bias=True):
        super(MyLinear, self).__init__()
        self.in_features = in_features
        self.out_features = out_features
        # Weight shape: (in_features, out_features)
        self.weight = Parameter(torch.Tensor(in_features, out_features))
        if bias:
            self.bias = Parameter(torch.Tensor(out_features))
        else:
            self.register_parameter('bias', None)
    def forward(self, input):
        # Matrix multiplication and optional bias addition
        if self.bias is None:
            return torch.mm(input, self.weight)
        else:
            return torch.mm(input, self.weight) + self.bias

7.4. Compare outputs of nn.Linear and custom MyLinear

This code compares the outputs of PyTorch’s nn.Linear and the custom MyLinear for the same input and parameters.

1
2
3
4
5
6
7
8

x = torch.from_numpy(np.random.randn(2, 3)).float()
linear1 = nn.Linear(3,4)
linear2 = MyLinear(3,4)
# Set the weight and bias of linear2 to match linear1's
linear2.weight.data = linear1.weight.data.transpose(1,0)
linear2.bias.data = linear1.bias.data
# Compare outputs for the same input
print(torch.eq(linear1(x), linear2(x)))

7.5. Residual linear layer example for ResNet-like skip connections

This class implements a residual linear layer, adding skip connections for better gradient flow.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20

class ResLinear(nn.Module):
    def __init__(self, in_features, out_features, activation=nn.ReLU()):
        super(ResLinear, self).__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.activation = activation
        self.linear = nn.Linear(in_features, out_features)
        # If input and output dimensions differ, add a projection layer
        if in_features != out_features:
            self.project_linear = nn.Linear(in_features, out_features)
    def forward(self, x):
        # Apply activation to linear output
        inner = self.activation(self.linear(x))
        # Use projection for skip connection if needed
        if self.in_features != self.out_features:
            skip = self.project_linear(x)
        else:
            skip = x
        # Add skip connection
        return inner + skip

7.6. Test standard and residual linear layers

This code tests both standard and residual linear layers for their output shapes.

1
2
3
4
5

x = torch.from_numpy(np.random.randn(2, 3)).float()
res1 = nn.Linear(3,3) # Standard linear layer
res2 = ResLinear(3,5) # Residual linear layer with projection
print(res1(x).size()) # Output shape for standard layer
print(res2(x).size()) # Output shape for residual layer

7.7. Putting things altogether, Sequential, Parameter updates

This model uses Sequential and custom layers, and includes methods for prediction and loss calculation.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24

class MyModel(nn.Module):
    def __init__(self, Linear=ResLinear):
        super(MyModel, self).__init__()
        # Build a sequence of layers for prediction
        self.predict_ = nn.Sequential(
            Linear(784, 328),
            nn.ReLU(),
            Linear(328, 328),
            nn.ReLU(),
            Linear(328, 10),
        )
        self.criterion = nn.CrossEntropyLoss()
    def predict_proba(self, x):
        # Softmax for probability output
        return F.softmax(x)
    def predict(self, x):
        # Argmax for class prediction
        return torch.max(self.predict_proba(x))[1]
    def loss(self, x, target):
        # Compute loss for training
        proba = self.predict_(x)
        return self.criterion(proba, target)
# CrossEntropyLoss expects pre-softmax output
# NLLLoss expects log-softmax output

7.8. Caveat: CrossEntropyLoss versus NLLLoss

  • CrossEntropyLoss takes in pre-softmax as input
  • NLLLoss takes in log-softmax as input
1
2
3
4
5
6

y = torch.Tensor(1,10).normal_() # Random output
t = torch.from_numpy(np.random.choice(10, size=1)) # Random target
loss1 = nn.CrossEntropyLoss()
loss2 = nn.NLLLoss()
print(loss1(y, t)) # CrossEntropyLoss on raw output
print(loss2(nn.LogSoftmax(dim=1)(y), t)) # NLLLoss on log-softmax output

7.9. Test model loss computation

1
2
3
4

x = torch.from_numpy(np.random.randn(64, 784)).float() # Random input
t = torch.from_numpy(np.random.choice(10, size=64)) # Random target
model = MyModel()
print(model.loss(x, t)) # Print loss value

7.10. Updating Parameters (Manually)

This code demonstrates manual parameter updates using gradient descent.

1
2
3
4
5
6
7
8
9
10
11
12

x = torch.from_numpy(np.random.randn(64, 784)).float() # Random input
t = torch.from_numpy(np.random.choice(10, size=64)) # Random target
model = MyModel()
lr = 0.1 # Learning rate
for i in range(10):
    loss = model.loss(x, t)
    loss.backward()
    for param in model.parameters():
        # Update parameters manually using gradient descent
        param.data.sub_(param.grad.data*lr)
        param.grad.data.zero_()
    print(param.grad) # Print gradients after update

7.11. Updating Parameters (torch.optim)

This code demonstrates parameter updates using PyTorch’s built-in optimizer.

1
2
3
4
5
6
7
8
9
10
11

x = torch.from_numpy(np.random.randn(64, 784)).float() # Random input
t = torch.from_numpy(np.random.choice(10, size=64)) # Random target
model = MyModel()
optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
for i in range(10):
    optimizer.zero_grad()
    loss = model.loss(x, t)
    loss.backward()
    optimizer.step()
    print(loss) # Print loss after optimizer step
# Momentum helps accelerate updates in the relevant direction

Resources

Project repository

GitHub Code: PyTorch Representation Learning Project

Machine Learning and Deep Learning, Deep Learning
PyTorch representation learning DCGAN generative models deep learning custom modules reinforcement learning
This post is licensed under CC BY 4.0 by the author.