chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,16 @@
|
||||
"""
|
||||
---
|
||||
title: Generative Adversarial Networks
|
||||
summary: >
|
||||
A set of PyTorch implementations/tutorials of GANs.
|
||||
---
|
||||
|
||||
# Generative Adversarial Networks
|
||||
|
||||
* [Original GAN](original/index.html)
|
||||
* [GAN with deep convolutional network](dcgan/index.html)
|
||||
* [Cycle GAN](cycle_gan/index.html)
|
||||
* [Wasserstein GAN](wasserstein/index.html)
|
||||
* [Wasserstein GAN with Gradient Penalty](wasserstein/gradient_penalty/index.html)
|
||||
* [StyleGAN 2](stylegan/index.html)
|
||||
"""
|
||||
@@ -0,0 +1,773 @@
|
||||
"""
|
||||
---
|
||||
title: Cycle GAN
|
||||
summary: >
|
||||
A simple PyTorch implementation/tutorial of Cycle GAN introduced in paper
|
||||
Unpaired Image-to-Image Translation using Cycle-Consistent Adversarial Networks.
|
||||
---
|
||||
|
||||
# Cycle GAN
|
||||
|
||||
This is a [PyTorch](https://pytorch.org) implementation/tutorial of the paper
|
||||
[Unpaired Image-to-Image Translation using Cycle-Consistent Adversarial Networks](https://arxiv.org/abs/1703.10593).
|
||||
|
||||
I've taken pieces of code from [eriklindernoren/PyTorch-GAN](https://github.com/eriklindernoren/PyTorch-GAN).
|
||||
It is a very good resource if you want to checkout other GAN variations too.
|
||||
|
||||
Cycle GAN does image-to-image translation.
|
||||
It trains a model to translate an image from given distribution to another, say, images of class A and B.
|
||||
Images of a certain distribution could be things like images of a certain style, or nature.
|
||||
The models do not need paired images between A and B.
|
||||
Just a set of images of each class is enough.
|
||||
This works very well on changing between image styles, lighting changes, pattern changes, etc.
|
||||
For example, changing summer to winter, painting style to photos, and horses to zebras.
|
||||
|
||||
Cycle GAN trains two generator models and two discriminator models.
|
||||
One generator translates images from A to B and the other from B to A.
|
||||
The discriminators test whether the generated images look real.
|
||||
|
||||
This file contains the model code as well as the training code.
|
||||
We also have a Google Colab notebook.
|
||||
|
||||
[](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/gan/cycle_gan/experiment.ipynb)
|
||||
"""
|
||||
|
||||
import itertools
|
||||
import random
|
||||
import zipfile
|
||||
from typing import Tuple
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torchvision.transforms as transforms
|
||||
from PIL import Image
|
||||
from torch.utils.data import DataLoader, Dataset
|
||||
from torchvision.transforms import InterpolationMode
|
||||
from torchvision.utils import make_grid
|
||||
|
||||
from labml import lab, tracker, experiment, monit
|
||||
from labml.configs import BaseConfigs
|
||||
from labml.utils.download import download_file
|
||||
from labml.utils.pytorch import get_modules
|
||||
from labml_nn.helpers.device import DeviceConfigs
|
||||
|
||||
|
||||
class GeneratorResNet(nn.Module):
|
||||
"""
|
||||
The generator is a residual network.
|
||||
"""
|
||||
|
||||
def __init__(self, input_channels: int, n_residual_blocks: int):
|
||||
super().__init__()
|
||||
# This first block runs a $7\times7$ convolution and maps the image to
|
||||
# a feature map.
|
||||
# The output feature map has the same height and width because we have
|
||||
# a padding of $3$.
|
||||
# Reflection padding is used because it gives better image quality at edges.
|
||||
#
|
||||
# `inplace=True` in `ReLU` saves a little bit of memory.
|
||||
out_features = 64
|
||||
layers = [
|
||||
nn.Conv2d(input_channels, out_features, kernel_size=7, padding=3, padding_mode='reflect'),
|
||||
nn.InstanceNorm2d(out_features),
|
||||
nn.ReLU(inplace=True),
|
||||
]
|
||||
in_features = out_features
|
||||
|
||||
# We down-sample with two $3 \times 3$ convolutions
|
||||
# with stride of 2
|
||||
for _ in range(2):
|
||||
out_features *= 2
|
||||
layers += [
|
||||
nn.Conv2d(in_features, out_features, kernel_size=3, stride=2, padding=1),
|
||||
nn.InstanceNorm2d(out_features),
|
||||
nn.ReLU(inplace=True),
|
||||
]
|
||||
in_features = out_features
|
||||
|
||||
# We take this through `n_residual_blocks`.
|
||||
# This module is defined below.
|
||||
for _ in range(n_residual_blocks):
|
||||
layers += [ResidualBlock(out_features)]
|
||||
|
||||
# Then the resulting feature map is up-sampled
|
||||
# to match the original image height and width.
|
||||
for _ in range(2):
|
||||
out_features //= 2
|
||||
layers += [
|
||||
nn.Upsample(scale_factor=2),
|
||||
nn.Conv2d(in_features, out_features, kernel_size=3, stride=1, padding=1),
|
||||
nn.InstanceNorm2d(out_features),
|
||||
nn.ReLU(inplace=True),
|
||||
]
|
||||
in_features = out_features
|
||||
|
||||
# Finally we map the feature map to an RGB image
|
||||
layers += [nn.Conv2d(out_features, input_channels, 7, padding=3, padding_mode='reflect'), nn.Tanh()]
|
||||
|
||||
# Create a sequential module with the layers
|
||||
self.layers = nn.Sequential(*layers)
|
||||
|
||||
# Initialize weights to $\mathcal{N}(0, 0.2)$
|
||||
self.apply(weights_init_normal)
|
||||
|
||||
def forward(self, x):
|
||||
return self.layers(x)
|
||||
|
||||
|
||||
class ResidualBlock(nn.Module):
|
||||
"""
|
||||
This is the residual block, with two convolution layers.
|
||||
"""
|
||||
|
||||
def __init__(self, in_features: int):
|
||||
super().__init__()
|
||||
self.block = nn.Sequential(
|
||||
nn.Conv2d(in_features, in_features, kernel_size=3, padding=1, padding_mode='reflect'),
|
||||
nn.InstanceNorm2d(in_features),
|
||||
nn.ReLU(inplace=True),
|
||||
nn.Conv2d(in_features, in_features, kernel_size=3, padding=1, padding_mode='reflect'),
|
||||
nn.InstanceNorm2d(in_features),
|
||||
nn.ReLU(inplace=True),
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor):
|
||||
return x + self.block(x)
|
||||
|
||||
|
||||
class Discriminator(nn.Module):
|
||||
"""
|
||||
This is the discriminator.
|
||||
"""
|
||||
|
||||
def __init__(self, input_shape: Tuple[int, int, int]):
|
||||
super().__init__()
|
||||
channels, height, width = input_shape
|
||||
|
||||
# Output of the discriminator is also a map of probabilities,
|
||||
# whether each region of the image is real or generated
|
||||
self.output_shape = (1, height // 2 ** 4, width // 2 ** 4)
|
||||
|
||||
self.layers = nn.Sequential(
|
||||
# Each of these blocks will shrink the height and width by a factor of 2
|
||||
DiscriminatorBlock(channels, 64, normalize=False),
|
||||
DiscriminatorBlock(64, 128),
|
||||
DiscriminatorBlock(128, 256),
|
||||
DiscriminatorBlock(256, 512),
|
||||
# Zero pad on top and left to keep the output height and width same
|
||||
# with the $4 \times 4$ kernel
|
||||
nn.ZeroPad2d((1, 0, 1, 0)),
|
||||
nn.Conv2d(512, 1, kernel_size=4, padding=1)
|
||||
)
|
||||
|
||||
# Initialize weights to $\mathcal{N}(0, 0.2)$
|
||||
self.apply(weights_init_normal)
|
||||
|
||||
def forward(self, img):
|
||||
return self.layers(img)
|
||||
|
||||
|
||||
class DiscriminatorBlock(nn.Module):
|
||||
"""
|
||||
This is the discriminator block module.
|
||||
It does a convolution, an optional normalization, and a leaky ReLU.
|
||||
|
||||
It shrinks the height and width of the input feature map by half.
|
||||
"""
|
||||
|
||||
def __init__(self, in_filters: int, out_filters: int, normalize: bool = True):
|
||||
super().__init__()
|
||||
layers = [nn.Conv2d(in_filters, out_filters, kernel_size=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 forward(self, x: torch.Tensor):
|
||||
return self.layers(x)
|
||||
|
||||
|
||||
def weights_init_normal(m):
|
||||
"""
|
||||
Initialize convolution layer weights to $\mathcal{N}(0, 0.2)$
|
||||
"""
|
||||
classname = m.__class__.__name__
|
||||
if classname.find("Conv") != -1:
|
||||
torch.nn.init.normal_(m.weight.data, 0.0, 0.02)
|
||||
|
||||
|
||||
def load_image(path: str):
|
||||
"""
|
||||
Load an image and change to RGB if in grey-scale.
|
||||
"""
|
||||
image = Image.open(path)
|
||||
if image.mode != 'RGB':
|
||||
image = Image.new("RGB", image.size).paste(image)
|
||||
|
||||
return image
|
||||
|
||||
|
||||
class ImageDataset(Dataset):
|
||||
"""
|
||||
### Dataset to load images
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def download(dataset_name: str):
|
||||
"""
|
||||
#### Download dataset and extract data
|
||||
"""
|
||||
# URL
|
||||
url = f'https://people.eecs.berkeley.edu/~taesung_park/CycleGAN/datasets/{dataset_name}.zip'
|
||||
# Download folder
|
||||
root = lab.get_data_path() / 'cycle_gan'
|
||||
if not root.exists():
|
||||
root.mkdir(parents=True)
|
||||
# Download destination
|
||||
archive = root / f'{dataset_name}.zip'
|
||||
# Download file (generally ~100MB)
|
||||
download_file(url, archive)
|
||||
# Extract the archive
|
||||
with zipfile.ZipFile(archive, 'r') as f:
|
||||
f.extractall(root)
|
||||
|
||||
def __init__(self, dataset_name: str, transforms_, mode: str):
|
||||
"""
|
||||
#### Initialize the dataset
|
||||
|
||||
* `dataset_name` is the name of the dataset
|
||||
* `transforms_` is the set of image transforms
|
||||
* `mode` is either `train` or `test`
|
||||
"""
|
||||
# Dataset path
|
||||
root = lab.get_data_path() / 'cycle_gan' / dataset_name
|
||||
# Download if missing
|
||||
if not root.exists():
|
||||
self.download(dataset_name)
|
||||
|
||||
# Image transforms
|
||||
self.transform = transforms.Compose(transforms_)
|
||||
|
||||
# Get image paths
|
||||
path_a = root / f'{mode}A'
|
||||
path_b = root / f'{mode}B'
|
||||
self.files_a = sorted(str(f) for f in path_a.iterdir())
|
||||
self.files_b = sorted(str(f) for f in path_b.iterdir())
|
||||
|
||||
def __getitem__(self, index):
|
||||
# Return a pair of images.
|
||||
# These pairs get batched together, and they do not act like pairs in training.
|
||||
# So it is kind of ok that we always keep giving the same pair.
|
||||
return {"x": self.transform(load_image(self.files_a[index % len(self.files_a)])),
|
||||
"y": self.transform(load_image(self.files_b[index % len(self.files_b)]))}
|
||||
|
||||
def __len__(self):
|
||||
# Number of images in the dataset
|
||||
return max(len(self.files_a), len(self.files_b))
|
||||
|
||||
|
||||
class ReplayBuffer:
|
||||
"""
|
||||
### Replay Buffer
|
||||
|
||||
Replay buffer is used to train the discriminator.
|
||||
Generated images are added to the replay buffer and sampled from it.
|
||||
|
||||
The replay buffer returns the newly added image with a probability of $0.5$.
|
||||
Otherwise, it sends an older generated image and replaces the older image
|
||||
with the newly generated image.
|
||||
|
||||
This is done to reduce model oscillation.
|
||||
"""
|
||||
|
||||
def __init__(self, max_size: int = 50):
|
||||
self.max_size = max_size
|
||||
self.data = []
|
||||
|
||||
def push_and_pop(self, data: torch.Tensor):
|
||||
"""Add/retrieve an image"""
|
||||
data = data.detach()
|
||||
res = []
|
||||
for element in data:
|
||||
if len(self.data) < self.max_size:
|
||||
self.data.append(element)
|
||||
res.append(element)
|
||||
else:
|
||||
if random.uniform(0, 1) > 0.5:
|
||||
i = random.randint(0, self.max_size - 1)
|
||||
res.append(self.data[i].clone())
|
||||
self.data[i] = element
|
||||
else:
|
||||
res.append(element)
|
||||
return torch.stack(res)
|
||||
|
||||
|
||||
class Configs(BaseConfigs):
|
||||
"""## Configurations"""
|
||||
|
||||
# `DeviceConfigs` will pick a GPU if available
|
||||
device: torch.device = DeviceConfigs()
|
||||
|
||||
# Hyper-parameters
|
||||
epochs: int = 200
|
||||
dataset_name: str = 'monet2photo'
|
||||
batch_size: int = 1
|
||||
|
||||
data_loader_workers = 8
|
||||
|
||||
learning_rate = 0.0002
|
||||
adam_betas = (0.5, 0.999)
|
||||
decay_start = 100
|
||||
|
||||
# The paper suggests using a least-squares loss instead of
|
||||
# negative log-likelihood, at it is found to be more stable.
|
||||
gan_loss = torch.nn.MSELoss()
|
||||
|
||||
# L1 loss is used for cycle loss and identity loss
|
||||
cycle_loss = torch.nn.L1Loss()
|
||||
identity_loss = torch.nn.L1Loss()
|
||||
|
||||
# Image dimensions
|
||||
img_height = 256
|
||||
img_width = 256
|
||||
img_channels = 3
|
||||
|
||||
# Number of residual blocks in the generator
|
||||
n_residual_blocks = 9
|
||||
|
||||
# Loss coefficients
|
||||
cyclic_loss_coefficient = 10.0
|
||||
identity_loss_coefficient = 5.
|
||||
|
||||
sample_interval = 500
|
||||
|
||||
# Models
|
||||
generator_xy: GeneratorResNet
|
||||
generator_yx: GeneratorResNet
|
||||
discriminator_x: Discriminator
|
||||
discriminator_y: Discriminator
|
||||
|
||||
# Optimizers
|
||||
generator_optimizer: torch.optim.Adam
|
||||
discriminator_optimizer: torch.optim.Adam
|
||||
|
||||
# Learning rate schedules
|
||||
generator_lr_scheduler: torch.optim.lr_scheduler.LambdaLR
|
||||
discriminator_lr_scheduler: torch.optim.lr_scheduler.LambdaLR
|
||||
|
||||
# Data loaders
|
||||
dataloader: DataLoader
|
||||
valid_dataloader: DataLoader
|
||||
|
||||
def sample_images(self, n: int):
|
||||
"""Generate samples from test set and save them"""
|
||||
batch = next(iter(self.valid_dataloader))
|
||||
self.generator_xy.eval()
|
||||
self.generator_yx.eval()
|
||||
with torch.no_grad():
|
||||
data_x, data_y = batch['x'].to(self.generator_xy.device), batch['y'].to(self.generator_yx.device)
|
||||
gen_y = self.generator_xy(data_x)
|
||||
gen_x = self.generator_yx(data_y)
|
||||
|
||||
# Arrange images along x-axis
|
||||
data_x = make_grid(data_x, nrow=5, normalize=True)
|
||||
data_y = make_grid(data_y, nrow=5, normalize=True)
|
||||
gen_x = make_grid(gen_x, nrow=5, normalize=True)
|
||||
gen_y = make_grid(gen_y, nrow=5, normalize=True)
|
||||
|
||||
# Arrange images along y-axis
|
||||
image_grid = torch.cat((data_x, gen_y, data_y, gen_x), 1)
|
||||
|
||||
# Show samples
|
||||
plot_image(image_grid)
|
||||
|
||||
def initialize(self):
|
||||
"""
|
||||
## Initialize models and data loaders
|
||||
"""
|
||||
input_shape = (self.img_channels, self.img_height, self.img_width)
|
||||
|
||||
# Create the models
|
||||
self.generator_xy = GeneratorResNet(self.img_channels, self.n_residual_blocks).to(self.device)
|
||||
self.generator_yx = GeneratorResNet(self.img_channels, self.n_residual_blocks).to(self.device)
|
||||
self.discriminator_x = Discriminator(input_shape).to(self.device)
|
||||
self.discriminator_y = Discriminator(input_shape).to(self.device)
|
||||
|
||||
# Create the optmizers
|
||||
self.generator_optimizer = torch.optim.Adam(
|
||||
itertools.chain(self.generator_xy.parameters(), self.generator_yx.parameters()),
|
||||
lr=self.learning_rate, betas=self.adam_betas)
|
||||
self.discriminator_optimizer = torch.optim.Adam(
|
||||
itertools.chain(self.discriminator_x.parameters(), self.discriminator_y.parameters()),
|
||||
lr=self.learning_rate, betas=self.adam_betas)
|
||||
|
||||
# Create the learning rate schedules.
|
||||
# The learning rate stars flat until `decay_start` epochs,
|
||||
# and then linearly reduce to $0$ at end of training.
|
||||
decay_epochs = self.epochs - self.decay_start
|
||||
self.generator_lr_scheduler = torch.optim.lr_scheduler.LambdaLR(
|
||||
self.generator_optimizer, lr_lambda=lambda e: 1.0 - max(0, e - self.decay_start) / decay_epochs)
|
||||
self.discriminator_lr_scheduler = torch.optim.lr_scheduler.LambdaLR(
|
||||
self.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), InterpolationMode.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
|
||||
self.dataloader = DataLoader(
|
||||
ImageDataset(self.dataset_name, transforms_, 'train'),
|
||||
batch_size=self.batch_size,
|
||||
shuffle=True,
|
||||
num_workers=self.data_loader_workers,
|
||||
)
|
||||
|
||||
# Validation data loader
|
||||
self.valid_dataloader = DataLoader(
|
||||
ImageDataset(self.dataset_name, transforms_, "test"),
|
||||
batch_size=5,
|
||||
shuffle=True,
|
||||
num_workers=self.data_loader_workers,
|
||||
)
|
||||
|
||||
def run(self):
|
||||
"""
|
||||
## Training
|
||||
|
||||
We aim to solve:
|
||||
$$G^{*}, F^{*} = \arg \min_{G,F} \max_{D_X, D_Y} \mathcal{L}(G, F, D_X, D_Y)$$
|
||||
|
||||
where,
|
||||
$G$ translates images from $X \rightarrow Y$,
|
||||
$F$ translates images from $Y \rightarrow X$,
|
||||
$D_X$ tests if images are from $X$ space,
|
||||
$D_Y$ tests if images are from $Y$ space, and
|
||||
|
||||
\begin{align}
|
||||
\mathcal{L}(G, F, D_X, D_Y)
|
||||
&= \mathcal{L}_{GAN}(G, D_Y, X, Y) \\
|
||||
&+ \mathcal{L}_{GAN}(F, D_X, Y, X) \\
|
||||
&+ \lambda_1 \mathcal{L}_{cyc}(G, F) \\
|
||||
&+ \lambda_2 \mathcal{L}_{identity}(G, F) \\
|
||||
\\
|
||||
\mathcal{L}_{GAN}(G, F, D_Y, X, Y)
|
||||
&= \mathbb{E}_{y \sim p_{data}(y)} \Big[log D_Y(y)\Big] \\
|
||||
&+ \mathbb{E}_{x \sim p_{data}(x)} \bigg[log\Big(1 - D_Y(G(x))\Big)\bigg] \\
|
||||
&+ \mathbb{E}_{x \sim p_{data}(x)} \Big[log D_X(x)\Big] \\
|
||||
&+ \mathbb{E}_{y \sim p_{data}(y)} \bigg[log\Big(1 - D_X(F(y))\Big)\bigg] \\
|
||||
\\
|
||||
\mathcal{L}_{cyc}(G, F)
|
||||
&= \mathbb{E}_{x \sim p_{data}(x)} \Big[\lVert F(G(x)) - x \lVert_1\Big] \\
|
||||
&+ \mathbb{E}_{y \sim p_{data}(y)} \Big[\lVert G(F(y)) - y \rVert_1\Big] \\
|
||||
\\
|
||||
\mathcal{L}_{identity}(G, F)
|
||||
&= \mathbb{E}_{x \sim p_{data}(x)} \Big[\lVert F(x) - x \lVert_1\Big] \\
|
||||
&+ \mathbb{E}_{y \sim p_{data}(y)} \Big[\lVert G(y) - y \rVert_1\Big] \\
|
||||
\end{align}
|
||||
|
||||
$\mathcal{L}_{GAN}$ is the generative adversarial loss from the original
|
||||
GAN paper.
|
||||
|
||||
$\mathcal{L}_{cyc}$ is the cyclic loss, where we try to get $F(G(x))$ to be similar to $x$,
|
||||
and $G(F(y))$ to be similar to $y$.
|
||||
Basically if the two generators (transformations) are applied in series it should give back the
|
||||
original image.
|
||||
This is the main contribution of this paper.
|
||||
It trains the generators to generate an image of the other distribution that is similar to
|
||||
the original image.
|
||||
Without this loss $G(x)$ could generate anything that's from the distribution of $Y$.
|
||||
Now it needs to generate something from the distribution of $Y$ but still has properties of $x$,
|
||||
so that $F(G(x)$ can re-generate something like $x$.
|
||||
|
||||
$\mathcal{L}_{cyc}$ is the identity loss.
|
||||
This was used to encourage the mapping to preserve color composition between
|
||||
the input and the output.
|
||||
|
||||
To solve $$G^*, F^*$$,
|
||||
discriminators $D_X$ and $D_Y$ should **ascend** on the gradient,
|
||||
|
||||
\begin{align}
|
||||
\nabla_{\theta_{D_X, D_Y}} \frac{1}{m} \sum_{i=1}^m
|
||||
&\Bigg[
|
||||
\log D_Y\Big(y^{(i)}\Big) \\
|
||||
&+ \log \Big(1 - D_Y\Big(G\Big(x^{(i)}\Big)\Big)\Big) \\
|
||||
&+ \log D_X\Big(x^{(i)}\Big) \\
|
||||
& +\log\Big(1 - D_X\Big(F\Big(y^{(i)}\Big)\Big)\Big)
|
||||
\Bigg]
|
||||
\end{align}
|
||||
|
||||
That is descend on *negative* log-likelihood loss.
|
||||
|
||||
In order to stabilize the training the negative log- likelihood objective
|
||||
was replaced by a least-squared loss -
|
||||
the least-squared error of discriminator, labelling real images with 1,
|
||||
and generated images with 0.
|
||||
So we want to descend on the gradient,
|
||||
|
||||
\begin{align}
|
||||
\nabla_{\theta_{D_X, D_Y}} \frac{1}{m} \sum_{i=1}^m
|
||||
&\Bigg[
|
||||
\bigg(D_Y\Big(y^{(i)}\Big) - 1\bigg)^2 \\
|
||||
&+ D_Y\Big(G\Big(x^{(i)}\Big)\Big)^2 \\
|
||||
&+ \bigg(D_X\Big(x^{(i)}\Big) - 1\bigg)^2 \\
|
||||
&+ D_X\Big(F\Big(y^{(i)}\Big)\Big)^2
|
||||
\Bigg]
|
||||
\end{align}
|
||||
|
||||
We use least-squares for generators also.
|
||||
The generators should *descend* on the gradient,
|
||||
|
||||
\begin{align}
|
||||
\nabla_{\theta_{F, G}} \frac{1}{m} \sum_{i=1}^m
|
||||
&\Bigg[
|
||||
\bigg(D_Y\Big(G\Big(x^{(i)}\Big)\Big) - 1\bigg)^2 \\
|
||||
&+ \bigg(D_X\Big(F\Big(y^{(i)}\Big)\Big) - 1\bigg)^2 \\
|
||||
&+ \mathcal{L}_{cyc}(G, F)
|
||||
+ \mathcal{L}_{identity}(G, F)
|
||||
\Bigg]
|
||||
\end{align}
|
||||
|
||||
We use `generator_xy` for $G$ and `generator_yx` for $F$.
|
||||
We use `discriminator_x` for $D_X$ and `discriminator_y` for $D_Y$.
|
||||
"""
|
||||
|
||||
# Replay buffers to keep generated samples
|
||||
gen_x_buffer = ReplayBuffer()
|
||||
gen_y_buffer = ReplayBuffer()
|
||||
|
||||
# Loop through epochs
|
||||
for epoch in monit.loop(self.epochs):
|
||||
# Loop through the dataset
|
||||
for i, batch in monit.enum('Train', self.dataloader):
|
||||
# Move images to the device
|
||||
data_x, data_y = batch['x'].to(self.device), batch['y'].to(self.device)
|
||||
|
||||
# true labels equal to $1$
|
||||
true_labels = torch.ones(data_x.size(0), *self.discriminator_x.output_shape,
|
||||
device=self.device, requires_grad=False)
|
||||
# false labels equal to $0$
|
||||
false_labels = torch.zeros(data_x.size(0), *self.discriminator_x.output_shape,
|
||||
device=self.device, requires_grad=False)
|
||||
|
||||
# Train the generators.
|
||||
# This returns the generated images.
|
||||
gen_x, gen_y = self.optimize_generators(data_x, data_y, true_labels)
|
||||
|
||||
# Train discriminators
|
||||
self.optimize_discriminator(data_x, data_y,
|
||||
gen_x_buffer.push_and_pop(gen_x), gen_y_buffer.push_and_pop(gen_y),
|
||||
true_labels, false_labels)
|
||||
|
||||
# Save training statistics and increment the global step counter
|
||||
tracker.save()
|
||||
tracker.add_global_step(max(len(data_x), len(data_y)))
|
||||
|
||||
# Save images at intervals
|
||||
batches_done = epoch * len(self.dataloader) + i
|
||||
if batches_done % self.sample_interval == 0:
|
||||
# Sample images
|
||||
self.sample_images(batches_done)
|
||||
|
||||
# Update learning rates
|
||||
self.generator_lr_scheduler.step()
|
||||
self.discriminator_lr_scheduler.step()
|
||||
# New line
|
||||
tracker.new_line()
|
||||
|
||||
def optimize_generators(self, data_x: torch.Tensor, data_y: torch.Tensor, true_labels: torch.Tensor):
|
||||
"""
|
||||
### Optimize the generators with identity, gan and cycle losses.
|
||||
"""
|
||||
|
||||
# Change to training mode
|
||||
self.generator_xy.train()
|
||||
self.generator_yx.train()
|
||||
|
||||
# Identity loss
|
||||
# $$\lVert F(G(x^{(i)})) - x^{(i)} \lVert_1\
|
||||
# \lVert G(F(y^{(i)})) - y^{(i)} \rVert_1$$
|
||||
loss_identity = (self.identity_loss(self.generator_yx(data_x), data_x) +
|
||||
self.identity_loss(self.generator_xy(data_y), data_y))
|
||||
|
||||
# Generate images $G(x)$ and $F(y)$
|
||||
gen_y = self.generator_xy(data_x)
|
||||
gen_x = self.generator_yx(data_y)
|
||||
|
||||
# GAN loss
|
||||
# $$\bigg(D_Y\Big(G\Big(x^{(i)}\Big)\Big) - 1\bigg)^2
|
||||
# + \bigg(D_X\Big(F\Big(y^{(i)}\Big)\Big) - 1\bigg)^2$$
|
||||
loss_gan = (self.gan_loss(self.discriminator_y(gen_y), true_labels) +
|
||||
self.gan_loss(self.discriminator_x(gen_x), true_labels))
|
||||
|
||||
# Cycle loss
|
||||
# $$
|
||||
# \lVert F(G(x^{(i)})) - x^{(i)} \lVert_1 +
|
||||
# \lVert G(F(y^{(i)})) - y^{(i)} \rVert_1
|
||||
# $$
|
||||
loss_cycle = (self.cycle_loss(self.generator_yx(gen_y), data_x) +
|
||||
self.cycle_loss(self.generator_xy(gen_x), data_y))
|
||||
|
||||
# Total loss
|
||||
loss_generator = (loss_gan +
|
||||
self.cyclic_loss_coefficient * loss_cycle +
|
||||
self.identity_loss_coefficient * loss_identity)
|
||||
|
||||
# Take a step in the optimizer
|
||||
self.generator_optimizer.zero_grad()
|
||||
loss_generator.backward()
|
||||
self.generator_optimizer.step()
|
||||
|
||||
# Log losses
|
||||
tracker.add({'loss.generator': loss_generator,
|
||||
'loss.generator.cycle': loss_cycle,
|
||||
'loss.generator.gan': loss_gan,
|
||||
'loss.generator.identity': loss_identity})
|
||||
|
||||
# Return generated images
|
||||
return gen_x, gen_y
|
||||
|
||||
def optimize_discriminator(self, data_x: torch.Tensor, data_y: torch.Tensor,
|
||||
gen_x: torch.Tensor, gen_y: torch.Tensor,
|
||||
true_labels: torch.Tensor, false_labels: torch.Tensor):
|
||||
"""
|
||||
### Optimize the discriminators with gan loss.
|
||||
"""
|
||||
|
||||
# GAN Loss
|
||||
#
|
||||
# \begin{align}
|
||||
# \bigg(D_Y\Big(y ^ {(i)}\Big) - 1\bigg) ^ 2
|
||||
# + D_Y\Big(G\Big(x ^ {(i)}\Big)\Big) ^ 2 + \\
|
||||
# \bigg(D_X\Big(x ^ {(i)}\Big) - 1\bigg) ^ 2
|
||||
# + D_X\Big(F\Big(y ^ {(i)}\Big)\Big) ^ 2
|
||||
# \end{align}
|
||||
loss_discriminator = (self.gan_loss(self.discriminator_x(data_x), true_labels) +
|
||||
self.gan_loss(self.discriminator_x(gen_x), false_labels) +
|
||||
self.gan_loss(self.discriminator_y(data_y), true_labels) +
|
||||
self.gan_loss(self.discriminator_y(gen_y), false_labels))
|
||||
|
||||
# Take a step in the optimizer
|
||||
self.discriminator_optimizer.zero_grad()
|
||||
loss_discriminator.backward()
|
||||
self.discriminator_optimizer.step()
|
||||
|
||||
# Log losses
|
||||
tracker.add({'loss.discriminator': loss_discriminator})
|
||||
|
||||
|
||||
def train():
|
||||
"""
|
||||
## Train Cycle GAN
|
||||
"""
|
||||
# Create configurations
|
||||
conf = Configs()
|
||||
# Create an experiment
|
||||
experiment.create(name='cycle_gan')
|
||||
# Calculate configurations.
|
||||
# It will calculate `conf.run` and all other configs required by it.
|
||||
experiment.configs(conf, {'dataset_name': 'summer2winter_yosemite'})
|
||||
conf.initialize()
|
||||
|
||||
# Register models for saving and loading.
|
||||
# `get_modules` gives a dictionary of `nn.Modules` in `conf`.
|
||||
# You can also specify a custom dictionary of models.
|
||||
experiment.add_pytorch_models(get_modules(conf))
|
||||
# Start and watch the experiment
|
||||
with experiment.start():
|
||||
# Run the training
|
||||
conf.run()
|
||||
|
||||
|
||||
def plot_image(img: torch.Tensor):
|
||||
"""
|
||||
### Plot an image with matplotlib
|
||||
"""
|
||||
from matplotlib import pyplot as plt
|
||||
|
||||
# Move tensor to CPU
|
||||
img = img.cpu()
|
||||
# Get min and max values of the image for normalization
|
||||
img_min, img_max = img.min(), img.max()
|
||||
# Scale image values to be [0...1]
|
||||
img = (img - img_min) / (img_max - img_min + 1e-5)
|
||||
# We have to change the order of dimensions to HWC.
|
||||
img = img.permute(1, 2, 0)
|
||||
# Show Image
|
||||
plt.imshow(img)
|
||||
# We don't need axes
|
||||
plt.axis('off')
|
||||
# Display
|
||||
plt.show()
|
||||
|
||||
|
||||
def evaluate():
|
||||
"""
|
||||
## Evaluate trained Cycle GAN
|
||||
"""
|
||||
# Set the run UUID from the training run
|
||||
trained_run_uuid = 'f73c1164184711eb9190b74249275441'
|
||||
# Create configs object
|
||||
conf = Configs()
|
||||
# Create experiment
|
||||
experiment.create(name='cycle_gan_inference')
|
||||
# Load hyper parameters set for training
|
||||
conf_dict = experiment.load_configs(trained_run_uuid)
|
||||
# Calculate configurations. We specify the generators `'generator_xy', 'generator_yx'`
|
||||
# so that it only loads those and their dependencies.
|
||||
# Configs like `device` and `img_channels` will be calculated, since these are required by
|
||||
# `generator_xy` and `generator_yx`.
|
||||
#
|
||||
# If you want other parameters like `dataset_name` you should specify them here.
|
||||
# If you specify nothing, all the configurations will be calculated, including data loaders.
|
||||
# Calculation of configurations and their dependencies will happen when you call `experiment.start`
|
||||
experiment.configs(conf, conf_dict)
|
||||
conf.initialize()
|
||||
|
||||
# Register models for saving and loading.
|
||||
# `get_modules` gives a dictionary of `nn.Modules` in `conf`.
|
||||
# You can also specify a custom dictionary of models.
|
||||
experiment.add_pytorch_models(get_modules(conf))
|
||||
# Specify which run to load from.
|
||||
# Loading will actually happen when you call `experiment.start`
|
||||
experiment.load(trained_run_uuid)
|
||||
|
||||
# Start the experiment
|
||||
with experiment.start():
|
||||
# Image transformations
|
||||
transforms_ = [
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
|
||||
]
|
||||
|
||||
# Load your own data. Here we try the test set.
|
||||
# I was trying with Yosemite photos, they look awesome.
|
||||
# You can use `conf.dataset_name`, if you specified `dataset_name` as something you wanted to be calculated
|
||||
# in the call to `experiment.configs`
|
||||
dataset = ImageDataset(conf.dataset_name, transforms_, 'train')
|
||||
# Get an image from dataset
|
||||
x_image = dataset[10]['x']
|
||||
# Display the image
|
||||
plot_image(x_image)
|
||||
|
||||
# Evaluation mode
|
||||
conf.generator_xy.eval()
|
||||
conf.generator_yx.eval()
|
||||
|
||||
# We don't need gradients
|
||||
with torch.no_grad():
|
||||
# Add batch dimension and move to the device we use
|
||||
data = x_image.unsqueeze(0).to(conf.device)
|
||||
generated_y = conf.generator_xy(data)
|
||||
|
||||
# Display the generated image.
|
||||
plot_image(generated_y[0].cpu())
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
train()
|
||||
# evaluate()
|
||||
@@ -0,0 +1,226 @@
|
||||
{
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"name": "Cycle GAN",
|
||||
"provenance": [],
|
||||
"collapsed_sections": [],
|
||||
"toc_visible": true
|
||||
},
|
||||
"kernelspec": {
|
||||
"name": "python3",
|
||||
"display_name": "Python 3"
|
||||
},
|
||||
"accelerator": "GPU"
|
||||
},
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "AYV_dMVDxyc2"
|
||||
},
|
||||
"source": [
|
||||
"[](https://github.com/labmlai/annotated_deep_learning_paper_implementations)\n",
|
||||
"[](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/gan/cycle_gan/experiment.ipynb)\n",
|
||||
"\n",
|
||||
"## Cycle GAN\n",
|
||||
"\n",
|
||||
"This is an experiment training Cycle GAN model."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "AahG_i2y5tY9"
|
||||
},
|
||||
"source": [
|
||||
"Install the `labml-nn` package"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "ZCzmCrAIVg0L",
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"outputId": "2fe2685f-731c-4c47-854e-a4f00e485281"
|
||||
},
|
||||
"source": [
|
||||
"!pip install labml-nn"
|
||||
],
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "SE2VUQ6L5zxI"
|
||||
},
|
||||
"source": [
|
||||
"Imports"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "0hJXx_g0wS2C"
|
||||
},
|
||||
"source": [
|
||||
"from labml import experiment\n",
|
||||
"from labml.utils.pytorch import get_modules\n",
|
||||
"from labml_nn.gan.cycle_gan import Configs"
|
||||
],
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "Lpggo0wM6qb-"
|
||||
},
|
||||
"source": [
|
||||
"Create an experiment"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "bFcr9k-l4cAg"
|
||||
},
|
||||
"source": [
|
||||
"experiment.create(name=\"cycle_gan\")"
|
||||
],
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "-OnHLi626tJt"
|
||||
},
|
||||
"source": [
|
||||
"Initialize configurations"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "Piz0c5f44hRo"
|
||||
},
|
||||
"source": [
|
||||
"conf = Configs()"
|
||||
],
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "wwMzCqpD6vkL"
|
||||
},
|
||||
"source": [
|
||||
"Set experiment configurations and assign a configurations dictionary to override configurations"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 17
|
||||
},
|
||||
"id": "e6hmQhTw4nks",
|
||||
"outputId": "4be767af-0ebd-4c35-8da0-0e532495e037"
|
||||
},
|
||||
"source": [
|
||||
"experiment.configs(conf, {'dataset_name': 'summer2winter_yosemite'})"
|
||||
],
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "DHyNvXfnzeWQ"
|
||||
},
|
||||
"source": [
|
||||
"Initialize"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 85
|
||||
},
|
||||
"id": "59ZeTv5SzcVe",
|
||||
"outputId": "55f4af22-b6df-4335-e4fb-d6d675e69b4e"
|
||||
},
|
||||
"source": [
|
||||
"conf.initialize()"
|
||||
],
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "EvI7MtgJ61w5"
|
||||
},
|
||||
"source": [
|
||||
"Set PyTorch models for loading and saving"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "GDlt7dp-5ALt"
|
||||
},
|
||||
"source": [
|
||||
"experiment.add_pytorch_models(get_modules(conf))"
|
||||
],
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "KJZRf8527GxL"
|
||||
},
|
||||
"source": [
|
||||
"Start the experiment and run the training loop."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 649
|
||||
},
|
||||
"id": "aIAWo7Fw5DR8",
|
||||
"outputId": "e3b02247-8ff9-47b5-8f52-49c9e3b8377f"
|
||||
},
|
||||
"source": [
|
||||
"with experiment.start():\n",
|
||||
" conf.run()"
|
||||
],
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "oBXXlP2b7XZO"
|
||||
},
|
||||
"source": [
|
||||
""
|
||||
],
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,4 @@
|
||||
# [Cycle GAN](https://nn.labml.ai/gan/cycle_gan/index.html)
|
||||
|
||||
This is a [PyTorch](https://pytorch.org) implementation/tutorial of the paper
|
||||
[Unpaired Image-to-Image Translation using Cycle-Consistent Adversarial Networks](https://arxiv.org/abs/1703.10593).
|
||||
@@ -0,0 +1,119 @@
|
||||
"""
|
||||
---
|
||||
title: Deep Convolutional Generative Adversarial Networks (DCGAN)
|
||||
summary: A simple PyTorch implementation/tutorial of Deep Convolutional Generative Adversarial Networks (DCGAN).
|
||||
---
|
||||
|
||||
# Deep Convolutional Generative Adversarial Networks (DCGAN)
|
||||
|
||||
This is a [PyTorch](https://pytorch.org) implementation of paper
|
||||
[Unsupervised Representation Learning with Deep Convolutional Generative Adversarial Networks](https://arxiv.org/abs/1511.06434).
|
||||
|
||||
This implementation is based on the [PyTorch DCGAN Tutorial](https://pytorch.org/tutorials/beginner/dcgan_faces_tutorial.html).
|
||||
"""
|
||||
|
||||
import torch.nn as nn
|
||||
|
||||
from labml import experiment
|
||||
from labml.configs import calculate
|
||||
from labml_nn.gan.original.experiment import Configs
|
||||
|
||||
|
||||
class Generator(nn.Module):
|
||||
"""
|
||||
### Convolutional Generator Network
|
||||
|
||||
This is similar to the de-convolutional network used for CelebA faces,
|
||||
but modified for MNIST images.
|
||||
|
||||
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
# The input is $1 \times 1$ with 100 channels
|
||||
self.layers = nn.Sequential(
|
||||
# This gives $3 \times 3$ output
|
||||
nn.ConvTranspose2d(100, 1024, 3, 1, 0, bias=False),
|
||||
nn.BatchNorm2d(1024),
|
||||
nn.ReLU(True),
|
||||
# This gives $7 \times 7$
|
||||
nn.ConvTranspose2d(1024, 512, 3, 2, 0, bias=False),
|
||||
nn.BatchNorm2d(512),
|
||||
nn.ReLU(True),
|
||||
# This gives $14 \times 14$
|
||||
nn.ConvTranspose2d(512, 256, 4, 2, 1, bias=False),
|
||||
nn.BatchNorm2d(256),
|
||||
nn.ReLU(True),
|
||||
# This gives $28 \times 28$
|
||||
nn.ConvTranspose2d(256, 1, 4, 2, 1, bias=False),
|
||||
nn.Tanh()
|
||||
)
|
||||
|
||||
self.apply(_weights_init)
|
||||
|
||||
def forward(self, x):
|
||||
# Change from shape `[batch_size, 100]` to `[batch_size, 100, 1, 1]`
|
||||
x = x.unsqueeze(-1).unsqueeze(-1)
|
||||
x = self.layers(x)
|
||||
return x
|
||||
|
||||
|
||||
class Discriminator(nn.Module):
|
||||
"""
|
||||
### Convolutional Discriminator Network
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
# The input is $28 \times 28$ with one channel
|
||||
self.layers = nn.Sequential(
|
||||
# This gives $14 \times 14$
|
||||
nn.Conv2d(1, 256, 4, 2, 1, bias=False),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
# This gives $7 \times 7$
|
||||
nn.Conv2d(256, 512, 4, 2, 1, bias=False),
|
||||
nn.BatchNorm2d(512),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
# This gives $3 \times 3$
|
||||
nn.Conv2d(512, 1024, 3, 2, 0, bias=False),
|
||||
nn.BatchNorm2d(1024),
|
||||
nn.LeakyReLU(0.2, inplace=True),
|
||||
# This gives $1 \times 1$
|
||||
nn.Conv2d(1024, 1, 3, 1, 0, bias=False),
|
||||
)
|
||||
self.apply(_weights_init)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.layers(x)
|
||||
return x.view(x.shape[0], -1)
|
||||
|
||||
|
||||
def _weights_init(m):
|
||||
classname = m.__class__.__name__
|
||||
if classname.find('Conv') != -1:
|
||||
nn.init.normal_(m.weight.data, 0.0, 0.02)
|
||||
elif classname.find('BatchNorm') != -1:
|
||||
nn.init.normal_(m.weight.data, 1.0, 0.02)
|
||||
nn.init.constant_(m.bias.data, 0)
|
||||
|
||||
|
||||
# We import the [simple gan experiment](../original/experiment.html) and change the
|
||||
# generator and discriminator networks
|
||||
calculate(Configs.generator, 'cnn', lambda c: Generator().to(c.device))
|
||||
calculate(Configs.discriminator, 'cnn', lambda c: Discriminator().to(c.device))
|
||||
|
||||
|
||||
def main():
|
||||
conf = Configs()
|
||||
experiment.create(name='mnist_dcgan')
|
||||
experiment.configs(conf,
|
||||
{'discriminator': 'cnn',
|
||||
'generator': 'cnn',
|
||||
'label_smoothing': 0.01})
|
||||
with experiment.start():
|
||||
conf.run()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,184 @@
|
||||
{
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"name": "Cycle GAN",
|
||||
"provenance": [],
|
||||
"collapsed_sections": [],
|
||||
"toc_visible": true
|
||||
},
|
||||
"kernelspec": {
|
||||
"name": "python3",
|
||||
"language": "python",
|
||||
"display_name": "Python 3"
|
||||
},
|
||||
"accelerator": "GPU"
|
||||
},
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "AYV_dMVDxyc2"
|
||||
},
|
||||
"source": [
|
||||
"[](https://github.com/labmlai/annotated_deep_learning_paper_implementations)\n",
|
||||
"[](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/gan/dcgan/experiment.ipynb)\n",
|
||||
"\n",
|
||||
"## DCGAN\n",
|
||||
"\n",
|
||||
"This is an experiment training DCGAN model."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "AahG_i2y5tY9"
|
||||
},
|
||||
"source": [
|
||||
"Install the `labml-nn` package"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "ZCzmCrAIVg0L",
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"outputId": "2fe2685f-731c-4c47-854e-a4f00e485281"
|
||||
},
|
||||
"source": [
|
||||
"!pip install labml-nn"
|
||||
],
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "SE2VUQ6L5zxI"
|
||||
},
|
||||
"source": [
|
||||
"Imports"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "0hJXx_g0wS2C"
|
||||
},
|
||||
"source": [
|
||||
"from labml import experiment\n",
|
||||
"from labml_nn.gan.dcgan import Configs"
|
||||
],
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "Lpggo0wM6qb-"
|
||||
},
|
||||
"source": [
|
||||
"Create an experiment"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "bFcr9k-l4cAg"
|
||||
},
|
||||
"source": [
|
||||
"experiment.create(name=\"mnist_dcgan\")"
|
||||
],
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "-OnHLi626tJt"
|
||||
},
|
||||
"source": [
|
||||
"Initialize configurations"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "Piz0c5f44hRo"
|
||||
},
|
||||
"source": [
|
||||
"conf = Configs()"
|
||||
],
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "wwMzCqpD6vkL"
|
||||
},
|
||||
"source": [
|
||||
"Set experiment configurations and assign a configurations dictionary to override configurations"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 17
|
||||
},
|
||||
"id": "e6hmQhTw4nks",
|
||||
"outputId": "4be767af-0ebd-4c35-8da0-0e532495e037"
|
||||
},
|
||||
"source": [
|
||||
"experiment.configs(conf,\n",
|
||||
" {'discriminator': 'cnn',\n",
|
||||
" 'generator': 'cnn',\n",
|
||||
" 'label_smoothing': 0.01})"
|
||||
],
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "KJZRf8527GxL"
|
||||
},
|
||||
"source": [
|
||||
"Start the experiment and run the training loop."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 649
|
||||
},
|
||||
"id": "aIAWo7Fw5DR8",
|
||||
"outputId": "e3b02247-8ff9-47b5-8f52-49c9e3b8377f"
|
||||
},
|
||||
"source": [
|
||||
"with experiment.start():\n",
|
||||
" conf.run()"
|
||||
],
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "oBXXlP2b7XZO"
|
||||
},
|
||||
"source": [
|
||||
""
|
||||
],
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,4 @@
|
||||
# [Deep Convolutional Generative Adversarial Networks - DCGAN](https://nn.labml.ai/gan/dcgan/index.html)
|
||||
|
||||
This is a [PyTorch](https://pytorch.org) implementation of paper
|
||||
[Unsupervised Representation Learning with Deep Convolutional Generative Adversarial Networks](https://arxiv.org/abs/1511.06434).
|
||||
@@ -0,0 +1,125 @@
|
||||
"""
|
||||
---
|
||||
title: Generative Adversarial Networks (GAN)
|
||||
summary: A simple PyTorch implementation/tutorial of Generative Adversarial Networks (GAN) loss functions.
|
||||
---
|
||||
|
||||
# Generative Adversarial Networks (GAN)
|
||||
|
||||
This is an implementation of
|
||||
[Generative Adversarial Networks](https://arxiv.org/abs/1406.2661).
|
||||
|
||||
The generator, $G(\pmb{z}; \theta_g)$ generates samples that match the
|
||||
distribution of data, while the discriminator, $D(\pmb{x}; \theta_g)$
|
||||
gives the probability that $\pmb{x}$ came from data rather than $G$.
|
||||
|
||||
We train $D$ and $G$ simultaneously on a two-player min-max game with value
|
||||
function $V(G, D)$.
|
||||
|
||||
$$\min_G \max_D V(D, G) =
|
||||
\mathop{\mathbb{E}}_{\pmb{x} \sim p_{data}(\pmb{x})}
|
||||
\big[\log D(\pmb{x})\big] +
|
||||
\mathop{\mathbb{E}}_{\pmb{z} \sim p_{\pmb{z}}(\pmb{z})}
|
||||
\big[\log (1 - D(G(\pmb{z}))\big]
|
||||
$$
|
||||
|
||||
$p_{data}(\pmb{x})$ is the probability distribution over data,
|
||||
whilst $p_{\pmb{z}}(\pmb{z})$ probability distribution of $\pmb{z}$, which is set to
|
||||
gaussian noise.
|
||||
|
||||
This file defines the loss functions. [Here](experiment.html) is an MNIST example
|
||||
with two multilayer perceptron for the generator and discriminator.
|
||||
"""
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.utils.data
|
||||
import torch.utils.data
|
||||
|
||||
|
||||
class DiscriminatorLogitsLoss(nn.Module):
|
||||
"""
|
||||
## Discriminator Loss
|
||||
|
||||
Discriminator should **ascend** on the gradient,
|
||||
|
||||
$$\nabla_{\theta_d} \frac{1}{m} \sum_{i=1}^m \Bigg[
|
||||
\log D\Big(\pmb{x}^{(i)}\Big) +
|
||||
\log \Big(1 - D\Big(G\Big(\pmb{z}^{(i)}\Big)\Big)\Big)
|
||||
\Bigg]$$
|
||||
|
||||
$m$ is the mini-batch size and $(i)$ is used to index samples in the mini-batch.
|
||||
$\pmb{x}$ are samples from $p_{data}$ and $\pmb{z}$ are samples from $p_z$.
|
||||
"""
|
||||
|
||||
def __init__(self, smoothing: float = 0.2):
|
||||
super().__init__()
|
||||
# We use PyTorch Binary Cross Entropy Loss, which is
|
||||
# $-\sum\Big[y \log(\hat{y}) + (1 - y) \log(1 - \hat{y})\Big]$,
|
||||
# where $y$ are the labels and $\hat{y}$ are the predictions.
|
||||
# *Note the negative sign*.
|
||||
# We use labels equal to $1$ for $\pmb{x}$ from $p_{data}$
|
||||
# and labels equal to $0$ for $\pmb{x}$ from $p_{G}.$
|
||||
# Then descending on the sum of these is the same as ascending on
|
||||
# the above gradient.
|
||||
#
|
||||
# `BCEWithLogitsLoss` combines softmax and binary cross entropy loss.
|
||||
self.loss_true = nn.BCEWithLogitsLoss()
|
||||
self.loss_false = nn.BCEWithLogitsLoss()
|
||||
|
||||
# We use label smoothing because it seems to work better in some cases
|
||||
self.smoothing = smoothing
|
||||
|
||||
# Labels are registered as buffered and persistence is set to `False`.
|
||||
self.register_buffer('labels_true', _create_labels(256, 1.0 - smoothing, 1.0), False)
|
||||
self.register_buffer('labels_false', _create_labels(256, 0.0, smoothing), False)
|
||||
|
||||
def forward(self, logits_true: torch.Tensor, logits_false: torch.Tensor):
|
||||
"""
|
||||
`logits_true` are logits from $D(\pmb{x}^{(i)})$ and
|
||||
`logits_false` are logits from $D(G(\pmb{z}^{(i)}))$
|
||||
"""
|
||||
if len(logits_true) > len(self.labels_true):
|
||||
self.register_buffer("labels_true",
|
||||
_create_labels(len(logits_true), 1.0 - self.smoothing, 1.0, logits_true.device), False)
|
||||
if len(logits_false) > len(self.labels_false):
|
||||
self.register_buffer("labels_false",
|
||||
_create_labels(len(logits_false), 0.0, self.smoothing, logits_false.device), False)
|
||||
|
||||
return (self.loss_true(logits_true, self.labels_true[:len(logits_true)]),
|
||||
self.loss_false(logits_false, self.labels_false[:len(logits_false)]))
|
||||
|
||||
|
||||
class GeneratorLogitsLoss(nn.Module):
|
||||
"""
|
||||
## Generator Loss
|
||||
|
||||
Generator should **descend** on the gradient,
|
||||
|
||||
$$\nabla_{\theta_g} \frac{1}{m} \sum_{i=1}^m \Bigg[
|
||||
\log \Big(1 - D\Big(G\Big(\pmb{z}^{(i)}\Big)\Big)\Big)
|
||||
\Bigg]$$
|
||||
"""
|
||||
|
||||
def __init__(self, smoothing: float = 0.2):
|
||||
super().__init__()
|
||||
self.loss_true = nn.BCEWithLogitsLoss()
|
||||
self.smoothing = smoothing
|
||||
# We use labels equal to $1$ for $\pmb{x}$ from $p_{G}.$
|
||||
# Then descending on this loss is the same as descending on
|
||||
# the above gradient.
|
||||
self.register_buffer('fake_labels', _create_labels(256, 1.0 - smoothing, 1.0), False)
|
||||
|
||||
def forward(self, logits: torch.Tensor):
|
||||
if len(logits) > len(self.fake_labels):
|
||||
self.register_buffer("fake_labels",
|
||||
_create_labels(len(logits), 1.0 - self.smoothing, 1.0, logits.device), False)
|
||||
|
||||
return self.loss_true(logits, self.fake_labels[:len(logits)])
|
||||
|
||||
|
||||
def _create_labels(n: int, r1: float, r2: float, device: torch.device = None):
|
||||
"""
|
||||
Create smoothed labels
|
||||
"""
|
||||
return torch.empty(n, 1, requires_grad=False, device=device).uniform_(r1, r2)
|
||||
@@ -0,0 +1,183 @@
|
||||
{
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"name": "Cycle GAN",
|
||||
"provenance": [],
|
||||
"collapsed_sections": [],
|
||||
"toc_visible": true
|
||||
},
|
||||
"kernelspec": {
|
||||
"name": "python3",
|
||||
"language": "python",
|
||||
"display_name": "Python 3"
|
||||
},
|
||||
"accelerator": "GPU"
|
||||
},
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "AYV_dMVDxyc2"
|
||||
},
|
||||
"source": [
|
||||
"[](https://github.com/labmlai/annotated_deep_learning_paper_implementations)\n",
|
||||
"[](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/gan/original/experiment.ipynb)\n",
|
||||
"\n",
|
||||
"## DCGAN\n",
|
||||
"\n",
|
||||
"This is an experiment training DCGAN model."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "AahG_i2y5tY9"
|
||||
},
|
||||
"source": [
|
||||
"Install the `labml-nn` package"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "ZCzmCrAIVg0L",
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"outputId": "2fe2685f-731c-4c47-854e-a4f00e485281"
|
||||
},
|
||||
"source": [
|
||||
"!pip install labml-nn"
|
||||
],
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "SE2VUQ6L5zxI"
|
||||
},
|
||||
"source": [
|
||||
"Imports"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "0hJXx_g0wS2C"
|
||||
},
|
||||
"source": [
|
||||
"\n",
|
||||
"from labml import experiment\n",
|
||||
"from labml_nn.gan.original.experiment import Configs"
|
||||
],
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "Lpggo0wM6qb-"
|
||||
},
|
||||
"source": [
|
||||
"Create an experiment"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "bFcr9k-l4cAg"
|
||||
},
|
||||
"source": [
|
||||
"experiment.create(name=\"mnist_gan\")"
|
||||
],
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "-OnHLi626tJt"
|
||||
},
|
||||
"source": [
|
||||
"Initialize configurations"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "Piz0c5f44hRo"
|
||||
},
|
||||
"source": [
|
||||
"conf = Configs()"
|
||||
],
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "wwMzCqpD6vkL"
|
||||
},
|
||||
"source": [
|
||||
"Set experiment configurations and assign a configurations dictionary to override configurations"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 17
|
||||
},
|
||||
"id": "e6hmQhTw4nks",
|
||||
"outputId": "4be767af-0ebd-4c35-8da0-0e532495e037"
|
||||
},
|
||||
"source": [
|
||||
"experiment.configs(conf,\n",
|
||||
" {'label_smoothing': 0.01})"
|
||||
],
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "KJZRf8527GxL"
|
||||
},
|
||||
"source": [
|
||||
"Start the experiment and run the training loop."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 649
|
||||
},
|
||||
"id": "aIAWo7Fw5DR8",
|
||||
"outputId": "e3b02247-8ff9-47b5-8f52-49c9e3b8377f"
|
||||
},
|
||||
"source": [
|
||||
"with experiment.start():\n",
|
||||
" conf.run()"
|
||||
],
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "oBXXlP2b7XZO"
|
||||
},
|
||||
"source": [
|
||||
""
|
||||
],
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,250 @@
|
||||
"""
|
||||
---
|
||||
title: Generative Adversarial Networks experiment with MNIST
|
||||
summary: This experiment generates MNIST images using multi-layer perceptron.
|
||||
---
|
||||
|
||||
# Generative Adversarial Networks experiment with MNIST
|
||||
"""
|
||||
|
||||
from typing import Any
|
||||
|
||||
from torchvision import transforms
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.utils.data
|
||||
from labml import tracker, monit, experiment
|
||||
from labml.configs import option, calculate
|
||||
from labml_nn.gan.original import DiscriminatorLogitsLoss, GeneratorLogitsLoss
|
||||
from labml_nn.helpers.datasets import MNISTConfigs
|
||||
from labml_nn.helpers.device import DeviceConfigs
|
||||
from labml_nn.helpers.optimizer import OptimizerConfigs
|
||||
from labml_nn.helpers.trainer import TrainValidConfigs, BatchIndex
|
||||
|
||||
|
||||
def weights_init(m):
|
||||
classname = m.__class__.__name__
|
||||
if classname.find('Linear') != -1:
|
||||
nn.init.normal_(m.weight.data, 0.0, 0.02)
|
||||
elif classname.find('BatchNorm') != -1:
|
||||
nn.init.normal_(m.weight.data, 1.0, 0.02)
|
||||
nn.init.constant_(m.bias.data, 0)
|
||||
|
||||
|
||||
class Generator(nn.Module):
|
||||
"""
|
||||
### Simple MLP Generator
|
||||
|
||||
This has three linear layers of increasing size with `LeakyReLU` activations.
|
||||
The final layer has a $tanh$ activation.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
layer_sizes = [256, 512, 1024]
|
||||
layers = []
|
||||
d_prev = 100
|
||||
for size in layer_sizes:
|
||||
layers = layers + [nn.Linear(d_prev, size), nn.LeakyReLU(0.2)]
|
||||
d_prev = size
|
||||
|
||||
self.layers = nn.Sequential(*layers, nn.Linear(d_prev, 28 * 28), nn.Tanh())
|
||||
|
||||
self.apply(weights_init)
|
||||
|
||||
def forward(self, x):
|
||||
return self.layers(x).view(x.shape[0], 1, 28, 28)
|
||||
|
||||
|
||||
class Discriminator(nn.Module):
|
||||
"""
|
||||
### Simple MLP Discriminator
|
||||
|
||||
This has three linear layers of decreasing size with `LeakyReLU` activations.
|
||||
The final layer has a single output that gives the logit of whether input
|
||||
is real or fake. You can get the probability by calculating the sigmoid of it.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
layer_sizes = [1024, 512, 256]
|
||||
layers = []
|
||||
d_prev = 28 * 28
|
||||
for size in layer_sizes:
|
||||
layers = layers + [nn.Linear(d_prev, size), nn.LeakyReLU(0.2)]
|
||||
d_prev = size
|
||||
|
||||
self.layers = nn.Sequential(*layers, nn.Linear(d_prev, 1))
|
||||
self.apply(weights_init)
|
||||
|
||||
def forward(self, x):
|
||||
return self.layers(x.view(x.shape[0], -1))
|
||||
|
||||
|
||||
class Configs(MNISTConfigs, TrainValidConfigs):
|
||||
"""
|
||||
## Configurations
|
||||
|
||||
This extends MNIST configurations to get the data loaders and Training and validation loop
|
||||
configurations to simplify our implementation.
|
||||
"""
|
||||
|
||||
device: torch.device = DeviceConfigs()
|
||||
dataset_transforms = 'mnist_gan_transforms'
|
||||
epochs: int = 10
|
||||
|
||||
is_save_models = True
|
||||
discriminator: nn.Module = 'mlp'
|
||||
generator: nn.Module = 'mlp'
|
||||
generator_optimizer: torch.optim.Adam
|
||||
discriminator_optimizer: torch.optim.Adam
|
||||
generator_loss: GeneratorLogitsLoss = 'original'
|
||||
discriminator_loss: DiscriminatorLogitsLoss = 'original'
|
||||
label_smoothing: float = 0.2
|
||||
discriminator_k: int = 1
|
||||
|
||||
def init(self):
|
||||
"""
|
||||
Initializations
|
||||
"""
|
||||
self.state_modules = []
|
||||
|
||||
tracker.set_scalar("loss.generator.*", True)
|
||||
tracker.set_scalar("loss.discriminator.*", True)
|
||||
tracker.set_image("generated", True, 1 / 100)
|
||||
|
||||
def sample_z(self, batch_size: int):
|
||||
"""
|
||||
$$z \sim p(z)$$
|
||||
"""
|
||||
return torch.randn(batch_size, 100, device=self.device)
|
||||
|
||||
def step(self, batch: Any, batch_idx: BatchIndex):
|
||||
"""
|
||||
Take a training step
|
||||
"""
|
||||
|
||||
# Set model states
|
||||
self.generator.train(self.mode.is_train)
|
||||
self.discriminator.train(self.mode.is_train)
|
||||
|
||||
# Get MNIST images
|
||||
data = batch[0].to(self.device)
|
||||
|
||||
# Increment step in training mode
|
||||
if self.mode.is_train:
|
||||
tracker.add_global_step(len(data))
|
||||
|
||||
# Train the discriminator
|
||||
with monit.section("discriminator"):
|
||||
# Get discriminator loss
|
||||
loss = self.calc_discriminator_loss(data)
|
||||
|
||||
# Train
|
||||
if self.mode.is_train:
|
||||
self.discriminator_optimizer.zero_grad()
|
||||
loss.backward()
|
||||
if batch_idx.is_last:
|
||||
tracker.add('discriminator', self.discriminator)
|
||||
self.discriminator_optimizer.step()
|
||||
|
||||
# Train the generator once in every `discriminator_k`
|
||||
if batch_idx.is_interval(self.discriminator_k):
|
||||
with monit.section("generator"):
|
||||
loss = self.calc_generator_loss(data.shape[0])
|
||||
|
||||
# Train
|
||||
if self.mode.is_train:
|
||||
self.generator_optimizer.zero_grad()
|
||||
loss.backward()
|
||||
if batch_idx.is_last:
|
||||
tracker.add('generator', self.generator)
|
||||
self.generator_optimizer.step()
|
||||
|
||||
tracker.save()
|
||||
|
||||
def calc_discriminator_loss(self, data):
|
||||
"""
|
||||
Calculate discriminator loss
|
||||
"""
|
||||
latent = self.sample_z(data.shape[0])
|
||||
logits_true = self.discriminator(data)
|
||||
logits_false = self.discriminator(self.generator(latent).detach())
|
||||
loss_true, loss_false = self.discriminator_loss(logits_true, logits_false)
|
||||
loss = loss_true + loss_false
|
||||
|
||||
# Log stuff
|
||||
tracker.add("loss.discriminator.true.", loss_true)
|
||||
tracker.add("loss.discriminator.false.", loss_false)
|
||||
tracker.add("loss.discriminator.", loss)
|
||||
|
||||
return loss
|
||||
|
||||
def calc_generator_loss(self, batch_size: int):
|
||||
"""
|
||||
Calculate generator loss
|
||||
"""
|
||||
latent = self.sample_z(batch_size)
|
||||
generated_images = self.generator(latent)
|
||||
logits = self.discriminator(generated_images)
|
||||
loss = self.generator_loss(logits)
|
||||
|
||||
# Log stuff
|
||||
tracker.add('generated', generated_images[0:6])
|
||||
tracker.add("loss.generator.", loss)
|
||||
|
||||
return loss
|
||||
|
||||
|
||||
@option(Configs.dataset_transforms)
|
||||
def mnist_gan_transforms():
|
||||
return transforms.Compose([
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize((0.5,), (0.5,))
|
||||
])
|
||||
|
||||
|
||||
@option(Configs.discriminator_optimizer)
|
||||
def _discriminator_optimizer(c: Configs):
|
||||
opt_conf = OptimizerConfigs()
|
||||
opt_conf.optimizer = 'Adam'
|
||||
opt_conf.parameters = c.discriminator.parameters()
|
||||
opt_conf.learning_rate = 2.5e-4
|
||||
# Setting exponent decay rate for first moment of gradient,
|
||||
# $\beta_1$ to `0.5` is important.
|
||||
# Default of `0.9` fails.
|
||||
opt_conf.betas = (0.5, 0.999)
|
||||
return opt_conf
|
||||
|
||||
|
||||
@option(Configs.generator_optimizer)
|
||||
def _generator_optimizer(c: Configs):
|
||||
opt_conf = OptimizerConfigs()
|
||||
opt_conf.optimizer = 'Adam'
|
||||
opt_conf.parameters = c.generator.parameters()
|
||||
opt_conf.learning_rate = 2.5e-4
|
||||
# Setting exponent decay rate for first moment of gradient,
|
||||
# $\beta_1$ to `0.5` is important.
|
||||
# Default of `0.9` fails.
|
||||
opt_conf.betas = (0.5, 0.999)
|
||||
return opt_conf
|
||||
|
||||
|
||||
calculate(Configs.generator, 'mlp', lambda c: Generator().to(c.device))
|
||||
calculate(Configs.discriminator, 'mlp', lambda c: Discriminator().to(c.device))
|
||||
calculate(Configs.generator_loss, 'original', lambda c: GeneratorLogitsLoss(c.label_smoothing).to(c.device))
|
||||
calculate(Configs.discriminator_loss, 'original', lambda c: DiscriminatorLogitsLoss(c.label_smoothing).to(c.device))
|
||||
|
||||
|
||||
def main():
|
||||
conf = Configs()
|
||||
experiment.create(name='mnist_gan', comment='test')
|
||||
experiment.configs(conf,
|
||||
{'label_smoothing': 0.01})
|
||||
with experiment.start():
|
||||
conf.run()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,4 @@
|
||||
# [Generative Adversarial Networks - GAN](https://nn.labml.ai/gan/original/index.html)
|
||||
|
||||
This is an annotated implementation of
|
||||
[Generative Adversarial Networks](https://arxiv.org/abs/1406.2661).
|
||||
@@ -0,0 +1,965 @@
|
||||
"""
|
||||
---
|
||||
title: StyleGAN 2
|
||||
summary: >
|
||||
An annotated PyTorch implementation of StyleGAN2.
|
||||
---
|
||||
|
||||
# StyleGAN 2
|
||||
|
||||
This is a [PyTorch](https://pytorch.org) implementation of the paper
|
||||
[Analyzing and Improving the Image Quality of StyleGAN](https://arxiv.org/abs/1912.04958)
|
||||
which introduces **StyleGAN 2**.
|
||||
StyleGAN 2 is an improvement over **StyleGAN** from the paper
|
||||
[A Style-Based Generator Architecture for Generative Adversarial Networks](https://arxiv.org/abs/1812.04948).
|
||||
And StyleGAN is based on **Progressive GAN** from the paper
|
||||
[Progressive Growing of GANs for Improved Quality, Stability, and Variation](https://arxiv.org/abs/1710.10196).
|
||||
All three papers are from the same authors from [NVIDIA AI](https://twitter.com/NVIDIAAI).
|
||||
|
||||
*Our implementation is a minimalistic StyleGAN 2 model training code.
|
||||
Only single GPU training is supported to keep the implementation simple.
|
||||
We managed to shrink it to keep it at less than 500 lines of code, including the training loop.*
|
||||
|
||||
**🏃 Here's the training code: [`experiment.py`](experiment.html).**
|
||||
|
||||
|
||||
|
||||
---*These are $64 \times 64$ images generated after training for about 80K steps.*---
|
||||
|
||||
|
||||
We'll first introduce the three papers at a high level.
|
||||
|
||||
## Generative Adversarial Networks
|
||||
|
||||
Generative adversarial networks have two components; the generator and the discriminator.
|
||||
The generator network takes a random latent vector ($z \in \mathcal{Z}$)
|
||||
and tries to generate a realistic image.
|
||||
The discriminator network tries to differentiate the real images from generated images.
|
||||
When we train the two networks together the generator starts generating images indistinguishable from real images.
|
||||
|
||||
## Progressive GAN
|
||||
|
||||
Progressive GAN generates high-resolution images ($1080 \times 1080$) of size.
|
||||
It does so by *progressively* increasing the image size.
|
||||
First, it trains a network that produces a $4 \times 4$ image, then $8 \times 8$ ,
|
||||
then an $16 \times 16$ image, and so on up to the desired image resolution.
|
||||
|
||||
At each resolution, the generator network produces an image in latent space which is converted into RGB,
|
||||
with a $1 \times 1$ convolution.
|
||||
When we progress from a lower resolution to a higher resolution
|
||||
(say from $4 \times 4$ to $8 \times 8$ ) we scale the latent image by $2\times$
|
||||
and add a new block (two $3 \times 3$ convolution layers)
|
||||
and a new $1 \times 1$ layer to get RGB.
|
||||
The transition is done smoothly by adding a residual connection to
|
||||
the $2\times$ scaled $4 \times 4$ RGB image.
|
||||
The weight of this residual connection is slowly reduced, to let the new block take over.
|
||||
|
||||
The discriminator is a mirror image of the generator network.
|
||||
The progressive growth of the discriminator is done similarly.
|
||||
|
||||
|
||||
|
||||
---*$2\times$ and $0.5\times$ denote feature map resolution scaling and scaling.
|
||||
$4\times4$, $8\times4$, ... denote feature map resolution at the generator or discriminator block.
|
||||
Each discriminator and generator block consists of 2 convolution layers with leaky ReLU activations.*---
|
||||
|
||||
They use **minibatch standard deviation** to increase variation and
|
||||
**equalized learning rate** which we discussed below in the implementation.
|
||||
They also use **pixel-wise normalization** where at each pixel the feature vector is normalized.
|
||||
They apply this to all the convolution layer outputs (except RGB).
|
||||
|
||||
|
||||
## StyleGAN
|
||||
|
||||
StyleGAN improves the generator of Progressive GAN keeping the discriminator architecture the same.
|
||||
|
||||
#### Mapping Network
|
||||
|
||||
It maps the random latent vector ($z \in \mathcal{Z}$)
|
||||
into a different latent space ($w \in \mathcal{W}$),
|
||||
with an 8-layer neural network.
|
||||
This gives an intermediate latent space $\mathcal{W}$
|
||||
where the factors of variations are more linear (disentangled).
|
||||
|
||||
#### AdaIN
|
||||
|
||||
Then $w$ is transformed into two vectors (**styles**) per layer,
|
||||
$i$, $y_i = (y_{s,i}, y_{b,i}) = f_{A_i}(w)$ and used for scaling and shifting (biasing)
|
||||
in each layer with $\text{AdaIN}$ operator (normalize and scale):
|
||||
$$\text{AdaIN}(x_i, y_i) = y_{s, i} \frac{x_i - \mu(x_i)}{\sigma(x_i)} + y_{b,i}$$
|
||||
|
||||
#### Style Mixing
|
||||
|
||||
To prevent the generator from assuming adjacent styles are correlated,
|
||||
they randomly use different styles for different blocks.
|
||||
That is, they sample two latent vectors $(z_1, z_2)$ and corresponding $(w_1, w_2)$ and
|
||||
use $w_1$ based styles for some blocks and $w_2$ based styles for some blacks randomly.
|
||||
|
||||
#### Stochastic Variation
|
||||
|
||||
Noise is made available to each block which helps the generator create more realistic images.
|
||||
Noise is scaled per channel by a learned weight.
|
||||
|
||||
#### Bilinear Up and Down Sampling
|
||||
|
||||
All the up and down-sampling operations are accompanied by bilinear smoothing.
|
||||
|
||||
|
||||
|
||||
---*$A$ denotes a linear layer.
|
||||
$B$ denotes a broadcast and scaling operation (noise is a single channel).
|
||||
StyleGAN also uses progressive growing like Progressive GAN.*---
|
||||
|
||||
## StyleGAN 2
|
||||
|
||||
StyleGAN 2 changes both the generator and the discriminator of StyleGAN.
|
||||
|
||||
#### Weight Modulation and Demodulation
|
||||
|
||||
They remove the $\text{AdaIN}$ operator and replace it with
|
||||
the weight modulation and demodulation step.
|
||||
This is supposed to improve what they call droplet artifacts that are present in generated images,
|
||||
which are caused by the normalization in $\text{AdaIN}$ operator.
|
||||
Style vector per layer is calculated from $w_i \in \mathcal{W}$ as $s_i = f_{A_i}(w_i)$.
|
||||
|
||||
Then the convolution weights $w$ are modulated as follows.
|
||||
($w$ here on refers to weights not intermediate latent space,
|
||||
we are sticking to the same notation as the paper.)
|
||||
|
||||
$$w'_{i, j, k} = s_i \cdot w_{i, j, k}$$
|
||||
Then it's demodulated by normalizing,
|
||||
$$w''_{i,j,k} = \frac{w'_{i,j,k}}{\sqrt{\sum_{i,k}{w'_{i, j, k}}^2 + \epsilon}}$$
|
||||
where $i$ is the input channel, $j$ is the output channel, and $k$ is the kernel index.
|
||||
|
||||
#### Path Length Regularization
|
||||
|
||||
Path length regularization encourages a fixed-size step in $\mathcal{W}$ to result in a non-zero,
|
||||
fixed-magnitude change in the generated image.
|
||||
|
||||
#### No Progressive Growing
|
||||
|
||||
StyleGAN2 uses residual connections (with down-sampling) in the discriminator and skip connections
|
||||
in the generator with up-sampling
|
||||
(the RGB outputs from each layer are added - no residual connections in feature maps).
|
||||
They show that with experiments that the contribution of low-resolution layers is higher
|
||||
at beginning of the training and then high-resolution layers take over.
|
||||
"""
|
||||
|
||||
import math
|
||||
from typing import Tuple, Optional, List
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import torch.utils.data
|
||||
from torch import nn
|
||||
|
||||
|
||||
class MappingNetwork(nn.Module):
|
||||
"""
|
||||
<a id="mapping_network"></a>
|
||||
|
||||
## Mapping Network
|
||||
|
||||
|
||||
|
||||
This is an MLP with 8 linear layers.
|
||||
The mapping network maps the latent vector $z \in \mathcal{W}$
|
||||
to an intermediate latent space $w \in \mathcal{W}$.
|
||||
$\mathcal{W}$ space will be disentangled from the image space
|
||||
where the factors of variation become more linear.
|
||||
"""
|
||||
|
||||
def __init__(self, features: int, n_layers: int):
|
||||
"""
|
||||
* `features` is the number of features in $z$ and $w$
|
||||
* `n_layers` is the number of layers in the mapping network.
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
# Create the MLP
|
||||
layers = []
|
||||
for i in range(n_layers):
|
||||
# [Equalized learning-rate linear layers](#equalized_linear)
|
||||
layers.append(EqualizedLinear(features, features))
|
||||
# Leaky Relu
|
||||
layers.append(nn.LeakyReLU(negative_slope=0.2, inplace=True))
|
||||
|
||||
self.net = nn.Sequential(*layers)
|
||||
|
||||
def forward(self, z: torch.Tensor):
|
||||
# Normalize $z$
|
||||
z = F.normalize(z, dim=1)
|
||||
# Map $z$ to $w$
|
||||
return self.net(z)
|
||||
|
||||
|
||||
class Generator(nn.Module):
|
||||
"""
|
||||
<a id="generator"></a>
|
||||
|
||||
## StyleGAN2 Generator
|
||||
|
||||
|
||||
|
||||
---*$A$ denotes a linear layer.
|
||||
$B$ denotes a broadcast and scaling operation (noise is a single channel).
|
||||
[`toRGB`](#to_rgb) also has a style modulation which is not shown in the diagram to keep it simple.*---
|
||||
|
||||
The generator starts with a learned constant.
|
||||
Then it has a series of blocks. The feature map resolution is doubled at each block
|
||||
Each block outputs an RGB image and they are scaled up and summed to get the final RGB image.
|
||||
"""
|
||||
|
||||
def __init__(self, log_resolution: int, d_latent: int, n_features: int = 32, max_features: int = 512):
|
||||
"""
|
||||
* `log_resolution` is the $\log_2$ of image resolution
|
||||
* `d_latent` is the dimensionality of $w$
|
||||
* `n_features` number of features in the convolution layer at the highest resolution (final block)
|
||||
* `max_features` maximum number of features in any generator block
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
# Calculate the number of features for each block
|
||||
#
|
||||
# Something like `[512, 512, 256, 128, 64, 32]`
|
||||
features = [min(max_features, n_features * (2 ** i)) for i in range(log_resolution - 2, -1, -1)]
|
||||
# Number of generator blocks
|
||||
self.n_blocks = len(features)
|
||||
|
||||
# Trainable $4 \times 4$ constant
|
||||
self.initial_constant = nn.Parameter(torch.randn((1, features[0], 4, 4)))
|
||||
|
||||
# First style block for $4 \times 4$ resolution and layer to get RGB
|
||||
self.style_block = StyleBlock(d_latent, features[0], features[0])
|
||||
self.to_rgb = ToRGB(d_latent, features[0])
|
||||
|
||||
# Generator blocks
|
||||
blocks = [GeneratorBlock(d_latent, features[i - 1], features[i]) for i in range(1, self.n_blocks)]
|
||||
self.blocks = nn.ModuleList(blocks)
|
||||
|
||||
# $2 \times$ up sampling layer. The feature space is up sampled
|
||||
# at each block
|
||||
self.up_sample = UpSample()
|
||||
|
||||
def forward(self, w: torch.Tensor, input_noise: List[Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]]):
|
||||
"""
|
||||
* `w` is $w$. In order to mix-styles (use different $w$ for different layers), we provide a separate
|
||||
$w$ for each [generator block](#generator_block). It has shape `[n_blocks, batch_size, d_latent]`.
|
||||
* `input_noise` is the noise for each block.
|
||||
It's a list of pairs of noise sensors because each block (except the initial) has two noise inputs
|
||||
after each convolution layer (see the diagram).
|
||||
"""
|
||||
|
||||
# Get batch size
|
||||
batch_size = w.shape[1]
|
||||
|
||||
# Expand the learned constant to match batch size
|
||||
x = self.initial_constant.expand(batch_size, -1, -1, -1)
|
||||
|
||||
# The first style block
|
||||
x = self.style_block(x, w[0], input_noise[0][1])
|
||||
# Get first rgb image
|
||||
rgb = self.to_rgb(x, w[0])
|
||||
|
||||
# Evaluate rest of the blocks
|
||||
for i in range(1, self.n_blocks):
|
||||
# Up sample the feature map
|
||||
x = self.up_sample(x)
|
||||
# Run it through the [generator block](#generator_block)
|
||||
x, rgb_new = self.blocks[i - 1](x, w[i], input_noise[i])
|
||||
# Up sample the RGB image and add to the rgb from the block
|
||||
rgb = self.up_sample(rgb) + rgb_new
|
||||
|
||||
# Return the final RGB image
|
||||
return rgb
|
||||
|
||||
|
||||
class GeneratorBlock(nn.Module):
|
||||
"""
|
||||
<a id="generator_block"></a>
|
||||
|
||||
### Generator Block
|
||||
|
||||
|
||||
|
||||
---*$A$ denotes a linear layer.
|
||||
$B$ denotes a broadcast and scaling operation (noise is a single channel).
|
||||
[`toRGB`](#to_rgb) also has a style modulation which is not shown in the diagram to keep it simple.*---
|
||||
|
||||
The generator block consists of two [style blocks](#style_block) ($3 \times 3$ convolutions with style modulation)
|
||||
and an RGB output.
|
||||
"""
|
||||
|
||||
def __init__(self, d_latent: int, in_features: int, out_features: int):
|
||||
"""
|
||||
* `d_latent` is the dimensionality of $w$
|
||||
* `in_features` is the number of features in the input feature map
|
||||
* `out_features` is the number of features in the output feature map
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
# First [style block](#style_block) changes the feature map size to `out_features`
|
||||
self.style_block1 = StyleBlock(d_latent, in_features, out_features)
|
||||
# Second [style block](#style_block)
|
||||
self.style_block2 = StyleBlock(d_latent, out_features, out_features)
|
||||
|
||||
# *toRGB* layer
|
||||
self.to_rgb = ToRGB(d_latent, out_features)
|
||||
|
||||
def forward(self, x: torch.Tensor, w: torch.Tensor, noise: Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]):
|
||||
"""
|
||||
* `x` is the input feature map of shape `[batch_size, in_features, height, width]`
|
||||
* `w` is $w$ with shape `[batch_size, d_latent]`
|
||||
* `noise` is a tuple of two noise tensors of shape `[batch_size, 1, height, width]`
|
||||
"""
|
||||
# First style block with first noise tensor.
|
||||
# The output is of shape `[batch_size, out_features, height, width]`
|
||||
x = self.style_block1(x, w, noise[0])
|
||||
# Second style block with second noise tensor.
|
||||
# The output is of shape `[batch_size, out_features, height, width]`
|
||||
x = self.style_block2(x, w, noise[1])
|
||||
|
||||
# Get RGB image
|
||||
rgb = self.to_rgb(x, w)
|
||||
|
||||
# Return feature map and rgb image
|
||||
return x, rgb
|
||||
|
||||
|
||||
class StyleBlock(nn.Module):
|
||||
"""
|
||||
<a id="style_block"></a>
|
||||
|
||||
### Style Block
|
||||
|
||||
|
||||
|
||||
---*$A$ denotes a linear layer.
|
||||
$B$ denotes a broadcast and scaling operation (noise is single channel).*---
|
||||
|
||||
Style block has a weight modulation convolution layer.
|
||||
"""
|
||||
|
||||
def __init__(self, d_latent: int, in_features: int, out_features: int):
|
||||
"""
|
||||
* `d_latent` is the dimensionality of $w$
|
||||
* `in_features` is the number of features in the input feature map
|
||||
* `out_features` is the number of features in the output feature map
|
||||
"""
|
||||
super().__init__()
|
||||
# Get style vector from $w$ (denoted by $A$ in the diagram) with
|
||||
# an [equalized learning-rate linear layer](#equalized_linear)
|
||||
self.to_style = EqualizedLinear(d_latent, in_features, bias=1.0)
|
||||
# Weight modulated convolution layer
|
||||
self.conv = Conv2dWeightModulate(in_features, out_features, kernel_size=3)
|
||||
# Noise scale
|
||||
self.scale_noise = nn.Parameter(torch.zeros(1))
|
||||
# Bias
|
||||
self.bias = nn.Parameter(torch.zeros(out_features))
|
||||
|
||||
# Activation function
|
||||
self.activation = nn.LeakyReLU(0.2, True)
|
||||
|
||||
def forward(self, x: torch.Tensor, w: torch.Tensor, noise: Optional[torch.Tensor]):
|
||||
"""
|
||||
* `x` is the input feature map of shape `[batch_size, in_features, height, width]`
|
||||
* `w` is $w$ with shape `[batch_size, d_latent]`
|
||||
* `noise` is a tensor of shape `[batch_size, 1, height, width]`
|
||||
"""
|
||||
# Get style vector $s$
|
||||
s = self.to_style(w)
|
||||
# Weight modulated convolution
|
||||
x = self.conv(x, s)
|
||||
# Scale and add noise
|
||||
if noise is not None:
|
||||
x = x + self.scale_noise[None, :, None, None] * noise
|
||||
# Add bias and evaluate activation function
|
||||
return self.activation(x + self.bias[None, :, None, None])
|
||||
|
||||
|
||||
class ToRGB(nn.Module):
|
||||
"""
|
||||
<a id="to_rgb"></a>
|
||||
|
||||
### To RGB
|
||||
|
||||
|
||||
|
||||
---*$A$ denotes a linear layer.*---
|
||||
|
||||
Generates an RGB image from a feature map using $1 \times 1$ convolution.
|
||||
"""
|
||||
|
||||
def __init__(self, d_latent: int, features: int):
|
||||
"""
|
||||
* `d_latent` is the dimensionality of $w$
|
||||
* `features` is the number of features in the feature map
|
||||
"""
|
||||
super().__init__()
|
||||
# Get style vector from $w$ (denoted by $A$ in the diagram) with
|
||||
# an [equalized learning-rate linear layer](#equalized_linear)
|
||||
self.to_style = EqualizedLinear(d_latent, features, bias=1.0)
|
||||
|
||||
# Weight modulated convolution layer without demodulation
|
||||
self.conv = Conv2dWeightModulate(features, 3, kernel_size=1, demodulate=False)
|
||||
# Bias
|
||||
self.bias = nn.Parameter(torch.zeros(3))
|
||||
# Activation function
|
||||
self.activation = nn.LeakyReLU(0.2, True)
|
||||
|
||||
def forward(self, x: torch.Tensor, w: torch.Tensor):
|
||||
"""
|
||||
* `x` is the input feature map of shape `[batch_size, in_features, height, width]`
|
||||
* `w` is $w$ with shape `[batch_size, d_latent]`
|
||||
"""
|
||||
# Get style vector $s$
|
||||
style = self.to_style(w)
|
||||
# Weight modulated convolution
|
||||
x = self.conv(x, style)
|
||||
# Add bias and evaluate activation function
|
||||
return self.activation(x + self.bias[None, :, None, None])
|
||||
|
||||
|
||||
class Conv2dWeightModulate(nn.Module):
|
||||
"""
|
||||
### Convolution with Weight Modulation and Demodulation
|
||||
|
||||
This layer scales the convolution weights by the style vector and demodulates by normalizing it.
|
||||
"""
|
||||
|
||||
def __init__(self, in_features: int, out_features: int, kernel_size: int,
|
||||
demodulate: float = True, eps: float = 1e-8):
|
||||
"""
|
||||
* `in_features` is the number of features in the input feature map
|
||||
* `out_features` is the number of features in the output feature map
|
||||
* `kernel_size` is the size of the convolution kernel
|
||||
* `demodulate` is flag whether to normalize weights by its standard deviation
|
||||
* `eps` is the $\epsilon$ for normalizing
|
||||
"""
|
||||
super().__init__()
|
||||
# Number of output features
|
||||
self.out_features = out_features
|
||||
# Whether to normalize weights
|
||||
self.demodulate = demodulate
|
||||
# Padding size
|
||||
self.padding = (kernel_size - 1) // 2
|
||||
|
||||
# [Weights parameter with equalized learning rate](#equalized_weight)
|
||||
self.weight = EqualizedWeight([out_features, in_features, kernel_size, kernel_size])
|
||||
# $\epsilon$
|
||||
self.eps = eps
|
||||
|
||||
def forward(self, x: torch.Tensor, s: torch.Tensor):
|
||||
"""
|
||||
* `x` is the input feature map of shape `[batch_size, in_features, height, width]`
|
||||
* `s` is style based scaling tensor of shape `[batch_size, in_features]`
|
||||
"""
|
||||
|
||||
# Get batch size, height and width
|
||||
b, _, h, w = x.shape
|
||||
|
||||
# Reshape the scales
|
||||
s = s[:, None, :, None, None]
|
||||
# Get [learning rate equalized weights](#equalized_weight)
|
||||
weights = self.weight()[None, :, :, :, :]
|
||||
# $$w`_{i,j,k} = s_i * w_{i,j,k}$$
|
||||
# where $i$ is the input channel, $j$ is the output channel, and $k$ is the kernel index.
|
||||
#
|
||||
# The result has shape `[batch_size, out_features, in_features, kernel_size, kernel_size]`
|
||||
weights = weights * s
|
||||
|
||||
# Demodulate
|
||||
if self.demodulate:
|
||||
# $$\sigma_j = \sqrt{\sum_{i,k} (w'_{i, j, k})^2 + \epsilon}$$
|
||||
sigma_inv = torch.rsqrt((weights ** 2).sum(dim=(2, 3, 4), keepdim=True) + self.eps)
|
||||
# $$w''_{i,j,k} = \frac{w'_{i,j,k}}{\sqrt{\sum_{i,k} (w'_{i, j, k})^2 + \epsilon}}$$
|
||||
weights = weights * sigma_inv
|
||||
|
||||
# Reshape `x`
|
||||
x = x.reshape(1, -1, h, w)
|
||||
|
||||
# Reshape weights
|
||||
_, _, *ws = weights.shape
|
||||
weights = weights.reshape(b * self.out_features, *ws)
|
||||
|
||||
# Use grouped convolution to efficiently calculate the convolution with sample wise kernel.
|
||||
# i.e. we have a different kernel (weights) for each sample in the batch
|
||||
x = F.conv2d(x, weights, padding=self.padding, groups=b)
|
||||
|
||||
# Reshape `x` to `[batch_size, out_features, height, width]` and return
|
||||
return x.reshape(-1, self.out_features, h, w)
|
||||
|
||||
|
||||
class Discriminator(nn.Module):
|
||||
"""
|
||||
<a id="discriminator"></a>
|
||||
|
||||
## StyleGAN 2 Discriminator
|
||||
|
||||
|
||||
|
||||
Discriminator first transforms the image to a feature map of the same resolution and then
|
||||
runs it through a series of blocks with residual connections.
|
||||
The resolution is down-sampled by $2 \times$ at each block while doubling the
|
||||
number of features.
|
||||
"""
|
||||
|
||||
def __init__(self, log_resolution: int, n_features: int = 64, max_features: int = 512):
|
||||
"""
|
||||
* `log_resolution` is the $\log_2$ of image resolution
|
||||
* `n_features` number of features in the convolution layer at the highest resolution (first block)
|
||||
* `max_features` maximum number of features in any generator block
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
# Layer to convert RGB image to a feature map with `n_features` number of features.
|
||||
self.from_rgb = nn.Sequential(
|
||||
EqualizedConv2d(3, n_features, 1),
|
||||
nn.LeakyReLU(0.2, True),
|
||||
)
|
||||
|
||||
# Calculate the number of features for each block.
|
||||
#
|
||||
# Something like `[64, 128, 256, 512, 512, 512]`.
|
||||
features = [min(max_features, n_features * (2 ** i)) for i in range(log_resolution - 1)]
|
||||
# Number of [discirminator blocks](#discriminator_block)
|
||||
n_blocks = len(features) - 1
|
||||
# Discriminator blocks
|
||||
blocks = [DiscriminatorBlock(features[i], features[i + 1]) for i in range(n_blocks)]
|
||||
self.blocks = nn.Sequential(*blocks)
|
||||
|
||||
# [Mini-batch Standard Deviation](#mini_batch_std_dev)
|
||||
self.std_dev = MiniBatchStdDev()
|
||||
# Number of features after adding the standard deviations map
|
||||
final_features = features[-1] + 1
|
||||
# Final $3 \times 3$ convolution layer
|
||||
self.conv = EqualizedConv2d(final_features, final_features, 3)
|
||||
# Final linear layer to get the classification
|
||||
self.final = EqualizedLinear(2 * 2 * final_features, 1)
|
||||
|
||||
def forward(self, x: torch.Tensor):
|
||||
"""
|
||||
* `x` is the input image of shape `[batch_size, 3, height, width]`
|
||||
"""
|
||||
|
||||
# Try to normalize the image (this is totally optional, but sped up the early training a little)
|
||||
x = x - 0.5
|
||||
# Convert from RGB
|
||||
x = self.from_rgb(x)
|
||||
# Run through the [discriminator blocks](#discriminator_block)
|
||||
x = self.blocks(x)
|
||||
|
||||
# Calculate and append [mini-batch standard deviation](#mini_batch_std_dev)
|
||||
x = self.std_dev(x)
|
||||
# $3 \times 3$ convolution
|
||||
x = self.conv(x)
|
||||
# Flatten
|
||||
x = x.reshape(x.shape[0], -1)
|
||||
# Return the classification score
|
||||
return self.final(x)
|
||||
|
||||
|
||||
class DiscriminatorBlock(nn.Module):
|
||||
"""
|
||||
<a id="discriminator_black"></a>
|
||||
|
||||
### Discriminator Block
|
||||
|
||||
|
||||
|
||||
Discriminator block consists of two $3 \times 3$ convolutions with a residual connection.
|
||||
"""
|
||||
|
||||
def __init__(self, in_features, out_features):
|
||||
"""
|
||||
* `in_features` is the number of features in the input feature map
|
||||
* `out_features` is the number of features in the output feature map
|
||||
"""
|
||||
super().__init__()
|
||||
# Down-sampling and $1 \times 1$ convolution layer for the residual connection
|
||||
self.residual = nn.Sequential(DownSample(),
|
||||
EqualizedConv2d(in_features, out_features, kernel_size=1))
|
||||
|
||||
# Two $3 \times 3$ convolutions
|
||||
self.block = nn.Sequential(
|
||||
EqualizedConv2d(in_features, in_features, kernel_size=3, padding=1),
|
||||
nn.LeakyReLU(0.2, True),
|
||||
EqualizedConv2d(in_features, out_features, kernel_size=3, padding=1),
|
||||
nn.LeakyReLU(0.2, True),
|
||||
)
|
||||
|
||||
# Down-sampling layer
|
||||
self.down_sample = DownSample()
|
||||
|
||||
# Scaling factor $\frac{1}{\sqrt 2}$ after adding the residual
|
||||
self.scale = 1 / math.sqrt(2)
|
||||
|
||||
def forward(self, x):
|
||||
# Get the residual connection
|
||||
residual = self.residual(x)
|
||||
|
||||
# Convolutions
|
||||
x = self.block(x)
|
||||
# Down-sample
|
||||
x = self.down_sample(x)
|
||||
|
||||
# Add the residual and scale
|
||||
return (x + residual) * self.scale
|
||||
|
||||
|
||||
class MiniBatchStdDev(nn.Module):
|
||||
"""
|
||||
<a id="mini_batch_std_dev"></a>
|
||||
|
||||
### Mini-batch Standard Deviation
|
||||
|
||||
Mini-batch standard deviation calculates the standard deviation
|
||||
across a mini-batch (or a subgroups within the mini-batch)
|
||||
for each feature in the feature map. Then it takes the mean of all
|
||||
the standard deviations and appends it to the feature map as one extra feature.
|
||||
"""
|
||||
|
||||
def __init__(self, group_size: int = 4):
|
||||
"""
|
||||
* `group_size` is the number of samples to calculate standard deviation across.
|
||||
"""
|
||||
super().__init__()
|
||||
self.group_size = group_size
|
||||
|
||||
def forward(self, x: torch.Tensor):
|
||||
"""
|
||||
* `x` is the feature map
|
||||
"""
|
||||
# Check if the batch size is divisible by the group size
|
||||
assert x.shape[0] % self.group_size == 0
|
||||
# Split the samples into groups of `group_size`, we flatten the feature map to a single dimension
|
||||
# since we want to calculate the standard deviation for each feature.
|
||||
grouped = x.view(self.group_size, -1)
|
||||
# Calculate the standard deviation for each feature among `group_size` samples
|
||||
#
|
||||
# \begin{align}
|
||||
# \mu_{i} &= \frac{1}{N} \sum_g x_{g,i} \\
|
||||
# \sigma_{i} &= \sqrt{\frac{1}{N} \sum_g (x_{g,i} - \mu_i)^2 + \epsilon}
|
||||
# \end{align}
|
||||
std = torch.sqrt(grouped.var(dim=0) + 1e-8)
|
||||
# Get the mean standard deviation
|
||||
std = std.mean().view(1, 1, 1, 1)
|
||||
# Expand the standard deviation to append to the feature map
|
||||
b, _, h, w = x.shape
|
||||
std = std.expand(b, -1, h, w)
|
||||
# Append (concatenate) the standard deviations to the feature map
|
||||
return torch.cat([x, std], dim=1)
|
||||
|
||||
|
||||
class DownSample(nn.Module):
|
||||
"""
|
||||
<a id="down_sample"></a>
|
||||
|
||||
### Down-sample
|
||||
|
||||
The down-sample operation [smoothens](#smooth) each feature channel and
|
||||
scale $2 \times$ using bilinear interpolation.
|
||||
This is based on the paper
|
||||
[Making Convolutional Networks Shift-Invariant Again](https://arxiv.org/abs/1904.11486).
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
# Smoothing layer
|
||||
self.smooth = Smooth()
|
||||
|
||||
def forward(self, x: torch.Tensor):
|
||||
# Smoothing or blurring
|
||||
x = self.smooth(x)
|
||||
# Scaled down
|
||||
return F.interpolate(x, (x.shape[2] // 2, x.shape[3] // 2), mode='bilinear', align_corners=False)
|
||||
|
||||
|
||||
class UpSample(nn.Module):
|
||||
"""
|
||||
<a id="up_sample"></a>
|
||||
|
||||
### Up-sample
|
||||
|
||||
The up-sample operation scales the image up by $2 \times$ and [smoothens](#smooth) each feature channel.
|
||||
This is based on the paper
|
||||
[Making Convolutional Networks Shift-Invariant Again](https://arxiv.org/abs/1904.11486).
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
# Up-sampling layer
|
||||
self.up_sample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False)
|
||||
# Smoothing layer
|
||||
self.smooth = Smooth()
|
||||
|
||||
def forward(self, x: torch.Tensor):
|
||||
# Up-sample and smoothen
|
||||
return self.smooth(self.up_sample(x))
|
||||
|
||||
|
||||
class Smooth(nn.Module):
|
||||
"""
|
||||
<a id="smooth"></a>
|
||||
|
||||
### Smoothing Layer
|
||||
|
||||
This layer blurs each channel
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
# Blurring kernel
|
||||
kernel = [[1, 2, 1],
|
||||
[2, 4, 2],
|
||||
[1, 2, 1]]
|
||||
# Convert the kernel to a PyTorch tensor
|
||||
kernel = torch.tensor([[kernel]], dtype=torch.float)
|
||||
# Normalize the kernel
|
||||
kernel /= kernel.sum()
|
||||
# Save kernel as a fixed parameter (no gradient updates)
|
||||
self.kernel = nn.Parameter(kernel, requires_grad=False)
|
||||
# Padding layer
|
||||
self.pad = nn.ReplicationPad2d(1)
|
||||
|
||||
def forward(self, x: torch.Tensor):
|
||||
# Get shape of the input feature map
|
||||
b, c, h, w = x.shape
|
||||
# Reshape for smoothening
|
||||
x = x.view(-1, 1, h, w)
|
||||
|
||||
# Add padding
|
||||
x = self.pad(x)
|
||||
|
||||
# Smoothen (blur) with the kernel
|
||||
x = F.conv2d(x, self.kernel)
|
||||
|
||||
# Reshape and return
|
||||
return x.view(b, c, h, w)
|
||||
|
||||
|
||||
class EqualizedLinear(nn.Module):
|
||||
"""
|
||||
<a id="equalized_linear"></a>
|
||||
|
||||
## Learning-rate Equalized Linear Layer
|
||||
|
||||
This uses [learning-rate equalized weights](#equalized_weights) for a linear layer.
|
||||
"""
|
||||
|
||||
def __init__(self, in_features: int, out_features: int, bias: float = 0.):
|
||||
"""
|
||||
* `in_features` is the number of features in the input feature map
|
||||
* `out_features` is the number of features in the output feature map
|
||||
* `bias` is the bias initialization constant
|
||||
"""
|
||||
|
||||
super().__init__()
|
||||
# [Learning-rate equalized weights](#equalized_weights)
|
||||
self.weight = EqualizedWeight([out_features, in_features])
|
||||
# Bias
|
||||
self.bias = nn.Parameter(torch.ones(out_features) * bias)
|
||||
|
||||
def forward(self, x: torch.Tensor):
|
||||
# Linear transformation
|
||||
return F.linear(x, self.weight(), bias=self.bias)
|
||||
|
||||
|
||||
class EqualizedConv2d(nn.Module):
|
||||
"""
|
||||
<a id="equalized_conv2d"></a>
|
||||
|
||||
## Learning-rate Equalized 2D Convolution Layer
|
||||
|
||||
This uses [learning-rate equalized weights](#equalized_weights) for a convolution layer.
|
||||
"""
|
||||
|
||||
def __init__(self, in_features: int, out_features: int,
|
||||
kernel_size: int, padding: int = 0):
|
||||
"""
|
||||
* `in_features` is the number of features in the input feature map
|
||||
* `out_features` is the number of features in the output feature map
|
||||
* `kernel_size` is the size of the convolution kernel
|
||||
* `padding` is the padding to be added on both sides of each size dimension
|
||||
"""
|
||||
super().__init__()
|
||||
# Padding size
|
||||
self.padding = padding
|
||||
# [Learning-rate equalized weights](#equalized_weights)
|
||||
self.weight = EqualizedWeight([out_features, in_features, kernel_size, kernel_size])
|
||||
# Bias
|
||||
self.bias = nn.Parameter(torch.ones(out_features))
|
||||
|
||||
def forward(self, x: torch.Tensor):
|
||||
# Convolution
|
||||
return F.conv2d(x, self.weight(), bias=self.bias, padding=self.padding)
|
||||
|
||||
|
||||
class EqualizedWeight(nn.Module):
|
||||
"""
|
||||
<a id="equalized_weight"></a>
|
||||
|
||||
## Learning-rate Equalized Weights Parameter
|
||||
|
||||
This is based on equalized learning rate introduced in the Progressive GAN paper.
|
||||
Instead of initializing weights at $\mathcal{N}(0,c)$ they initialize weights
|
||||
to $\mathcal{N}(0, 1)$ and then multiply them by $c$ when using it.
|
||||
$$w_i = c \hat{w}_i$$
|
||||
|
||||
The gradients on stored parameters $\hat{w}$ get multiplied by $c$ but this doesn't have
|
||||
an affect since optimizers such as Adam normalize them by a running mean of the squared gradients.
|
||||
|
||||
The optimizer updates on $\hat{w}$ are proportionate to the learning rate $\lambda$.
|
||||
But the effective weights $w$ get updated proportionately to $c \lambda$.
|
||||
Without equalized learning rate, the effective weights will get updated proportionately to just $\lambda$.
|
||||
|
||||
So we are effectively scaling the learning rate by $c$ for these weight parameters.
|
||||
"""
|
||||
|
||||
def __init__(self, shape: List[int]):
|
||||
"""
|
||||
* `shape` is the shape of the weight parameter
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
# He initialization constant
|
||||
self.c = 1 / math.sqrt(np.prod(shape[1:]))
|
||||
# Initialize the weights with $\mathcal{N}(0, 1)$
|
||||
self.weight = nn.Parameter(torch.randn(shape))
|
||||
# Weight multiplication coefficient
|
||||
|
||||
def forward(self):
|
||||
# Multiply the weights by $c$ and return
|
||||
return self.weight * self.c
|
||||
|
||||
|
||||
class GradientPenalty(nn.Module):
|
||||
"""
|
||||
<a id="gradient_penalty"></a>
|
||||
|
||||
## Gradient Penalty
|
||||
|
||||
This is the $R_1$ regularization penality from the paper
|
||||
[Which Training Methods for GANs do actually Converge?](https://arxiv.org/abs/1801.04406).
|
||||
|
||||
$$R_1(\psi) = \frac{\gamma}{2} \mathbb{E}_{p_\mathcal{D}(x)}
|
||||
\Big[\Vert \nabla_x D_\psi(x)^2 \Vert\Big]$$
|
||||
|
||||
That is we try to reduce the L2 norm of gradients of the discriminator with
|
||||
respect to images, for real images ($P_\mathcal{D}$).
|
||||
"""
|
||||
|
||||
def forward(self, x: torch.Tensor, d: torch.Tensor):
|
||||
"""
|
||||
* `x` is $x \sim \mathcal{D}$
|
||||
* `d` is $D(x)$
|
||||
"""
|
||||
|
||||
# Get batch size
|
||||
batch_size = x.shape[0]
|
||||
|
||||
# Calculate gradients of $D(x)$ with respect to $x$.
|
||||
# `grad_outputs` is set to $1$ since we want the gradients of $D(x)$,
|
||||
# and we need to create and retain graph since we have to compute gradients
|
||||
# with respect to weight on this loss.
|
||||
gradients, *_ = torch.autograd.grad(outputs=d,
|
||||
inputs=x,
|
||||
grad_outputs=d.new_ones(d.shape),
|
||||
create_graph=True)
|
||||
|
||||
# Reshape gradients to calculate the norm
|
||||
gradients = gradients.reshape(batch_size, -1)
|
||||
# Calculate the norm $\Vert \nabla_{x} D(x)^2 \Vert$
|
||||
norm = gradients.norm(2, dim=-1)
|
||||
# Return the loss $\Vert \nabla_x D_\psi(x)^2 \Vert$
|
||||
return torch.mean(norm ** 2)
|
||||
|
||||
|
||||
class PathLengthPenalty(nn.Module):
|
||||
"""
|
||||
<a id="path_length_penalty"></a>
|
||||
|
||||
## Path Length Penalty
|
||||
|
||||
This regularization encourages a fixed-size step in $w$ to result in a fixed-magnitude
|
||||
change in the image.
|
||||
|
||||
$$\mathbb{E}_{w \sim f(z), y \sim \mathcal{N}(0, \mathbf{I})}
|
||||
\Big(\Vert \mathbf{J}^\top_{w} y \Vert_2 - a \Big)^2$$
|
||||
|
||||
where $\mathbf{J}_w$ is the Jacobian
|
||||
$\mathbf{J}_w = \frac{\partial g}{\partial w}$,
|
||||
$w$ are sampled from $w \in \mathcal{W}$ from the mapping network, and
|
||||
$y$ are images with noise $\mathcal{N}(0, \mathbf{I})$.
|
||||
|
||||
$a$ is the exponential moving average of $\Vert \mathbf{J}^\top_{w} y \Vert_2$
|
||||
as the training progresses.
|
||||
|
||||
$\mathbf{J}^\top_{w} y$ is calculated without explicitly calculating the Jacobian using
|
||||
$$\mathbf{J}^\top_{w} y = \nabla_w \big(g(w) \cdot y \big)$$
|
||||
"""
|
||||
|
||||
def __init__(self, beta: float):
|
||||
"""
|
||||
* `beta` is the constant $\beta$ used to calculate the exponential moving average $a$
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
# $\beta$
|
||||
self.beta = beta
|
||||
# Number of steps calculated $N$
|
||||
self.steps = nn.Parameter(torch.tensor(0.), requires_grad=False)
|
||||
# Exponential sum of $\mathbf{J}^\top_{w} y$
|
||||
# $$\sum^N_{i=1} \beta^{(N - i)}[\mathbf{J}^\top_{w} y]_i$$
|
||||
# where $[\mathbf{J}^\top_{w} y]_i$ is the value of it at $i$-th step of training
|
||||
self.exp_sum_a = nn.Parameter(torch.tensor(0.), requires_grad=False)
|
||||
|
||||
def forward(self, w: torch.Tensor, x: torch.Tensor):
|
||||
"""
|
||||
* `w` is the batch of $w$ of shape `[batch_size, d_latent]`
|
||||
* `x` are the generated images of shape `[batch_size, 3, height, width]`
|
||||
"""
|
||||
|
||||
# Get the device
|
||||
device = x.device
|
||||
# Get number of pixels
|
||||
image_size = x.shape[2] * x.shape[3]
|
||||
# Calculate $y \in \mathcal{N}(0, \mathbf{I})$
|
||||
y = torch.randn(x.shape, device=device)
|
||||
# Calculate $\big(g(w) \cdot y \big)$ and normalize by the square root of image size.
|
||||
# This is scaling is not mentioned in the paper but was present in
|
||||
# [their implementation](https://github.com/NVlabs/stylegan2/blob/master/training/loss.py#L167).
|
||||
output = (x * y).sum() / math.sqrt(image_size)
|
||||
|
||||
# Calculate gradients to get $\mathbf{J}^\top_{w} y$
|
||||
gradients, *_ = torch.autograd.grad(outputs=output,
|
||||
inputs=w,
|
||||
grad_outputs=torch.ones(output.shape, device=device),
|
||||
create_graph=True)
|
||||
|
||||
# Calculate L2-norm of $\mathbf{J}^\top_{w} y$
|
||||
norm = (gradients ** 2).sum(dim=2).mean(dim=1).sqrt()
|
||||
|
||||
# Regularize after first step
|
||||
if self.steps > 0:
|
||||
# Calculate $a$
|
||||
# $$\frac{1}{1 - \beta^N} \sum^N_{i=1} \beta^{(N - i)}[\mathbf{J}^\top_{w} y]_i$$
|
||||
a = self.exp_sum_a / (1 - self.beta ** self.steps)
|
||||
# Calculate the penalty
|
||||
# $$\mathbb{E}_{w \sim f(z), y \sim \mathcal{N}(0, \mathbf{I})}
|
||||
# \Big(\Vert \mathbf{J}^\top_{w} y \Vert_2 - a \Big)^2$$
|
||||
loss = torch.mean((norm - a) ** 2)
|
||||
else:
|
||||
# Return a dummy loss if we can't calculate $a$
|
||||
loss = norm.new_tensor(0)
|
||||
|
||||
# Calculate the mean of $\Vert \mathbf{J}^\top_{w} y \Vert_2$
|
||||
mean = norm.mean().detach()
|
||||
# Update exponential sum
|
||||
self.exp_sum_a.mul_(self.beta).add_(mean, alpha=1 - self.beta)
|
||||
# Increment $N$
|
||||
self.steps.add_(1.)
|
||||
|
||||
# Return the penalty
|
||||
return loss
|
||||
@@ -0,0 +1,459 @@
|
||||
"""
|
||||
---
|
||||
title: StyleGAN 2 Model Training
|
||||
summary: >
|
||||
An annotated PyTorch implementation of StyleGAN2 model training code.
|
||||
---
|
||||
|
||||
# [StyleGAN 2](index.html) Model Training
|
||||
|
||||
This is the training code for [StyleGAN 2](index.html) model.
|
||||
|
||||
|
||||
|
||||
---*These are $64 \times 64$ images generated after training for about 80K steps.*---
|
||||
|
||||
*Our implementation is a minimalistic StyleGAN 2 model training code.
|
||||
Only single GPU training is supported to keep the implementation simple.
|
||||
We managed to shrink it to keep it at less than 500 lines of code, including the training loop.*
|
||||
|
||||
*Without DDP (distributed data parallel) and multi-gpu training it will not be possible to train the model
|
||||
for large resolutions (128+).
|
||||
If you want training code with fp16 and DDP take a look at
|
||||
[lucidrains/stylegan2-pytorch](https://github.com/lucidrains/stylegan2-pytorch).*
|
||||
|
||||
We trained this on [CelebA-HQ dataset](https://github.com/tkarras/progressive_growing_of_gans).
|
||||
You can find the download instruction in this
|
||||
[discussion on fast.ai](https://forums.fast.ai/t/download-celeba-hq-dataset/45873/3).
|
||||
Save the images inside [`data/stylegan` folder](#dataset_path).
|
||||
"""
|
||||
|
||||
import math
|
||||
from pathlib import Path
|
||||
from typing import Iterator, Tuple
|
||||
|
||||
import torchvision
|
||||
from PIL import Image
|
||||
|
||||
import torch
|
||||
import torch.utils.data
|
||||
from labml import tracker, lab, monit, experiment
|
||||
from labml.configs import BaseConfigs
|
||||
from labml_nn.gan.stylegan import Discriminator, Generator, MappingNetwork, GradientPenalty, PathLengthPenalty
|
||||
from labml_nn.gan.wasserstein import DiscriminatorLoss, GeneratorLoss
|
||||
from labml_nn.helpers.device import DeviceConfigs
|
||||
from labml_nn.helpers.trainer import ModeState
|
||||
from labml_nn.utils import cycle_dataloader
|
||||
|
||||
|
||||
class Dataset(torch.utils.data.Dataset):
|
||||
"""
|
||||
## Dataset
|
||||
|
||||
This loads the training dataset and resize it to the give image size.
|
||||
"""
|
||||
|
||||
def __init__(self, path: str, image_size: int):
|
||||
"""
|
||||
* `path` path to the folder containing the images
|
||||
* `image_size` size of the image
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
# Get the paths of all `jpg` files
|
||||
self.paths = [p for p in Path(path).glob(f'**/*.jpg')]
|
||||
|
||||
# Transformation
|
||||
self.transform = torchvision.transforms.Compose([
|
||||
# Resize the image
|
||||
torchvision.transforms.Resize(image_size),
|
||||
# Convert to PyTorch tensor
|
||||
torchvision.transforms.ToTensor(),
|
||||
])
|
||||
|
||||
def __len__(self):
|
||||
"""Number of images"""
|
||||
return len(self.paths)
|
||||
|
||||
def __getitem__(self, index):
|
||||
"""Get the the `index`-th image"""
|
||||
path = self.paths[index]
|
||||
img = Image.open(path)
|
||||
return self.transform(img)
|
||||
|
||||
|
||||
class Configs(BaseConfigs):
|
||||
"""
|
||||
## Configurations
|
||||
"""
|
||||
|
||||
# Device to train the model on.
|
||||
# [`DeviceConfigs`](../../helpers/device.html)
|
||||
# picks up an available CUDA device or defaults to CPU.
|
||||
device: torch.device = DeviceConfigs()
|
||||
|
||||
# [StyleGAN2 Discriminator](index.html#discriminator)
|
||||
discriminator: Discriminator
|
||||
# [StyleGAN2 Generator](index.html#generator)
|
||||
generator: Generator
|
||||
# [Mapping network](index.html#mapping_network)
|
||||
mapping_network: MappingNetwork
|
||||
|
||||
# Discriminator and generator loss functions.
|
||||
# We use [Wasserstein loss](../wasserstein/index.html)
|
||||
discriminator_loss: DiscriminatorLoss
|
||||
generator_loss: GeneratorLoss
|
||||
|
||||
# Optimizers
|
||||
generator_optimizer: torch.optim.Adam
|
||||
discriminator_optimizer: torch.optim.Adam
|
||||
mapping_network_optimizer: torch.optim.Adam
|
||||
|
||||
# [Gradient Penalty Regularization Loss](index.html#gradient_penalty)
|
||||
gradient_penalty = GradientPenalty()
|
||||
# Gradient penalty coefficient $\gamma$
|
||||
gradient_penalty_coefficient: float = 10.
|
||||
|
||||
# [Path length penalty](index.html#path_length_penalty)
|
||||
path_length_penalty: PathLengthPenalty
|
||||
|
||||
# Data loader
|
||||
loader: Iterator
|
||||
|
||||
# Batch size
|
||||
batch_size: int = 32
|
||||
# Dimensionality of $z$ and $w$
|
||||
d_latent: int = 512
|
||||
# Height/width of the image
|
||||
image_size: int = 32
|
||||
# Number of layers in the mapping network
|
||||
mapping_network_layers: int = 8
|
||||
# Generator & Discriminator learning rate
|
||||
learning_rate: float = 1e-3
|
||||
# Mapping network learning rate ($100 \times$ lower than the others)
|
||||
mapping_network_learning_rate: float = 1e-5
|
||||
# Number of steps to accumulate gradients on. Use this to increase the effective batch size.
|
||||
gradient_accumulate_steps: int = 1
|
||||
# $\beta_1$ and $\beta_2$ for Adam optimizer
|
||||
adam_betas: Tuple[float, float] = (0.0, 0.99)
|
||||
# Probability of mixing styles
|
||||
style_mixing_prob: float = 0.9
|
||||
|
||||
# Total number of training steps
|
||||
training_steps: int = 150_000
|
||||
|
||||
# Number of blocks in the generator (calculated based on image resolution)
|
||||
n_gen_blocks: int
|
||||
|
||||
# ### Lazy regularization
|
||||
# Instead of calculating the regularization losses, the paper proposes lazy regularization
|
||||
# where the regularization terms are calculated once in a while.
|
||||
# This improves the training efficiency a lot.
|
||||
|
||||
# The interval at which to compute gradient penalty
|
||||
lazy_gradient_penalty_interval: int = 4
|
||||
# Path length penalty calculation interval
|
||||
lazy_path_penalty_interval: int = 32
|
||||
# Skip calculating path length penalty during the initial phase of training
|
||||
lazy_path_penalty_after: int = 5_000
|
||||
|
||||
# How often to log generated images
|
||||
log_generated_interval: int = 500
|
||||
# How often to save model checkpoints
|
||||
save_checkpoint_interval: int = 2_000
|
||||
|
||||
# Training mode state for logging activations
|
||||
mode: ModeState
|
||||
|
||||
# <a id="dataset_path"></a>
|
||||
# We trained this on [CelebA-HQ dataset](https://github.com/tkarras/progressive_growing_of_gans).
|
||||
# You can find the download instruction in this
|
||||
# [discussion on fast.ai](https://forums.fast.ai/t/download-celeba-hq-dataset/45873/3).
|
||||
# Save the images inside `data/stylegan` folder.
|
||||
dataset_path: str = str(lab.get_data_path() / 'stylegan2')
|
||||
|
||||
def init(self):
|
||||
"""
|
||||
### Initialize
|
||||
"""
|
||||
# Create dataset
|
||||
dataset = Dataset(self.dataset_path, self.image_size)
|
||||
# Create data loader
|
||||
dataloader = torch.utils.data.DataLoader(dataset, batch_size=self.batch_size, num_workers=8,
|
||||
shuffle=True, drop_last=True, pin_memory=True)
|
||||
# Continuous [cyclic loader](../../utils.html#cycle_dataloader)
|
||||
self.loader = cycle_dataloader(dataloader)
|
||||
|
||||
# $\log_2$ of image resolution
|
||||
log_resolution = int(math.log2(self.image_size))
|
||||
|
||||
# Create discriminator and generator
|
||||
self.discriminator = Discriminator(log_resolution).to(self.device)
|
||||
self.generator = Generator(log_resolution, self.d_latent).to(self.device)
|
||||
# Get number of generator blocks for creating style and noise inputs
|
||||
self.n_gen_blocks = self.generator.n_blocks
|
||||
# Create mapping network
|
||||
self.mapping_network = MappingNetwork(self.d_latent, self.mapping_network_layers).to(self.device)
|
||||
# Create path length penalty loss
|
||||
self.path_length_penalty = PathLengthPenalty(0.99).to(self.device)
|
||||
|
||||
# Discriminator and generator losses
|
||||
self.discriminator_loss = DiscriminatorLoss().to(self.device)
|
||||
self.generator_loss = GeneratorLoss().to(self.device)
|
||||
|
||||
# Create optimizers
|
||||
self.discriminator_optimizer = torch.optim.Adam(
|
||||
self.discriminator.parameters(),
|
||||
lr=self.learning_rate, betas=self.adam_betas
|
||||
)
|
||||
self.generator_optimizer = torch.optim.Adam(
|
||||
self.generator.parameters(),
|
||||
lr=self.learning_rate, betas=self.adam_betas
|
||||
)
|
||||
self.mapping_network_optimizer = torch.optim.Adam(
|
||||
self.mapping_network.parameters(),
|
||||
lr=self.mapping_network_learning_rate, betas=self.adam_betas
|
||||
)
|
||||
|
||||
# Set tracker configurations
|
||||
tracker.set_image("generated", True)
|
||||
|
||||
def get_w(self, batch_size: int):
|
||||
"""
|
||||
### Sample $w$
|
||||
|
||||
This samples $z$ randomly and get $w$ from the mapping network.
|
||||
|
||||
We also apply style mixing sometimes where we generate two latent variables
|
||||
$z_1$ and $z_2$ and get corresponding $w_1$ and $w_2$.
|
||||
Then we randomly sample a cross-over point and apply $w_1$ to
|
||||
the generator blocks before the cross-over point and
|
||||
$w_2$ to the blocks after.
|
||||
"""
|
||||
|
||||
# Mix styles
|
||||
if torch.rand(()).item() < self.style_mixing_prob:
|
||||
# Random cross-over point
|
||||
cross_over_point = int(torch.rand(()).item() * self.n_gen_blocks)
|
||||
# Sample $z_1$ and $z_2$
|
||||
z2 = torch.randn(batch_size, self.d_latent).to(self.device)
|
||||
z1 = torch.randn(batch_size, self.d_latent).to(self.device)
|
||||
# Get $w_1$ and $w_2$
|
||||
w1 = self.mapping_network(z1)
|
||||
w2 = self.mapping_network(z2)
|
||||
# Expand $w_1$ and $w_2$ for the generator blocks and concatenate
|
||||
w1 = w1[None, :, :].expand(cross_over_point, -1, -1)
|
||||
w2 = w2[None, :, :].expand(self.n_gen_blocks - cross_over_point, -1, -1)
|
||||
return torch.cat((w1, w2), dim=0)
|
||||
# Without mixing
|
||||
else:
|
||||
# Sample $z$ and $z$
|
||||
z = torch.randn(batch_size, self.d_latent).to(self.device)
|
||||
# Get $w$ and $w$
|
||||
w = self.mapping_network(z)
|
||||
# Expand $w$ for the generator blocks
|
||||
return w[None, :, :].expand(self.n_gen_blocks, -1, -1)
|
||||
|
||||
def get_noise(self, batch_size: int):
|
||||
"""
|
||||
### Generate noise
|
||||
|
||||
This generates noise for each [generator block](index.html#generator_block)
|
||||
"""
|
||||
# List to store noise
|
||||
noise = []
|
||||
# Noise resolution starts from $4$
|
||||
resolution = 4
|
||||
|
||||
# Generate noise for each generator block
|
||||
for i in range(self.n_gen_blocks):
|
||||
# The first block has only one $3 \times 3$ convolution
|
||||
if i == 0:
|
||||
n1 = None
|
||||
# Generate noise to add after the first convolution layer
|
||||
else:
|
||||
n1 = torch.randn(batch_size, 1, resolution, resolution, device=self.device)
|
||||
# Generate noise to add after the second convolution layer
|
||||
n2 = torch.randn(batch_size, 1, resolution, resolution, device=self.device)
|
||||
|
||||
# Add noise tensors to the list
|
||||
noise.append((n1, n2))
|
||||
|
||||
# Next block has $2 \times$ resolution
|
||||
resolution *= 2
|
||||
|
||||
# Return noise tensors
|
||||
return noise
|
||||
|
||||
def generate_images(self, batch_size: int):
|
||||
"""
|
||||
### Generate images
|
||||
|
||||
This generate images using the generator
|
||||
"""
|
||||
|
||||
# Get $w$
|
||||
w = self.get_w(batch_size)
|
||||
# Get noise
|
||||
noise = self.get_noise(batch_size)
|
||||
|
||||
# Generate images
|
||||
images = self.generator(w, noise)
|
||||
|
||||
# Return images and $w$
|
||||
return images, w
|
||||
|
||||
def step(self, idx: int):
|
||||
"""
|
||||
### Training Step
|
||||
"""
|
||||
|
||||
# Train the discriminator
|
||||
with monit.section('Discriminator'):
|
||||
# Reset gradients
|
||||
self.discriminator_optimizer.zero_grad()
|
||||
|
||||
# Accumulate gradients for `gradient_accumulate_steps`
|
||||
for i in range(self.gradient_accumulate_steps):
|
||||
# Sample images from generator
|
||||
generated_images, _ = self.generate_images(self.batch_size)
|
||||
# Discriminator classification for generated images
|
||||
fake_output = self.discriminator(generated_images.detach())
|
||||
|
||||
# Get real images from the data loader
|
||||
real_images = next(self.loader).to(self.device)
|
||||
# We need to calculate gradients w.r.t. real images for gradient penalty
|
||||
if (idx + 1) % self.lazy_gradient_penalty_interval == 0:
|
||||
real_images.requires_grad_()
|
||||
# Discriminator classification for real images
|
||||
real_output = self.discriminator(real_images)
|
||||
|
||||
# Get discriminator loss
|
||||
real_loss, fake_loss = self.discriminator_loss(real_output, fake_output)
|
||||
disc_loss = real_loss + fake_loss
|
||||
|
||||
# Add gradient penalty
|
||||
if (idx + 1) % self.lazy_gradient_penalty_interval == 0:
|
||||
# Calculate and log gradient penalty
|
||||
gp = self.gradient_penalty(real_images, real_output)
|
||||
tracker.add('loss.gp', gp)
|
||||
# Multiply by coefficient and add gradient penalty
|
||||
disc_loss = disc_loss + 0.5 * self.gradient_penalty_coefficient * gp * self.lazy_gradient_penalty_interval
|
||||
|
||||
# Compute gradients
|
||||
disc_loss.backward()
|
||||
|
||||
# Log discriminator loss
|
||||
tracker.add('loss.discriminator', disc_loss)
|
||||
|
||||
if (idx + 1) % self.log_generated_interval == 0:
|
||||
# Log discriminator model parameters occasionally
|
||||
tracker.add('discriminator', self.discriminator)
|
||||
|
||||
# Clip gradients for stabilization
|
||||
torch.nn.utils.clip_grad_norm_(self.discriminator.parameters(), max_norm=1.0)
|
||||
# Take optimizer step
|
||||
self.discriminator_optimizer.step()
|
||||
|
||||
# Train the generator
|
||||
with monit.section('Generator'):
|
||||
# Reset gradients
|
||||
self.generator_optimizer.zero_grad()
|
||||
self.mapping_network_optimizer.zero_grad()
|
||||
|
||||
# Accumulate gradients for `gradient_accumulate_steps`
|
||||
for i in range(self.gradient_accumulate_steps):
|
||||
# Sample images from generator
|
||||
generated_images, w = self.generate_images(self.batch_size)
|
||||
# Discriminator classification for generated images
|
||||
fake_output = self.discriminator(generated_images)
|
||||
|
||||
# Get generator loss
|
||||
gen_loss = self.generator_loss(fake_output)
|
||||
|
||||
# Add path length penalty
|
||||
if idx > self.lazy_path_penalty_after and (idx + 1) % self.lazy_path_penalty_interval == 0:
|
||||
# Calculate path length penalty
|
||||
plp = self.path_length_penalty(w, generated_images)
|
||||
# Ignore if `nan`
|
||||
if not torch.isnan(plp):
|
||||
tracker.add('loss.plp', plp)
|
||||
gen_loss = gen_loss + plp
|
||||
|
||||
# Calculate gradients
|
||||
gen_loss.backward()
|
||||
|
||||
# Log generator loss
|
||||
tracker.add('loss.generator', gen_loss)
|
||||
|
||||
if (idx + 1) % self.log_generated_interval == 0:
|
||||
# Log discriminator model parameters occasionally
|
||||
tracker.add('generator', self.generator)
|
||||
tracker.add('mapping_network', self.mapping_network)
|
||||
|
||||
# Clip gradients for stabilization
|
||||
torch.nn.utils.clip_grad_norm_(self.generator.parameters(), max_norm=1.0)
|
||||
torch.nn.utils.clip_grad_norm_(self.mapping_network.parameters(), max_norm=1.0)
|
||||
|
||||
# Take optimizer step
|
||||
self.generator_optimizer.step()
|
||||
self.mapping_network_optimizer.step()
|
||||
|
||||
# Log generated images
|
||||
if (idx + 1) % self.log_generated_interval == 0:
|
||||
tracker.add('generated', torch.cat([generated_images[:6], real_images[:3]], dim=0))
|
||||
# Save model checkpoints
|
||||
if (idx + 1) % self.save_checkpoint_interval == 0:
|
||||
# Save checkpoint
|
||||
pass
|
||||
|
||||
# Flush tracker
|
||||
tracker.save()
|
||||
|
||||
def train(self):
|
||||
"""
|
||||
## Train model
|
||||
"""
|
||||
|
||||
# Loop for `training_steps`
|
||||
for i in monit.loop(self.training_steps):
|
||||
# Take a training step
|
||||
self.step(i)
|
||||
#
|
||||
if (i + 1) % self.log_generated_interval == 0:
|
||||
tracker.new_line()
|
||||
|
||||
|
||||
def main():
|
||||
"""
|
||||
### Train StyleGAN2
|
||||
"""
|
||||
|
||||
# Create an experiment
|
||||
experiment.create(name='stylegan2')
|
||||
# Create configurations object
|
||||
configs = Configs()
|
||||
|
||||
# Set configurations and override some
|
||||
experiment.configs(configs, {
|
||||
'device.cuda_device': 0,
|
||||
'image_size': 64,
|
||||
'log_generated_interval': 200
|
||||
})
|
||||
|
||||
# Initialize
|
||||
configs.init()
|
||||
# Set models for saving and loading
|
||||
experiment.add_pytorch_models(mapping_network=configs.mapping_network,
|
||||
generator=configs.generator,
|
||||
discriminator=configs.discriminator)
|
||||
|
||||
# Start the experiment
|
||||
with experiment.start():
|
||||
# Run the training loop
|
||||
configs.train()
|
||||
|
||||
|
||||
#
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,10 @@
|
||||
# [StyleGAN 2](https://nn.labml.ai/gan/stylegan/index.html)
|
||||
|
||||
This is a [PyTorch](https://pytorch.org) implementation of the paper
|
||||
[Analyzing and Improving the Image Quality of StyleGAN](https://arxiv.org/abs/1912.04958)
|
||||
which introduces **StyleGAN2**.
|
||||
StyleGAN 2 is an improvement over **StyleGAN** from the paper
|
||||
[A Style-Based Generator Architecture for Generative Adversarial Networks](https://arxiv.org/abs/1812.04948).
|
||||
And StyleGAN is based on **Progressive GAN** from the paper
|
||||
[Progressive Growing of GANs for Improved Quality, Stability, and Variation](https://arxiv.org/abs/1710.10196).
|
||||
All three papers are from the same authors from [NVIDIA AI](https://twitter.com/NVIDIAAI).
|
||||
@@ -0,0 +1,133 @@
|
||||
r"""
|
||||
---
|
||||
title: Wasserstein GAN (WGAN)
|
||||
summary: A simple PyTorch implementation/tutorial of Wasserstein Generative Adversarial Networks (WGAN) loss functions.
|
||||
---
|
||||
|
||||
# Wasserstein GAN (WGAN)
|
||||
|
||||
This is an implementation of
|
||||
[Wasserstein GAN](https://arxiv.org/abs/1701.07875).
|
||||
|
||||
The original GAN loss is based on Jensen-Shannon (JS) divergence
|
||||
between the real distribution $\mathbb{P}_r$ and generated distribution $\mathbb{P}_g$.
|
||||
The Wasserstein GAN is based on Earth Mover distance between these distributions.
|
||||
|
||||
$$
|
||||
W(\mathbb{P}_r, \mathbb{P}_g) =
|
||||
\underset{\gamma \in \Pi(\mathbb{P}_r, \mathbb{P}_g)} {\mathrm{inf}}
|
||||
\mathbb{E}_{(x,y) \sim \gamma}
|
||||
\Vert x - y \Vert
|
||||
$$
|
||||
|
||||
$\Pi(\mathbb{P}_r, \mathbb{P}_g)$ is the set of all joint distributions, whose
|
||||
marginal probabilities are $\gamma(x, y)$.
|
||||
|
||||
$\mathbb{E}_{(x,y) \sim \gamma} \Vert x - y \Vert$ is the earth mover distance for
|
||||
a given joint distribution ($x$ and $y$ are probabilities).
|
||||
|
||||
So $W(\mathbb{P}_r, \mathbb{P}_g)$ is equal to the least earth mover distance for
|
||||
any joint distribution between the real distribution $\mathbb{P}_r$ and generated distribution $\mathbb{P}_g$.
|
||||
|
||||
The paper shows that Jensen-Shannon (JS) divergence and other measures for the difference between two probability
|
||||
distributions are not smooth. And therefore if we are doing gradient descent on one of the probability
|
||||
distributions (parameterized) it will not converge.
|
||||
|
||||
Based on Kantorovich-Rubinstein duality,
|
||||
$$
|
||||
W(\mathbb{P}_r, \mathbb{P}_g) =
|
||||
\underset{\Vert f \Vert_L \le 1} {\mathrm{sup}}
|
||||
\mathbb{E}_{x \sim \mathbb{P}_r} [f(x)]- \mathbb{E}_{x \sim \mathbb{P}_g} [f(x)]
|
||||
$$
|
||||
|
||||
where $\Vert f \Vert_L \le 1$ are all 1-Lipschitz functions.
|
||||
|
||||
That is, it is equal to the greatest difference
|
||||
$$\mathbb{E}_{x \sim \mathbb{P}_r} [f(x)] - \mathbb{E}_{x \sim \mathbb{P}_g} [f(x)]$$
|
||||
among all 1-Lipschitz functions.
|
||||
|
||||
For $K$-Lipschitz functions,
|
||||
$$
|
||||
W(\mathbb{P}_r, \mathbb{P}_g) =
|
||||
\underset{\Vert f \Vert_L \le K} {\mathrm{sup}}
|
||||
\mathbb{E}_{x \sim \mathbb{P}_r} \Bigg[\frac{1}{K} f(x) \Bigg]
|
||||
- \mathbb{E}_{x \sim \mathbb{P}_g} \Bigg[\frac{1}{K} f(x) \Bigg]
|
||||
$$
|
||||
|
||||
If all $K$-Lipschitz functions can be represented as $f_w$ where $f$ is parameterized by
|
||||
$w \in \mathcal{W}$,
|
||||
|
||||
$$
|
||||
K \cdot W(\mathbb{P}_r, \mathbb{P}_g) =
|
||||
\max_{w \in \mathcal{W}}
|
||||
\mathbb{E}_{x \sim \mathbb{P}_r} [f_w(x)]- \mathbb{E}_{x \sim \mathbb{P}_g} [f_w(x)]
|
||||
$$
|
||||
|
||||
If $(\mathbb{P}_{g})$ is represented by a generator $$g_\theta (z)$$ and $z$ is from a known
|
||||
distribution $z \sim p(z)$,
|
||||
|
||||
$$
|
||||
K \cdot W(\mathbb{P}_r, \mathbb{P}_\theta) =
|
||||
\max_{w \in \mathcal{W}}
|
||||
\mathbb{E}_{x \sim \mathbb{P}_r} [f_w(x)]- \mathbb{E}_{z \sim p(z)} [f_w(g_\theta(z))]
|
||||
$$
|
||||
|
||||
Now to converge $g_\theta$ with $\mathbb{P}_{r}$ we can gradient descent on $\theta$
|
||||
to minimize above formula.
|
||||
|
||||
Similarly we can find $\max_{w \in \mathcal{W}}$ by ascending on $w$,
|
||||
while keeping $K$ bounded. *One way to keep $K$ bounded is to clip all weights in the neural
|
||||
network that defines $f$ clipped within a range.*
|
||||
|
||||
Here is the code to try this on a [simple MNIST generation experiment](experiment.html).
|
||||
|
||||
[](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/gan/wasserstein/experiment.ipynb)
|
||||
"""
|
||||
|
||||
import torch.utils.data
|
||||
from torch import nn
|
||||
from torch.nn import functional as F
|
||||
|
||||
|
||||
class DiscriminatorLoss(nn.Module):
|
||||
"""
|
||||
## Discriminator Loss
|
||||
|
||||
We want to find $w$ to maximize
|
||||
$$\mathbb{E}_{x \sim \mathbb{P}_r} [f_w(x)]- \mathbb{E}_{z \sim p(z)} [f_w(g_\theta(z))]$$,
|
||||
so we minimize,
|
||||
$$-\frac{1}{m} \sum_{i=1}^m f_w \big(x^{(i)} \big) +
|
||||
\frac{1}{m} \sum_{i=1}^m f_w \big( g_\theta(z^{(i)}) \big)$$
|
||||
"""
|
||||
|
||||
def forward(self, f_real: torch.Tensor, f_fake: torch.Tensor):
|
||||
"""
|
||||
* `f_real` is $f_w(x)$
|
||||
* `f_fake` is $f_w(g_\theta(z))$
|
||||
|
||||
This returns the a tuple with losses for $f_w(x)$ and $f_w(g_\theta(z))$,
|
||||
which are later added.
|
||||
They are kept separate for logging.
|
||||
"""
|
||||
|
||||
# We use ReLUs to clip the loss to keep $f \in [-1, +1]$ range.
|
||||
return F.relu(1 - f_real).mean(), F.relu(1 + f_fake).mean()
|
||||
|
||||
|
||||
class GeneratorLoss(nn.Module):
|
||||
"""
|
||||
## Generator Loss
|
||||
|
||||
We want to find $\theta$ to minimize
|
||||
$$\mathbb{E}_{x \sim \mathbb{P}_r} [f_w(x)]- \mathbb{E}_{z \sim p(z)} [f_w(g_\theta(z))]$$
|
||||
The first component is independent of $\theta$,
|
||||
so we minimize,
|
||||
$$-\frac{1}{m} \sum_{i=1}^m f_w \big( g_\theta(z^{(i)}) \big)$$
|
||||
|
||||
"""
|
||||
|
||||
def forward(self, f_fake: torch.Tensor):
|
||||
"""
|
||||
* `f_fake` is $f_w(g_\theta(z))$
|
||||
"""
|
||||
return -f_fake.mean()
|
||||
@@ -0,0 +1,189 @@
|
||||
{
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"name": "Cycle GAN",
|
||||
"provenance": [],
|
||||
"collapsed_sections": [],
|
||||
"toc_visible": true
|
||||
},
|
||||
"kernelspec": {
|
||||
"name": "python3",
|
||||
"language": "python",
|
||||
"display_name": "Python 3"
|
||||
},
|
||||
"accelerator": "GPU"
|
||||
},
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "AYV_dMVDxyc2"
|
||||
},
|
||||
"source": [
|
||||
"[](https://github.com/labmlai/annotated_deep_learning_paper_implementations)\n",
|
||||
"[](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/gan/wasserstein/experiment.ipynb)\n",
|
||||
"\n",
|
||||
"## DCGAN\n",
|
||||
"\n",
|
||||
"This is an experiment training DCGAN model."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "AahG_i2y5tY9"
|
||||
},
|
||||
"source": [
|
||||
"Install the `labml-nn` package"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "ZCzmCrAIVg0L",
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"outputId": "2fe2685f-731c-4c47-854e-a4f00e485281"
|
||||
},
|
||||
"source": [
|
||||
"!pip install labml-nn"
|
||||
],
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "SE2VUQ6L5zxI"
|
||||
},
|
||||
"source": [
|
||||
"Imports"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "0hJXx_g0wS2C"
|
||||
},
|
||||
"source": [
|
||||
"\n",
|
||||
"from labml import experiment\n",
|
||||
"from labml_nn.gan.wasserstein.experiment import Configs"
|
||||
],
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "Lpggo0wM6qb-"
|
||||
},
|
||||
"source": [
|
||||
"Create an experiment"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "bFcr9k-l4cAg"
|
||||
},
|
||||
"source": [
|
||||
"experiment.create(name=\"mnist_wgan\")"
|
||||
],
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "-OnHLi626tJt"
|
||||
},
|
||||
"source": [
|
||||
"Initialize configurations"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "Piz0c5f44hRo"
|
||||
},
|
||||
"source": [
|
||||
"conf = Configs()"
|
||||
],
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "wwMzCqpD6vkL"
|
||||
},
|
||||
"source": [
|
||||
"Set experiment configurations and assign a configurations dictionary to override configurations"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 17
|
||||
},
|
||||
"id": "e6hmQhTw4nks",
|
||||
"outputId": "4be767af-0ebd-4c35-8da0-0e532495e037"
|
||||
},
|
||||
"source": [
|
||||
"experiment.configs(conf,\n",
|
||||
" {\n",
|
||||
" 'discriminator': 'cnn',\n",
|
||||
" 'generator': 'cnn',\n",
|
||||
" 'label_smoothing': 0.01,\n",
|
||||
" 'generator_loss': 'wasserstein',\n",
|
||||
" 'discriminator_loss': 'wasserstein',\n",
|
||||
" })"
|
||||
],
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "KJZRf8527GxL"
|
||||
},
|
||||
"source": [
|
||||
"Start the experiment and run the training loop."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 649
|
||||
},
|
||||
"id": "aIAWo7Fw5DR8",
|
||||
"outputId": "e3b02247-8ff9-47b5-8f52-49c9e3b8377f"
|
||||
},
|
||||
"source": [
|
||||
"with experiment.start():\n",
|
||||
" conf.run()"
|
||||
],
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "oBXXlP2b7XZO"
|
||||
},
|
||||
"source": [
|
||||
""
|
||||
],
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,44 @@
|
||||
"""
|
||||
---
|
||||
title: WGAN experiment with MNIST
|
||||
summary: This experiment generates MNIST images using convolutional neural network.
|
||||
---
|
||||
|
||||
# WGAN experiment with MNIST
|
||||
"""
|
||||
from labml import experiment
|
||||
|
||||
from labml.configs import calculate
|
||||
# Import configurations from [DCGAN experiment](../dcgan/index.html)
|
||||
from labml_nn.gan.dcgan import Configs
|
||||
|
||||
# Import [Wasserstein GAN losses](./index.html)
|
||||
from labml_nn.gan.wasserstein import GeneratorLoss, DiscriminatorLoss
|
||||
|
||||
# Set configurations options for Wasserstein GAN losses
|
||||
calculate(Configs.generator_loss, 'wasserstein', lambda c: GeneratorLoss())
|
||||
calculate(Configs.discriminator_loss, 'wasserstein', lambda c: DiscriminatorLoss())
|
||||
|
||||
|
||||
def main():
|
||||
# Create configs object
|
||||
conf = Configs()
|
||||
# Create experiment
|
||||
experiment.create(name='mnist_wassertein_dcgan', comment='test')
|
||||
# Override configurations
|
||||
experiment.configs(conf,
|
||||
{
|
||||
'discriminator': 'cnn',
|
||||
'generator': 'cnn',
|
||||
'label_smoothing': 0.01,
|
||||
'generator_loss': 'wasserstein',
|
||||
'discriminator_loss': 'wasserstein',
|
||||
})
|
||||
|
||||
# Start the experiment and run training loop
|
||||
with experiment.start():
|
||||
conf.run()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,83 @@
|
||||
r"""
|
||||
---
|
||||
title: Gradient Penalty for Wasserstein GAN (WGAN-GP)
|
||||
summary: >
|
||||
An annotated PyTorch implementation/tutorial of
|
||||
Improved Training of Wasserstein GANs.
|
||||
---
|
||||
|
||||
# Gradient Penalty for Wasserstein GAN (WGAN-GP)
|
||||
|
||||
This is an implementation of
|
||||
[Improved Training of Wasserstein GANs](https://arxiv.org/abs/1704.00028).
|
||||
|
||||
[WGAN](../index.html) suggests clipping weights to enforce Lipschitz constraint
|
||||
on the discriminator network (critic).
|
||||
This and other weight constraints like L2 norm clipping, weight normalization,
|
||||
L1, L2 weight decay have problems:
|
||||
|
||||
1. Limiting the capacity of the discriminator
|
||||
2. Exploding and vanishing gradients (without [Batch Normalization](../../../normalization/batch_norm/index.html)).
|
||||
|
||||
The paper [Improved Training of Wasserstein GANs](https://arxiv.org/abs/1704.00028)
|
||||
proposal a better way to improve Lipschitz constraint, a gradient penalty.
|
||||
|
||||
$$\mathcal{L}_{GP} = \lambda \underset{\hat{x} \sim \mathbb{P}_{\hat{x}}}{\mathbb{E}}
|
||||
\Big[ \big(\Vert \nabla_{\hat{x}} D(\hat{x}) \Vert_2 - 1\big)^2 \Big]
|
||||
$$
|
||||
|
||||
where $\lambda$ is the penalty weight and
|
||||
|
||||
\begin{align}
|
||||
x &\sim \mathbb{P}_r \\
|
||||
z &\sim p(z) \\
|
||||
\epsilon &\sim U[0,1] \\
|
||||
\tilde{x} &\leftarrow G_\theta (z) \\
|
||||
\hat{x} &\leftarrow \epsilon x + (1 - \epsilon) \tilde{x}
|
||||
\end{align}
|
||||
|
||||
That is we try to keep the gradient norm $\Vert \nabla_{\hat{x}} D(\hat{x}) \Vert_2$ close to $1$.
|
||||
|
||||
In this implementation we set $\epsilon = 1$.
|
||||
|
||||
Here is the [code for an experiment](experiment.html) that uses gradient penalty.
|
||||
"""
|
||||
|
||||
import torch
|
||||
import torch.autograd
|
||||
|
||||
from torch import nn
|
||||
|
||||
|
||||
class GradientPenalty(nn.Module):
|
||||
"""
|
||||
## Gradient Penalty
|
||||
"""
|
||||
|
||||
def forward(self, x: torch.Tensor, f: torch.Tensor):
|
||||
"""
|
||||
* `x` is $x \sim \mathbb{P}_r$
|
||||
* `f` is $D(x)$
|
||||
|
||||
$\hat{x} \leftarrow x$
|
||||
since we set $\epsilon = 1$ for this implementation.
|
||||
"""
|
||||
|
||||
# Get batch size
|
||||
batch_size = x.shape[0]
|
||||
|
||||
# Calculate gradients of $D(x)$ with respect to $x$.
|
||||
# `grad_outputs` is set to ones since we want the gradients of $D(x)$,
|
||||
# and we need to create and retain graph since we have to compute gradients
|
||||
# with respect to weight on this loss.
|
||||
gradients, *_ = torch.autograd.grad(outputs=f,
|
||||
inputs=x,
|
||||
grad_outputs=f.new_ones(f.shape),
|
||||
create_graph=True)
|
||||
|
||||
# Reshape gradients to calculate the norm
|
||||
gradients = gradients.reshape(batch_size, -1)
|
||||
# Calculate the norm $\Vert \nabla_{\hat{x}} D(\hat{x}) \Vert_2$
|
||||
norm = gradients.norm(2, dim=-1)
|
||||
# Return the loss $\big(\Vert \nabla_{\hat{x}} D(\hat{x}) \Vert_2 - 1\big)^2$
|
||||
return torch.mean((norm - 1) ** 2)
|
||||
@@ -0,0 +1,86 @@
|
||||
"""
|
||||
---
|
||||
title: WGAN-GP experiment with MNIST
|
||||
summary: This experiment generates MNIST images using convolutional neural network.
|
||||
---
|
||||
|
||||
# WGAN-GP experiment with MNIST
|
||||
"""
|
||||
|
||||
import torch
|
||||
|
||||
from labml import experiment, tracker
|
||||
# Import configurations from [Wasserstein experiment](../experiment.html)
|
||||
from labml_nn.gan.wasserstein.experiment import Configs as OriginalConfigs
|
||||
#
|
||||
from labml_nn.gan.wasserstein.gradient_penalty import GradientPenalty
|
||||
|
||||
|
||||
class Configs(OriginalConfigs):
|
||||
"""
|
||||
## Configuration class
|
||||
|
||||
We extend [original GAN implementation](../../original/experiment.html) and override the discriminator (critic) loss
|
||||
calculation to include gradient penalty.
|
||||
"""
|
||||
|
||||
# Gradient penalty coefficient $\lambda$
|
||||
gradient_penalty_coefficient: float = 10.0
|
||||
#
|
||||
gradient_penalty = GradientPenalty()
|
||||
|
||||
def calc_discriminator_loss(self, data: torch.Tensor):
|
||||
"""
|
||||
This overrides the original discriminator loss calculation and
|
||||
includes gradient penalty.
|
||||
"""
|
||||
# Require gradients on $x$ to calculate gradient penalty
|
||||
data.requires_grad_()
|
||||
# Sample $z \sim p(z)$
|
||||
latent = self.sample_z(data.shape[0])
|
||||
# $D(x)$
|
||||
f_real = self.discriminator(data)
|
||||
# $D(G_\theta(z))$
|
||||
f_fake = self.discriminator(self.generator(latent).detach())
|
||||
# Get discriminator losses
|
||||
loss_true, loss_false = self.discriminator_loss(f_real, f_fake)
|
||||
# Calculate gradient penalties in training mode
|
||||
if self.mode.is_train:
|
||||
gradient_penalty = self.gradient_penalty(data, f_real)
|
||||
tracker.add("loss.gp.", gradient_penalty)
|
||||
loss = loss_true + loss_false + self.gradient_penalty_coefficient * gradient_penalty
|
||||
# Skip gradient penalty otherwise
|
||||
else:
|
||||
loss = loss_true + loss_false
|
||||
|
||||
# Log stuff
|
||||
tracker.add("loss.discriminator.true.", loss_true)
|
||||
tracker.add("loss.discriminator.false.", loss_false)
|
||||
tracker.add("loss.discriminator.", loss)
|
||||
|
||||
return loss
|
||||
|
||||
|
||||
def main():
|
||||
# Create configs object
|
||||
conf = Configs()
|
||||
# Create experiment
|
||||
experiment.create(name='mnist_wassertein_gp_dcgan')
|
||||
# Override configurations
|
||||
experiment.configs(conf,
|
||||
{
|
||||
'discriminator': 'cnn',
|
||||
'generator': 'cnn',
|
||||
'label_smoothing': 0.01,
|
||||
'generator_loss': 'wasserstein',
|
||||
'discriminator_loss': 'wasserstein',
|
||||
'discriminator_k': 5,
|
||||
})
|
||||
|
||||
# Start the experiment and run training loop
|
||||
with experiment.start():
|
||||
conf.run()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,16 @@
|
||||
# [Gradient Penalty for Wasserstein GAN (WGAN-GP)](https://nn.labml.ai/gan/wasserstein/gradient_penalty/index.html)
|
||||
|
||||
This is an implementation of
|
||||
[Improved Training of Wasserstein GANs](https://arxiv.org/abs/1704.00028).
|
||||
|
||||
[WGAN](https://nn.labml.ai/gan/wasserstein/index.html) suggests
|
||||
clipping weights to enforce Lipschitz constraint
|
||||
on the discriminator network (critic).
|
||||
This and other weight constraints like L2 norm clipping, weight normalization,
|
||||
L1, L2 weight decay have problems:
|
||||
|
||||
1. Limiting the capacity of the discriminator
|
||||
2. Exploding and vanishing gradients (without [Batch Normalization](https://nn.labml.ai/normalization/batch_norm/index.html)).
|
||||
|
||||
The paper [Improved Training of Wasserstein GANs](https://arxiv.org/abs/1704.00028)
|
||||
proposal a better way to improve Lipschitz constraint, a gradient penalty.
|
||||
@@ -0,0 +1,4 @@
|
||||
# [Wasserstein GAN - WGAN](https://nn.labml.ai/gan/wasserstein/index.html)
|
||||
|
||||
This is an implementation of
|
||||
[Wasserstein GAN](https://arxiv.org/abs/1701.07875).
|
||||
Reference in New Issue
Block a user