Save the output of conditional GAN for each class ~ With cGAN implementation by PyTorch ~

This is the second post. Last time implemented DCGAN with PyTorch and made it possible to save output images one by one.

This time, we will implement an improved conditional GAN (conditional GAN) so that the output of GAN can be controlled. At the same time, as in the last time, we will be able to save the output images one by one.

Purpose

Implement conditional GAN and save output one by one

conditional GAN Conditional GAN allows you to explicitly separate the generated images. This was made possible by training using the label information of the teacher data during training. The paper is here

From the following paper 180B6B55-C45F-40F9-8863-D5A7B5E1D19D.png It's like learning by adding class label information to both Generator and Discriminator inputs. It seems that the input format changes a little, but the basic structure of GAN does not change.

Implementation

Let's move on to implementation. This time, we will implement conditional GAN based on DCGAN that was implemented last time.

Execution environment

Google Colaboratory

Module import & save destination settings

First from the module import

import argparse
import os
import numpy as np

import torchvision.transforms as transforms
from torchvision.utils import save_image

from torch.utils.data import DataLoader
from torchvision import datasets

import torch.nn as nn
import torch.nn.functional as F
import torch

img_save_path = 'images-C_dcgan'
os.makedirs(img_save_path, exist_ok=True)

Command line & default value settings

It's almost the same as last time. The subtle change is that the generated image size is 32x32 instead of the MNIST default of 28x28.

parser = argparse.ArgumentParser()
parser.add_argument('--n_epochs', type=int, default=200, help='number of epochs of training')
parser.add_argument('--batch_size', type=int, default=64, help='size of the batches')
parser.add_argument('--lr', type=float, default=0.0002, help='adam: learning rate')
parser.add_argument('--beta1', type=float, default=0.5, help='adam: decay of first order momentum of gradient')
parser.add_argument('--beta2', type=float, default=0.999, help='adam: decay of second order momentum of gradient')
parser.add_argument('--n_cpu', type=int, default=8, help='number of cpu threads to use during batch generation')
parser.add_argument('--latent_dim', type=int, default=100, help='dimensionality of the latent space')
parser.add_argument('--n_classes', type=int, default=10, help='number of classes for dataset')
parser.add_argument('--img_size', type=int, default=32, help='size of each image dimension')
parser.add_argument('--channels', type=int, default=1, help='number of image channels')
parser.add_argument('--sample_interval', type=int, default=400, help='interval between image sampling')
args = parser.parse_args()
#args for google colab=parser.parse_args(args=[])
print(args)

C,H,W = args.channels, args.img_size, args.img_size

Weight setting

def weights_init_normal(m):
    classname = m.__class__.__name__
    if classname.find('Conv') != -1:
        torch.nn.init.normal(m.weight, 0.0, 0.02)
    elif classname.find('BatchNorm2d') != -1:
        torch.nn.init.normal(m.weight, 1.0, 0.02)
        torch.nn.init.constant(m.bias, 0.0)

Generator Let's define the Generator. Generate with cat Combine the image information and the label information to generate.

class Generator(nn.Module):
    # initializers
    def __init__(self, d=128):
        super(Generator, self).__init__()
        self.deconv1_1 = nn.ConvTranspose2d(100, d*2, 4, 1, 0)
        self.deconv1_1_bn = nn.BatchNorm2d(d*2)
        self.deconv1_2 = nn.ConvTranspose2d(10, d*2, 4, 1, 0)
        self.deconv1_2_bn = nn.BatchNorm2d(d*2)
        self.deconv2 = nn.ConvTranspose2d(d*4, d*2, 4, 2, 1)
        self.deconv2_bn = nn.BatchNorm2d(d*2)
        self.deconv3 = nn.ConvTranspose2d(d*2, d, 4, 2, 1)
        self.deconv3_bn = nn.BatchNorm2d(d)
        self.deconv4 = nn.ConvTranspose2d(d, C, 4, 2, 1)


    # forward method
    def forward(self, input, label):
        x = F.relu(self.deconv1_1_bn(self.deconv1_1(input)))
        y = F.relu(self.deconv1_2_bn(self.deconv1_2(label)))
        x = torch.cat([x, y], 1)
        x = F.relu(self.deconv2_bn(self.deconv2(x)))
        x = F.relu(self.deconv3_bn(self.deconv3(x)))
        x = torch.tanh(self.deconv4(x))
        return x

Last time, I implemented Generator with Upsampling + Conv2d. This time, we are implementing using ConvTranspose2d instead of the previous method. This difference is summarized in this article, so please have a look if you are interested.

Discriminator The definition of Discriminator. Label information is also attached with cat here.


