Generating Cartoon Avatars with Deep Convolutional Generative Adversarial Networks

Understanding DCGAN Architecture

Deep Convolutional Generative Adversarial Networks (DCGAN) represent an evolution of the original GAN architecture. The primary distinction lies in the incorporation of convolutional layers in both the discriminator and generator networks. This architectural enhancement enables more effective learning of hierarchical representations from image data.

The DCGAN framework was first introduced by Radford et al. in their paper "Unsupervised Representation Learning With Deep Convolutional Generative Adversarial Networks." In this implementation, the discriminator employs a series of convolutional layers, each followed by batch normalization and LeakyReLU activation functions. The network accepts 3x64x64 images as input and outputs a probability score indicating whether the input image is real or generated. Conversely, the generator utilizes transposed convolutional layers, batch normalization, and ReLU activations to transform a latent vector z into a 3x64x64 RGB image.

Dataset Preparation and Processing

First, we'll download the dataset required for training our DCGAN model:

from download import download

dataset_url = "https://download.mindspore.cn/dataset/Faces/faces.zip"
dataset_path = download(dataset_url, "./faces", kind="zip", replace=True)

The downloaded dataset has the following directory structure:

./faces/faces
├── 0.jpg
├── 1.jpg
├── 2.jpg
├── 3.jpg
├── 4.jpg
    ...
├── 70169.jpg
└── 70170.jpg

Data Processing Pipeline

Let's define the necessary hyperparameters for our training process:

batch_size = 128          # Batch size for training
image_dimension = 64       # Spatial size of training images
color_channels = 3         # Number of color channels in images
latent_vector_size = 100   # Dimensionality of the latent vector
generator_feature_maps = 64  # Size of feature maps in generator
discriminator_feature_maps = 64  # Size of feature maps in discriminator
training_epochs = 3       # Number of training epochs
learning_rate = 0.0002    # Learning rate for optimization
beta1_param = 0.5         # Beta1 hyperparameter for Adam optimizer

Next, we'll implement the data loading and preprocessing pipeline:

import numpy as np
import mindspore.dataset as ds
import mindspore.dataset.vision as vision

def prepare_dataset(dataset_path):
    """Load and preprocess the dataset"""
    dataset = ds.ImageFolderDataset(dataset_path,
                                    num_parallel_workers=4,
                                    shuffle=True,
                                    decode=True)

    # Define data augmentation transformations
    transformations = [
        vision.Resize(image_dimension),
        vision.CenterCrop(image_dimension),
        vision.HWC2CHW(),
        lambda x: ((x / 255).astype("float32"))
    ]

    # Apply transformations
    dataset = dataset.project('image')
    dataset = dataset.map(transformations, 'image')

    # Create batches
    dataset = dataset.batch(batch_size)
    return dataset

dataset = prepare_dataset('./faces')

import matplotlib.pyplot as plt

def visualize_sample_data(data):
    """Display a subset of training samples"""
    plt.figure(figsize=(10, 3), dpi=140)
    for i, image in enumerate(data[0][:30], 1):
        plt.subplot(3, 10, i)
        plt.axis("off")
        plt.imshow(image.transpose(1, 2, 0))
    plt.show()

sample_batch = next(dataset.create_tuple_iterator(output_numpy=True))
visualize_sample_data(sample_batch)

Network Architecture Implementation

After preparing the data, we can construct our DCGAN network. According to the DCGAN paper, all model weights should be randomly initialized from a normal distribution with mean=0 and standard deviation=0.02.

Generator Network

The generator network's function is to map a latent vector z to the data space, creating RGB images of the same dimensions as the real images. This is achieved through a series of transposed convolutional layers, each paired with batch normalization and ReLU activation functions. The final output passes through a tanh function to constrain values to the range [-1, 1].

The generator structure is influenced by the hyperparameters defined earlier: latent_vector_size determines the dimensionality of the input vector, generator_feature_maps affects the size of feature maps propagated through the network, and color_channels specifies the number of channels in the output image.

Here's the implementation of our generator network:

import mindspore as ms
from mindspore import nn, ops
from mindspore.common.initializer import Normal

# Initialize weights with normal distribution
weight_initializer = Normal(mean=0, sigma=0.02)
gamma_initializer = Normal(mean=1, sigma=0.02)

