chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -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).
|
||||
Reference in New Issue
Block a user