class Discriminator(nn.Module):
    # initializers
    def __init__(self, d=128):
        super(Discriminator, self).__init__()
        self.conv1_1 = nn.Conv2d(C, d//2, 4, 2, 1)
        self.conv1_2 = nn.Conv2d(10, d//2, 4, 2, 1)
        self.conv2 = nn.Conv2d(d, d*2, 4, 2, 1)
        self.conv2_bn = nn.BatchNorm2d(d*2)
        self.conv3 = nn.Conv2d(d*2, d*4, 4, 2, 1)
        self.conv3_bn = nn.BatchNorm2d(d*4)
        self.conv4 = nn.Conv2d(d * 4, 1, 4, 1, 0)

    def forward(self, input, label):
        x = F.leaky_relu(self.conv1_1(input), 0.2)
        y = F.leaky_relu(self.conv1_2(label), 0.2)
        x = torch.cat([x, y], 1)
        x = F.leaky_relu(self.conv2_bn(self.conv2(x)), 0.2)
        x = F.leaky_relu(self.conv3_bn(self.conv3(x)), 0.2)
        x = F.sigmoid(self.conv4(x))
        return x

Loss function and network settings

Define the Loss function, initialize the weights, initialize the Generator / Discriminator, and set the Optimizer.


# Loss function
adversarial_loss = torch.nn.BCELoss()

# Initialize Generator and discriminator
generator = Generator()
discriminator = Discriminator()

if torch.cuda.is_available():
    generator.cuda()
    discriminator.cuda()
    adversarial_loss.cuda()

# Initialize weights
generator.apply(weights_init_normal)
discriminator.apply(weights_init_normal)

# Optimizers
optimizer_G = torch.optim.Adam(generator.parameters(), lr=args.lr, betas=(args.beta1, args.beta2))
optimizer_D = torch.optim.Adam(discriminator.parameters(), lr=args.lr, betas=(args.beta1, args.beta2))

Creating a Dataloader

We will create a Dataloader. This time, the image is generated with a size of 32 * 32, so the MNIST image is resized in the pre-processing part of the image.


# Configure data loader
os.makedirs('./data', exist_ok=True)
dataloader = torch.utils.data.DataLoader(
    datasets.MNIST('./data', train=True, download=True,
                   transform=transforms.Compose([
                       transforms.Resize(args.img_size),
                       transforms.ToTensor(),
                       transforms.Normalize([0.5,], [0.5,])
                   ])),
    batch_size=args.batch_size, shuffle=True, drop_last=True)
print('the data is ok')

Training Training of GAN.


for epoch in range(1, args.n_epochs+1):
    for i, (imgs, labels) in enumerate(dataloader):

        Batch_Size = args.batch_size
        N_Class = args.n_classes
        img_size = args.img_size
        # Adversarial ground truths
        valid = torch.ones(Batch_Size).cuda()
        fake = torch.zeros(Batch_Size).cuda()

        # Configure input
        real_imgs = imgs.type(torch.FloatTensor).cuda()

        real_y = torch.zeros(Batch_Size, N_Class)
        real_y = real_y.scatter_(1, labels.view(Batch_Size, 1), 1).view(Batch_Size, N_Class, 1, 1).contiguous()
        real_y = real_y.expand(-1, -1, img_size, img_size).cuda()

        # Sample noise and labels as generator input
        noise = torch.randn((Batch_Size, args.latent_dim,1,1)).cuda()
        gen_labels = (torch.rand(Batch_Size, 1) * N_Class).type(torch.LongTensor)
        gen_y = torch.zeros(Batch_Size, N_Class)
        gen_y = gen_y.scatter_(1, gen_labels.view(Batch_Size, 1), 1).view(Batch_Size, N_Class,1,1).cuda()
        # ---------------------
        #  Train Discriminator
        # ---------------------
        optimizer_D.zero_grad()
        # Loss for real images
        d_real_loss = adversarial_loss(discriminator(real_imgs, real_y).squeeze(), valid)
        # Loss for fake images
        gen_imgs = generator(noise, gen_y)
        gen_y_for_D = gen_y.view(Batch_Size, N_Class, 1, 1).contiguous().expand(-1, -1, img_size, img_size)

        d_fake_loss = adversarial_loss(discriminator(gen_imgs.detach(),gen_y_for_D).squeeze(), fake)
        # Total discriminator loss
        d_loss = (d_real_loss + d_fake_loss)
        d_loss.backward()
        optimizer_D.step()

        # -----------------
        #  Train Generator
        # -----------------

        optimizer_G.zero_grad()

        g_loss = adversarial_loss(discriminator(gen_imgs,gen_y_for_D).squeeze(), valid)
        g_loss.backward()
        optimizer_G.step()


        print ("[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f]" % (epoch, args.n_epochs, i, len(dataloader),
                                                            d_loss.data.cpu(), g_loss.data.cpu()))

        batches_done = epoch * len(dataloader) + i
        if epoch % 20 == 0:
            noise = torch.FloatTensor(np.random.normal(0, 1, (N_Class**2, args.latent_dim,1,1))).cuda()
            #fixed labels
            y_ = torch.LongTensor(np.array([num for num in range(N_Class)])).view(N_Class,1).expand(-1,N_Class).contiguous()
            y_fixed = torch.zeros(N_Class**2, N_Class)
            y_fixed = y_fixed.scatter_(1,y_.view(N_Class**2,1),1).view(N_Class**2, N_Class,1,1).cuda()

            with torch.no_grad():
                gen_imgs = generator(noise, y_fixed).view(-1,C,H,W)

            save_image(gen_imgs.data, img_save_path + '/epoch:%d.png' % epoch, nrow=N_Class, normalize=True) 

Execution result

The execution result is as follows. 20-19600.png You can see that the generated images are neatly arranged for each class. Conditional GAN allows you to control the images generated in this way.

Generate and save images for each class

Like last time, we will be able to save images one by one.


if epoch % 20 == 0:
    noise = torch.FloatTensor(np.random.normal(0, 1, (N_Class**2, args.latent_dim,1,1))).cuda()
    #fixed labels
    y_ = torch.LongTensor(np.array([num for num in range(N_Class)])).view(N_Class,1).expand(-1,N_Class).contiguous()
    y_fixed = torch.zeros(N_Class**2, N_Class)
    y_fixed = y_fixed.scatter_(1,y_.view(N_Class**2,1),1).view(N_Class**2, N_Class,1,1).cuda()

    with torch.no_grad():
        gen_imgs = generator(noise, y_fixed).view(-1,C,H,W)

    save_image(gen_imgs.data, img_save_path + '/epoch:%d.png' % epoch, nrow=N_Class, normalize=True)

Here part


if epoch % 20 == 0:
    for l in range(10): #Save 10 sheets for each class
        noise = torch.FloatTensor(np.random.normal(0, 1, (N_Class**2, args.latent_dim,1,1))).cuda()
        #fixed labels
        y_ = torch.LongTensor(np.array([num for num in range(N_Class)])).view(N_Class,1).expand(-1,N_Class).contiguous()
        y_fixed = torch.zeros(N_Class**2, N_Class)
        y_fixed = y_fixed.scatter_(1,y_.view(N_Class**2,1),1).view(N_Class**2, N_Class,1,1).cuda()

        for m in range()
            with torch.no_grad():
                gen_imgs = generator(noise, y_fixed).view(-1,C,H,W)

            save_gen_imgs = gen_imgs[10*i]
            save_image(save_gen_imgs, img_save_path + '/epochs:%d/%d/epoch:%d-%d_%d.png' % (epoch, i, epoch,i, j), normalize=True)

Change it like this. If you want to do this, you need to change the directory structure for saving images.

images-C_dcgan
├── epochs:20
│   ├── 0
│   ├── 1
│   ├── 2
│   ├── 3
│   ├── 4
│   ├── 5
│   ├── 6
│   ├── 7
│   ├── 8
│   └── 9
│     .
│     .
│     .
│
└── epochs:200
    ├── 0
    ├── 1
    ├── 2
    ├── 3
    ├── 4
    ├── 5
    ├── 6
    ├── 7
    ├── 8
    └── 9

There are 0 to 9 directories every 20 epochs. It's easier to create at once using ʻos.makedirs`. Images are now saved for each class.

Summary

This time, we implemented conditional GAN following DCGAN so that the generated images can be saved one by one. This time, we implemented conditional GAN by adding label information to both Generator and Discriminator inputs, which is the simplest. Currently, the de facto standards for implementing conditional GAN are technologies such as Projection Discriminator and Conditional Batch Normalization. I don't understand much about the technology around here, so if I have a chance, I would like to study while implementing it.

Recommended Posts

Save the output of conditional GAN for each class ~ With cGAN implementation by PyTorch ~
Save the output of GAN one by one ~ With the implementation of GAN by PyTorch ~
Add the attribute of the object of the class with the for statement
Generation of junk character MNIST (KMNIST) with cGAN (conditional GAN)
Summary of basic implementation by PyTorch
Output csv with different number of digits for each column with numpy
The third night of the loop with for
The second night of the loop with for
Build API server for checking the operation of front implementation with python3 and Flask
Output the specified table of Oracle database in Python to Excel for each file
Gradually display the output of the command executed by subprocess.Popen
Make Fatjar by changing the main class with Gradle
Get the sum of each of multiple columns with awk
Display the integrated temperature for each field with Z-GIS
A python implementation of the Bayesian linear regression class
Preparing the execution environment of PyTorch with Docker November 2019
Date the name of the decomposed PDF for each page
[For beginners] Quantify the similarity of sentences with TF-IDF