class Generator(nn.Cell):
    """DCGAN Generator Network"""

    def __init__(self):
        super(Generator, self).__init__()
        self.network = nn.SequentialCell(
            nn.Conv2dTranspose(latent_vector_size, generator_feature_maps * 8, 4, 1, 'valid', weight_init=weight_initializer),
            nn.BatchNorm2d(generator_feature_maps * 8, gamma_init=gamma_initializer),
            nn.ReLU(),
            nn.Conv2dTranspose(generator_feature_maps * 8, generator_feature_maps * 4, 4, 2, 'pad', 1, weight_init=weight_initializer),
            nn.BatchNorm2d(generator_feature_maps * 4, gamma_init=gamma_initializer),
            nn.ReLU(),
            nn.Conv2dTranspose(generator_feature_maps * 4, generator_feature_maps * 2, 4, 2, 'pad', 1, weight_init=weight_initializer),
            nn.BatchNorm2d(generator_feature_maps * 2, gamma_init=gamma_initializer),
            nn.ReLU(),
            nn.Conv2dTranspose(generator_feature_maps * 2, generator_feature_maps, 4, 2, 'pad', 1, weight_init=weight_initializer),
            nn.BatchNorm2d(generator_feature_maps, gamma_init=gamma_initializer),
            nn.ReLU(),
            nn.Conv2dTranspose(generator_feature_maps, color_channels, 4, 2, 'pad', 1, weight_init=weight_initializer),
            nn.Tanh()
        )

    def construct(self, input_vector):
        return self.network(input_vector)

generator_model = Generator()

Discriminator Network

As previously mentioned, the discriminator is a binary classification network that outputs the probability of an image being real. It processes input through a series of convolutional layers, batch normalization, and LeakyReLU activations, with a final Sigmoid activation function to produce the probability score.

The DCGAN paper suggests using convolutional layers rather than pooling operations for downsampling, as this allows the network to learn its own pooling features.

Here's the implementation of our discriminator network:

class Discriminator(nn.Cell):
    """DCGAN Discriminator Network"""

    def __init__(self):
        super(Discriminator, self).__init__()
        self.network = nn.SequentialCell(
            nn.Conv2d(color_channels, discriminator_feature_maps, 4, 2, 'pad', 1, weight_init=weight_initializer),
            nn.LeakyReLU(0.2),
            nn.Conv2d(discriminator_feature_maps, discriminator_feature_maps * 2, 4, 2, 'pad', 1, weight_init=weight_initializer),
            nn.BatchNorm2d(discriminator_feature_maps * 2, gamma_init=gamma_initializer),
            nn.LeakyReLU(0.2),
            nn.Conv2d(discriminator_feature_maps * 2, discriminator_feature_maps * 4, 4, 2, 'pad', 1, weight_init=weight_initializer),
            nn.BatchNorm2d(discriminator_feature_maps * 4, gamma_init=gamma_initializer),
            nn.LeakyReLU(0.2),
            nn.Conv2d(discriminator_feature_maps * 4, discriminator_feature_maps * 8, 4, 2, 'pad', 1, weight_init=weight_initializer),
            nn.BatchNorm2d(discriminator_feature_maps * 8, gamma_init=gamma_initializer),
            nn.LeakyReLU(0.2),
            nn.Conv2d(discriminator_feature_maps * 8, 1, 4, 1, 'valid', weight_init=weight_initializer),
        )
        self.output_layer = nn.Sigmoid()

    def construct(self, input_image):
        output = self.network(input_image)
        output = output.reshape(output.shape[0], -1)
        return self.output_layer(output)

discriminator_model = Discriminator()

Model Training Process

Loss Function Definition

After defining the generator and discriminator networks, we'll use the binary cross-entropy loss function (BCELoss) from MindSpore for our training process:

# Define the loss function
adversarial_loss = nn.BCELoss(reduction='mean')

Optimizer Configuration

We'll set up two separate optimizers, one for the discriminator and one for the generator. Both will use the Adam optimizer with a learning rate of 0.0002 and beta1 parameter of 0.5:

# Configure optimizers for generator and discriminator
discriminator_optimizer = nn.Adam(discriminator_model.trainable_params(), learning_rate=learning_rate, beta1=beta1_param)
generator_optimizer = nn.Adam(generator_model.trainable_params(), learning_rate=learning_rate, beta1=beta1_param)
generator_optimizer.update_parameters_name('optim_g.')
discriminator_optimizer.update_parameters_name('optim_d.')

Training Implementation

The training process consists of two main parts: training the discriminator and training the generator.

• Discriminator Training: The objective is to maximize the probability of correctly classifying real images as real and generated images as fake. This involves maximizing log(D(x)) + log(1 - D(G(z))).
• Generator Training: As described in the DCGAN paper, we aim to train the generator to produce better fake images by minimizing log(1 - D(G(z))).

Here's the implementation of our training logic:

def generator_step(real_images, real_labels):
    # Sample noise as input to the generator
    noise_vector = ops.standard_normal((real_images.shape[0], latent_vector_size, 1, 1))

    # Generate a batch of images
    generated_images = generator_model(noise_vector)

    # Calculate generator loss based on its ability to fool the discriminator
    generator_loss = adversarial_loss(discriminator_model(generated_images), real_labels)

    return generator_loss, generated_images

