annotated_deep_learning_paper_implementations/labml_nn/gan/cycle_gan.py

"""
Download datasets from https://people.eecs.berkeley.edu/~taesung_park/CycleGAN/datasets/[DATASET NAME].zip
and extract them into labml_nn/data/cycle_gan/[DATASET NAME]

I've taken pieces of code from https://github.com/eriklindernoren/PyTorch-GAN
"""

import itertools
import random
from pathlib import PurePath, Path

import torch
import torch.nn as nn
import torchvision.transforms as transforms
from PIL import Image
from torch.utils.data import DataLoader
from torch.utils.data import Dataset
from torchvision.utils import make_grid
from torchvision.utils import save_image

from labml import lab, tracker, experiment
from labml_helpers.device import DeviceConfigs
from labml_helpers.module import Module
from labml_helpers.training_loop import TrainingLoopConfigs


class ReplayBuffer:
    def __init__(self, max_size: int = 50):
        self.max_size = max_size
        self.data = []

    def push_and_pop(self, data):
        to_return = []
        for element in data:
            if len(self.data) < self.max_size:
                self.data.append(element)
                to_return.append(element)
            else:
                if random.uniform(0, 1) > 0.5:
                    i = random.randint(0, self.max_size - 1)
                    to_return.append(self.data[i].clone())
                    self.data[i] = element
                else:
                    to_return.append(element)
        return torch.stack(to_return)


def load_image(path: str):
    image = Image.open(path)
    if image.mode != 'RGB':
        image = Image.new("RGB", image.size).pase(image)

    return image


class ImageDataset(Dataset):
    def __init__(self, root: PurePath, transforms_, unaligned: bool, mode: str):
        root = Path(root)
        self.transform = transforms.Compose(transforms_)
        self.unaligned = unaligned

        self.files_A = sorted(str(f) for f in (root / f'{mode}A').iterdir())
        self.files_B = sorted(str(f) for f in (root / f'{mode}B').iterdir())

    def __getitem__(self, index):
        return {"a": self.transform(load_image(self.files_A[index % len(self.files_A)])),
                "b": self.transform(load_image(self.files_B[index % len(self.files_B)]))}

    def __len__(self):
        return max(len(self.files_A), len(self.files_B))


def weights_init_normal(m):
    classname = m.__class__.__name__
    if classname.find("Conv") != -1:
        torch.nn.init.normal_(m.weight.data, 0.0, 0.02)
    elif classname.find("BatchNorm2d") != -1:
        torch.nn.init.normal_(m.weight.data, 1.0, 0.02)


class ResidualBlock(Module):
    def __init__(self, in_features: int):
        super().__init__()
        self.block = nn.Sequential(
            nn.ReflectionPad2d(1),
            nn.Conv2d(in_features, in_features, 3),
            nn.InstanceNorm2d(in_features),
            nn.ReLU(inplace=True),
            nn.ReflectionPad2d(1),
            nn.Conv2d(in_features, in_features, 3),
            nn.InstanceNorm2d(in_features),
        )

    def __call__(self, x: torch.Tensor):
        return x + self.block(x)


class GeneratorResNet(Module):
    def __init__(self, input_shape, num_residual_blocks):
        super().__init__()
        channels = input_shape[0]

        # Initial convolution block
        out_features = 64
        layers = [
            nn.ReflectionPad2d(channels),
            nn.Conv2d(channels, out_features, 7),
            nn.InstanceNorm2d(out_features),
            nn.ReLU(inplace=True),
        ]
        in_features = out_features

        # Downsampling
        for _ in range(2):
            out_features *= 2
            layers += [
                nn.Conv2d(in_features, out_features, 3, stride=2, padding=1),
                nn.InstanceNorm2d(out_features),
                nn.ReLU(inplace=True),
            ]
            in_features = out_features

        # Residual blocks
        for _ in range(num_residual_blocks):
            layers += [ResidualBlock(out_features)]

        # Upsampling
        for _ in range(2):
            out_features //= 2
            layers += [
                nn.Upsample(scale_factor=2),
                nn.Conv2d(in_features, out_features, 3, stride=1, padding=1),
                nn.InstanceNorm2d(out_features),
                nn.ReLU(inplace=True),
            ]
            in_features = out_features

        # Output layer
        layers += [nn.ReflectionPad2d(channels), nn.Conv2d(out_features, channels, 7), nn.Tanh()]

        self.layers = nn.Sequential(*layers)

        self.apply(weights_init_normal)

    def __call__(self, x):
        return self.layers(x)