def discriminator_step(real_images, generated_images, real_labels, fake_labels):
    # Calculate discriminator's ability to classify real samples correctly
    real_loss = adversarial_loss(discriminator_model(real_images), real_labels)
    fake_loss = adversarial_loss(discriminator_model(generated_images), fake_labels)
    discriminator_loss = (real_loss + fake_loss) / 2
    return discriminator_loss

# Create gradient functions
generator_gradient_fn = ms.value_and_grad(generator_step, None,
                                        generator_optimizer.parameters,
                                        has_aux=True)
discriminator_gradient_fn = ms.value_and_grad(discriminator_step, None,
                                            discriminator_optimizer.parameters)

@ms.jit
def training_step(images):
    real_labels = ops.ones((images.shape[0], 1), mindspore.float32)
    fake_labels = ops.zeros((images.shape[0], 1), mindspore.float32)

    (generator_loss, generated_images), generator_gradients = generator_gradient_fn(images, real_labels)
    generator_optimizer(generator_gradients)
    discriminator_loss, discriminator_gradients = discriminator_gradient_fn(images, generated_images, real_labels, fake_labels)
    discriminator_optimizer(discriminator_gradients)

    return generator_loss, discriminator_loss, generated_images

We'll now train our network by iterating through the dataset for the specified number of epochs. We'll collect the losses from both the generator and discriminator every 50 iterations to visualize the training progress:

import mindspore

# Lists to store losses for visualization
generator_losses = []
discriminator_losses = []
generated_images_list = []

dataset_size = dataset.get_dataset_size()
for epoch in range(training_epochs):
    generator_model.set_train()
    discriminator_model.set_train()
    
    # Process data for each epoch
    for i, (images, ) in enumerate(dataset.create_tuple_iterator()):
        g_loss, d_loss, gen_imgs = training_step(images)
        if i % 100 == 0 or i == dataset_size - 1:
            # Print training progress
            print('[%2d/%d][%3d/%d]   Loss_D:%7.4f  Loss_G:%7.4f' % (
                epoch + 1, training_epochs, i + 1, dataset_size, d_loss.asnumpy(), g_loss.asnumpy()))
        discriminator_losses.append(d_loss.asnumpy())
        generator_losses.append(g_loss.asnumpy())

    # Generate images after each epoch for visualization
    generator_model.set_train(False)
    fixed_noise = ops.standard_normal((batch_size, latent_vector_size, 1, 1))
    generated_batch = generator_model(fixed_noise)
    generated_images_list.append(generated_batch.transpose(0, 2, 3, 1).asnumpy())

    # Save model parameters
    mindspore.save_checkpoint(generator_model, "./generator_model.ckpt")
    mindspore.save_checkpoint(discriminator_model, "./discriminator_model.ckpt")

Training Results Visualization

Let's plot the loss functions of both the generator and discriminator during training:

plt.figure(figsize=(10, 5))
plt.title("Generator and Discriminator Loss During Training")
plt.plot(generator_losses, label="Generator", color='blue')
plt.plot(discriminator_losses, label="Discriminator", color='orange')
plt.xlabel("Iterations")
plt.ylabel("Loss")
plt.legend()
plt.show()

Next, we'll create an animation to visualize the progression of generated images throughout training:

import matplotlib.pyplot as plt
import matplotlib.animation as animation

def create_animation(image_sequence):
    animation_frames = []
    figure = plt.figure(figsize=(8, 3), dpi=120)
    for epoch in range(len(image_sequence)):
        images = []
        for i in range(3):
            row = np.concatenate((image_sequence[epoch][i * 8:(i + 1) * 8]), axis=1)
            images.append(row)
        img = np.clip(np.concatenate((images[:]), axis=0), 0, 1)
        plt.axis("off")
        animation_frames.append([plt.imshow(img)])

    animation_obj = animation.ArtistAnimation(figure, animation_frames, interval=1000, repeat_delay=1000, blit=True)
    animation_obj.save('./dcgan_animation.gif', writer='pillow', fps=1)

create_animation(generated_images_list)

As shown in the animation, the quality of generated images improves as training progresses. If we increase the number of training epochs to 50 or more, the generated cartoon avatars become increasingly similar to those in the training dataset.

Finally, let's load the trained generator model to generate new images:

# Load pre-trained model parameters
mindspore.load_checkpoint("./generator_model.ckpt", generator_model)

# Generate images using the trained model
fixed_noise = ops.standard_normal((batch_size, latent_vector_size, 1, 1))
generated_images = generator_model(fixed_noise).transpose(0, 2, 3, 1).asnumpy()

# Display generated images
figure = plt.figure(figsize=(8, 3), dpi=120)
images = []
for i in range(3):
    images.append(np.concatenate((generated_images[i * 8:(i + 1) * 8]), axis=1))
img = np.clip(np.concatenate((images[:]), axis=0), 0, 1)
plt.axis("off")
plt.imshow(img)
plt.show()