class DiscriminatorBlock(Module):
    def __init__(self, in_filters, out_filters, normalize=True):
        super().__init__()
        layers = [nn.Conv2d(in_filters, out_filters, 4, stride=2, padding=1)]
        if normalize:
            layers.append(nn.InstanceNorm2d(out_filters))
        layers.append(nn.LeakyReLU(0.2, inplace=True))
        self.layers = nn.Sequential(*layers)

    def __call__(self, x: torch.Tensor):
        return self.layers(x)


class Discriminator(Module):
    def __init__(self, input_shape):
        super().__init__()
        channels, height, width = input_shape

        # Calculate output shape of image discriminator (PatchGAN)
        self.output_shape = (1, height // 2 ** 4, width // 2 ** 4)

        self.model = nn.Sequential(
            DiscriminatorBlock(channels, 64, normalize=False),
            DiscriminatorBlock(64, 128),
            DiscriminatorBlock(128, 256),
            DiscriminatorBlock(256, 512),
            nn.ZeroPad2d((1, 0, 1, 0)),
            nn.Conv2d(512, 1, 4, padding=1)
        )

        self.apply(weights_init_normal)

    def forward(self, img):
        return self.model(img)


def sample_images(n, dataset_name, valid_dataloader, generator_ab, generator_ba):
    """Saves a generated sample from the test set"""
    batch = next(iter(valid_dataloader))
    generator_ab.eval()
    generator_ba.eval()
    with torch.no_grad():
        real_a, real_b = batch['a'].to(generator_ab.device), batch['b'].to(generator_ba.device)
        fake_b = generator_ab(real_a)
        fake_a = generator_ba(real_b)

        # Arange images along x-axis
        real_a = make_grid(real_a, nrow=5, normalize=True)
        real_b = make_grid(real_b, nrow=5, normalize=True)
        fake_a = make_grid(fake_a, nrow=5, normalize=True)
        fake_b = make_grid(fake_b, nrow=5, normalize=True)

        # arange images along y-axis
        image_grid = torch.cat((real_a, fake_b, real_b, fake_a), 1)
    save_image(image_grid, "images/%s/%s.png" % (dataset_name, n), normalize=False)


class Configs(TrainingLoopConfigs):
    device: torch.device = DeviceConfigs()
    loop_count: int = 200
    dataset_name: str = 'monet2photo'
    batch_size: int = 1

    data_loader_workers = 8
    is_save_models = True

    learning_rate = 0.0002
    adam_betas = (0.5, 0.999)
    decay_start = 100

    identity_loss = torch.nn.L1Loss()
    cycle_loss = torch.nn.L1Loss()
    gan_loss = torch.nn.MSELoss()

    batch_step = 'cycle_gan_batch_step'

    img_height = 256
    img_width = 256
    img_channels = 3

    n_residual_blocks = 9

    cyclic_loss_coefficient = 10.0
    identity_loss_coefficient = 5.

    sample_interval = 100

    def run(self):
        images_path = lab.get_data_path() / 'cycle_gan' / self.dataset_name

        input_shape = (self.img_channels, self.img_height, self.img_width)
        generator_ab = GeneratorResNet(input_shape, self.n_residual_blocks).to(self.device)
        generator_ba = GeneratorResNet(input_shape, self.n_residual_blocks).to(self.device)
        discriminator_a = Discriminator(input_shape).to(self.device)
        discriminator_b = Discriminator(input_shape).to(self.device)

        generator_optimizer = torch.optim.Adam(
            itertools.chain(generator_ab.parameters(), generator_ba.parameters()),
            lr=self.learning_rate, betas=self.adam_betas)
        discriminator_optimizer = torch.optim.Adam(
            itertools.chain(discriminator_a.parameters(), discriminator_b.parameters()),
            lr=self.learning_rate, betas=self.adam_betas)

        decay_epochs = self.loop_count - self.decay_start
        generator_lr_scheduler = torch.optim.lr_scheduler.LambdaLR(
            generator_optimizer, lr_lambda=lambda e: 1.0 - max(0, e - self.decay_start) / decay_epochs)
        discriminator_lr_scheduler = torch.optim.lr_scheduler.LambdaLR(
            discriminator_optimizer, lr_lambda=lambda e: 1.0 - max(0, e - self.decay_start) / decay_epochs)

        # Image transformations
        transforms_ = [
            transforms.Resize(int(self.img_height * 1.12), Image.BICUBIC),
            transforms.RandomCrop((self.img_height, self.img_width)),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
        ]

        # Training data loader
        dataloader = DataLoader(
            ImageDataset(images_path, transforms_, True, 'train'),
            batch_size=self.batch_size,
            shuffle=True,
            num_workers=self.data_loader_workers,
        )
        # Test data loader
        valid_dataloader = DataLoader(
            ImageDataset(images_path, transforms_, True, "test"),
            batch_size=5,
            shuffle=True,
            num_workers=self.data_loader_workers,
        )

        # Buffers of previously generated samples
        fake_a_buffer = ReplayBuffer()
        fake_b_buffer = ReplayBuffer()

        for epoch in self.training_loop:
            for i, batch in enumerate(dataloader):
                # Set model input
                real_a, real_b = batch['a'].to(self.device), batch['b'].to(self.device)

                # adversarial ground truths
                valid = torch.ones(real_a.size(0), *discriminator_a.output_shape,
                                   device=self.device, requires_grad=False)
                fake = torch.zeros(real_a.size(0), *discriminator_a.output_shape,
                                   device=self.device, requires_grad=False)

                #  Train generators
                generator_ab.train()
                generator_ba.train()

                # Identity loss
                loss_identity = self.identity_loss(generator_ba(real_a), real_a) + \
                                self.identity_loss(generator_ab(real_b), real_b)

                # GAN loss
                fake_b = generator_ab(real_a)
                fake_a = generator_ba(real_b)

                loss_gan = self.gan_loss(discriminator_b(fake_b), valid) + \
                           self.gan_loss(discriminator_a(fake_a), valid)

                loss_cycle = self.cycle_loss(generator_ba(fake_b), real_a) + \
                             self.cycle_loss(generator_ab(fake_a), real_b)

                # Total loss
                loss_generator = (loss_gan + self.cyclic_loss_coefficient * loss_cycle
                                  + self.identity_loss_coefficient * loss_identity)

                generator_optimizer.zero_grad()
                loss_generator.backward()
                generator_optimizer.step()

                #  Train discriminators
                fake_a_replay = fake_a_buffer.push_and_pop(fake_a)
                fake_b_replay = fake_b_buffer.push_and_pop(fake_b)
                loss_discriminator = self.gan_loss(discriminator_a(real_a), valid) + \
                                     self.gan_loss(discriminator_a(fake_a_replay.detach()), fake) + \
                                     self.gan_loss(discriminator_b(real_b), valid) + \
                                     self.gan_loss(discriminator_b(fake_b_replay.detach()), fake)

                discriminator_optimizer.zero_grad()
                loss_discriminator.backward()
                discriminator_optimizer.step()

                tracker.save({'loss.generator': loss_generator,
                              'loss.discriminator': loss_discriminator,
                              'loss.generator.cycle': loss_cycle,
                              'loss.generator.gan': loss_gan,
                              'loss.generator.identity': loss_identity})

                # If at sample interval save image
                batches_done = epoch * len(dataloader) + i
                if batches_done % self.sample_interval == 0:
                    sample_images(batches_done, self.dataset_name, valid_dataloader, generator_ab, generator_ba)

                tracker.add_global_step(max(len(real_a), len(real_b)))

            # Update learning rates
            generator_lr_scheduler.step()
            discriminator_lr_scheduler.step()


def main():
    conf = Configs()
    experiment.create(name='cycle_gan')
    experiment.configs(conf, 'run')
    with experiment.start():
        conf.run()


if __name__ == '__main__':
    main()