chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,155 @@
|
||||
"""
|
||||
# [Annotated Research Paper Implementations: Transformers, StyleGAN, Stable Diffusion, DDPM/DDIM, LayerNorm, Nucleus Sampling and more](index.html)
|
||||
|
||||
This is a collection of simple PyTorch implementations of
|
||||
neural networks and related algorithms.
|
||||
[These implementations](https://github.com/labmlai/annotated_deep_learning_paper_implementations) are documented with explanations,
|
||||
and the [website](index.html)
|
||||
renders these as side-by-side formatted notes.
|
||||
We believe these would help you understand these algorithms better.
|
||||
|
||||
|
||||
|
||||
We are actively maintaining this repo and adding new
|
||||
implementations.
|
||||
[](https://twitter.com/labmlai) for updates.
|
||||
|
||||
## Translations
|
||||
|
||||
### **[English (original)](https://nn.labml.ai)**
|
||||
### **[Chinese (translated)](https://nn.labml.ai/zh/)**
|
||||
### **[Japanese (translated)](https://nn.labml.ai/ja/)**
|
||||
|
||||
## Paper Implementations
|
||||
|
||||
#### ✨ [Transformers](transformers/index.html)
|
||||
|
||||
* [JAX implementation](transformers/jax_transformer/index.html)
|
||||
* [Multi-headed attention](transformers/mha.html)
|
||||
* [Triton Flash Attention](transformers/flash/index.html)
|
||||
* [Transformer building blocks](transformers/models.html)
|
||||
* [Transformer XL](transformers/xl/index.html)
|
||||
* [Relative multi-headed attention](transformers/xl/relative_mha.html)
|
||||
* [Rotary Positional Embeddings (RoPE)](transformers/rope/index.html)
|
||||
* [Attention with Linear Biases (ALiBi)](transformers/alibi/index.html)
|
||||
* [RETRO](transformers/retro/index.html)
|
||||
* [Compressive Transformer](transformers/compressive/index.html)
|
||||
* [GPT Architecture](transformers/gpt/index.html)
|
||||
* [GLU Variants](transformers/glu_variants/simple.html)
|
||||
* [kNN-LM: Generalization through Memorization](transformers/knn/index.html)
|
||||
* [Feedback Transformer](transformers/feedback/index.html)
|
||||
* [Switch Transformer](transformers/switch/index.html)
|
||||
* [Fast Weights Transformer](transformers/fast_weights/index.html)
|
||||
* [FNet](transformers/fnet/index.html)
|
||||
* [Attention Free Transformer](transformers/aft/index.html)
|
||||
* [Masked Language Model](transformers/mlm/index.html)
|
||||
* [MLP-Mixer: An all-MLP Architecture for Vision](transformers/mlp_mixer/index.html)
|
||||
* [Pay Attention to MLPs (gMLP)](transformers/gmlp/index.html)
|
||||
* [Vision Transformer (ViT)](transformers/vit/index.html)
|
||||
* [Primer EZ](transformers/primer_ez/index.html)
|
||||
* [Hourglass](transformers/hour_glass/index.html)
|
||||
|
||||
#### ✨ [Low-Rank Adaptation (LoRA)](lora/index.html)
|
||||
|
||||
#### ✨ [Eleuther GPT-NeoX](neox/index.html)
|
||||
* [Generate on a 48GB GPU](neox/samples/generate.html)
|
||||
* [Finetune on two 48GB GPUs](neox/samples/finetune.html)
|
||||
* [LLM.int8()](neox/utils/llm_int8.html)
|
||||
|
||||
#### ✨ [Diffusion models](diffusion/index.html)
|
||||
|
||||
* [Denoising Diffusion Probabilistic Models (DDPM)](diffusion/ddpm/index.html)
|
||||
* [Denoising Diffusion Implicit Models (DDIM)](diffusion/stable_diffusion/sampler/ddim.html)
|
||||
* [Latent Diffusion Models](diffusion/stable_diffusion/latent_diffusion.html)
|
||||
* [Stable Diffusion](diffusion/stable_diffusion/index.html)
|
||||
|
||||
#### ✨ [Generative Adversarial Networks](gan/index.html)
|
||||
* [Original GAN](gan/original/index.html)
|
||||
* [GAN with deep convolutional network](gan/dcgan/index.html)
|
||||
* [Cycle GAN](gan/cycle_gan/index.html)
|
||||
* [Wasserstein GAN](gan/wasserstein/index.html)
|
||||
* [Wasserstein GAN with Gradient Penalty](gan/wasserstein/gradient_penalty/index.html)
|
||||
* [StyleGAN 2](gan/stylegan/index.html)
|
||||
|
||||
#### ✨ [Recurrent Highway Networks](recurrent_highway_networks/index.html)
|
||||
|
||||
#### ✨ [LSTM](lstm/index.html)
|
||||
|
||||
#### ✨ [HyperNetworks - HyperLSTM](hypernetworks/hyper_lstm.html)
|
||||
|
||||
#### ✨ [ResNet](resnet/index.html)
|
||||
|
||||
#### ✨ [ConvMixer](conv_mixer/index.html)
|
||||
|
||||
#### ✨ [Capsule Networks](capsule_networks/index.html)
|
||||
|
||||
#### ✨ [U-Net](unet/index.html)
|
||||
|
||||
#### ✨ [Sketch RNN](sketch_rnn/index.html)
|
||||
|
||||
#### ✨ Graph Neural Networks
|
||||
|
||||
* [Graph Attention Networks (GAT)](graphs/gat/index.html)
|
||||
* [Graph Attention Networks v2 (GATv2)](graphs/gatv2/index.html)
|
||||
|
||||
#### ✨ [Reinforcement Learning](rl/index.html)
|
||||
* [Proximal Policy Optimization](rl/ppo/index.html) with
|
||||
[Generalized Advantage Estimation](rl/ppo/gae.html)
|
||||
* [Deep Q Networks](rl/dqn/index.html) with
|
||||
with [Dueling Network](rl/dqn/model.html),
|
||||
[Prioritized Replay](rl/dqn/replay_buffer.html)
|
||||
and Double Q Network.
|
||||
|
||||
#### ✨ [Counterfactual Regret Minimization (CFR)](cfr/index.html)
|
||||
|
||||
Solving games with incomplete information such as poker with CFR.
|
||||
|
||||
* [Kuhn Poker](cfr/kuhn/index.html)
|
||||
|
||||
#### ✨ [Optimizers](optimizers/index.html)
|
||||
* [Adam](optimizers/adam.html)
|
||||
* [AMSGrad](optimizers/amsgrad.html)
|
||||
* [Adam Optimizer with warmup](optimizers/adam_warmup.html)
|
||||
* [Noam Optimizer](optimizers/noam.html)
|
||||
* [Rectified Adam Optimizer](optimizers/radam.html)
|
||||
* [AdaBelief Optimizer](optimizers/ada_belief.html)
|
||||
* [Sophia-G Optimizer](optimizers/sophia.html)
|
||||
|
||||
#### ✨ [Normalization Layers](normalization/index.html)
|
||||
* [Batch Normalization](normalization/batch_norm/index.html)
|
||||
* [Layer Normalization](normalization/layer_norm/index.html)
|
||||
* [Instance Normalization](normalization/instance_norm/index.html)
|
||||
* [Group Normalization](normalization/group_norm/index.html)
|
||||
* [Weight Standardization](normalization/weight_standardization/index.html)
|
||||
* [Batch-Channel Normalization](normalization/batch_channel_norm/index.html)
|
||||
* [DeepNorm](normalization/deep_norm/index.html)
|
||||
|
||||
#### ✨ [Distillation](distillation/index.html)
|
||||
|
||||
#### ✨ [Adaptive Computation](adaptive_computation/index.html)
|
||||
|
||||
* [PonderNet](adaptive_computation/ponder_net/index.html)
|
||||
|
||||
#### ✨ [Uncertainty](uncertainty/index.html)
|
||||
|
||||
* [Evidential Deep Learning to Quantify Classification Uncertainty](uncertainty/evidence/index.html)
|
||||
|
||||
#### ✨ [Activations](activations/index.html)
|
||||
|
||||
* [Fuzzy Tiling Activations](activations/fta/index.html)
|
||||
|
||||
#### ✨ [Language Model Sampling Techniques](sampling/index.html)
|
||||
* [Greedy Sampling](sampling/greedy.html)
|
||||
* [Temperature Sampling](sampling/temperature.html)
|
||||
* [Top-k Sampling](sampling/top_k.html)
|
||||
* [Nucleus Sampling](sampling/nucleus.html)
|
||||
|
||||
#### ✨ [Scalable Training/Inference](scaling/index.html)
|
||||
* [Zero3 memory optimizations](scaling/zero3/index.html)
|
||||
|
||||
### Installation
|
||||
|
||||
```bash
|
||||
pip install labml-nn
|
||||
```
|
||||
"""
|
||||
@@ -0,0 +1,14 @@
|
||||
"""
|
||||
---
|
||||
title: Neural Network Activation Functions
|
||||
summary: >
|
||||
A set of PyTorch implementations/tutorials related to neural network activations
|
||||
---
|
||||
|
||||
# Neural Networks Activations
|
||||
|
||||
* [Fuzzy Tiling Activations](fta/index.html)
|
||||
* 🚧 [Swish](swish/index.html)
|
||||
"""
|
||||
|
||||
from .swish import Swish
|
||||
@@ -0,0 +1,131 @@
|
||||
"""
|
||||
---
|
||||
title: Fuzzy Tiling Activations
|
||||
summary: >
|
||||
PyTorch implementation and tutorial of Fuzzy Tiling Activations from the
|
||||
paper Fuzzy Tiling Activations: A Simple Approach to Learning Sparse Representations Online.
|
||||
---
|
||||
|
||||
# Fuzzy Tiling Activations (FTA)
|
||||
|
||||
[](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/activations/fta/experiment.ipynb)
|
||||
|
||||
This is a [PyTorch](https://pytorch.org) implementation/tutorial of
|
||||
[Fuzzy Tiling Activations: A Simple Approach to Learning Sparse Representations Online](https://arxiv.org/abs/1911.08068).
|
||||
|
||||
Fuzzy tiling activations are a form of sparse activations based on binning.
|
||||
|
||||
Binning is classification of a scalar value into a bin based on intervals.
|
||||
One problem with binning is that it gives zero gradients for most values (except at the boundary of bins).
|
||||
The other is that binning loses precision if the bin intervals are large.
|
||||
|
||||
FTA overcomes these disadvantages.
|
||||
Instead of hard boundaries like in Tiling Activations, FTA uses soft boundaries
|
||||
between bins.
|
||||
This gives non-zero gradients for all or a wide range of values.
|
||||
And also doesn't lose precision since it's captured in partial values.
|
||||
|
||||
#### Tiling Activations
|
||||
|
||||
$\mathbf{c}$ is the tiling vector,
|
||||
|
||||
$$\mathbf{c} = (l, l + \delta, l + 2 \delta, \dots, u - 2 \delta, u - \delta)$$
|
||||
|
||||
where $[l, u]$ is the input range, $\delta$ is the bin size, and $u - l$ is divisible by $\delta$.
|
||||
|
||||
Tiling activation is,
|
||||
|
||||
$$\phi(z) = 1 - I_+ \big( \max(\mathbf{c} - z, 0) + \max(z - \delta - \mathbf{c}) \big)$$
|
||||
|
||||
where $I_+(\cdot)$ is the indicator function which gives $1$ if the input is positive and $0$ otherwise.
|
||||
|
||||
Note that tiling activation gives zero gradients because it has hard boundaries.
|
||||
|
||||
#### Fuzzy Tiling Activations
|
||||
|
||||
The fuzzy indicator function,
|
||||
|
||||
$$I_{\eta,+}(x) = I_+(\eta - x) x + I_+ (x - \eta)$$
|
||||
|
||||
which increases linearly from $0$ to $1$ when $0 \le x \lt \eta$
|
||||
and is equal to $1$ for $\eta \le x$.
|
||||
$\eta$ is a hyper-parameter.
|
||||
|
||||
FTA uses this to create soft boundaries between bins.
|
||||
|
||||
$$\phi_\eta(z) = 1 - I_{\eta,+} \big( \max(\mathbf{c} - z, 0) + \max(z - \delta - \mathbf{c}, 0) \big)$$
|
||||
|
||||
[Here's a simple experiment](experiment.html) that uses FTA in a transformer.
|
||||
"""
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
|
||||
class FTA(nn.Module):
|
||||
"""
|
||||
### Fuzzy Tiling Activations (FTA)
|
||||
"""
|
||||
|
||||
def __init__(self, lower_limit: float, upper_limit: float, delta: float, eta: float):
|
||||
"""
|
||||
:param lower_limit: is the lower limit $l$
|
||||
:param upper_limit: is the upper limit $u$
|
||||
:param delta: is the bin size $\delta$
|
||||
:param eta: is the parameter $\eta$ that detemines the softness of the boundaries.
|
||||
"""
|
||||
super().__init__()
|
||||
# Initialize tiling vector
|
||||
# $$\mathbf{c} = (l, l + \delta, l + 2 \delta, \dots, u - 2 \delta, u - \delta)$$
|
||||
self.c = nn.Parameter(torch.arange(lower_limit, upper_limit, delta), requires_grad=False)
|
||||
# The input vector expands by a factor equal to the number of bins $\frac{u - l}{\delta}$
|
||||
self.expansion_factor = len(self.c)
|
||||
# $\delta$
|
||||
self.delta = delta
|
||||
# $\eta$
|
||||
self.eta = eta
|
||||
|
||||
def fuzzy_i_plus(self, x: torch.Tensor):
|
||||
"""
|
||||
#### Fuzzy indicator function
|
||||
|
||||
$$I_{\eta,+}(x) = I_+(\eta - x) x + I_+ (x - \eta)$$
|
||||
"""
|
||||
return (x <= self.eta) * x + (x > self.eta)
|
||||
|
||||
def forward(self, z: torch.Tensor):
|
||||
# Add another dimension of size $1$.
|
||||
# We will expand this into bins.
|
||||
z = z.view(*z.shape, 1)
|
||||
|
||||
# $$\phi_\eta(z) = 1 - I_{\eta,+} \big( \max(\mathbf{c} - z, 0) + \max(z - \delta - \mathbf{c}, 0) \big)$$
|
||||
z = 1. - self.fuzzy_i_plus(torch.clip(self.c - z, min=0.) + torch.clip(z - self.delta - self.c, min=0.))
|
||||
|
||||
# Reshape back to original number of dimensions.
|
||||
# The last dimension size gets expanded by the number of bins, $\frac{u - l}{\delta}$.
|
||||
return z.view(*z.shape[:-2], -1)
|
||||
|
||||
|
||||
def _test():
|
||||
"""
|
||||
#### Code to test the FTA module
|
||||
"""
|
||||
from labml.logger import inspect
|
||||
|
||||
# Initialize
|
||||
a = FTA(-10, 10, 2., 0.5)
|
||||
# Print $\mathbf{c}$
|
||||
inspect(a.c)
|
||||
# Print number of bins $\frac{u - l}{\delta}$
|
||||
inspect(a.expansion_factor)
|
||||
|
||||
# Input $z$
|
||||
z = torch.tensor([1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9., 10., 11.])
|
||||
# Print $z$
|
||||
inspect(z)
|
||||
# Print $\phi_\eta(z)$
|
||||
inspect(a(z))
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
_test()
|
||||
@@ -0,0 +1,269 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "AYV_dMVDxyc2",
|
||||
"pycharm": {
|
||||
"name": "#%% md\n"
|
||||
}
|
||||
},
|
||||
"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/activations/fta/experiment.ipynb)\n",
|
||||
"\n",
|
||||
"## [Fuzzy Tiling Activations](https://nn.labml.ai/activations/fta/index.html)\n",
|
||||
"\n",
|
||||
"Here we train a transformer that uses [Fuzzy Tiling Activation](https://nn.labml.ai/activations/fta/index.html) in the\n",
|
||||
"[Feed-Forward Network](https://nn.labml.ai/transformers/feed_forward.html).\n",
|
||||
"We use it for a language model and train it on Tiny Shakespeare dataset\n",
|
||||
"for demonstration.\n",
|
||||
"However, this is probably not the ideal task for FTA, and we\n",
|
||||
"believe FTA is more suitable for modeling data with continuous variables."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "AahG_i2y5tY9",
|
||||
"pycharm": {
|
||||
"name": "#%% md\n"
|
||||
}
|
||||
},
|
||||
"source": [
|
||||
"### Install the packages"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "ZCzmCrAIVg0L",
|
||||
"outputId": "cf107fb2-4d50-4c67-af34-367624553421",
|
||||
"pycharm": {
|
||||
"name": "#%%\n"
|
||||
}
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"!pip install labml-nn --quiet"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "SE2VUQ6L5zxI",
|
||||
"pycharm": {
|
||||
"name": "#%% md\n"
|
||||
}
|
||||
},
|
||||
"source": [
|
||||
"### Imports"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"import torch\n",
|
||||
"import torch.nn as nn\n",
|
||||
"\n",
|
||||
"from labml import experiment\n",
|
||||
"from labml.configs import option\n",
|
||||
"from labml_nn.activations.fta.experiment import Configs"
|
||||
],
|
||||
"metadata": {
|
||||
"collapsed": false,
|
||||
"pycharm": {
|
||||
"name": "#%%\n"
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"### Create an experiment"
|
||||
],
|
||||
"metadata": {
|
||||
"collapsed": false,
|
||||
"pycharm": {
|
||||
"name": "#%% md\n"
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"experiment.create(name=\"fta\", writers={'screen'})"
|
||||
],
|
||||
"metadata": {
|
||||
"collapsed": false,
|
||||
"pycharm": {
|
||||
"name": "#%%\n"
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"### Configurations"
|
||||
],
|
||||
"metadata": {
|
||||
"collapsed": false,
|
||||
"pycharm": {
|
||||
"name": "#%% md\n"
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"conf = Configs()"
|
||||
],
|
||||
"metadata": {
|
||||
"collapsed": false,
|
||||
"pycharm": {
|
||||
"name": "#%%\n"
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"Set experiment configurations and assign a configurations dictionary to override configurations"
|
||||
],
|
||||
"metadata": {
|
||||
"collapsed": false,
|
||||
"pycharm": {
|
||||
"name": "#%% md\n"
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"experiment.configs(conf, {\n",
|
||||
" 'tokenizer': 'character',\n",
|
||||
" 'prompt_separator': '',\n",
|
||||
" 'prompt': 'It is ',\n",
|
||||
" 'text': 'tiny_shakespeare',\n",
|
||||
"\n",
|
||||
" 'seq_len': 256,\n",
|
||||
" 'epochs': 32,\n",
|
||||
" 'batch_size': 16,\n",
|
||||
" 'inner_iterations': 10,\n",
|
||||
"\n",
|
||||
" 'optimizer.optimizer': 'Adam',\n",
|
||||
" 'optimizer.learning_rate': 3e-4,\n",
|
||||
"})"
|
||||
],
|
||||
"metadata": {
|
||||
"collapsed": false,
|
||||
"pycharm": {
|
||||
"name": "#%%\n"
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "EvI7MtgJ61w5",
|
||||
"pycharm": {
|
||||
"name": "#%% md\n"
|
||||
}
|
||||
},
|
||||
"source": [
|
||||
"Set PyTorch models for loading and saving"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 255
|
||||
},
|
||||
"id": "GDlt7dp-5ALt",
|
||||
"outputId": "e7548e8f-c541-4618-dc5a-1597cae42003",
|
||||
"pycharm": {
|
||||
"name": "#%%\n"
|
||||
}
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"experiment.add_pytorch_models({'model': conf.model})"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "KJZRf8527GxL",
|
||||
"pycharm": {
|
||||
"name": "#%% md\n"
|
||||
}
|
||||
},
|
||||
"source": [
|
||||
"### Start the experiment and run the training loop."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 1000
|
||||
},
|
||||
"id": "aIAWo7Fw5DR8",
|
||||
"outputId": "db979785-bfe3-4eda-d3eb-8ccbe61053e5",
|
||||
"pycharm": {
|
||||
"name": "#%%\n"
|
||||
}
|
||||
},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# Start the experiment\n",
|
||||
"with experiment.start():\n",
|
||||
" conf.run()"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"accelerator": "GPU",
|
||||
"colab": {
|
||||
"collapsed_sections": [],
|
||||
"name": "FTA",
|
||||
"provenance": []
|
||||
},
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3 (ipykernel)",
|
||||
"language": "python",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"codemirror_mode": {
|
||||
"name": "ipython",
|
||||
"version": 3
|
||||
},
|
||||
"file_extension": ".py",
|
||||
"mimetype": "text/x-python",
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython3",
|
||||
"version": "3.7.11"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 4
|
||||
}
|
||||
@@ -0,0 +1,218 @@
|
||||
"""
|
||||
---
|
||||
title: Fuzzy Tiling Activation Experiment
|
||||
summary: >
|
||||
Training a transformer with FTA in FFN on Tiny Shakespeare.
|
||||
---
|
||||
|
||||
# [Fuzzy Tiling Activation](index.html) Experiment
|
||||
|
||||
[](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/activations/fta/experiment.ipynb)
|
||||
|
||||
Here we train a transformer that uses [Fuzzy Tiling Activation](index.html) in the
|
||||
[Feed-Forward Network](../../transformers/feed_forward.html).
|
||||
We use it for a language model and train it on Tiny Shakespeare dataset
|
||||
for demonstration.
|
||||
|
||||
However, this is probably not the ideal task for FTA, and we
|
||||
believe FTA is more suitable for modeling data with continuous variables.
|
||||
"""
|
||||
|
||||
import copy
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from labml import experiment
|
||||
from labml.configs import option
|
||||
from labml_nn.activations.fta import FTA
|
||||
from labml_nn.experiments.nlp_autoregression import NLPAutoRegressionConfigs
|
||||
from labml_nn.transformers import MultiHeadAttention, TransformerLayer
|
||||
from labml_nn.transformers.utils import subsequent_mask
|
||||
|
||||
|
||||
class FeedForwardFTA(nn.Module):
|
||||
"""
|
||||
## FFN module with [FTA](index.html) activation
|
||||
"""
|
||||
|
||||
def __init__(self, d_model: int, d_ff: int,
|
||||
activation: FTA,
|
||||
dropout: float = 0.1):
|
||||
"""
|
||||
* `d_model` is the number of features in a token embedding
|
||||
* `d_ff` is the number of features in the hidden layer of the FFN
|
||||
* `activation` is FTA activation module
|
||||
* `dropout` is dropout probability for the hidden layer
|
||||
"""
|
||||
super().__init__()
|
||||
# Layer one parameterized by weight $W_1$ and bias $b_1$
|
||||
self.layer1 = nn.Linear(d_model, d_ff)
|
||||
# Layer two parameterized by weight $W_1$ and bias $b_1$
|
||||
self.layer2 = nn.Linear(d_ff * activation.expansion_factor, d_model)
|
||||
# Hidden layer dropout
|
||||
self.dropout = nn.Dropout(dropout)
|
||||
# Activation function $f$
|
||||
self.activation = activation
|
||||
|
||||
def forward(self, x: torch.Tensor):
|
||||
# $f(x W_1 + b_1)$
|
||||
x = self.activation(self.layer1(x))
|
||||
# Apply dropout
|
||||
x = self.dropout(x)
|
||||
#
|
||||
return self.layer2(x)
|
||||
|
||||
|
||||
class AutoregressiveTransformer(nn.Module):
|
||||
"""
|
||||
## Auto-Regressive model
|
||||
|
||||
This is an autoregressive transformer model that uses Feed-Forward Networks with
|
||||
(Fuzzy Tiling Activations)(index.html).
|
||||
"""
|
||||
|
||||
def __init__(self, n_tokens: int, d_model: int, n_layers: int, layer: TransformerLayer):
|
||||
"""
|
||||
:param n_tokens: is the number of tokens in the vocabulary
|
||||
:param d_model: is the embedding size
|
||||
:param n_layers: is the number of transformer layers
|
||||
:param layer: is the layer. We use `n_layers` copies of this for the transformer.
|
||||
"""
|
||||
super().__init__()
|
||||
# Transformer with `n_layers` layers
|
||||
self.transformer_layers = nn.ModuleList([copy.deepcopy(layer) for _ in range(n_layers)])
|
||||
|
||||
# Token embedding layer
|
||||
self.emb = nn.Embedding(n_tokens, d_model)
|
||||
# Readout layer
|
||||
self.readout = nn.Linear(d_model, n_tokens)
|
||||
|
||||
# The mask will be initialized on the first call
|
||||
self.mask = None
|
||||
|
||||
def forward(self, x: torch.Tensor):
|
||||
"""
|
||||
:param x: are the input tokens of shape `[seq_len, batch_size]`
|
||||
"""
|
||||
# Create auto-regressive mask
|
||||
if self.mask is None or self.mask.size(0) != len(x):
|
||||
# Subsequent mask, will mask out tokens from seeing future tokens
|
||||
self.mask = subsequent_mask(len(x)).to(x.device)
|
||||
|
||||
# Get the token embeddings
|
||||
x = self.emb(x)
|
||||
# Transformer encoder
|
||||
for layer in self.transformer_layers:
|
||||
x = layer(x=x, mask=self.mask)
|
||||
# Get logits
|
||||
x = self.readout(x)
|
||||
|
||||
# Return results
|
||||
return x, None
|
||||
|
||||
|
||||
class Configs(NLPAutoRegressionConfigs):
|
||||
"""
|
||||
## Configurations
|
||||
|
||||
This inherits from
|
||||
[`NLPAutoRegressionConfigs`](../../experiments/nlp_autoregression.html#NLPAutoRegressionConfigs)
|
||||
"""
|
||||
|
||||
# Model
|
||||
model: AutoregressiveTransformer
|
||||
|
||||
# Number of layers
|
||||
n_layers: int = 4
|
||||
|
||||
# $\alpha$ and $\beta$ for DeepNorm
|
||||
deep_norm_alpha: float
|
||||
deep_norm_beta: float
|
||||
|
||||
# Number of heads in the attention
|
||||
n_heads: int = 4
|
||||
# Embedding size
|
||||
d_model: int = 256
|
||||
# Size of each attention head
|
||||
d_k: int = 16
|
||||
# Feed forward layer size
|
||||
d_ff: int = 256
|
||||
|
||||
# FTA
|
||||
fta_lower_limit: float = -1.
|
||||
fta_upper_limit: float = +1.
|
||||
fta_delta: float = 0.2
|
||||
fta_eta: float = 0.05
|
||||
|
||||
|
||||
@option(Configs.model)
|
||||
def _model(c: Configs):
|
||||
"""
|
||||
#### Initialize the model
|
||||
"""
|
||||
|
||||
# Create FTA activation module
|
||||
fta = FTA(c.fta_lower_limit, c.fta_upper_limit, c.fta_delta, c.fta_eta)
|
||||
# Create the transformer.
|
||||
# We re-use [`TransformerLayer`](../../transformers/models.html#TransformerLayer) and
|
||||
# [`MultiHeadAttention`](../../transformers/mha.html) implementations.
|
||||
m = AutoregressiveTransformer(c.n_tokens, c.d_model, c.n_layers,
|
||||
TransformerLayer(d_model=c.d_model,
|
||||
feed_forward=FeedForwardFTA(d_model=c.d_model,
|
||||
d_ff=c.d_ff,
|
||||
activation=fta,
|
||||
dropout=0.1),
|
||||
self_attn=MultiHeadAttention(c.n_heads, c.d_model,
|
||||
dropout_prob=0.0),
|
||||
dropout_prob=0.0))
|
||||
|
||||
# Move to the device
|
||||
return m.to(c.device)
|
||||
|
||||
|
||||
def main():
|
||||
"""
|
||||
#### Create and run the experiment
|
||||
"""
|
||||
# Create experiment
|
||||
experiment.create(name="fta", writers={'screen', 'labml'})
|
||||
# Create configs
|
||||
conf = Configs()
|
||||
# Override configurations
|
||||
experiment.configs(conf, {
|
||||
# Use character level tokenizer
|
||||
'tokenizer': 'character',
|
||||
# Prompt separator is blank
|
||||
'prompt_separator': '',
|
||||
# Starting prompt for sampling
|
||||
'prompt': 'It is ',
|
||||
# Use Tiny Shakespeare dataset
|
||||
'text': 'tiny_shakespeare',
|
||||
|
||||
# Use a context size of $256$
|
||||
'seq_len': 256,
|
||||
# Train for 32 epochs
|
||||
'epochs': 32,
|
||||
# Batch size $16$
|
||||
'batch_size': 16,
|
||||
# Switch between training and validation for $10$ times per epoch
|
||||
'inner_iterations': 10,
|
||||
|
||||
# Adam optimizer with no warmup
|
||||
'optimizer.optimizer': 'Adam',
|
||||
'optimizer.learning_rate': 3e-4,
|
||||
})
|
||||
|
||||
# Set model(s) for saving and loading
|
||||
experiment.add_pytorch_models({'model': conf.model})
|
||||
|
||||
# Start the experiment
|
||||
with experiment.start():
|
||||
# Run training
|
||||
conf.run()
|
||||
|
||||
|
||||
#
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,12 @@
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
|
||||
|
||||
class Swish(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.sigmoid = nn.Sigmoid()
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
return x * self.sigmoid(x)
|
||||
@@ -0,0 +1,15 @@
|
||||
"""
|
||||
---
|
||||
title: Neural Networks with Adaptive Computation
|
||||
summary: >
|
||||
A set of PyTorch implementations/tutorials related to adaptive computation
|
||||
---
|
||||
|
||||
# Neural Networks with Adaptive Computation
|
||||
|
||||
These are neural network architectures that change the computation complexity based on the
|
||||
complexity of the input sample.
|
||||
|
||||
* 🚧 TODO: Adaptive Computation Time for Recurrent Neural Networks
|
||||
* [PonderNet: Learning to Ponder](ponder_net/index.html)
|
||||
"""
|
||||
@@ -0,0 +1,62 @@
|
||||
"""
|
||||
---
|
||||
title: "Parity Task"
|
||||
summary: >
|
||||
This creates data for Parity Task from the paper Adaptive Computation Time
|
||||
for Recurrent Neural Networks
|
||||
---
|
||||
|
||||
# Parity Task
|
||||
|
||||
This creates data for Parity Task from the paper
|
||||
[Adaptive Computation Time for Recurrent Neural Networks](https://arxiv.org/abs/1603.08983).
|
||||
|
||||
The input of the parity task is a vector with $0$'s $1$'s and $-1$'s.
|
||||
The output is the parity of $1$'s - one if there is an odd number of $1$'s and zero otherwise.
|
||||
The input is generated by making a random number of elements in the vector either $1$ or $-1$'s.
|
||||
"""
|
||||
|
||||
from typing import Tuple
|
||||
|
||||
import torch
|
||||
from torch.utils.data import Dataset
|
||||
|
||||
|
||||
class ParityDataset(Dataset):
|
||||
"""
|
||||
### Parity dataset
|
||||
"""
|
||||
|
||||
def __init__(self, n_samples: int, n_elems: int = 64):
|
||||
"""
|
||||
* `n_samples` is the number of samples
|
||||
* `n_elems` is the number of elements in the input vector
|
||||
"""
|
||||
self.n_samples = n_samples
|
||||
self.n_elems = n_elems
|
||||
|
||||
def __len__(self):
|
||||
"""
|
||||
Size of the dataset
|
||||
"""
|
||||
return self.n_samples
|
||||
|
||||
def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Generate a sample
|
||||
"""
|
||||
|
||||
# Empty vector
|
||||
x = torch.zeros((self.n_elems,))
|
||||
# Number of non-zero elements - a random number between $1$ and total number of elements
|
||||
n_non_zero = torch.randint(1, self.n_elems + 1, (1,)).item()
|
||||
# Fill non-zero elements with $1$'s and $-1$'s
|
||||
x[:n_non_zero] = torch.randint(0, 2, (n_non_zero,)) * 2 - 1
|
||||
# Randomly permute the elements
|
||||
x = x[torch.randperm(self.n_elems)]
|
||||
|
||||
# The parity
|
||||
y = (x == 1.).sum() % 2
|
||||
|
||||
#
|
||||
return x, y
|
||||
@@ -0,0 +1,267 @@
|
||||
"""
|
||||
---
|
||||
title: "PonderNet: Learning to Ponder"
|
||||
summary: >
|
||||
A PyTorch implementation/tutorial of PonderNet: Learning to Ponder.
|
||||
---
|
||||
|
||||
# PonderNet: Learning to Ponder
|
||||
|
||||
This is a [PyTorch](https://pytorch.org) implementation of the paper
|
||||
[PonderNet: Learning to Ponder](https://arxiv.org/abs/2107.05407).
|
||||
|
||||
PonderNet adapts the computation based on the input.
|
||||
It changes the number of steps to take on a recurrent network based on the input.
|
||||
PonderNet learns this with end-to-end gradient descent.
|
||||
|
||||
PonderNet has a step function of the form
|
||||
|
||||
$$\hat{y}_n, h_{n+1}, \lambda_n = s(x, h_n)$$
|
||||
|
||||
where $x$ is the input, $h_n$ is the state, $\hat{y}_n$ is the prediction at step $n$,
|
||||
and $\lambda_n$ is the probability of halting (stopping) at current step.
|
||||
|
||||
$s$ can be any neural network (e.g. LSTM, MLP, GRU, Attention layer).
|
||||
|
||||
The unconditioned probability of halting at step $n$ is then,
|
||||
|
||||
$$p_n = \lambda_n \prod_{j=1}^{n-1} (1 - \lambda_j)$$
|
||||
|
||||
That is the probability of not being halted at any of the previous steps and halting at step $n$.
|
||||
|
||||
During inference, we halt by sampling based on the halting probability $\lambda_n$
|
||||
and get the prediction at the halting layer $\hat{y}_n$ as the final output.
|
||||
|
||||
During training, we get the predictions from all the layers and calculate the losses for each of them.
|
||||
And then take the weighted average of the losses based on the probabilities of getting halted at each layer
|
||||
$p_n$.
|
||||
|
||||
The step function is applied to a maximum number of steps donated by $N$.
|
||||
|
||||
The overall loss of PonderNet is
|
||||
|
||||
\begin{align}
|
||||
L &= L_{Rec} + \beta L_{Reg} \\
|
||||
L_{Rec} &= \sum_{n=1}^N p_n \mathcal{L}(y, \hat{y}_n) \\
|
||||
L_{Reg} &= \mathop{KL} \Big(p_n \Vert p_G(\lambda_p) \Big)
|
||||
\end{align}
|
||||
|
||||
$\mathcal{L}$ is the normal loss function between target $y$ and prediction $\hat{y}_n$.
|
||||
|
||||
$\mathop{KL}$ is the [Kullback–Leibler divergence](https://en.wikipedia.org/wiki/Kullback%E2%80%93Leibler_divergence).
|
||||
|
||||
$p_G$ is the [Geometric distribution](https://en.wikipedia.org/wiki/Geometric_distribution) parameterized by
|
||||
$\lambda_p$. *$\lambda_p$ has nothing to do with $\lambda_n$; we are just sticking to same notation as the paper*.
|
||||
$$Pr_{p_G(\lambda_p)}(X = k) = (1 - \lambda_p)^k \lambda_p$$.
|
||||
|
||||
The regularization loss biases the network towards taking $\frac{1}{\lambda_p}$ steps and incentivizes
|
||||
non-zero probabilities for all steps; i.e. promotes exploration.
|
||||
|
||||
Here is the [training code `experiment.py`](experiment.html) to train a PonderNet on [Parity Task](../parity.html).
|
||||
"""
|
||||
|
||||
from typing import Tuple
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
|
||||
|
||||
class ParityPonderGRU(nn.Module):
|
||||
"""
|
||||
## PonderNet with GRU for Parity Task
|
||||
|
||||
This is a simple model that uses a [GRU Cell](https://pytorch.org/docs/stable/generated/torch.nn.GRUCell.html)
|
||||
as the step function.
|
||||
|
||||
This model is for the [Parity Task](../parity.html) where the input is a vector of `n_elems`.
|
||||
Each element of the vector is either `0`, `1` or `-1` and the output is the parity
|
||||
- a binary value that is true if the number of `1`s is odd and false otherwise.
|
||||
|
||||
The prediction of the model is the log probability of the parity being $1$.
|
||||
"""
|
||||
|
||||
def __init__(self, n_elems: int, n_hidden: int, max_steps: int):
|
||||
"""
|
||||
* `n_elems` is the number of elements in the input vector
|
||||
* `n_hidden` is the state vector size of the GRU
|
||||
* `max_steps` is the maximum number of steps $N$
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
self.max_steps = max_steps
|
||||
self.n_hidden = n_hidden
|
||||
|
||||
# GRU
|
||||
# $$h_{n+1} = s_h(x, h_n)$$
|
||||
self.gru = nn.GRUCell(n_elems, n_hidden)
|
||||
# $$\hat{y}_n = s_y(h_n)$$
|
||||
# We could use a layer that takes the concatenation of $h$ and $x$ as input
|
||||
# but we went with this for simplicity.
|
||||
self.output_layer = nn.Linear(n_hidden, 1)
|
||||
# $$\lambda_n = s_\lambda(h_n)$$
|
||||
self.lambda_layer = nn.Linear(n_hidden, 1)
|
||||
self.lambda_prob = nn.Sigmoid()
|
||||
# An option to set during inference so that computation is actually halted at inference time
|
||||
self.is_halt = False
|
||||
|
||||
def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
* `x` is the input of shape `[batch_size, n_elems]`
|
||||
|
||||
This outputs a tuple of four tensors:
|
||||
|
||||
1. $p_1 \dots p_N$ in a tensor of shape `[N, batch_size]`
|
||||
2. $\hat{y}_1 \dots \hat{y}_N$ in a tensor of shape `[N, batch_size]` - the log probabilities of the parity being $1$
|
||||
3. $p_m$ of shape `[batch_size]`
|
||||
4. $\hat{y}_m$ of shape `[batch_size]` where the computation was halted at step $m$
|
||||
"""
|
||||
|
||||
#
|
||||
batch_size = x.shape[0]
|
||||
|
||||
# We get initial state $h_1 = s_h(x)$
|
||||
h = x.new_zeros((x.shape[0], self.n_hidden))
|
||||
h = self.gru(x, h)
|
||||
|
||||
# Lists to store $p_1 \dots p_N$ and $\hat{y}_1 \dots \hat{y}_N$
|
||||
p = []
|
||||
y = []
|
||||
# $\prod_{j=1}^{n-1} (1 - \lambda_j)$
|
||||
un_halted_prob = h.new_ones((batch_size,))
|
||||
|
||||
# A vector to maintain which samples has halted computation
|
||||
halted = h.new_zeros((batch_size,))
|
||||
# $p_m$ and $\hat{y}_m$ where the computation was halted at step $m$
|
||||
p_m = h.new_zeros((batch_size,))
|
||||
y_m = h.new_zeros((batch_size,))
|
||||
|
||||
# Iterate for $N$ steps
|
||||
for n in range(1, self.max_steps + 1):
|
||||
# The halting probability $\lambda_N = 1$ for the last step
|
||||
if n == self.max_steps:
|
||||
lambda_n = h.new_ones(h.shape[0])
|
||||
# $\lambda_n = s_\lambda(h_n)$
|
||||
else:
|
||||
lambda_n = self.lambda_prob(self.lambda_layer(h))[:, 0]
|
||||
# $\hat{y}_n = s_y(h_n)$
|
||||
y_n = self.output_layer(h)[:, 0]
|
||||
|
||||
# $$p_n = \lambda_n \prod_{j=1}^{n-1} (1 - \lambda_j)$$
|
||||
p_n = un_halted_prob * lambda_n
|
||||
# Update $\prod_{j=1}^{n-1} (1 - \lambda_j)$
|
||||
un_halted_prob = un_halted_prob * (1 - lambda_n)
|
||||
|
||||
# Halt based on halting probability $\lambda_n$
|
||||
halt = torch.bernoulli(lambda_n) * (1 - halted)
|
||||
|
||||
# Collect $p_n$ and $\hat{y}_n$
|
||||
p.append(p_n)
|
||||
y.append(y_n)
|
||||
|
||||
# Update $p_m$ and $\hat{y}_m$ based on what was halted at current step $n$
|
||||
p_m = p_m * (1 - halt) + p_n * halt
|
||||
y_m = y_m * (1 - halt) + y_n * halt
|
||||
|
||||
# Update halted samples
|
||||
halted = halted + halt
|
||||
# Get next state $h_{n+1} = s_h(x, h_n)$
|
||||
h = self.gru(x, h)
|
||||
|
||||
# Stop the computation if all samples have halted
|
||||
if self.is_halt and halted.sum() == batch_size:
|
||||
break
|
||||
|
||||
#
|
||||
return torch.stack(p), torch.stack(y), p_m, y_m
|
||||
|
||||
|
||||
class ReconstructionLoss(nn.Module):
|
||||
"""
|
||||
## Reconstruction loss
|
||||
|
||||
$$L_{Rec} = \sum_{n=1}^N p_n \mathcal{L}(y, \hat{y}_n)$$
|
||||
|
||||
$\mathcal{L}$ is the normal loss function between target $y$ and prediction $\hat{y}_n$.
|
||||
"""
|
||||
|
||||
def __init__(self, loss_func: nn.Module):
|
||||
"""
|
||||
* `loss_func` is the loss function $\mathcal{L}$
|
||||
"""
|
||||
super().__init__()
|
||||
self.loss_func = loss_func
|
||||
|
||||
def forward(self, p: torch.Tensor, y_hat: torch.Tensor, y: torch.Tensor):
|
||||
"""
|
||||
* `p` is $p_1 \dots p_N$ in a tensor of shape `[N, batch_size]`
|
||||
* `y_hat` is $\hat{y}_1 \dots \hat{y}_N$ in a tensor of shape `[N, batch_size, ...]`
|
||||
* `y` is the target of shape `[batch_size, ...]`
|
||||
"""
|
||||
|
||||
# The total $\sum_{n=1}^N p_n \mathcal{L}(y, \hat{y}_n)$
|
||||
total_loss = p.new_tensor(0.)
|
||||
# Iterate upto $N$
|
||||
for n in range(p.shape[0]):
|
||||
# $p_n \mathcal{L}(y, \hat{y}_n)$ for each sample and the mean of them
|
||||
loss = (p[n] * self.loss_func(y_hat[n], y)).mean()
|
||||
# Add to total loss
|
||||
total_loss = total_loss + loss
|
||||
|
||||
#
|
||||
return total_loss
|
||||
|
||||
|
||||
class RegularizationLoss(nn.Module):
|
||||
"""
|
||||
## Regularization loss
|
||||
|
||||
$$L_{Reg} = \mathop{KL} \Big(p_n \Vert p_G(\lambda_p) \Big)$$
|
||||
|
||||
$\mathop{KL}$ is the [Kullback–Leibler divergence](https://en.wikipedia.org/wiki/Kullback%E2%80%93Leibler_divergence).
|
||||
|
||||
$p_G$ is the [Geometric distribution](https://en.wikipedia.org/wiki/Geometric_distribution) parameterized by
|
||||
$\lambda_p$. *$\lambda_p$ has nothing to do with $\lambda_n$; we are just sticking to same notation as the paper*.
|
||||
$$Pr_{p_G(\lambda_p)}(X = k) = (1 - \lambda_p)^k \lambda_p$$.
|
||||
|
||||
The regularization loss biases the network towards taking $\frac{1}{\lambda_p}$ steps and incentivies non-zero probabilities
|
||||
for all steps; i.e. promotes exploration.
|
||||
"""
|
||||
|
||||
def __init__(self, lambda_p: float, max_steps: int = 1_000):
|
||||
"""
|
||||
* `lambda_p` is $\lambda_p$ - the success probability of geometric distribution
|
||||
* `max_steps` is the highest $N$; we use this to pre-compute $p_G(\lambda_p)$
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
# Empty vector to calculate $p_G(\lambda_p)$
|
||||
p_g = torch.zeros((max_steps,))
|
||||
# $(1 - \lambda_p)^k$
|
||||
not_halted = 1.
|
||||
# Iterate upto `max_steps`
|
||||
for k in range(max_steps):
|
||||
# $$Pr_{p_G(\lambda_p)}(X = k) = (1 - \lambda_p)^k \lambda_p$$
|
||||
p_g[k] = not_halted * lambda_p
|
||||
# Update $(1 - \lambda_p)^k$
|
||||
not_halted = not_halted * (1 - lambda_p)
|
||||
|
||||
# Save $Pr_{p_G(\lambda_p)}$
|
||||
self.p_g = nn.Parameter(p_g, requires_grad=False)
|
||||
|
||||
# KL-divergence loss
|
||||
self.kl_div = nn.KLDivLoss(reduction='batchmean')
|
||||
|
||||
def forward(self, p: torch.Tensor):
|
||||
"""
|
||||
* `p` is $p_1 \dots p_N$ in a tensor of shape `[N, batch_size]`
|
||||
"""
|
||||
# Transpose `p` to `[batch_size, N]`
|
||||
p = p.transpose(0, 1)
|
||||
# Get $Pr_{p_G(\lambda_p)}$ upto $N$ and expand it across the batch dimension
|
||||
p_g = self.p_g[None, :p.shape[1]].expand_as(p)
|
||||
|
||||
# Calculate the KL-divergence.
|
||||
# *The [PyTorch KL-divergence](https://pytorch.org/docs/stable/generated/torch.nn.KLDivLoss.html)
|
||||
# implementation accepts log probabilities.*
|
||||
return self.kl_div(p.log(), p_g)
|
||||
@@ -0,0 +1,161 @@
|
||||
"""
|
||||
---
|
||||
title: "PonderNet Parity Task Experiment"
|
||||
summary: >
|
||||
This trains is a PonderNet on Parity Task
|
||||
---
|
||||
|
||||
# [PonderNet](index.html) [Parity Task](../parity.html) Experiment
|
||||
|
||||
This trains a [PonderNet](index.html) on [Parity Task](../parity.html).
|
||||
"""
|
||||
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.utils.data import DataLoader
|
||||
|
||||
from labml import tracker, experiment
|
||||
from labml_nn.helpers.metrics import AccuracyDirect
|
||||
from labml_nn.helpers.trainer import SimpleTrainValidConfigs, BatchIndex
|
||||
from labml_nn.adaptive_computation.parity import ParityDataset
|
||||
from labml_nn.adaptive_computation.ponder_net import ParityPonderGRU, ReconstructionLoss, RegularizationLoss
|
||||
|
||||
|
||||
class Configs(SimpleTrainValidConfigs):
|
||||
"""
|
||||
Configurations with a
|
||||
[simple training loop](../../helpers/trainer.html)
|
||||
"""
|
||||
|
||||
# Number of epochs
|
||||
epochs: int = 100
|
||||
# Number of batches per epoch
|
||||
n_batches: int = 500
|
||||
# Batch size
|
||||
batch_size: int = 128
|
||||
|
||||
# Model
|
||||
model: ParityPonderGRU
|
||||
|
||||
# $L_{Rec}$
|
||||
loss_rec: ReconstructionLoss
|
||||
# $L_{Reg}$
|
||||
loss_reg: RegularizationLoss
|
||||
|
||||
# The number of elements in the input vector.
|
||||
# *We keep it low for demonstration; otherwise, training takes a lot of time.
|
||||
# Although the parity task seems simple, figuring out the pattern by looking at samples
|
||||
# is quite hard.*
|
||||
n_elems: int = 8
|
||||
# Number of units in the hidden layer (state)
|
||||
n_hidden: int = 64
|
||||
# Maximum number of steps $N$
|
||||
max_steps: int = 20
|
||||
|
||||
# $\lambda_p$ for the geometric distribution $p_G(\lambda_p)$
|
||||
lambda_p: float = 0.2
|
||||
# Regularization loss $L_{Reg}$ coefficient $\beta$
|
||||
beta: float = 0.01
|
||||
|
||||
# Gradient clipping by norm
|
||||
grad_norm_clip: float = 1.0
|
||||
|
||||
# Training and validation loaders
|
||||
train_loader: DataLoader
|
||||
valid_loader: DataLoader
|
||||
|
||||
# Accuracy calculator
|
||||
accuracy = AccuracyDirect()
|
||||
|
||||
def init(self):
|
||||
# Print indicators to screen
|
||||
tracker.set_scalar('loss.*', True)
|
||||
tracker.set_scalar('loss_reg.*', True)
|
||||
tracker.set_scalar('accuracy.*', True)
|
||||
tracker.set_scalar('steps.*', True)
|
||||
|
||||
# We need to set the metrics to calculate them for the epoch for training and validation
|
||||
self.state_modules = [self.accuracy]
|
||||
|
||||
# Initialize the model
|
||||
self.model = ParityPonderGRU(self.n_elems, self.n_hidden, self.max_steps).to(self.device)
|
||||
# $L_{Rec}$
|
||||
self.loss_rec = ReconstructionLoss(nn.BCEWithLogitsLoss(reduction='none')).to(self.device)
|
||||
# $L_{Reg}$
|
||||
self.loss_reg = RegularizationLoss(self.lambda_p, self.max_steps).to(self.device)
|
||||
|
||||
# Training and validation loaders
|
||||
self.train_loader = DataLoader(ParityDataset(self.batch_size * self.n_batches, self.n_elems),
|
||||
batch_size=self.batch_size)
|
||||
self.valid_loader = DataLoader(ParityDataset(self.batch_size * 32, self.n_elems),
|
||||
batch_size=self.batch_size)
|
||||
|
||||
def step(self, batch: Any, batch_idx: BatchIndex):
|
||||
"""
|
||||
This method gets called by the trainer for each batch
|
||||
"""
|
||||
# Set the model mode
|
||||
self.model.train(self.mode.is_train)
|
||||
|
||||
# Get the input and labels and move them to the model's device
|
||||
data, target = batch[0].to(self.device), batch[1].to(self.device)
|
||||
|
||||
# Increment step in training mode
|
||||
if self.mode.is_train:
|
||||
tracker.add_global_step(len(data))
|
||||
|
||||
# Run the model
|
||||
p, y_hat, p_sampled, y_hat_sampled = self.model(data)
|
||||
|
||||
# Calculate the reconstruction loss
|
||||
loss_rec = self.loss_rec(p, y_hat, target.to(torch.float))
|
||||
tracker.add("loss.", loss_rec)
|
||||
|
||||
# Calculate the regularization loss
|
||||
loss_reg = self.loss_reg(p)
|
||||
tracker.add("loss_reg.", loss_reg)
|
||||
|
||||
# $L = L_{Rec} + \beta L_{Reg}$
|
||||
loss = loss_rec + self.beta * loss_reg
|
||||
|
||||
# Calculate the expected number of steps taken
|
||||
steps = torch.arange(1, p.shape[0] + 1, device=p.device)
|
||||
expected_steps = (p * steps[:, None]).sum(dim=0)
|
||||
tracker.add("steps.", expected_steps)
|
||||
|
||||
# Call accuracy metric
|
||||
self.accuracy(y_hat_sampled > 0, target)
|
||||
|
||||
if self.mode.is_train:
|
||||
# Compute gradients
|
||||
loss.backward()
|
||||
# Clip gradients
|
||||
torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=self.grad_norm_clip)
|
||||
# Optimizer
|
||||
self.optimizer.step()
|
||||
# Clear gradients
|
||||
self.optimizer.zero_grad()
|
||||
#
|
||||
tracker.save()
|
||||
|
||||
|
||||
def main():
|
||||
"""
|
||||
Run the experiment
|
||||
"""
|
||||
experiment.create(name='ponder_net')
|
||||
|
||||
conf = Configs()
|
||||
experiment.configs(conf, {
|
||||
'optimizer.optimizer': 'Adam',
|
||||
'optimizer.learning_rate': 0.0003,
|
||||
})
|
||||
|
||||
with experiment.start():
|
||||
conf.run()
|
||||
|
||||
#
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,8 @@
|
||||
# [PonderNet: Learning to Ponder](https://nn.labml.ai/adaptive_computation/ponder_net/index.html)
|
||||
|
||||
This is a [PyTorch](https://pytorch.org) implementation of the paper
|
||||
[PonderNet: Learning to Ponder](https://arxiv.org/abs/2107.05407).
|
||||
|
||||
PonderNet adapts the computation based on the input.
|
||||
It changes the number of steps to take on a recurrent network based on the input.
|
||||
PonderNet learns this with end-to-end gradient descent.
|
||||
@@ -0,0 +1,7 @@
|
||||
# [Neural Networks with Adaptive Computation](https://nn.labml.ai/adaptive_computation/index.html)
|
||||
|
||||
These are neural network architectures that change the computation complexity based on the
|
||||
complexity of the input sample.
|
||||
|
||||
* 🚧 TODO: Adaptive Computation Time for Recurrent Neural Networks
|
||||
* [PonderNet: Learning to Ponder](https://nn.labml.ai/adaptive_computation/ponder_net/index.html)
|
||||
@@ -0,0 +1,185 @@
|
||||
"""
|
||||
---
|
||||
title: Capsule Networks
|
||||
summary: >
|
||||
PyTorch implementation and tutorial of Capsule Networks.
|
||||
Capsule network is a neural network architecture that embeds features
|
||||
as capsules and routes them with a voting mechanism to next layer of capsules.
|
||||
---
|
||||
|
||||
# Capsule Networks
|
||||
|
||||
This is a [PyTorch](https://pytorch.org) implementation/tutorial of
|
||||
[Dynamic Routing Between Capsules](https://arxiv.org/abs/1710.09829).
|
||||
|
||||
Capsule network is a neural network architecture that embeds features
|
||||
as capsules and routes them with a voting mechanism to next layer of capsules.
|
||||
|
||||
Unlike in other implementations of models, we've included a sample, because
|
||||
it is difficult to understand some concepts with just the modules.
|
||||
[This is the annotated code for a model that uses capsules to classify MNIST dataset](mnist.html)
|
||||
|
||||
This file holds the implementations of the core modules of Capsule Networks.
|
||||
|
||||
I used [jindongwang/Pytorch-CapsuleNet](https://github.com/jindongwang/Pytorch-CapsuleNet) to clarify some
|
||||
confusions I had with the paper.
|
||||
|
||||
Here's a notebook for training a Capsule Network on MNIST dataset.
|
||||
|
||||
[](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/capsule_networks/mnist.ipynb)
|
||||
"""
|
||||
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import torch.utils.data
|
||||
|
||||
|
||||
class Squash(nn.Module):
|
||||
"""
|
||||
## Squash
|
||||
|
||||
This is **squashing** function from paper, given by equation $(1)$.
|
||||
|
||||
$$\mathbf{v}_j = \frac{{\lVert \mathbf{s}_j \rVert}^2}{1 + {\lVert \mathbf{s}_j \rVert}^2}
|
||||
\frac{\mathbf{s}_j}{\lVert \mathbf{s}_j \rVert}$$
|
||||
|
||||
$\frac{\mathbf{s}_j}{\lVert \mathbf{s}_j \rVert}$
|
||||
normalizes the length of all the capsules, whilst
|
||||
$\frac{{\lVert \mathbf{s}_j \rVert}^2}{1 + {\lVert \mathbf{s}_j \rVert}^2}$
|
||||
shrinks the capsules that have a length smaller than one .
|
||||
"""
|
||||
|
||||
def __init__(self, epsilon=1e-8):
|
||||
super().__init__()
|
||||
self.epsilon = epsilon
|
||||
|
||||
def forward(self, s: torch.Tensor):
|
||||
"""
|
||||
The shape of `s` is `[batch_size, n_capsules, n_features]`
|
||||
"""
|
||||
|
||||
# ${\lVert \mathbf{s}_j \rVert}^2$
|
||||
s2 = (s ** 2).sum(dim=-1, keepdims=True)
|
||||
|
||||
# We add an epsilon when calculating $\lVert \mathbf{s}_j \rVert$ to make sure it doesn't become zero.
|
||||
# If this becomes zero it starts giving out `nan` values and training fails.
|
||||
# $$\mathbf{v}_j = \frac{{\lVert \mathbf{s}_j \rVert}^2}{1 + {\lVert \mathbf{s}_j \rVert}^2}
|
||||
# \frac{\mathbf{s}_j}{\sqrt{{\lVert \mathbf{s}_j \rVert}^2 + \epsilon}}$$
|
||||
return (s2 / (1 + s2)) * (s / torch.sqrt(s2 + self.epsilon))
|
||||
|
||||
|
||||
class Router(nn.Module):
|
||||
"""
|
||||
## Routing Algorithm
|
||||
|
||||
This is the routing mechanism described in the paper.
|
||||
You can use multiple routing layers in your models.
|
||||
|
||||
This combines calculating $\mathbf{s}_j$ for this layer and
|
||||
the routing algorithm described in *Procedure 1*.
|
||||
"""
|
||||
|
||||
def __init__(self, in_caps: int, out_caps: int, in_d: int, out_d: int, iterations: int):
|
||||
"""
|
||||
`in_caps` is the number of capsules, and `in_d` is the number of features per capsule from the layer below.
|
||||
`out_caps` and `out_d` are the same for this layer.
|
||||
|
||||
`iterations` is the number of routing iterations, symbolized by $r$ in the paper.
|
||||
"""
|
||||
super().__init__()
|
||||
self.in_caps = in_caps
|
||||
self.out_caps = out_caps
|
||||
self.iterations = iterations
|
||||
self.softmax = nn.Softmax(dim=1)
|
||||
self.squash = Squash()
|
||||
|
||||
# This is the weight matrix $\mathbf{W}_{ij}$. It maps each capsule in the
|
||||
# lower layer to each capsule in this layer
|
||||
self.weight = nn.Parameter(torch.randn(in_caps, out_caps, in_d, out_d), requires_grad=True)
|
||||
|
||||
def forward(self, u: torch.Tensor):
|
||||
"""
|
||||
The shape of `u` is `[batch_size, n_capsules, n_features]`.
|
||||
These are the capsules from the lower layer.
|
||||
"""
|
||||
|
||||
# $$\hat{\mathbf{u}}_{j|i} = \mathbf{W}_{ij} \mathbf{u}_i$$
|
||||
# Here $j$ is used to index capsules in this layer, whilst $i$ is
|
||||
# used to index capsules in the layer below (previous).
|
||||
u_hat = torch.einsum('ijnm,bin->bijm', self.weight, u)
|
||||
|
||||
# Initial logits $b_{ij}$ are the log prior probabilities that capsule $i$
|
||||
# should be coupled with $j$.
|
||||
# We initialize these at zero
|
||||
b = u.new_zeros(u.shape[0], self.in_caps, self.out_caps)
|
||||
|
||||
v = None
|
||||
|
||||
# Iterate
|
||||
for i in range(self.iterations):
|
||||
# routing softmax $$c_{ij} = \frac{\exp({b_{ij}})}{\sum_k\exp({b_{ik}})}$$
|
||||
c = self.softmax(b)
|
||||
# $$\mathbf{s}_j = \sum_i{c_{ij} \hat{\mathbf{u}}_{j|i}}$$
|
||||
s = torch.einsum('bij,bijm->bjm', c, u_hat)
|
||||
# $$\mathbf{v}_j = squash(\mathbf{s}_j)$$
|
||||
v = self.squash(s)
|
||||
# $$a_{ij} = \mathbf{v}_j \cdot \hat{\mathbf{u}}_{j|i}$$
|
||||
a = torch.einsum('bjm,bijm->bij', v, u_hat)
|
||||
# $$b_{ij} \gets b_{ij} + \mathbf{v}_j \cdot \hat{\mathbf{u}}_{j|i}$$
|
||||
b = b + a
|
||||
|
||||
return v
|
||||
|
||||
|
||||
class MarginLoss(nn.Module):
|
||||
"""
|
||||
## Margin loss for class existence
|
||||
|
||||
A separate margin loss is used for each output capsule and the total loss is the sum of them.
|
||||
The length of each output capsule is the probability that class is present in the input.
|
||||
|
||||
Loss for each output capsule or class $k$ is,
|
||||
$$\mathcal{L}_k = T_k \max(0, m^{+} - \lVert\mathbf{v}_k\rVert)^2 +
|
||||
\lambda (1 - T_k) \max(0, \lVert\mathbf{v}_k\rVert - m^{-})^2$$
|
||||
|
||||
$T_k$ is $1$ if the class $k$ is present and $0$ otherwise.
|
||||
The first component of the loss is $0$ when the class is not present,
|
||||
and the second component is $0$ if the class is present.
|
||||
The $\max(0, x)$ is used to avoid predictions going to extremes.
|
||||
$m^{+}$ is set to be $0.9$ and $m^{-}$ to be $0.1$ in the paper.
|
||||
|
||||
The $\lambda$ down-weighting is used to stop the length of all capsules from
|
||||
falling during the initial phase of training.
|
||||
"""
|
||||
|
||||
def __init__(self, *, n_labels: int, lambda_: float = 0.5, m_positive: float = 0.9, m_negative: float = 0.1):
|
||||
super().__init__()
|
||||
|
||||
self.m_negative = m_negative
|
||||
self.m_positive = m_positive
|
||||
self.lambda_ = lambda_
|
||||
self.n_labels = n_labels
|
||||
|
||||
def forward(self, v: torch.Tensor, labels: torch.Tensor):
|
||||
"""
|
||||
`v`, $\mathbf{v}_j$ are the squashed output capsules.
|
||||
This has shape `[batch_size, n_labels, n_features]`; that is, there is a capsule for each label.
|
||||
|
||||
`labels` are the labels, and has shape `[batch_size]`.
|
||||
"""
|
||||
# $$\lVert \mathbf{v}_j \rVert$$
|
||||
v_norm = torch.sqrt((v ** 2).sum(dim=-1))
|
||||
|
||||
# $$\mathcal{L}$$
|
||||
# `labels` is one-hot encoded labels of shape `[batch_size, n_labels]`
|
||||
labels = torch.eye(self.n_labels, device=labels.device)[labels]
|
||||
|
||||
# $$\mathcal{L}_k = T_k \max(0, m^{+} - \lVert\mathbf{v}_k\rVert)^2 +
|
||||
# \lambda (1 - T_k) \max(0, \lVert\mathbf{v}_k\rVert - m^{-})^2$$
|
||||
# `loss` has shape `[batch_size, n_labels]`. We have parallelized the computation
|
||||
# of $\mathcal{L}_k$ for for all $k$.
|
||||
loss = labels * F.relu(self.m_positive - v_norm) + \
|
||||
self.lambda_ * (1.0 - labels) * F.relu(v_norm - self.m_negative)
|
||||
|
||||
# $$\sum_k \mathcal{L}_k$$
|
||||
return loss.sum(dim=-1).mean()
|
||||
@@ -0,0 +1,208 @@
|
||||
{
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"name": "Capsule Networks",
|
||||
"provenance": [],
|
||||
"collapsed_sections": []
|
||||
},
|
||||
"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/capsule_networks/mnist.ipynb) \n",
|
||||
"\n",
|
||||
"## Training a Capsule Network to classify MNIST digits\n",
|
||||
"\n",
|
||||
"This is an experiment to train a Capsule Network to classify MNIST digits using PyTorch."
|
||||
]
|
||||
},
|
||||
{
|
||||
"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": "7ab15f72-c99f-4097-ecd2-5740ee9ed61c"
|
||||
},
|
||||
"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": [
|
||||
"import torch\n",
|
||||
"\n",
|
||||
"from labml import experiment\n",
|
||||
"from labml_nn.capsule_networks.mnist 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=\"capsule_networks\")"
|
||||
],
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "-OnHLi626tJt"
|
||||
},
|
||||
"source": [
|
||||
"Initialize [Capsule Network configurations](https://nn.labml.ai/capsule_networks/mnist.html)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"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": "ebefa8fa-93d2-4131-db95-e27f15aa3aa0"
|
||||
},
|
||||
"source": [
|
||||
"experiment.configs(conf, {'optimizer.optimizer': 'Adam',\n",
|
||||
" 'optimizer.learning_rate': 1e-3,\n",
|
||||
" 'inner_iterations': 5})"
|
||||
],
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "EvI7MtgJ61w5"
|
||||
},
|
||||
"source": [
|
||||
"Set PyTorch models for loading and saving"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 102
|
||||
},
|
||||
"id": "GDlt7dp-5ALt",
|
||||
"outputId": "9701092b-c88a-4687-c90e-b193c369e59e"
|
||||
},
|
||||
"source": [
|
||||
"experiment.add_pytorch_models({'model': conf.model})"
|
||||
],
|
||||
"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": 646
|
||||
},
|
||||
"id": "aIAWo7Fw5DR8",
|
||||
"outputId": "5ddbfce3-91f8-4506-e483-1640cb5a14b3"
|
||||
},
|
||||
"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,174 @@
|
||||
"""
|
||||
---
|
||||
title: Classify MNIST digits with Capsule Networks
|
||||
summary: Code for training Capsule Networks on MNIST dataset
|
||||
---
|
||||
|
||||
# Classify MNIST digits with Capsule Networks
|
||||
|
||||
This is an annotated PyTorch code to classify MNIST digits with PyTorch.
|
||||
|
||||
This paper implements the experiment described in paper
|
||||
[Dynamic Routing Between Capsules](https://arxiv.org/abs/1710.09829).
|
||||
"""
|
||||
from typing import Any
|
||||
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import torch.utils.data
|
||||
from labml import experiment, tracker
|
||||
from labml.configs import option
|
||||
from labml_nn.capsule_networks import Squash, Router, MarginLoss
|
||||
from labml_nn.helpers.datasets import MNISTConfigs
|
||||
from labml_nn.helpers.metrics import AccuracyDirect
|
||||
from labml_nn.helpers.trainer import SimpleTrainValidConfigs, BatchIndex
|
||||
|
||||
|
||||
class MNISTCapsuleNetworkModel(nn.Module):
|
||||
"""
|
||||
## Model for classifying MNIST digits
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
# First convolution layer has $256$, $9 \times 9$ convolution kernels
|
||||
self.conv1 = nn.Conv2d(in_channels=1, out_channels=256, kernel_size=9, stride=1)
|
||||
# The second layer (Primary Capsules) s a convolutional capsule layer with $32$ channels
|
||||
# of convolutional $8D$ capsules ($8$ features per capsule).
|
||||
# That is, each primary capsule contains 8 convolutional units with a 9 × 9 kernel and a stride of 2.
|
||||
# In order to implement this we create a convolutional layer with $32 \times 8$ channels and
|
||||
# reshape and permutate its output to get the capsules of $8$ features each.
|
||||
self.conv2 = nn.Conv2d(in_channels=256, out_channels=32 * 8, kernel_size=9, stride=2, padding=0)
|
||||
self.squash = Squash()
|
||||
|
||||
# Routing layer gets the $32 \times 6 \times 6$ primary capsules and produces $10$ capsules.
|
||||
# Each of the primary capsules have $8$ features, while output capsules (Digit Capsules)
|
||||
# have $16$ features.
|
||||
# The routing algorithm iterates $3$ times.
|
||||
self.digit_capsules = Router(32 * 6 * 6, 10, 8, 16, 3)
|
||||
|
||||
# This is the decoder mentioned in the paper.
|
||||
# It takes the outputs of the $10$ digit capsules, each with $16$ features to reproduce the
|
||||
# image. It goes through linear layers of sizes $512$ and $1024$ with $ReLU$ activations.
|
||||
self.decoder = nn.Sequential(
|
||||
nn.Linear(16 * 10, 512),
|
||||
nn.ReLU(),
|
||||
nn.Linear(512, 1024),
|
||||
nn.ReLU(),
|
||||
nn.Linear(1024, 784),
|
||||
nn.Sigmoid()
|
||||
)
|
||||
|
||||
def forward(self, data: torch.Tensor):
|
||||
"""
|
||||
`data` are the MNIST images, with shape `[batch_size, 1, 28, 28]`
|
||||
"""
|
||||
# Pass through the first convolution layer.
|
||||
# Output of this layer has shape `[batch_size, 256, 20, 20]`
|
||||
x = F.relu(self.conv1(data))
|
||||
# Pass through the second convolution layer.
|
||||
# Output of this has shape `[batch_size, 32 * 8, 6, 6]`.
|
||||
# *Note that this layer has a stride length of $2$*.
|
||||
x = self.conv2(x)
|
||||
|
||||
# Resize and permutate to get the capsules
|
||||
caps = x.view(x.shape[0], 8, 32 * 6 * 6).permute(0, 2, 1)
|
||||
# Squash the capsules
|
||||
caps = self.squash(caps)
|
||||
# Take them through the router to get digit capsules.
|
||||
# This has shape `[batch_size, 10, 16]`.
|
||||
caps = self.digit_capsules(caps)
|
||||
|
||||
# Get masks for reconstructioon
|
||||
with torch.no_grad():
|
||||
# The prediction by the capsule network is the capsule with longest length
|
||||
pred = (caps ** 2).sum(-1).argmax(-1)
|
||||
# Create a mask to maskout all the other capsules
|
||||
mask = torch.eye(10, device=data.device)[pred]
|
||||
|
||||
# Mask the digit capsules to get only the capsule that made the prediction and
|
||||
# take it through decoder to get reconstruction
|
||||
reconstructions = self.decoder((caps * mask[:, :, None]).view(x.shape[0], -1))
|
||||
# Reshape the reconstruction to match the image dimensions
|
||||
reconstructions = reconstructions.view(-1, 1, 28, 28)
|
||||
|
||||
return caps, reconstructions, pred
|
||||
|
||||
|
||||
class Configs(MNISTConfigs, SimpleTrainValidConfigs):
|
||||
"""
|
||||
Configurations with MNIST data and Train & Validation setup
|
||||
"""
|
||||
epochs: int = 10
|
||||
model: nn.Module = 'capsule_network_model'
|
||||
reconstruction_loss = nn.MSELoss()
|
||||
margin_loss = MarginLoss(n_labels=10)
|
||||
accuracy = AccuracyDirect()
|
||||
|
||||
def init(self):
|
||||
# Print losses and accuracy to screen
|
||||
tracker.set_scalar('loss.*', True)
|
||||
tracker.set_scalar('accuracy.*', True)
|
||||
|
||||
# We need to set the metrics to calculate them for the epoch for training and validation
|
||||
self.state_modules = [self.accuracy]
|
||||
|
||||
def step(self, batch: Any, batch_idx: BatchIndex):
|
||||
"""
|
||||
This method gets called by the trainer
|
||||
"""
|
||||
# Set the model mode
|
||||
self.model.train(self.mode.is_train)
|
||||
|
||||
# Get the images and labels and move them to the model's device
|
||||
data, target = batch[0].to(self.device), batch[1].to(self.device)
|
||||
|
||||
# Increment step in training mode
|
||||
if self.mode.is_train:
|
||||
tracker.add_global_step(len(data))
|
||||
|
||||
# Run the model
|
||||
caps, reconstructions, pred = self.model(data)
|
||||
|
||||
# Calculate the total loss
|
||||
loss = self.margin_loss(caps, target) + 0.0005 * self.reconstruction_loss(reconstructions, data)
|
||||
tracker.add("loss.", loss)
|
||||
|
||||
# Call accuracy metric
|
||||
self.accuracy(pred, target)
|
||||
|
||||
if self.mode.is_train:
|
||||
loss.backward()
|
||||
|
||||
self.optimizer.step()
|
||||
# Log parameters and gradients
|
||||
if batch_idx.is_last:
|
||||
tracker.add('model', self.model)
|
||||
self.optimizer.zero_grad()
|
||||
|
||||
tracker.save()
|
||||
|
||||
|
||||
@option(Configs.model)
|
||||
def capsule_network_model(c: Configs):
|
||||
"""Set the model"""
|
||||
return MNISTCapsuleNetworkModel().to(c.device)
|
||||
|
||||
|
||||
def main():
|
||||
"""
|
||||
Run the experiment
|
||||
"""
|
||||
experiment.create(name='capsule_network_mnist')
|
||||
conf = Configs()
|
||||
experiment.configs(conf, {'optimizer.optimizer': 'Adam',
|
||||
'optimizer.learning_rate': 1e-3})
|
||||
|
||||
experiment.add_pytorch_models({'model': conf.model})
|
||||
|
||||
with experiment.start():
|
||||
conf.run()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,20 @@
|
||||
# [Capsule Networks](https://nn.labml.ai/capsule_networks/index.html)
|
||||
|
||||
This is a [PyTorch](https://pytorch.org) implementation/tutorial of
|
||||
[Dynamic Routing Between Capsules](https://arxiv.org/abs/1710.09829).
|
||||
|
||||
Capsule network is a neural network architecture that embeds features
|
||||
as capsules and routes them with a voting mechanism to next layer of capsules.
|
||||
|
||||
Unlike in other implementations of models, we've included a sample, because
|
||||
it is difficult to understand some concepts with just the modules.
|
||||
[This is the annotated code for a model that uses capsules to classify MNIST dataset](mnist.html)
|
||||
|
||||
This file holds the implementations of the core modules of Capsule Networks.
|
||||
|
||||
I used [jindongwang/Pytorch-CapsuleNet](https://github.com/jindongwang/Pytorch-CapsuleNet) to clarify some
|
||||
confusions I had with the paper.
|
||||
|
||||
Here's a notebook for training a Capsule Network on MNIST dataset.
|
||||
|
||||
[](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/capsule_networks/mnist.ipynb)
|
||||
@@ -0,0 +1,753 @@
|
||||
"""
|
||||
---
|
||||
title: Regret Minimization in Games with Incomplete Information (CFR)
|
||||
summary: >
|
||||
This is an annotated implementation/tutorial of Regret Minimization in Games with Incomplete Information
|
||||
---
|
||||
|
||||
# Regret Minimization in Games with Incomplete Information (CFR)
|
||||
|
||||
The paper
|
||||
[Regret Minimization in Games with Incomplete Information](http://martin.zinkevich.org/publications/regretpoker.pdf)
|
||||
introduces counterfactual regret and how minimizing counterfactual regret through self-play
|
||||
can be used to reach Nash equilibrium.
|
||||
The algorithm is called Counterfactual Regret Minimization (**CFR**).
|
||||
|
||||
The paper
|
||||
[Monte Carlo Sampling for Regret Minimization in Extensive Games](http://mlanctot.info/files/papers/nips09mccfr.pdf)
|
||||
introduces Monte Carlo Counterfactual Regret Minimization (**MCCFR**),
|
||||
where we sample from the game tree and estimate the regrets.
|
||||
|
||||
We tried to keep our Python implementation easy-to-understand like a tutorial.
|
||||
We run it on [a very simple imperfect information game called Kuhn poker](kuhn/index.html).
|
||||
|
||||
[](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/cfr/kuhn/experiment.ipynb)
|
||||
|
||||
[](https://twitter.com/labmlai/status/1407186002255380484)
|
||||
Twitter thread
|
||||
|
||||
## Introduction
|
||||
|
||||
We implement Monte Carlo Counterfactual Regret Minimization (MCCFR) with chance sampling (CS).
|
||||
It iteratively, explores part of the game tree by trying all player actions,
|
||||
but sampling chance events.
|
||||
Chance events are things like dealing cards; they are kept sampled once per iteration.
|
||||
Then it calculates, for each action, the *regret* of following the current strategy instead of taking that action.
|
||||
Then it updates the strategy based on these regrets for the next iteration, using regret matching.
|
||||
Finally, it computes the average of the strategies throughout the iterations,
|
||||
which is very close to the Nash equilibrium if we ran enough iterations.
|
||||
|
||||
We will first introduce the mathematical notation and theory.
|
||||
|
||||
### Player
|
||||
|
||||
A player is denoted by $i \in N$, where $N$ is the set of players.
|
||||
|
||||
### [History](#History)
|
||||
|
||||
History $h \in H$ is a sequence of actions including chance events,
|
||||
and $H$ is the set of all histories.
|
||||
|
||||
$Z \subseteq H$ is the set of terminal histories (game over).
|
||||
|
||||
### Action
|
||||
|
||||
Action $a$, $A(h) = {a: (h, a) \in H}$ where $h \in H$ is a non-terminal [history](#History).
|
||||
|
||||
### [Information Set $I_i$](#InfoSet)
|
||||
|
||||
**Information set** $I_i \in \mathcal{I}_i$ for player $i$
|
||||
is similar to a history $h \in H$
|
||||
but only contains the actions visible to player $i$.
|
||||
That is, the history $h$ will contain actions/events such as cards dealt to the
|
||||
opposing player while $I_i$ will not have them.
|
||||
|
||||
$\mathcal{I}_i$ is known as the **information partition** of player $i$.
|
||||
|
||||
$h \in I$ is the set of all histories that belong to a given information set;
|
||||
i.e. all those histories look the same in the eye of the player.
|
||||
|
||||
<a id="Strategy"></a>
|
||||
|
||||
### Strategy
|
||||
|
||||
**Strategy of player** $i$, $\sigma_i \in \Sigma_i$ is a distribution over actions $A(I_i)$,
|
||||
where $\Sigma_i$ is the set of all strategies for player $i$.
|
||||
Strategy on $t$-th iteration is denoted by $\sigma^t_i$.
|
||||
|
||||
Strategy is defined as a probability for taking an action $a$ in for a given information set $I$,
|
||||
|
||||
$$\sigma_i(I)(a)$$
|
||||
|
||||
$\sigma$ is the **strategy profile** which consists of strategies of all players
|
||||
$\sigma_1, \sigma_2, \ldots$
|
||||
|
||||
$\sigma_{-i}$ is strategies of all players except $\sigma_i$
|
||||
|
||||
<a id="HistoryProbability"></a>
|
||||
|
||||
### Probability of History
|
||||
|
||||
$\pi^\sigma(h)$ is the probability of reaching the history $h$ with strategy profile $\sigma$.
|
||||
$\pi^\sigma(h)_{-i}$ is the probability of reaching $h$ without player $i$'s contribution;
|
||||
i.e. player $i$ took the actions to follow $h$ with a probability of $1$.
|
||||
|
||||
$\pi^\sigma(h)_{i}$ is the probability of reaching $h$ with only player $i$'s contribution.
|
||||
That is,
|
||||
$$\pi^\sigma(h) = \pi^\sigma(h)_{i} \pi^\sigma(h)_{-i}$$
|
||||
|
||||
Probability of reaching a information set $I$ is,
|
||||
$$\pi^\sigma(I) = \sum_{h \in I} \pi^\sigma(h)$$
|
||||
|
||||
### Utility (Pay off)
|
||||
|
||||
The [terminal utility](#terminal_utility) is the utility (or pay off)
|
||||
of a player $i$ for a terminal history $h$.
|
||||
|
||||
$$u_i(h)$$ where $h \in Z$
|
||||
|
||||
$u_i(\sigma)$ is the expected utility (payoff) for player $i$ with strategy profile $\sigma$.
|
||||
|
||||
$$u_i(\sigma) = \sum_{h \in Z} u_i(h) \pi^\sigma(h)$$
|
||||
|
||||
<a id="NashEquilibrium"></a>
|
||||
|
||||
### Nash Equilibrium
|
||||
|
||||
Nash equilibrium is a state where none of the players can increase their expected utility (or payoff)
|
||||
by changing their strategy alone.
|
||||
|
||||
For two players, Nash equilibrium is a [strategy profile](#Strategy) where
|
||||
|
||||
\begin{align}
|
||||
u_1(\sigma) &\ge \max_{\sigma'_1 \in \Sigma_1} u_1(\sigma'_1, \sigma_2) \\
|
||||
u_2(\sigma) &\ge \max_{\sigma'_2 \in \Sigma_2} u_1(\sigma_1, \sigma'_2) \\
|
||||
\end{align}
|
||||
|
||||
$\epsilon$-Nash equilibrium is,
|
||||
|
||||
\begin{align}
|
||||
u_1(\sigma) + \epsilon &\ge \max_{\sigma'_1 \in \Sigma_1} u_1(\sigma'_1, \sigma_2) \\
|
||||
u_2(\sigma) + \epsilon &\ge \max_{\sigma'_2 \in \Sigma_2} u_1(\sigma_1, \sigma'_2) \\
|
||||
\end{align}
|
||||
|
||||
### Regret Minimization
|
||||
|
||||
Regret is the utility (or pay off) that the player didn't get because
|
||||
she didn't follow the optimal strategy or took the best action.
|
||||
|
||||
Average overall regret for Player $i$ is the average regret of not following the
|
||||
optimal strategy in all $T$ rounds of iterations.
|
||||
|
||||
$$R^T_i = \frac{1}{T} \max_{\sigma^*_i \in \Sigma_i} \sum_{t=1}^T
|
||||
\Big( u_i(\sigma^*_i, \sigma^t_{-i}) - u_i(\sigma^t) \Big)$$
|
||||
|
||||
where $\sigma^t$ is the strategy profile of all players in iteration $t$,
|
||||
and
|
||||
|
||||
$$(\sigma^*_i, \sigma^t_{-i})$$
|
||||
|
||||
is the strategy profile $\sigma^t$ with player $i$'s strategy
|
||||
replaced with $\sigma^*_i$.
|
||||
|
||||
The average strategy is the average of strategies followed in each round,
|
||||
for all $I \in \mathcal{I}, a \in A(I)$
|
||||
|
||||
$$\textcolor{cyan}{\bar{\sigma}^T_i(I)(a)} =
|
||||
\frac{\sum_{t=1}^T \pi_i^{\sigma^t}(I)\textcolor{lightgreen}{\sigma^t(I)(a)}}{\sum_{t=1}^T \pi_i^{\sigma^t}(I)}$$
|
||||
|
||||
That is the mean regret of not playing with the optimal strategy.
|
||||
|
||||
If $R^T_i < \epsilon$ for all players then $\bar{\sigma}^T_i(I)(a)$ is a
|
||||
$2\epsilon$-Nash equilibrium.
|
||||
|
||||
\begin{align}
|
||||
R^T_i &< \epsilon \\
|
||||
R^T_i &= \frac{1}{T} \max_{\sigma^*_i \in \Sigma_i} \sum_{t=1}^T
|
||||
\Big( u_i(\sigma^*_i, \sigma^t_{-i}) - u_i(\sigma^t) \Big) \\
|
||||
&= \frac{1}{T} \max_{\sigma^*_i \in \Sigma_i} \sum_{t=1}^T u_i(\sigma^*_i, \sigma^t_{-i})
|
||||
- \frac{1}{T} \sum_{t=1}^T u_i(\sigma^t) < \epsilon
|
||||
\end{align}
|
||||
|
||||
Since $u_1 = -u_2$ because it's a zero-sum game, we can add $R^T_1$ and $R^T_i$ and the
|
||||
second term will cancel out.
|
||||
|
||||
\begin{align}
|
||||
2\epsilon &>
|
||||
\frac{1}{T} \max_{\sigma^*_1 \in \Sigma_1} \sum_{t=1}^T u_1(\sigma^*_1, \sigma^t_{-1}) +
|
||||
\frac{1}{T} \max_{\sigma^*_2 \in \Sigma_2} \sum_{t=1}^T u_2(\sigma^*_2, \sigma^t_{-2})
|
||||
\end{align}
|
||||
|
||||
The average of utilities over a set of strategies is equal to the utility of the average strategy.
|
||||
|
||||
$$\frac{1}{T} \sum_{t=1}^T u_i(\sigma^t) = u_i(\bar{\sigma}^T)$$
|
||||
|
||||
Therefore,
|
||||
|
||||
\begin{align}
|
||||
2\epsilon &>
|
||||
\max_{\sigma^*_1 \in \Sigma_1} u_1(\sigma^*_1, \bar{\sigma}^T_{-1}) +
|
||||
\max_{\sigma^*_2 \in \Sigma_2} u_2(\sigma^*_2, \bar{\sigma}^T_{-2})
|
||||
\end{align}
|
||||
|
||||
From the definition of $\max$,
|
||||
$$\max_{\sigma^*_2 \in \Sigma_2} u_2(\sigma^*_2, \bar{\sigma}^T_{-2}) \ge u_2(\bar{\sigma}^T)
|
||||
= -u_1(\bar{\sigma}^T)$$
|
||||
|
||||
Then,
|
||||
|
||||
\begin{align}
|
||||
2\epsilon &>
|
||||
\max_{\sigma^*_1 \in \Sigma_1} u_1(\sigma^*_1, \bar{\sigma}^T_{-1}) +
|
||||
-u_1(\bar{\sigma}^T) \\
|
||||
u_1(\bar{\sigma}^T) + 2\epsilon &> \max_{\sigma^*_1 \in \Sigma_1} u_1(\sigma^*_1, \bar{\sigma}^T_{-1})
|
||||
\end{align}
|
||||
|
||||
This is $2\epsilon$-Nash equilibrium.
|
||||
You can similarly prove for games with more than 2 players.
|
||||
|
||||
So we need to minimize $R^T_i$ to get close to a Nash equilibrium.
|
||||
|
||||
<a id="CounterfactualRegret"></a>
|
||||
|
||||
### Counterfactual regret
|
||||
|
||||
**Counterfactual value** $\textcolor{pink}{v_i(\sigma, I)}$ is the expected utility for player $i$ if
|
||||
if player $i$ tried to reach $I$ (took the actions leading to $I$ with a probability of $1$).
|
||||
|
||||
$$\textcolor{pink}{v_i(\sigma, I)} = \sum_{z \in Z_I} \pi^\sigma_{-i}(z[I]) \pi^\sigma(z[I], z) u_i(z)$$
|
||||
|
||||
where $Z_I$ is the set of terminal histories reachable from $I$,
|
||||
and $z[I]$ is the prefix of $z$ up to $I$.
|
||||
$\pi^\sigma(z[I], z)$ is the probability of reaching z from $z[I]$.
|
||||
|
||||
**Immediate counterfactual regret** is,
|
||||
|
||||
$$R^T_{i,imm}(I) = \max_{a \in A{I}} R^T_{i,imm}(I, a)$$
|
||||
|
||||
where
|
||||
|
||||
$$R^T_{i,imm}(I) = \frac{1}{T} \sum_{t=1}^T
|
||||
\Big(
|
||||
\textcolor{pink}{v_i(\sigma^t |_{I \rightarrow a}, I)} - \textcolor{pink}{v_i(\sigma^t, I)}
|
||||
\Big)$$
|
||||
|
||||
where $\sigma |_{I \rightarrow a}$ is the strategy profile $\sigma$ with the modification
|
||||
of always taking action $a$ at information set $I$.
|
||||
|
||||
The [paper](http://martin.zinkevich.org/publications/regretpoker.pdf) proves that (Theorem 3),
|
||||
|
||||
$$R^T_i \le \sum_{I \in \mathcal{I}} R^{T,+}_{i,imm}(I)$$
|
||||
where $$R^{T,+}_{i,imm}(I) = \max(R^T_{i,imm}(I), 0)$$
|
||||
|
||||
<a id="RegretMatching"></a>
|
||||
|
||||
### Regret Matching
|
||||
|
||||
The strategy is calculated using regret matching.
|
||||
|
||||
The regret for each information set and action pair $\textcolor{orange}{R^T_i(I, a)}$ is maintained,
|
||||
|
||||
\begin{align}
|
||||
\textcolor{coral}{r^t_i(I, a)} &=
|
||||
\textcolor{pink}{v_i(\sigma^t |_{I \rightarrow a}, I)} - \textcolor{pink}{v_i(\sigma^t, I)}
|
||||
\\
|
||||
\textcolor{orange}{R^T_i(I, a)} &=
|
||||
\frac{1}{T} \sum_{t=1}^T \textcolor{coral}{r^t_i(I, a)}
|
||||
\end{align}
|
||||
|
||||
and the strategy is calculated with regret matching,
|
||||
|
||||
\begin{align}
|
||||
\textcolor{lightgreen}{\sigma_i^{T+1}(I)(a)} =
|
||||
\begin{cases}
|
||||
\frac{\textcolor{orange}{R^{T,+}_i(I, a)}}{\sum_{a'\in A(I)}\textcolor{orange}{R^{T,+}_i(I, a')}},
|
||||
& \text{if} \sum_{a'\in A(I)}\textcolor{orange}{R^{T,+}_i(I, a')} \gt 0 \\
|
||||
\frac{1}{\lvert A(I) \rvert},
|
||||
& \text{otherwise}
|
||||
\end{cases}
|
||||
\end{align}
|
||||
|
||||
where $\textcolor{orange}{R^{T,+}_i(I, a)} = \max \Big(\textcolor{orange}{R^T_i(I, a)}, 0 \Big)$
|
||||
|
||||
The paper
|
||||
The paper
|
||||
[Regret Minimization in Games with Incomplete Information](http://martin.zinkevich.org/publications/regretpoker.pdf)
|
||||
proves that if the strategy is selected according to above equation
|
||||
$R^T_i$ gets smaller proportionate to $\frac{1}{\sqrt T}$, and
|
||||
therefore reaches $\epsilon$-[Nash equilibrium](#NashEquilibrium).
|
||||
|
||||
<a id="MCCFR"></a>
|
||||
|
||||
### Monte Carlo CFR (MCCFR)
|
||||
|
||||
Computing $\textcolor{coral}{r^t_i(I, a)}$ requires expanding the full game tree
|
||||
on each iteration.
|
||||
|
||||
The paper
|
||||
[Monte Carlo Sampling for Regret Minimization in Extensive Games](http://mlanctot.info/files/papers/nips09mccfr.pdf)
|
||||
shows we can sample from the game tree and estimate the regrets.
|
||||
|
||||
$\mathcal{Q} = {Q_1, \ldots, Q_r}$ is a set of subsets of $Z$ ($Q_j \subseteq Z$) where
|
||||
we look at only a single block $Q_j$ in an iteration.
|
||||
Union of all subsets spans $Z$ ($Q_1 \cap \ldots \cap Q_r = Z$).
|
||||
$q_j$ is the probability of picking block $Q_j$.
|
||||
|
||||
$q(z)$ is the probability of picking $z$ in current iteration; i.e. $q(z) = \sum_{j:z \in Q_j} q_j$ -
|
||||
the sum of $q_j$ where $z \in Q_j$.
|
||||
|
||||
Then we get **sampled counterfactual value** fro block $j$,
|
||||
|
||||
$$\textcolor{pink}{\tilde{v}(\sigma, I|j)} =
|
||||
\sum_{z \in Q_j} \frac{1}{q(z)}
|
||||
\pi^\sigma_{-i}(z[I]) \pi^\sigma(z[I], z) u_i(z)$$
|
||||
|
||||
The paper shows that
|
||||
|
||||
$$\mathbb{E}_{j \sim q_j} \Big[ \textcolor{pink}{\tilde{v}(\sigma, I|j)} \Big]
|
||||
= \textcolor{pink}{v_i(\sigma, I)}$$
|
||||
|
||||
with a simple proof.
|
||||
|
||||
Therefore we can sample a part of the game tree and calculate the regrets.
|
||||
We calculate an estimate of regrets
|
||||
|
||||
$$
|
||||
\textcolor{coral}{\tilde{r}^t_i(I, a)} =
|
||||
\textcolor{pink}{\tilde{v}_i(\sigma^t |_{I \rightarrow a}, I)} - \textcolor{pink}{\tilde{v}_i(\sigma^t, I)}
|
||||
$$
|
||||
|
||||
And use that to update $\textcolor{orange}{R^T_i(I, a)}$ and calculate
|
||||
the strategy $\textcolor{lightgreen}{\sigma_i^{T+1}(I)(a)}$ on each iteration.
|
||||
Finally, we calculate the overall average strategy $\textcolor{cyan}{\bar{\sigma}^T_i(I)(a)}$.
|
||||
|
||||
Here is a [Kuhn Poker](kuhn/index.html) implementation to try CFR on Kuhn Poker.
|
||||
|
||||
*Let's dive into the code!*
|
||||
"""
|
||||
from typing import NewType, Dict, List, Callable, cast
|
||||
|
||||
from labml import monit, tracker, logger, experiment
|
||||
from labml.configs import BaseConfigs, option
|
||||
|
||||
# A player $i \in N$ where $N$ is the set of players
|
||||
Player = NewType('Player', int)
|
||||
# Action $a$, $A(h) = {a: (h, a) \in H}$ where $h \in H$ is a non-terminal [history](#History)
|
||||
Action = NewType('Action', str)
|
||||
|
||||
|
||||
class History:
|
||||
"""
|
||||
<a id="History"></a>
|
||||
|
||||
## History
|
||||
|
||||
History $h \in H$ is a sequence of actions including chance events,
|
||||
and $H$ is the set of all histories.
|
||||
|
||||
This class should be extended with game specific logic.
|
||||
"""
|
||||
|
||||
def is_terminal(self):
|
||||
"""
|
||||
Whether it's a terminal history; i.e. game over.
|
||||
$h \in Z$
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
def terminal_utility(self, i: Player) -> float:
|
||||
"""
|
||||
<a id="terminal_utility"></a>
|
||||
Utility of player $i$ for a terminal history.
|
||||
$u_i(h)$ where $h \in Z$
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
def player(self) -> Player:
|
||||
"""
|
||||
Get current player, denoted by $P(h)$, where $P$ is known as **Player function**.
|
||||
|
||||
If $P(h) = c$ it means that current event is a chance $c$ event.
|
||||
Something like dealing cards, or opening common cards in poker.
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
def is_chance(self) -> bool:
|
||||
"""
|
||||
Whether the next step is a chance step; something like dealing a new card.
|
||||
$P(h) = c$
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
def sample_chance(self) -> Action:
|
||||
"""
|
||||
Sample a chance when $P(h) = c$.
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
def __add__(self, action: Action):
|
||||
"""
|
||||
Add an action to the history.
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
def info_set_key(self) -> str:
|
||||
"""
|
||||
Get [information set](#InfoSet) for the current player
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def new_info_set(self) -> 'InfoSet':
|
||||
"""
|
||||
Create a new [information set](#InfoSet) for the current player
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
def __repr__(self):
|
||||
"""
|
||||
Human readable representation
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
|
||||
class InfoSet:
|
||||
"""
|
||||
<a id="InfoSet"></a>
|
||||
|
||||
## Information Set $I_i$
|
||||
"""
|
||||
|
||||
# Unique key identifying the information set
|
||||
key: str
|
||||
# $\sigma_i$, the [strategy](#Strategy) of player $i$
|
||||
strategy: Dict[Action, float]
|
||||
# Total regret of not taking each action $A(I_i)$,
|
||||
#
|
||||
# \begin{align}
|
||||
# \textcolor{coral}{\tilde{r}^t_i(I, a)} &=
|
||||
# \textcolor{pink}{\tilde{v}_i(\sigma^t |_{I \rightarrow a}, I)} -
|
||||
# \textcolor{pink}{\tilde{v}_i(\sigma^t, I)}
|
||||
# \\
|
||||
# \textcolor{orange}{R^T_i(I, a)} &=
|
||||
# \frac{1}{T} \sum_{t=1}^T \textcolor{coral}{\tilde{r}^t_i(I, a)}
|
||||
# \end{align}
|
||||
#
|
||||
# We maintain $T \textcolor{orange}{R^T_i(I, a)}$ instead of $\textcolor{orange}{R^T_i(I, a)}$
|
||||
# since $\frac{1}{T}$ term cancels out anyway when computing strategy
|
||||
# $\textcolor{lightgreen}{\sigma_i^{T+1}(I)(a)}$
|
||||
regret: Dict[Action, float]
|
||||
# We maintain the cumulative strategy
|
||||
# $$\sum_{t=1}^T \pi_i^{\sigma^t}(I)\textcolor{lightgreen}{\sigma^t(I)(a)}$$
|
||||
# to compute overall average strategy
|
||||
#
|
||||
# $$\textcolor{cyan}{\bar{\sigma}^T_i(I)(a)} =
|
||||
# \frac{\sum_{t=1}^T \pi_i^{\sigma^t}(I)\textcolor{lightgreen}{\sigma^t(I)(a)}}{\sum_{t=1}^T \pi_i^{\sigma^t}(I)}$$
|
||||
cumulative_strategy: Dict[Action, float]
|
||||
|
||||
def __init__(self, key: str):
|
||||
"""
|
||||
Initialize
|
||||
"""
|
||||
self.key = key
|
||||
self.regret = {a: 0 for a in self.actions()}
|
||||
self.cumulative_strategy = {a: 0 for a in self.actions()}
|
||||
self.calculate_strategy()
|
||||
|
||||
def actions(self) -> List[Action]:
|
||||
"""
|
||||
Actions $A(I_i)$
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
@staticmethod
|
||||
def from_dict(data: Dict[str, any]) -> 'InfoSet':
|
||||
"""
|
||||
Load information set from a saved dictionary
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
def to_dict(self):
|
||||
"""
|
||||
Save the information set to a dictionary
|
||||
"""
|
||||
return {
|
||||
'key': self.key,
|
||||
'regret': self.regret,
|
||||
'average_strategy': self.cumulative_strategy,
|
||||
}
|
||||
|
||||
def load_dict(self, data: Dict[str, any]):
|
||||
"""
|
||||
Load data from a saved dictionary
|
||||
"""
|
||||
self.regret = data['regret']
|
||||
self.cumulative_strategy = data['average_strategy']
|
||||
self.calculate_strategy()
|
||||
|
||||
def calculate_strategy(self):
|
||||
"""
|
||||
## Calculate strategy
|
||||
|
||||
Calculate current strategy using [regret matching](#RegretMatching).
|
||||
|
||||
\begin{align}
|
||||
\textcolor{lightgreen}{\sigma_i^{T+1}(I)(a)} =
|
||||
\begin{cases}
|
||||
\frac{\textcolor{orange}{R^{T,+}_i(I, a)}}{\sum_{a'\in A(I)}\textcolor{orange}{R^{T,+}_i(I, a')}},
|
||||
& \text{if} \sum_{a'\in A(I)}\textcolor{orange}{R^{T,+}_i(I, a')} \gt 0 \\
|
||||
\frac{1}{\lvert A(I) \rvert},
|
||||
& \text{otherwise}
|
||||
\end{cases}
|
||||
\end{align}
|
||||
|
||||
where $\textcolor{orange}{R^{T,+}_i(I, a)} = \max \Big(\textcolor{orange}{R^T_i(I, a)}, 0 \Big)$
|
||||
"""
|
||||
# $$\textcolor{orange}{R^{T,+}_i(I, a)} = \max \Big(\textcolor{orange}{R^T_i(I, a)}, 0 \Big)$$
|
||||
regret = {a: max(r, 0) for a, r in self.regret.items()}
|
||||
# $$\sum_{a'\in A(I)}\textcolor{orange}{R^{T,+}_i(I, a')}$$
|
||||
regret_sum = sum(regret.values())
|
||||
# if $\sum_{a'\in A(I)}\textcolor{orange}{R^{T,+}_i(I, a')} \gt 0$,
|
||||
if regret_sum > 0:
|
||||
# $$\textcolor{lightgreen}{\sigma_i^{T+1}(I)(a)} =
|
||||
# \frac{\textcolor{orange}{R^{T,+}_i(I, a)}}{\sum_{a'\in A(I)}\textcolor{orange}{R^{T,+}_i(I, a')}}$$
|
||||
self.strategy = {a: r / regret_sum for a, r in regret.items()}
|
||||
# Otherwise,
|
||||
else:
|
||||
# $\lvert A(I) \rvert$
|
||||
count = len(list(a for a in self.regret))
|
||||
# $$\textcolor{lightgreen}{\sigma_i^{T+1}(I)(a)} =
|
||||
# \frac{1}{\lvert A(I) \rvert}$$
|
||||
self.strategy = {a: 1 / count for a, r in regret.items()}
|
||||
|
||||
def get_average_strategy(self):
|
||||
"""
|
||||
## Get average strategy
|
||||
|
||||
$$\textcolor{cyan}{\bar{\sigma}^T_i(I)(a)} =
|
||||
\frac{\sum_{t=1}^T \pi_i^{\sigma^t}(I)\textcolor{lightgreen}{\sigma^t(I)(a)}}
|
||||
{\sum_{t=1}^T \pi_i^{\sigma^t}(I)}$$
|
||||
"""
|
||||
# $$\sum_{t=1}^T \pi_i^{\sigma^t}(I) \textcolor{lightgreen}{\sigma^t(I)(a)}$$
|
||||
cum_strategy = {a: self.cumulative_strategy.get(a, 0.) for a in self.actions()}
|
||||
# $$\sum_{t=1}^T \pi_i^{\sigma^t}(I) =
|
||||
# \sum_{a \in A(I)} \sum_{t=1}^T
|
||||
# \pi_i^{\sigma^t}(I)\textcolor{lightgreen}{\sigma^t(I)(a)}$$
|
||||
strategy_sum = sum(cum_strategy.values())
|
||||
# If $\sum_{t=1}^T \pi_i^{\sigma^t}(I) > 0$,
|
||||
if strategy_sum > 0:
|
||||
# $$\textcolor{cyan}{\bar{\sigma}^T_i(I)(a)} =
|
||||
# \frac{\sum_{t=1}^T \pi_i^{\sigma^t}(I)\textcolor{lightgreen}{\sigma^t(I)(a)}}
|
||||
# {\sum_{t=1}^T \pi_i^{\sigma^t}(I)}$$
|
||||
return {a: s / strategy_sum for a, s in cum_strategy.items()}
|
||||
# Otherwise,
|
||||
else:
|
||||
# $\lvert A(I) \rvert$
|
||||
count = len(list(a for a in cum_strategy))
|
||||
# $$\textcolor{cyan}{\bar{\sigma}^T_i(I)(a)} =
|
||||
# \frac{1}{\lvert A(I) \rvert}$$
|
||||
return {a: 1 / count for a, r in cum_strategy.items()}
|
||||
|
||||
def __repr__(self):
|
||||
"""
|
||||
Human readable representation
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
|
||||
class CFR:
|
||||
"""
|
||||
## Counterfactual Regret Minimization (CFR) Algorithm
|
||||
|
||||
We do chance sampling (**CS**) where all the chance events (nodes) are sampled and
|
||||
all other events (nodes) are explored.
|
||||
|
||||
We can ignore the term $q(z)$ since it's the same for all terminal histories
|
||||
since we are doing chance sampling and it cancels out when calculating
|
||||
strategy (common in numerator and denominator).
|
||||
"""
|
||||
|
||||
# $\mathcal{I}$ set of all information sets.
|
||||
info_sets: Dict[str, InfoSet]
|
||||
|
||||
def __init__(self, *,
|
||||
create_new_history: Callable[[], History],
|
||||
epochs: int,
|
||||
n_players: int = 2):
|
||||
"""
|
||||
* `create_new_history` creates a new empty history
|
||||
* `epochs` is the number of iterations to train on $T$
|
||||
* `n_players` is the number of players
|
||||
"""
|
||||
self.n_players = n_players
|
||||
self.epochs = epochs
|
||||
self.create_new_history = create_new_history
|
||||
# A dictionary for $\mathcal{I}$ set of all information sets
|
||||
self.info_sets = {}
|
||||
# Tracker for analytics
|
||||
self.tracker = InfoSetTracker()
|
||||
|
||||
def _get_info_set(self, h: History):
|
||||
"""
|
||||
Returns the information set $I$ of the current player for a given history $h$
|
||||
"""
|
||||
info_set_key = h.info_set_key()
|
||||
if info_set_key not in self.info_sets:
|
||||
self.info_sets[info_set_key] = h.new_info_set()
|
||||
return self.info_sets[info_set_key]
|
||||
|
||||
def walk_tree(self, h: History, i: Player, pi_i: float, pi_neg_i: float) -> float:
|
||||
"""
|
||||
### Walk Tree
|
||||
|
||||
This function walks the game tree.
|
||||
|
||||
* `h` is the current history $h$
|
||||
* `i` is the player $i$ that we are computing regrets of
|
||||
* [`pi_i`](#HistoryProbability) is
|
||||
$\pi^{\sigma^t}_i(h)$
|
||||
* [`pi_neg_i`](#HistoryProbability) is
|
||||
$\pi^{\sigma^t}_{-i}(h)$
|
||||
|
||||
It returns the expected utility, for the history $h$
|
||||
$$\sum_{z \in Z_h} \pi^\sigma(h, z) u_i(z)$$
|
||||
where $Z_h$ is the set of terminal histories with prefix $h$
|
||||
|
||||
While walking the tee it updates the total regrets $\textcolor{orange}{R^T_i(I, a)}$.
|
||||
"""
|
||||
|
||||
# If it's a terminal history $h \in Z$ return the terminal utility $u_i(h)$.
|
||||
if h.is_terminal():
|
||||
return h.terminal_utility(i)
|
||||
# If it's a chance event $P(h) = c$ sample a and go to next step.
|
||||
elif h.is_chance():
|
||||
a = h.sample_chance()
|
||||
return self.walk_tree(h + a, i, pi_i, pi_neg_i)
|
||||
|
||||
# Get current player's information set for $h$
|
||||
I = self._get_info_set(h)
|
||||
# To store $\sum_{z \in Z_h} \pi^\sigma(h, z) u_i(z)$
|
||||
v = 0
|
||||
# To store
|
||||
# $$\sum_{z \in Z_h} \pi^{\sigma^t |_{I \rightarrow a}}(h, z) u_i(z)$$
|
||||
# for each action $a \in A(h)$
|
||||
va = {}
|
||||
|
||||
# Iterate through all actions
|
||||
for a in I.actions():
|
||||
# If the current player is $i$,
|
||||
if i == h.player():
|
||||
# \begin{align}
|
||||
# \pi^{\sigma^t}_i(h + a) &= \pi^{\sigma^t}_i(h) \sigma^t_i(I)(a) \\
|
||||
# \pi^{\sigma^t}_{-i}(h + a) &= \pi^{\sigma^t}_{-i}(h)
|
||||
# \end{align}
|
||||
va[a] = self.walk_tree(h + a, i, pi_i * I.strategy[a], pi_neg_i)
|
||||
# Otherwise,
|
||||
else:
|
||||
# \begin{align}
|
||||
# \pi^{\sigma^t}_i(h + a) &= \pi^{\sigma^t}_i(h) \\
|
||||
# \pi^{\sigma^t}_{-i}(h + a) &= \pi^{\sigma^t}_{-i}(h) * \sigma^t_i(I)(a)
|
||||
# \end{align}
|
||||
va[a] = self.walk_tree(h + a, i, pi_i, pi_neg_i * I.strategy[a])
|
||||
# $$\sum_{z \in Z_h} \pi^\sigma(h, z) u_i(z) =
|
||||
# \sum_{a \in A(I)} \Bigg[ \sigma^t_i(I)(a)
|
||||
# \sum_{z \in Z_h} \pi^{\sigma^t |_{I \rightarrow a}}(h, z) u_i(z)
|
||||
# \Bigg]$$
|
||||
v = v + I.strategy[a] * va[a]
|
||||
|
||||
# If the current player is $i$,
|
||||
# update the cumulative strategies and total regrets
|
||||
if h.player() == i:
|
||||
# Update cumulative strategies
|
||||
# $$\sum_{t=1}^T \pi_i^{\sigma^t}(I)\textcolor{lightgreen}{\sigma^t(I)(a)}
|
||||
# = \sum_{t=1}^T \Big[ \sum_{h \in I} \pi_i^{\sigma^t}(h)
|
||||
# \textcolor{lightgreen}{\sigma^t(I)(a)} \Big]$$
|
||||
for a in I.actions():
|
||||
I.cumulative_strategy[a] = I.cumulative_strategy[a] + pi_i * I.strategy[a]
|
||||
# \begin{align}
|
||||
# \textcolor{coral}{\tilde{r}^t_i(I, a)} &=
|
||||
# \textcolor{pink}{\tilde{v}_i(\sigma^t |_{I \rightarrow a}, I)} -
|
||||
# \textcolor{pink}{\tilde{v}_i(\sigma^t, I)} \\
|
||||
# &=
|
||||
# \pi^{\sigma^t}_{-i} (h) \Big(
|
||||
# \sum_{z \in Z_h} \pi^{\sigma^t |_{I \rightarrow a}}(h, z) u_i(z) -
|
||||
# \sum_{z \in Z_h} \pi^\sigma(h, z) u_i(z)
|
||||
# \Big) \\
|
||||
# T \textcolor{orange}{R^T_i(I, a)} &=
|
||||
# \sum_{t=1}^T \textcolor{coral}{\tilde{r}^t_i(I, a)}
|
||||
# \end{align}
|
||||
for a in I.actions():
|
||||
I.regret[a] += pi_neg_i * (va[a] - v)
|
||||
|
||||
# Update the strategy $\textcolor{lightgreen}{\sigma^t(I)(a)}$
|
||||
I.calculate_strategy()
|
||||
|
||||
# Return the expected utility for player $i$,
|
||||
# $$\sum_{z \in Z_h} \pi^\sigma(h, z) u_i(z)$$
|
||||
return v
|
||||
|
||||
def iterate(self):
|
||||
"""
|
||||
### Iteratively update $\textcolor{lightgreen}{\sigma^t(I)(a)}$
|
||||
|
||||
This updates the strategies for $T$ iterations.
|
||||
"""
|
||||
|
||||
# Loop for `epochs` times
|
||||
for t in monit.iterate('Train', self.epochs):
|
||||
# Walk tree and update regrets for each player
|
||||
for i in range(self.n_players):
|
||||
self.walk_tree(self.create_new_history(), cast(Player, i), 1, 1)
|
||||
|
||||
# Track data for analytics
|
||||
tracker.add_global_step()
|
||||
self.tracker(self.info_sets)
|
||||
tracker.save()
|
||||
|
||||
# Print the information sets
|
||||
logger.inspect(self.info_sets)
|
||||
|
||||
|
||||
class InfoSetTracker:
|
||||
"""
|
||||
### Information set tracker
|
||||
|
||||
This is a small helper class to track data from information sets
|
||||
"""
|
||||
def __init__(self):
|
||||
"""
|
||||
Set tracking indicators
|
||||
"""
|
||||
tracker.set_histogram(f'strategy.*')
|
||||
tracker.set_histogram(f'average_strategy.*')
|
||||
tracker.set_histogram(f'regret.*')
|
||||
|
||||
def __call__(self, info_sets: Dict[str, InfoSet]):
|
||||
"""
|
||||
Track the data from all information sets
|
||||
"""
|
||||
for I in info_sets.values():
|
||||
avg_strategy = I.get_average_strategy()
|
||||
for a in I.actions():
|
||||
tracker.add({
|
||||
f'strategy.{I.key}.{a}': I.strategy[a],
|
||||
f'average_strategy.{I.key}.{a}': avg_strategy[a],
|
||||
f'regret.{I.key}.{a}': I.regret[a],
|
||||
})
|
||||
|
||||
|
||||
class CFRConfigs(BaseConfigs):
|
||||
"""
|
||||
### Configurable CFR module
|
||||
"""
|
||||
create_new_history: Callable[[], History]
|
||||
epochs: int = 1_00_000
|
||||
cfr: CFR = 'simple_cfr'
|
||||
|
||||
|
||||
@option(CFRConfigs.cfr)
|
||||
def simple_cfr(c: CFRConfigs):
|
||||
"""
|
||||
Initialize **CFR** algorithm
|
||||
"""
|
||||
return CFR(create_new_history=c.create_new_history,
|
||||
epochs=c.epochs)
|
||||
@@ -0,0 +1,66 @@
|
||||
from typing import List
|
||||
|
||||
import altair as alt
|
||||
import numpy as np
|
||||
|
||||
from labml import analytics
|
||||
from labml.analytics import IndicatorCollection
|
||||
|
||||
|
||||
def calculate_percentages(means: List[np.ndarray], names: List[List[str]]):
|
||||
normalized = []
|
||||
|
||||
for i in range(len(means)):
|
||||
total = np.zeros_like(means[i])
|
||||
for j, n in enumerate(names):
|
||||
if n[-1][:-1] == names[i][-1][:-1]:
|
||||
total += means[j]
|
||||
normalized.append(means[i] / (total + np.finfo(float).eps))
|
||||
|
||||
return normalized
|
||||
|
||||
|
||||
def plot_infosets(indicators: IndicatorCollection, *,
|
||||
is_normalize: bool = True,
|
||||
width: int = 600,
|
||||
height: int = 300):
|
||||
data, names = analytics.indicator_data(indicators)
|
||||
step = [d[:, 0] for d in data]
|
||||
means = [d[:, 5] for d in data]
|
||||
|
||||
if is_normalize:
|
||||
normalized = calculate_percentages(means, names)
|
||||
else:
|
||||
normalized = means
|
||||
|
||||
common = names[0][-1]
|
||||
for i, n in enumerate(names):
|
||||
n = n[-1]
|
||||
if len(n) < len(common):
|
||||
common = common[:len(n)]
|
||||
for j in range(len(common)):
|
||||
if common[j] != n[j]:
|
||||
common = common[:j]
|
||||
break
|
||||
|
||||
table = []
|
||||
for i, n in enumerate(names):
|
||||
for j, v in zip(step[i], normalized[i]):
|
||||
table.append({
|
||||
'series': n[-1][len(common):],
|
||||
'step': j,
|
||||
'value': v
|
||||
})
|
||||
|
||||
table = alt.Data(values=table)
|
||||
|
||||
selection = alt.selection_multi(fields=['series'], bind='legend')
|
||||
|
||||
return alt.Chart(table).mark_line().encode(
|
||||
alt.X('step:Q'),
|
||||
alt.Y('value:Q'),
|
||||
alt.Color('series:N', scale=alt.Scale(scheme='tableau20')),
|
||||
opacity=alt.condition(selection, alt.value(1), alt.value(0.0001))
|
||||
).add_selection(
|
||||
selection
|
||||
).properties(width=width, height=height)
|
||||
@@ -0,0 +1,27 @@
|
||||
import json
|
||||
import pathlib
|
||||
from typing import Dict
|
||||
|
||||
from labml import experiment
|
||||
from labml_nn.cfr import InfoSet
|
||||
|
||||
|
||||
class InfoSetSaver(experiment.ModelSaver):
|
||||
def __init__(self, infosets: Dict[str, InfoSet]):
|
||||
self.infosets = infosets
|
||||
|
||||
def save(self, checkpoint_path: pathlib.Path) -> any:
|
||||
data = {key: infoset.to_dict() for key, infoset in self.infosets.items()}
|
||||
file_name = f"infosets.json"
|
||||
|
||||
with open(str(checkpoint_path / file_name), 'w') as f:
|
||||
f.write(json.dumps(data))
|
||||
|
||||
return file_name
|
||||
|
||||
def load(self, checkpoint_path: pathlib.Path, file_name: str):
|
||||
with open(str(checkpoint_path / file_name), 'w') as f:
|
||||
data = json.loads(f.read())
|
||||
|
||||
for key, d in data.items():
|
||||
self.infosets[key] = InfoSet.from_dict(d)
|
||||
@@ -0,0 +1,248 @@
|
||||
"""
|
||||
---
|
||||
title: CFR on Kuhn Poker
|
||||
summary: >
|
||||
This is an annotated implementation/tutorial of CFR on Kuhn Poker
|
||||
---
|
||||
|
||||
# [Counterfactual Regret Minimization (CFR)](../index.html) on Kuhn Poker
|
||||
|
||||
This applies [Counterfactual Regret Minimization (CFR)](../index.html) to Kuhn poker.
|
||||
|
||||
[Kuhn Poker](https://en.wikipedia.org/wiki/Kuhn_poker) is a two player 3-card betting game.
|
||||
The players are dealt one card each out of Ace, King and Queen (no suits).
|
||||
There are only three cards in the pack so one card is left out.
|
||||
Ace beats King and Queen and King beats Queen - just like in normal ranking of cards.
|
||||
|
||||
Both players ante $1$ chip (blindly bet $1$ chip).
|
||||
After looking at the cards, the first player can either pass or bet $1$ chip.
|
||||
If first player passes, the the player with higher card wins the pot.
|
||||
If first player bets, the second play can bet (i.e. call) $1$ chip or pass (i.e. fold).
|
||||
If the second player bets and the player with the higher card wins the pot.
|
||||
If the second player passes (i.e. folds) the first player gets the pot.
|
||||
This game is played repeatedly and a good strategy will optimize for the long term utility (or winnings).
|
||||
|
||||
Here's some example games:
|
||||
|
||||
* `KAp` - Player 1 gets K. Player 2 gets A. Player 1 passes. Player 2 doesn't get a betting chance and Player 2 wins the pot of $2$ chips.
|
||||
* `QKbp` - Player 1 gets Q. Player 2 gets K. Player 1 bets a chip. Player 2 passes (folds). Player 1 gets the pot of $4$ because Player 2 folded.
|
||||
* `QAbb` - Player 1 gets Q. Player 2 gets A. Player 1 bets a chip. Player 2 also bets (calls). Player 2 wins the pot of $4$.
|
||||
|
||||
He we extend the `InfoSet` class and `History` class defined in [`__init__.py`](../index.html)
|
||||
with Kuhn Poker specifics.
|
||||
|
||||
[](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/cfr/kuhn/experiment.ipynb)
|
||||
"""
|
||||
|
||||
from typing import List, cast, Dict
|
||||
|
||||
import numpy as np
|
||||
|
||||
from labml import experiment
|
||||
from labml.configs import option
|
||||
from labml_nn.cfr import History as _History, InfoSet as _InfoSet, Action, Player, CFRConfigs
|
||||
|
||||
# Kuhn poker actions are pass (`p`) or bet (`b`)
|
||||
ACTIONS = cast(List[Action], ['p', 'b'])
|
||||
# The three cards in play are Ace, King and Queen
|
||||
CHANCES = cast(List[Action], ['A', 'K', 'Q'])
|
||||
# There are two players
|
||||
PLAYERS = cast(List[Player], [0, 1])
|
||||
|
||||
|
||||
class InfoSet(_InfoSet):
|
||||
"""
|
||||
## [Information set](../index.html#InfoSet)
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def from_dict(data: Dict[str, any]) -> 'InfoSet':
|
||||
"""Does not support save/load"""
|
||||
pass
|
||||
|
||||
def actions(self) -> List[Action]:
|
||||
"""
|
||||
Return the list of actions. Terminal states are handled by `History` class.
|
||||
"""
|
||||
return ACTIONS
|
||||
|
||||
def __repr__(self):
|
||||
"""
|
||||
Human readable string representation - it gives the betting probability
|
||||
"""
|
||||
total = sum(self.cumulative_strategy.values())
|
||||
total = max(total, 1e-6)
|
||||
bet = self.cumulative_strategy[cast(Action, 'b')] / total
|
||||
return f'{bet * 100: .1f}%'
|
||||
|
||||
|
||||
class History(_History):
|
||||
"""
|
||||
## [History](../index.html#History)
|
||||
|
||||
This defines when a game ends, calculates the utility and sample chance events (dealing cards).
|
||||
|
||||
The history is stored in a string:
|
||||
|
||||
* First two characters are the cards dealt to player 1 and player 2
|
||||
* The third character is the action by the first player
|
||||
* Fourth character is the action by the second player
|
||||
"""
|
||||
|
||||
# History
|
||||
history: str
|
||||
|
||||
def __init__(self, history: str = ''):
|
||||
"""
|
||||
Initialize with a given history string
|
||||
"""
|
||||
self.history = history
|
||||
|
||||
def is_terminal(self):
|
||||
"""
|
||||
Whether the history is terminal (game over).
|
||||
"""
|
||||
# Players are yet to take actions
|
||||
if len(self.history) <= 2:
|
||||
return False
|
||||
# Last player to play passed (game over)
|
||||
elif self.history[-1] == 'p':
|
||||
return True
|
||||
# Both players called (bet) (game over)
|
||||
elif self.history[-2:] == 'bb':
|
||||
return True
|
||||
# Any other combination
|
||||
else:
|
||||
return False
|
||||
|
||||
def _terminal_utility_p1(self) -> float:
|
||||
"""
|
||||
Calculate the terminal utility for player $1$, $u_1(z)$
|
||||
"""
|
||||
# $+1$ if Player 1 has a better card and $-1$ otherwise
|
||||
winner = -1 + 2 * (self.history[0] < self.history[1])
|
||||
|
||||
# Second player passed
|
||||
if self.history[-2:] == 'bp':
|
||||
return 1
|
||||
# Both players called, the player with better card wins $2$ chips
|
||||
elif self.history[-2:] == 'bb':
|
||||
return winner * 2
|
||||
# First player passed, the player with better card wins $1$ chip
|
||||
elif self.history[-1] == 'p':
|
||||
return winner
|
||||
# History is non-terminal
|
||||
else:
|
||||
raise RuntimeError()
|
||||
|
||||
def terminal_utility(self, i: Player) -> float:
|
||||
"""
|
||||
Get the terminal utility for player $i$
|
||||
"""
|
||||
# If $i$ is Player 1
|
||||
if i == PLAYERS[0]:
|
||||
return self._terminal_utility_p1()
|
||||
# Otherwise, $u_2(z) = -u_1(z)$
|
||||
else:
|
||||
return -1 * self._terminal_utility_p1()
|
||||
|
||||
def is_chance(self) -> bool:
|
||||
"""
|
||||
The first two events are card dealing; i.e. chance events
|
||||
"""
|
||||
return len(self.history) < 2
|
||||
|
||||
def __add__(self, other: Action):
|
||||
"""
|
||||
Add an action to the history and return a new history
|
||||
"""
|
||||
return History(self.history + other)
|
||||
|
||||
def player(self) -> Player:
|
||||
"""
|
||||
Current player
|
||||
"""
|
||||
return cast(Player, len(self.history) % 2)
|
||||
|
||||
def sample_chance(self) -> Action:
|
||||
"""
|
||||
Sample a chance action
|
||||
"""
|
||||
while True:
|
||||
# Randomly pick a card
|
||||
r = np.random.randint(len(CHANCES))
|
||||
chance = CHANCES[r]
|
||||
# See if the card was dealt before
|
||||
for c in self.history:
|
||||
if c == chance:
|
||||
chance = None
|
||||
break
|
||||
|
||||
# Return the card if it was not dealt before
|
||||
if chance is not None:
|
||||
return cast(Action, chance)
|
||||
|
||||
def __repr__(self):
|
||||
"""
|
||||
Human readable representation
|
||||
"""
|
||||
return repr(self.history)
|
||||
|
||||
def info_set_key(self) -> str:
|
||||
"""
|
||||
Information set key for the current history.
|
||||
This is a string of actions only visible to the current player.
|
||||
"""
|
||||
# Get current player
|
||||
i = self.player()
|
||||
# Current player sees her card and the betting actions
|
||||
return self.history[i] + self.history[2:]
|
||||
|
||||
def new_info_set(self) -> InfoSet:
|
||||
# Create a new information set object
|
||||
return InfoSet(self.info_set_key())
|
||||
|
||||
|
||||
def create_new_history():
|
||||
"""A function to create an empty history object"""
|
||||
return History()
|
||||
|
||||
|
||||
class Configs(CFRConfigs):
|
||||
"""
|
||||
Configurations extends the CFR configurations class
|
||||
"""
|
||||
pass
|
||||
|
||||
|
||||
@option(Configs.create_new_history)
|
||||
def _cnh():
|
||||
"""
|
||||
Set the `create_new_history` method for Kuhn Poker
|
||||
"""
|
||||
return create_new_history
|
||||
|
||||
|
||||
def main():
|
||||
"""
|
||||
### Run the experiment
|
||||
"""
|
||||
|
||||
# Create an experiment, we only write tracking information to `sqlite` to speed things up.
|
||||
# Since the algorithm iterates fast and we track data on each iteration, writing to
|
||||
# other destinations such as Tensorboard can be relatively time consuming.
|
||||
# SQLite is enough for our analytics.
|
||||
experiment.create(name='kuhn_poker', writers={'sqlite'})
|
||||
# Initialize configuration
|
||||
conf = Configs()
|
||||
# Load configuration
|
||||
experiment.configs(conf)
|
||||
# Start the experiment
|
||||
with experiment.start():
|
||||
# Start iterating
|
||||
conf.cfr.iterate()
|
||||
|
||||
|
||||
#
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,250 @@
|
||||
{
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0,
|
||||
"metadata": {
|
||||
"accelerator": "GPU",
|
||||
"colab": {
|
||||
"name": "Counterfactual Regret Minimization (CFR) on Kuhn Poker",
|
||||
"provenance": [],
|
||||
"collapsed_sections": []
|
||||
},
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3",
|
||||
"language": "python",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"codemirror_mode": {
|
||||
"name": "ipython",
|
||||
"version": 3
|
||||
},
|
||||
"file_extension": ".py",
|
||||
"mimetype": "text/x-python",
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython3",
|
||||
"version": "3.7.5"
|
||||
}
|
||||
},
|
||||
"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/cfr/kuhn/experiment.ipynb) \n",
|
||||
"\n",
|
||||
"## [Counterfactual Regret Minimization (CFR)](https://nn.labml.ai/cfr/index.html) on Kuhn Poker\n",
|
||||
"\n",
|
||||
"This is an experiment learning to play Kuhn Poker with Counterfactual Regret Minimization CFR algorithm."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "AahG_i2y5tY9"
|
||||
},
|
||||
"source": [
|
||||
"Install the `labml-nn` package"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "ZCzmCrAIVg0L"
|
||||
},
|
||||
"source": [
|
||||
"%%capture\n",
|
||||
"!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, analytics\n",
|
||||
"from labml_nn.cfr.analytics import plot_infosets\n",
|
||||
"from labml_nn.cfr.kuhn import Configs\n",
|
||||
"from labml_nn.cfr.infoset_saver import InfoSetSaver"
|
||||
],
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "Lpggo0wM6qb-"
|
||||
},
|
||||
"source": [
|
||||
"Create an experiment, we only write tracking information to `sqlite` to speed things up.\n",
|
||||
"Since the algorithm iterates fast and we track data on each iteration, writing to\n",
|
||||
"other destinations such as Tensorboard can be relatively time consuming.\n",
|
||||
"SQLite is enough for our analytics."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "bFcr9k-l4cAg"
|
||||
},
|
||||
"source": [
|
||||
"experiment.create(name='kuhn_poker', writers={'sqlite'})"
|
||||
],
|
||||
"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": "e20b5ea3-605b-4bcc-c9de-0da93b6c7f32"
|
||||
},
|
||||
"source": [
|
||||
"experiment.configs(conf, {'epochs': 1_000_000})"
|
||||
],
|
||||
"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": 187
|
||||
},
|
||||
"id": "aIAWo7Fw5DR8",
|
||||
"outputId": "18cd2384-d6c0-42a3-feae-5a309563bb33"
|
||||
},
|
||||
"source": [
|
||||
"# Start the experiment\n",
|
||||
"with experiment.start():\n",
|
||||
" conf.cfr.iterate()"
|
||||
],
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "oBXXlP2b7XZO"
|
||||
},
|
||||
"source": [
|
||||
"inds = analytics.runs(experiment.get_uuid())"
|
||||
],
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "RJ0L4XH2Y8g4"
|
||||
},
|
||||
"source": [
|
||||
"# dir(inds)"
|
||||
],
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "htumVaMnY8g4",
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 568
|
||||
},
|
||||
"outputId": "735a944d-1b96-49e8-97b6-64317ea515b1"
|
||||
},
|
||||
"source": [
|
||||
"plot_infosets(inds['average_strategy.*'], width=600, height=500).display()"
|
||||
],
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "yTDu8JKdY8g4",
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 690
|
||||
},
|
||||
"outputId": "6047dae2-095e-4323-ee91-f49573ad509c"
|
||||
},
|
||||
"source": [
|
||||
"analytics.scatter(inds.average_strategy_Q_b, inds.average_strategy_Kb_b,\n",
|
||||
" width=400, height=400)"
|
||||
],
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "GnI67bbLY8g5"
|
||||
},
|
||||
"source": [
|
||||
""
|
||||
],
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,211 @@
|
||||
"""
|
||||
---
|
||||
title: Patches Are All You Need? (ConvMixer)
|
||||
summary: >
|
||||
A PyTorch implementation/tutorial of the paper
|
||||
"Patches Are All You Need?"
|
||||
---
|
||||
|
||||
# Patches Are All You Need? (ConvMixer)
|
||||
|
||||
This is a [PyTorch](https://pytorch.org) implementation of the paper
|
||||
[Patches Are All You Need?](https://arxiv.org/abs/2201.09792).
|
||||
|
||||
|
||||
|
||||
ConvMixer is Similar to [MLP-Mixer](../transformers/mlp_mixer/index.html).
|
||||
MLP-Mixer separates mixing of spatial and channel dimensions, by applying an MLP across spatial dimension
|
||||
and then an MLP across the channel dimension
|
||||
(spatial MLP replaces the [ViT](../transformers/vit/index.html) attention
|
||||
and channel MLP is the [FFN](../transformers/feed_forward.html) of ViT).
|
||||
|
||||
ConvMixer uses a $1 \times 1$ convolution for channel mixing and a
|
||||
depth-wise convolution for spatial mixing.
|
||||
Since it's a convolution instead of a full MLP across the space, it mixes only the nearby batches in
|
||||
contrast to ViT or MLP-Mixer.
|
||||
Also, the MLP-mixer uses MLPs of two layers for each mixing and ConvMixer uses a single layer for each mixing.
|
||||
|
||||
The paper recommends removing the residual connection across the channel mixing (point-wise convolution)
|
||||
and having only a residual connection over the spatial mixing (depth-wise convolution).
|
||||
They also use [Batch normalization](../normalization/batch_norm/index.html) instead
|
||||
of [Layer normalization](../normalization/layer_norm/index.html).
|
||||
|
||||
Here's [an experiment](experiment.html) that trains ConvMixer on CIFAR-10.
|
||||
"""
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from labml_nn.utils import clone_module_list
|
||||
|
||||
|
||||
class ConvMixerLayer(nn.Module):
|
||||
"""
|
||||
<a id="ConvMixerLayer"></a>
|
||||
|
||||
## ConvMixer layer
|
||||
|
||||
This is a single ConvMixer layer. The model will have a series of these.
|
||||
"""
|
||||
|
||||
def __init__(self, d_model: int, kernel_size: int):
|
||||
"""
|
||||
* `d_model` is the number of channels in patch embeddings, $h$
|
||||
* `kernel_size` is the size of the kernel of spatial convolution, $k$
|
||||
"""
|
||||
super().__init__()
|
||||
# Depth-wise convolution is separate convolution for each channel.
|
||||
# We do this with a convolution layer with the number of groups equal to the number of channels.
|
||||
# So that each channel is it's own group.
|
||||
self.depth_wise_conv = nn.Conv2d(d_model, d_model,
|
||||
kernel_size=kernel_size,
|
||||
groups=d_model,
|
||||
padding=(kernel_size - 1) // 2)
|
||||
# Activation after depth-wise convolution
|
||||
self.act1 = nn.GELU()
|
||||
# Normalization after depth-wise convolution
|
||||
self.norm1 = nn.BatchNorm2d(d_model)
|
||||
|
||||
# Point-wise convolution is a $1 \times 1$ convolution.
|
||||
# i.e. a linear transformation of patch embeddings
|
||||
self.point_wise_conv = nn.Conv2d(d_model, d_model, kernel_size=1)
|
||||
# Activation after point-wise convolution
|
||||
self.act2 = nn.GELU()
|
||||
# Normalization after point-wise convolution
|
||||
self.norm2 = nn.BatchNorm2d(d_model)
|
||||
|
||||
def forward(self, x: torch.Tensor):
|
||||
# For the residual connection around the depth-wise convolution
|
||||
residual = x
|
||||
|
||||
# Depth-wise convolution, activation and normalization
|
||||
x = self.depth_wise_conv(x)
|
||||
x = self.act1(x)
|
||||
x = self.norm1(x)
|
||||
|
||||
# Add residual connection
|
||||
x += residual
|
||||
|
||||
# Point-wise convolution, activation and normalization
|
||||
x = self.point_wise_conv(x)
|
||||
x = self.act2(x)
|
||||
x = self.norm2(x)
|
||||
|
||||
#
|
||||
return x
|
||||
|
||||
|
||||
class PatchEmbeddings(nn.Module):
|
||||
"""
|
||||
<a id="PatchEmbeddings"></a>
|
||||
|
||||
## Get patch embeddings
|
||||
|
||||
This splits the image into patches of size $p \times p$ and gives an embedding for each patch.
|
||||
"""
|
||||
|
||||
def __init__(self, d_model: int, patch_size: int, in_channels: int):
|
||||
"""
|
||||
* `d_model` is the number of channels in patch embeddings $h$
|
||||
* `patch_size` is the size of the patch, $p$
|
||||
* `in_channels` is the number of channels in the input image (3 for rgb)
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
# We create a convolution layer with a kernel size and and stride length equal to patch size.
|
||||
# This is equivalent to splitting the image into patches and doing a linear
|
||||
# transformation on each patch.
|
||||
self.conv = nn.Conv2d(in_channels, d_model, kernel_size=patch_size, stride=patch_size)
|
||||
# Activation function
|
||||
self.act = nn.GELU()
|
||||
# Batch normalization
|
||||
self.norm = nn.BatchNorm2d(d_model)
|
||||
|
||||
def forward(self, x: torch.Tensor):
|
||||
"""
|
||||
* `x` is the input image of shape `[batch_size, channels, height, width]`
|
||||
"""
|
||||
# Apply convolution layer
|
||||
x = self.conv(x)
|
||||
# Activation and normalization
|
||||
x = self.act(x)
|
||||
x = self.norm(x)
|
||||
|
||||
#
|
||||
return x
|
||||
|
||||
|
||||
class ClassificationHead(nn.Module):
|
||||
"""
|
||||
<a id="ClassificationHead"></a>
|
||||
|
||||
## Classification Head
|
||||
|
||||
They do average pooling (taking the mean of all patch embeddings) and a final linear transformation
|
||||
to predict the log-probabilities of the image classes.
|
||||
"""
|
||||
|
||||
def __init__(self, d_model: int, n_classes: int):
|
||||
"""
|
||||
* `d_model` is the number of channels in patch embeddings, $h$
|
||||
* `n_classes` is the number of classes in the classification task
|
||||
"""
|
||||
super().__init__()
|
||||
# Average Pool
|
||||
self.pool = nn.AdaptiveAvgPool2d((1, 1))
|
||||
# Linear layer
|
||||
self.linear = nn.Linear(d_model, n_classes)
|
||||
|
||||
def forward(self, x: torch.Tensor):
|
||||
# Average pooling
|
||||
x = self.pool(x)
|
||||
# Get the embedding, `x` will have shape `[batch_size, d_model, 1, 1]`
|
||||
x = x[:, :, 0, 0]
|
||||
# Linear layer
|
||||
x = self.linear(x)
|
||||
|
||||
#
|
||||
return x
|
||||
|
||||
|
||||
class ConvMixer(nn.Module):
|
||||
"""
|
||||
## ConvMixer
|
||||
|
||||
This combines the patch embeddings block, a number of ConvMixer layers and a classification head.
|
||||
"""
|
||||
|
||||
def __init__(self, conv_mixer_layer: ConvMixerLayer, n_layers: int,
|
||||
patch_emb: PatchEmbeddings,
|
||||
classification: ClassificationHead):
|
||||
"""
|
||||
* `conv_mixer_layer` is a copy of a single [ConvMixer layer](#ConvMixerLayer).
|
||||
We make copies of it to make ConvMixer with `n_layers`.
|
||||
* `n_layers` is the number of ConvMixer layers (or depth), $d$.
|
||||
* `patch_emb` is the [patch embeddings layer](#PatchEmbeddings).
|
||||
* `classification` is the [classification head](#ClassificationHead).
|
||||
"""
|
||||
super().__init__()
|
||||
# Patch embeddings
|
||||
self.patch_emb = patch_emb
|
||||
# Classification head
|
||||
self.classification = classification
|
||||
# Make copies of the [ConvMixer layer](#ConvMixerLayer)
|
||||
self.conv_mixer_layers = clone_module_list(conv_mixer_layer, n_layers)
|
||||
|
||||
def forward(self, x: torch.Tensor):
|
||||
"""
|
||||
* `x` is the input image of shape `[batch_size, channels, height, width]`
|
||||
"""
|
||||
# Get patch embeddings. This gives a tensor of shape `[batch_size, d_model, height / patch_size, width / patch_size]`.
|
||||
x = self.patch_emb(x)
|
||||
|
||||
# Pass through [ConvMixer layers](#ConvMixerLayer)
|
||||
for layer in self.conv_mixer_layers:
|
||||
x = layer(x)
|
||||
|
||||
# Classification head, to get logits
|
||||
x = self.classification(x)
|
||||
|
||||
#
|
||||
return x
|
||||
@@ -0,0 +1,85 @@
|
||||
"""
|
||||
---
|
||||
title: Train ConvMixer on CIFAR 10
|
||||
summary: >
|
||||
Train ConvMixer on CIFAR 10
|
||||
---
|
||||
|
||||
# Train a [ConvMixer](index.html) on CIFAR 10
|
||||
|
||||
This script trains a ConvMixer on CIFAR 10 dataset.
|
||||
|
||||
This is not an attempt to reproduce the results of the paper.
|
||||
The paper uses image augmentations
|
||||
present in [PyTorch Image Models (timm)](https://github.com/rwightman/pytorch-image-models)
|
||||
for training. We haven't done this for simplicity - which causes our validation accuracy to drop.
|
||||
"""
|
||||
|
||||
from labml import experiment
|
||||
from labml.configs import option
|
||||
from labml_nn.experiments.cifar10 import CIFAR10Configs
|
||||
|
||||
|
||||
class Configs(CIFAR10Configs):
|
||||
"""
|
||||
## Configurations
|
||||
|
||||
We use [`CIFAR10Configs`](../experiments/cifar10.html) which defines all the
|
||||
dataset related configurations, optimizer, and a training loop.
|
||||
"""
|
||||
|
||||
# Size of a patch, $p$
|
||||
patch_size: int = 2
|
||||
# Number of channels in patch embeddings, $h$
|
||||
d_model: int = 256
|
||||
# Number of [ConvMixer layers](#ConvMixerLayer) or depth, $d$
|
||||
n_layers: int = 8
|
||||
# Kernel size of the depth-wise convolution, $k$
|
||||
kernel_size: int = 7
|
||||
# Number of classes in the task
|
||||
n_classes: int = 10
|
||||
|
||||
|
||||
@option(Configs.model)
|
||||
def _conv_mixer(c: Configs):
|
||||
"""
|
||||
### Create model
|
||||
"""
|
||||
from labml_nn.conv_mixer import ConvMixerLayer, ConvMixer, ClassificationHead, PatchEmbeddings
|
||||
|
||||
# Create ConvMixer
|
||||
return ConvMixer(ConvMixerLayer(c.d_model, c.kernel_size), c.n_layers,
|
||||
PatchEmbeddings(c.d_model, c.patch_size, 3),
|
||||
ClassificationHead(c.d_model, c.n_classes)).to(c.device)
|
||||
|
||||
|
||||
def main():
|
||||
# Create experiment
|
||||
experiment.create(name='ConvMixer', comment='cifar10')
|
||||
# Create configurations
|
||||
conf = Configs()
|
||||
# Load configurations
|
||||
experiment.configs(conf, {
|
||||
# Optimizer
|
||||
'optimizer.optimizer': 'Adam',
|
||||
'optimizer.learning_rate': 2.5e-4,
|
||||
|
||||
# Training epochs and batch size
|
||||
'epochs': 150,
|
||||
'train_batch_size': 64,
|
||||
|
||||
# Simple image augmentations
|
||||
'train_dataset': 'cifar10_train_augmented',
|
||||
# Do not augment images for validation
|
||||
'valid_dataset': 'cifar10_valid_no_augment',
|
||||
})
|
||||
# Set model for saving/loading
|
||||
experiment.add_pytorch_models({'model': conf.model})
|
||||
# Start the experiment and run the training loop
|
||||
with experiment.start():
|
||||
conf.run()
|
||||
|
||||
|
||||
#
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,23 @@
|
||||
# [Patches Are All You Need?](https://nn.labml.ai/conv_mixer/index.html)
|
||||
|
||||
This is a [PyTorch](https://pytorch.org) implementation of the paper
|
||||
[Patches Are All You Need?](https://arxiv.org/abs/2201.09792).
|
||||
|
||||
ConvMixer is Similar to [MLP-Mixer](https://nn.labml.ai/transformers/mlp_mixer/index.html).
|
||||
MLP-Mixer separates mixing of spatial and channel dimensions, by applying an MLP across spatial dimension
|
||||
and then an MLP across the channel dimension
|
||||
(spatial MLP replaces the [ViT](https://nn.labml.ai/transformers/vit/index.html) attention
|
||||
and channel MLP is the [FFN](https://nn.labml.ai/transformers/feed_forward.html) of ViT).
|
||||
|
||||
ConvMixer uses a 1x1 convolution for channel mixing and a
|
||||
depth-wise convolution for spatial mixing.
|
||||
Since it's a convolution instead of a full MLP across the space, it mixes only the nearby batches in
|
||||
contrast to ViT or MLP-Mixer.
|
||||
Also, the MLP-mixer uses MLPs of two layers for each mixing and ConvMixer uses a single layer for each mixing.
|
||||
|
||||
The paper recommends removing the residual connection across the channel mixing (point-wise convolution)
|
||||
and having only a residual connection over the spatial mixing (depth-wise convolution).
|
||||
They also use [Batch normalization](https://nn.labml.ai/normalization/batch_norm/index.html) instead
|
||||
of [Layer normalization](../normalization/layer_norm/index.html).
|
||||
|
||||
Here's [an experiment](https://nn.labml.ai/conv_mixer/experiment.html) that trains ConvMixer on CIFAR-10.
|
||||
@@ -0,0 +1,14 @@
|
||||
"""
|
||||
---
|
||||
title: Diffusion models
|
||||
summary: >
|
||||
A set of PyTorch implementations/tutorials of diffusion models.
|
||||
---
|
||||
|
||||
# Diffusion models
|
||||
|
||||
* [Denoising Diffusion Probabilistic Models (DDPM)](ddpm/index.html)
|
||||
* [Stable Diffusion](stable_diffusion/index.html)
|
||||
* [Latent Diffusion Model](stable_diffusion/latent_diffusion.html)
|
||||
* [Denoising Diffusion Implicit Models (DDIM) Sampling](stable_diffusion/sampler/ddim.html)
|
||||
"""
|
||||
@@ -0,0 +1,287 @@
|
||||
"""
|
||||
---
|
||||
title: Denoising Diffusion Probabilistic Models (DDPM)
|
||||
summary: >
|
||||
PyTorch implementation and tutorial of the paper
|
||||
Denoising Diffusion Probabilistic Models (DDPM).
|
||||
---
|
||||
|
||||
# Denoising Diffusion Probabilistic Models (DDPM)
|
||||
|
||||
[](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/diffusion/ddpm/experiment.ipynb)
|
||||
|
||||
This is a [PyTorch](https://pytorch.org) implementation/tutorial of the paper
|
||||
[Denoising Diffusion Probabilistic Models](https://arxiv.org/abs/2006.11239).
|
||||
|
||||
In simple terms, we get an image from data and add noise step by step.
|
||||
Then We train a model to predict that noise at each step and use the model to
|
||||
generate images.
|
||||
|
||||
The following definitions and derivations show how this works.
|
||||
For details please refer to [the paper](https://arxiv.org/abs/2006.11239).
|
||||
|
||||
## Forward Process
|
||||
|
||||
The forward process adds noise to the data $x_0 \sim q(x_0)$, for $T$ timesteps.
|
||||
|
||||
\begin{align}
|
||||
q(x_t | x_{t-1}) = \mathcal{N}\big(x_t; \sqrt{1- \beta_t} x_{t-1}, \beta_t \mathbf{I}\big) \\
|
||||
q(x_{1:T} | x_0) = \prod_{t = 1}^{T} q(x_t | x_{t-1})
|
||||
\end{align}
|
||||
|
||||
where $\beta_1, \dots, \beta_T$ is the variance schedule.
|
||||
|
||||
We can sample $x_t$ at any timestep $t$ with,
|
||||
|
||||
\begin{align}
|
||||
q(x_t|x_0) &= \mathcal{N} \Big(x_t; \sqrt{\bar\alpha_t} x_0, (1-\bar\alpha_t) \mathbf{I} \Big)
|
||||
\end{align}
|
||||
|
||||
where $\alpha_t = 1 - \beta_t$ and $\bar\alpha_t = \prod_{s=1}^t \alpha_s$
|
||||
|
||||
## Reverse Process
|
||||
|
||||
The reverse process removes noise starting at $p(x_T) = \mathcal{N}(x_T; \mathbf{0}, \mathbf{I})$
|
||||
for $T$ time steps.
|
||||
|
||||
\begin{align}
|
||||
\textcolor{lightgreen}{p_\theta}(x_{t-1} | x_t) &= \mathcal{N}\big(x_{t-1};
|
||||
\textcolor{lightgreen}{\mu_\theta}(x_t, t), \textcolor{lightgreen}{\Sigma_\theta}(x_t, t)\big) \\
|
||||
\textcolor{lightgreen}{p_\theta}(x_{0:T}) &= \textcolor{lightgreen}{p_\theta}(x_T) \prod_{t = 1}^{T} \textcolor{lightgreen}{p_\theta}(x_{t-1} | x_t) \\
|
||||
\textcolor{lightgreen}{p_\theta}(x_0) &= \int \textcolor{lightgreen}{p_\theta}(x_{0:T}) dx_{1:T}
|
||||
\end{align}
|
||||
|
||||
$\textcolor{lightgreen}\theta$ are the parameters we train.
|
||||
|
||||
## Loss
|
||||
|
||||
We optimize the ELBO (from Jenson's inequality) on the negative log likelihood.
|
||||
|
||||
\begin{align}
|
||||
\mathbb{E}[-\log \textcolor{lightgreen}{p_\theta}(x_0)]
|
||||
&\le \mathbb{E}_q [ -\log \frac{\textcolor{lightgreen}{p_\theta}(x_{0:T})}{q(x_{1:T}|x_0)} ] \\
|
||||
&=L
|
||||
\end{align}
|
||||
|
||||
The loss can be rewritten as follows.
|
||||
|
||||
\begin{align}
|
||||
L
|
||||
&= \mathbb{E}_q [ -\log \frac{\textcolor{lightgreen}{p_\theta}(x_{0:T})}{q(x_{1:T}|x_0)} ] \\
|
||||
&= \mathbb{E}_q [ -\log p(x_T) - \sum_{t=1}^T \log \frac{\textcolor{lightgreen}{p_\theta}(x_{t-1}|x_t)}{q(x_t|x_{t-1})} ] \\
|
||||
&= \mathbb{E}_q [
|
||||
-\log \frac{p(x_T)}{q(x_T|x_0)}
|
||||
-\sum_{t=2}^T \log \frac{\textcolor{lightgreen}{p_\theta}(x_{t-1}|x_t)}{q(x_{t-1}|x_t,x_0)}
|
||||
-\log \textcolor{lightgreen}{p_\theta}(x_0|x_1)] \\
|
||||
&= \mathbb{E}_q [
|
||||
D_{KL}(q(x_T|x_0) \Vert p(x_T))
|
||||
+\sum_{t=2}^T D_{KL}(q(x_{t-1}|x_t,x_0) \Vert \textcolor{lightgreen}{p_\theta}(x_{t-1}|x_t))
|
||||
-\log \textcolor{lightgreen}{p_\theta}(x_0|x_1)]
|
||||
\end{align}
|
||||
|
||||
$D_{KL}(q(x_T|x_0) \Vert p(x_T))$ is constant since we keep $\beta_1, \dots, \beta_T$ constant.
|
||||
|
||||
### Computing $L_{t-1} = D_{KL}(q(x_{t-1}|x_t,x_0) \Vert \textcolor{lightgreen}{p_\theta}(x_{t-1}|x_t))$
|
||||
|
||||
The forward process posterior conditioned by $x_0$ is,
|
||||
|
||||
\begin{align}
|
||||
q(x_{t-1}|x_t, x_0) &= \mathcal{N} \Big(x_{t-1}; \tilde\mu_t(x_t, x_0), \tilde\beta_t \mathbf{I} \Big) \\
|
||||
\tilde\mu_t(x_t, x_0) &= \frac{\sqrt{\bar\alpha_{t-1}}\beta_t}{1 - \bar\alpha_t}x_0
|
||||
+ \frac{\sqrt{\alpha_t}(1 - \bar\alpha_{t-1})}{1-\bar\alpha_t}x_t \\
|
||||
\tilde\beta_t &= \frac{1 - \bar\alpha_{t-1}}{1 - \bar\alpha_t} \beta_t
|
||||
\end{align}
|
||||
|
||||
The paper sets $\textcolor{lightgreen}{\Sigma_\theta}(x_t, t) = \sigma_t^2 \mathbf{I}$ where $\sigma_t^2$ is set to constants
|
||||
$\beta_t$ or $\tilde\beta_t$.
|
||||
|
||||
Then,
|
||||
$$\textcolor{lightgreen}{p_\theta}(x_{t-1} | x_t) = \mathcal{N}\big(x_{t-1}; \textcolor{lightgreen}{\mu_\theta}(x_t, t), \sigma_t^2 \mathbf{I} \big)$$
|
||||
|
||||
For given noise $\epsilon \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$ using $q(x_t|x_0)$
|
||||
|
||||
\begin{align}
|
||||
x_t(x_0, \epsilon) &= \sqrt{\bar\alpha_t} x_0 + \sqrt{1-\bar\alpha_t}\epsilon \\
|
||||
x_0 &= \frac{1}{\sqrt{\bar\alpha_t}} \Big(x_t(x_0, \epsilon) - \sqrt{1-\bar\alpha_t}\epsilon\Big)
|
||||
\end{align}
|
||||
|
||||
This gives,
|
||||
|
||||
\begin{align}
|
||||
L_{t-1}
|
||||
&= D_{KL}(q(x_{t-1}|x_t,x_0) \Vert \textcolor{lightgreen}{p_\theta}(x_{t-1}|x_t)) \\
|
||||
&= \mathbb{E}_q \Bigg[ \frac{1}{2\sigma_t^2}
|
||||
\Big \Vert \tilde\mu(x_t, x_0) - \textcolor{lightgreen}{\mu_\theta}(x_t, t) \Big \Vert^2 \Bigg] \\
|
||||
&= \mathbb{E}_{x_0, \epsilon} \Bigg[ \frac{1}{2\sigma_t^2}
|
||||
\bigg\Vert \frac{1}{\sqrt{\alpha_t}} \Big(
|
||||
x_t(x_0, \epsilon) - \frac{\beta_t}{\sqrt{1 - \bar\alpha_t}} \epsilon
|
||||
\Big) - \textcolor{lightgreen}{\mu_\theta}(x_t(x_0, \epsilon), t) \bigg\Vert^2 \Bigg] \\
|
||||
\end{align}
|
||||
|
||||
Re-parameterizing with a model to predict noise
|
||||
|
||||
\begin{align}
|
||||
\textcolor{lightgreen}{\mu_\theta}(x_t, t) &= \tilde\mu \bigg(x_t,
|
||||
\frac{1}{\sqrt{\bar\alpha_t}} \Big(x_t -
|
||||
\sqrt{1-\bar\alpha_t}\textcolor{lightgreen}{\epsilon_\theta}(x_t, t) \Big) \bigg) \\
|
||||
&= \frac{1}{\sqrt{\alpha_t}} \Big(x_t -
|
||||
\frac{\beta_t}{\sqrt{1-\bar\alpha_t}}\textcolor{lightgreen}{\epsilon_\theta}(x_t, t) \Big)
|
||||
\end{align}
|
||||
|
||||
where $\epsilon_\theta$ is a learned function that predicts $\epsilon$ given $(x_t, t)$.
|
||||
|
||||
This gives,
|
||||
|
||||
\begin{align}
|
||||
L_{t-1}
|
||||
&= \mathbb{E}_{x_0, \epsilon} \Bigg[ \frac{\beta_t^2}{2\sigma_t^2 \alpha_t (1 - \bar\alpha_t)}
|
||||
\Big\Vert
|
||||
\epsilon - \textcolor{lightgreen}{\epsilon_\theta}(\sqrt{\bar\alpha_t} x_0 + \sqrt{1-\bar\alpha_t}\epsilon, t)
|
||||
\Big\Vert^2 \Bigg]
|
||||
\end{align}
|
||||
|
||||
That is, we are training to predict the noise.
|
||||
|
||||
### Simplified loss
|
||||
|
||||
$$L_{\text{simple}}(\theta) = \mathbb{E}_{t,x_0, \epsilon} \Bigg[ \bigg\Vert
|
||||
\epsilon - \textcolor{lightgreen}{\epsilon_\theta}(\sqrt{\bar\alpha_t} x_0 + \sqrt{1-\bar\alpha_t}\epsilon, t)
|
||||
\bigg\Vert^2 \Bigg]$$
|
||||
|
||||
This minimizes $-\log \textcolor{lightgreen}{p_\theta}(x_0|x_1)$ when $t=1$ and $L_{t-1}$ for $t\gt1$ discarding the
|
||||
weighting in $L_{t-1}$. Discarding the weights $\frac{\beta_t^2}{2\sigma_t^2 \alpha_t (1 - \bar\alpha_t)}$
|
||||
increase the weight given to higher $t$ (which have higher noise levels), therefore increasing the sample quality.
|
||||
|
||||
This file implements the loss calculation and a basic sampling method that we use to generate images during
|
||||
training.
|
||||
|
||||
Here is the [UNet model](unet.html) that gives $\textcolor{lightgreen}{\epsilon_\theta}(x_t, t)$ and
|
||||
[training code](experiment.html).
|
||||
[This file](evaluate.html) can generate samples and interpolations from a trained model.
|
||||
"""
|
||||
from typing import Tuple, Optional
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import torch.utils.data
|
||||
from torch import nn
|
||||
|
||||
from labml_nn.diffusion.ddpm.utils import gather
|
||||
|
||||
|
||||
class DenoiseDiffusion:
|
||||
"""
|
||||
## Denoise Diffusion
|
||||
"""
|
||||
|
||||
def __init__(self, eps_model: nn.Module, n_steps: int, device: torch.device):
|
||||
"""
|
||||
* `eps_model` is $\textcolor{lightgreen}{\epsilon_\theta}(x_t, t)$ model
|
||||
* `n_steps` is $t$
|
||||
* `device` is the device to place constants on
|
||||
"""
|
||||
super().__init__()
|
||||
self.eps_model = eps_model
|
||||
|
||||
# Create $\beta_1, \dots, \beta_T$ linearly increasing variance schedule
|
||||
self.beta = torch.linspace(0.0001, 0.02, n_steps).to(device)
|
||||
|
||||
# $\alpha_t = 1 - \beta_t$
|
||||
self.alpha = 1. - self.beta
|
||||
# $\bar\alpha_t = \prod_{s=1}^t \alpha_s$
|
||||
self.alpha_bar = torch.cumprod(self.alpha, dim=0)
|
||||
# $T$
|
||||
self.n_steps = n_steps
|
||||
# $\sigma^2 = \beta$
|
||||
self.sigma2 = self.beta
|
||||
|
||||
def q_xt_x0(self, x0: torch.Tensor, t: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
#### Get $q(x_t|x_0)$ distribution
|
||||
|
||||
\begin{align}
|
||||
q(x_t|x_0) &= \mathcal{N} \Big(x_t; \sqrt{\bar\alpha_t} x_0, (1-\bar\alpha_t) \mathbf{I} \Big)
|
||||
\end{align}
|
||||
"""
|
||||
|
||||
# [gather](utils.html) $\alpha_t$ and compute $\sqrt{\bar\alpha_t} x_0$
|
||||
mean = gather(self.alpha_bar, t) ** 0.5 * x0
|
||||
# $(1-\bar\alpha_t) \mathbf{I}$
|
||||
var = 1 - gather(self.alpha_bar, t)
|
||||
#
|
||||
return mean, var
|
||||
|
||||
def q_sample(self, x0: torch.Tensor, t: torch.Tensor, eps: Optional[torch.Tensor] = None):
|
||||
"""
|
||||
#### Sample from $q(x_t|x_0)$
|
||||
|
||||
\begin{align}
|
||||
q(x_t|x_0) &= \mathcal{N} \Big(x_t; \sqrt{\bar\alpha_t} x_0, (1-\bar\alpha_t) \mathbf{I} \Big)
|
||||
\end{align}
|
||||
"""
|
||||
|
||||
# $\epsilon \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$
|
||||
if eps is None:
|
||||
eps = torch.randn_like(x0)
|
||||
|
||||
# get $q(x_t|x_0)$
|
||||
mean, var = self.q_xt_x0(x0, t)
|
||||
# Sample from $q(x_t|x_0)$
|
||||
return mean + (var ** 0.5) * eps
|
||||
|
||||
def p_sample(self, xt: torch.Tensor, t: torch.Tensor):
|
||||
"""
|
||||
#### Sample from $\textcolor{lightgreen}{p_\theta}(x_{t-1}|x_t)$
|
||||
|
||||
\begin{align}
|
||||
\textcolor{lightgreen}{p_\theta}(x_{t-1} | x_t) &= \mathcal{N}\big(x_{t-1};
|
||||
\textcolor{lightgreen}{\mu_\theta}(x_t, t), \sigma_t^2 \mathbf{I} \big) \\
|
||||
\textcolor{lightgreen}{\mu_\theta}(x_t, t)
|
||||
&= \frac{1}{\sqrt{\alpha_t}} \Big(x_t -
|
||||
\frac{\beta_t}{\sqrt{1-\bar\alpha_t}}\textcolor{lightgreen}{\epsilon_\theta}(x_t, t) \Big)
|
||||
\end{align}
|
||||
"""
|
||||
|
||||
# $\textcolor{lightgreen}{\epsilon_\theta}(x_t, t)$
|
||||
eps_theta = self.eps_model(xt, t)
|
||||
# [gather](utils.html) $\bar\alpha_t$
|
||||
alpha_bar = gather(self.alpha_bar, t)
|
||||
# $\alpha_t$
|
||||
alpha = gather(self.alpha, t)
|
||||
# $\frac{\beta}{\sqrt{1-\bar\alpha_t}}$
|
||||
eps_coef = (1 - alpha) / (1 - alpha_bar) ** .5
|
||||
# $$\frac{1}{\sqrt{\alpha_t}} \Big(x_t -
|
||||
# \frac{\beta_t}{\sqrt{1-\bar\alpha_t}}\textcolor{lightgreen}{\epsilon_\theta}(x_t, t) \Big)$$
|
||||
mean = 1 / (alpha ** 0.5) * (xt - eps_coef * eps_theta)
|
||||
# $\sigma^2$
|
||||
var = gather(self.sigma2, t)
|
||||
|
||||
# $\epsilon \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$
|
||||
eps = torch.randn(xt.shape, device=xt.device)
|
||||
# Sample
|
||||
return mean + (var ** .5) * eps
|
||||
|
||||
def loss(self, x0: torch.Tensor, noise: Optional[torch.Tensor] = None):
|
||||
"""
|
||||
#### Simplified Loss
|
||||
|
||||
$$L_{\text{simple}}(\theta) = \mathbb{E}_{t,x_0, \epsilon} \Bigg[ \bigg\Vert
|
||||
\epsilon - \textcolor{lightgreen}{\epsilon_\theta}(\sqrt{\bar\alpha_t} x_0 + \sqrt{1-\bar\alpha_t}\epsilon, t)
|
||||
\bigg\Vert^2 \Bigg]$$
|
||||
"""
|
||||
# Get batch size
|
||||
batch_size = x0.shape[0]
|
||||
# Get random $t$ for each sample in the batch
|
||||
t = torch.randint(0, self.n_steps, (batch_size,), device=x0.device, dtype=torch.long)
|
||||
|
||||
# $\epsilon \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$
|
||||
if noise is None:
|
||||
noise = torch.randn_like(x0)
|
||||
|
||||
# Sample $x_t$ for $q(x_t|x_0)$
|
||||
xt = self.q_sample(x0, t, eps=noise)
|
||||
# Get $\textcolor{lightgreen}{\epsilon_\theta}(\sqrt{\bar\alpha_t} x_0 + \sqrt{1-\bar\alpha_t}\epsilon, t)$
|
||||
eps_theta = self.eps_model(xt, t)
|
||||
|
||||
# MSE loss
|
||||
return F.mse_loss(noise, eps_theta)
|
||||
@@ -0,0 +1,328 @@
|
||||
"""
|
||||
---
|
||||
title: Denoising Diffusion Probabilistic Models (DDPM) evaluation/sampling
|
||||
summary: >
|
||||
Code to generate samples from a trained
|
||||
Denoising Diffusion Probabilistic Model.
|
||||
---
|
||||
|
||||
# [Denoising Diffusion Probabilistic Models (DDPM)](index.html) evaluation/sampling
|
||||
|
||||
This is the code to generate images and create interpolations between given images.
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from matplotlib import pyplot as plt
|
||||
from torchvision.transforms.functional import to_pil_image, resize
|
||||
|
||||
from labml import experiment, monit
|
||||
from labml_nn.diffusion.ddpm import DenoiseDiffusion, gather
|
||||
from labml_nn.diffusion.ddpm.experiment import Configs
|
||||
|
||||
|
||||
class Sampler:
|
||||
"""
|
||||
## Sampler class
|
||||
"""
|
||||
|
||||
def __init__(self, diffusion: DenoiseDiffusion, image_channels: int, image_size: int, device: torch.device):
|
||||
"""
|
||||
* `diffusion` is the `DenoiseDiffusion` instance
|
||||
* `image_channels` is the number of channels in the image
|
||||
* `image_size` is the image size
|
||||
* `device` is the device of the model
|
||||
"""
|
||||
self.device = device
|
||||
self.image_size = image_size
|
||||
self.image_channels = image_channels
|
||||
self.diffusion = diffusion
|
||||
|
||||
# $T$
|
||||
self.n_steps = diffusion.n_steps
|
||||
# $\textcolor{lightgreen}{\epsilon_\theta}(x_t, t)$
|
||||
self.eps_model = diffusion.eps_model
|
||||
# $\beta_t$
|
||||
self.beta = diffusion.beta
|
||||
# $\alpha_t$
|
||||
self.alpha = diffusion.alpha
|
||||
# $\bar\alpha_t$
|
||||
self.alpha_bar = diffusion.alpha_bar
|
||||
# $\bar\alpha_{t-1}$
|
||||
alpha_bar_tm1 = torch.cat([self.alpha_bar.new_ones((1,)), self.alpha_bar[:-1]])
|
||||
|
||||
# To calculate
|
||||
#
|
||||
# \begin{align}
|
||||
# q(x_{t-1}|x_t, x_0) &= \mathcal{N} \Big(x_{t-1}; \tilde\mu_t(x_t, x_0), \tilde\beta_t \mathbf{I} \Big) \\
|
||||
# \tilde\mu_t(x_t, x_0) &= \frac{\sqrt{\bar\alpha_{t-1}}\beta_t}{1 - \bar\alpha_t}x_0
|
||||
# + \frac{\sqrt{\alpha_t}(1 - \bar\alpha_{t-1})}{1-\bar\alpha_t}x_t \\
|
||||
# \tilde\beta_t &= \frac{1 - \bar\alpha_{t-1}}{1 - \bar\alpha_t} \beta_t
|
||||
# \end{align}
|
||||
|
||||
# $$\tilde\beta_t = \frac{1 - \bar\alpha_{t-1}}{1 - \bar\alpha_t} \beta_t$$
|
||||
self.beta_tilde = self.beta * (1 - alpha_bar_tm1) / (1 - self.alpha_bar)
|
||||
# $$\frac{\sqrt{\bar\alpha_{t-1}}\beta_t}{1 - \bar\alpha_t}$$
|
||||
self.mu_tilde_coef1 = self.beta * (alpha_bar_tm1 ** 0.5) / (1 - self.alpha_bar)
|
||||
# $$\frac{\sqrt{\alpha_t}(1 - \bar\alpha_{t-1}}{1-\bar\alpha_t}$$
|
||||
self.mu_tilde_coef2 = (self.alpha ** 0.5) * (1 - alpha_bar_tm1) / (1 - self.alpha_bar)
|
||||
# $\sigma^2 = \beta$
|
||||
self.sigma2 = self.beta
|
||||
|
||||
def show_image(self, img, title=""):
|
||||
"""Helper function to display an image"""
|
||||
img = img.clip(0, 1)
|
||||
img = img.cpu().numpy()
|
||||
plt.imshow(img.transpose(1, 2, 0))
|
||||
plt.title(title)
|
||||
plt.show()
|
||||
|
||||
def make_video(self, frames, path="video.mp4"):
|
||||
"""Helper function to create a video"""
|
||||
import imageio
|
||||
# 20 second video
|
||||
writer = imageio.get_writer(path, fps=len(frames) // 20)
|
||||
# Add each image
|
||||
for f in frames:
|
||||
f = f.clip(0, 1)
|
||||
f = to_pil_image(resize(f, [368, 368]))
|
||||
writer.append_data(np.array(f))
|
||||
#
|
||||
writer.close()
|
||||
|
||||
def sample_animation(self, n_frames: int = 1000, create_video: bool = True):
|
||||
"""
|
||||
#### Sample an image step-by-step using $\textcolor{lightgreen}{p_\theta}(x_{t-1}|x_t)$
|
||||
|
||||
We sample an image step-by-step using $\textcolor{lightgreen}{p_\theta}(x_{t-1}|x_t)$ and at each step
|
||||
show the estimate
|
||||
$$x_0 \approx \hat{x}_0 = \frac{1}{\sqrt{\bar\alpha}}
|
||||
\Big( x_t - \sqrt{1 - \bar\alpha_t} \textcolor{lightgreen}{\epsilon_\theta}(x_t, t) \Big)$$
|
||||
"""
|
||||
|
||||
# $x_T \sim p(x_T) = \mathcal{N}(x_T; \mathbf{0}, \mathbf{I})$
|
||||
xt = torch.randn([1, self.image_channels, self.image_size, self.image_size], device=self.device)
|
||||
|
||||
# Interval to log $\hat{x}_0$
|
||||
interval = self.n_steps // n_frames
|
||||
# Frames for video
|
||||
frames = []
|
||||
# Sample $T$ steps
|
||||
for t_inv in monit.iterate('Denoise', self.n_steps):
|
||||
# $t$
|
||||
t_ = self.n_steps - t_inv - 1
|
||||
# $t$ in a tensor
|
||||
t = xt.new_full((1,), t_, dtype=torch.long)
|
||||
# $\textcolor{lightgreen}{\epsilon_\theta}(x_t, t)$
|
||||
eps_theta = self.eps_model(xt, t)
|
||||
if t_ % interval == 0:
|
||||
# Get $\hat{x}_0$ and add to frames
|
||||
x0 = self.p_x0(xt, t, eps_theta)
|
||||
frames.append(x0[0])
|
||||
if not create_video:
|
||||
self.show_image(x0[0], f"{t_}")
|
||||
# Sample from $\textcolor{lightgreen}{p_\theta}(x_{t-1}|x_t)$
|
||||
xt = self.p_sample(xt, t, eps_theta)
|
||||
|
||||
# Make video
|
||||
if create_video:
|
||||
self.make_video(frames)
|
||||
|
||||
def interpolate(self, x1: torch.Tensor, x2: torch.Tensor, lambda_: float, t_: int = 100):
|
||||
"""
|
||||
#### Interpolate two images $x_0$ and $x'_0$
|
||||
|
||||
We get $x_t \sim q(x_t|x_0)$ and $x'_t \sim q(x'_t|x_0)$.
|
||||
|
||||
Then interpolate to
|
||||
$$\bar{x}_t = (1 - \lambda)x_t + \lambda x'_0$$
|
||||
|
||||
Then get
|
||||
$$\bar{x}_0 \sim \textcolor{lightgreen}{p_\theta}(x_0|\bar{x}_t)$$
|
||||
|
||||
* `x1` is $x_0$
|
||||
* `x2` is $x'_0$
|
||||
* `lambda_` is $\lambda$
|
||||
* `t_` is $t$
|
||||
"""
|
||||
|
||||
# Number of samples
|
||||
n_samples = x1.shape[0]
|
||||
# $t$ tensor
|
||||
t = torch.full((n_samples,), t_, device=self.device)
|
||||
# $$\bar{x}_t = (1 - \lambda)x_t + \lambda x'_0$$
|
||||
xt = (1 - lambda_) * self.diffusion.q_sample(x1, t) + lambda_ * self.diffusion.q_sample(x2, t)
|
||||
|
||||
# $$\bar{x}_0 \sim \textcolor{lightgreen}{p_\theta}(x_0|\bar{x}_t)$$
|
||||
return self._sample_x0(xt, t_)
|
||||
|
||||
def interpolate_animate(self, x1: torch.Tensor, x2: torch.Tensor, n_frames: int = 100, t_: int = 100,
|
||||
create_video=True):
|
||||
"""
|
||||
#### Interpolate two images $x_0$ and $x'_0$ and make a video
|
||||
|
||||
* `x1` is $x_0$
|
||||
* `x2` is $x'_0$
|
||||
* `n_frames` is the number of frames for the image
|
||||
* `t_` is $t$
|
||||
* `create_video` specifies whether to make a video or to show each frame
|
||||
"""
|
||||
|
||||
# Show original images
|
||||
self.show_image(x1, "x1")
|
||||
self.show_image(x2, "x2")
|
||||
# Add batch dimension
|
||||
x1 = x1[None, :, :, :]
|
||||
x2 = x2[None, :, :, :]
|
||||
# $t$ tensor
|
||||
t = torch.full((1,), t_, device=self.device)
|
||||
# $x_t \sim q(x_t|x_0)$
|
||||
x1t = self.diffusion.q_sample(x1, t)
|
||||
# $x'_t \sim q(x'_t|x_0)$
|
||||
x2t = self.diffusion.q_sample(x2, t)
|
||||
|
||||
frames = []
|
||||
# Get frames with different $\lambda$
|
||||
for i in monit.iterate('Interpolate', n_frames + 1, is_children_silent=True):
|
||||
# $\lambda$
|
||||
lambda_ = i / n_frames
|
||||
# $$\bar{x}_t = (1 - \lambda)x_t + \lambda x'_0$$
|
||||
xt = (1 - lambda_) * x1t + lambda_ * x2t
|
||||
# $$\bar{x}_0 \sim \textcolor{lightgreen}{p_\theta}(x_0|\bar{x}_t)$$
|
||||
x0 = self._sample_x0(xt, t_)
|
||||
# Add to frames
|
||||
frames.append(x0[0])
|
||||
# Show frame
|
||||
if not create_video:
|
||||
self.show_image(x0[0], f"{lambda_ :.2f}")
|
||||
|
||||
# Make video
|
||||
if create_video:
|
||||
self.make_video(frames)
|
||||
|
||||
def _sample_x0(self, xt: torch.Tensor, n_steps: int):
|
||||
"""
|
||||
#### Sample an image using $\textcolor{lightgreen}{p_\theta}(x_{t-1}|x_t)$
|
||||
|
||||
* `xt` is $x_t$
|
||||
* `n_steps` is $t$
|
||||
"""
|
||||
|
||||
# Number of sampels
|
||||
n_samples = xt.shape[0]
|
||||
# Iterate until $t$ steps
|
||||
for t_ in monit.iterate('Denoise', n_steps):
|
||||
t = n_steps - t_ - 1
|
||||
# Sample from $\textcolor{lightgreen}{p_\theta}(x_{t-1}|x_t)$
|
||||
xt = self.diffusion.p_sample(xt, xt.new_full((n_samples,), t, dtype=torch.long))
|
||||
|
||||
# Return $x_0$
|
||||
return xt
|
||||
|
||||
def sample(self, n_samples: int = 16):
|
||||
"""
|
||||
#### Generate images
|
||||
"""
|
||||
# $x_T \sim p(x_T) = \mathcal{N}(x_T; \mathbf{0}, \mathbf{I})$
|
||||
xt = torch.randn([n_samples, self.image_channels, self.image_size, self.image_size], device=self.device)
|
||||
|
||||
# $$x_0 \sim \textcolor{lightgreen}{p_\theta}(x_0|x_t)$$
|
||||
x0 = self._sample_x0(xt, self.n_steps)
|
||||
|
||||
# Show images
|
||||
for i in range(n_samples):
|
||||
self.show_image(x0[i])
|
||||
|
||||
def p_sample(self, xt: torch.Tensor, t: torch.Tensor, eps_theta: torch.Tensor):
|
||||
"""
|
||||
#### Sample from $\textcolor{lightgreen}{p_\theta}(x_{t-1}|x_t)$
|
||||
|
||||
\begin{align}
|
||||
\textcolor{lightgreen}{p_\theta}(x_{t-1} | x_t) &= \mathcal{N}\big(x_{t-1};
|
||||
\textcolor{lightgreen}{\mu_\theta}(x_t, t), \sigma_t^2 \mathbf{I} \big) \\
|
||||
\textcolor{lightgreen}{\mu_\theta}(x_t, t)
|
||||
&= \frac{1}{\sqrt{\alpha_t}} \Big(x_t -
|
||||
\frac{\beta_t}{\sqrt{1-\bar\alpha_t}}\textcolor{lightgreen}{\epsilon_\theta}(x_t, t) \Big)
|
||||
\end{align}
|
||||
"""
|
||||
# [gather](utils.html) $\bar\alpha_t$
|
||||
alpha_bar = gather(self.alpha_bar, t)
|
||||
# $\alpha_t$
|
||||
alpha = gather(self.alpha, t)
|
||||
# $\frac{\beta}{\sqrt{1-\bar\alpha_t}}$
|
||||
eps_coef = (1 - alpha) / (1 - alpha_bar) ** .5
|
||||
# $$\frac{1}{\sqrt{\alpha_t}} \Big(x_t -
|
||||
# \frac{\beta_t}{\sqrt{1-\bar\alpha_t}}\textcolor{lightgreen}{\epsilon_\theta}(x_t, t) \Big)$$
|
||||
mean = 1 / (alpha ** 0.5) * (xt - eps_coef * eps_theta)
|
||||
# $\sigma^2$
|
||||
var = gather(self.sigma2, t)
|
||||
|
||||
# $\epsilon \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$
|
||||
eps = torch.randn(xt.shape, device=xt.device)
|
||||
# Sample
|
||||
return mean + (var ** .5) * eps
|
||||
|
||||
def p_x0(self, xt: torch.Tensor, t: torch.Tensor, eps: torch.Tensor):
|
||||
"""
|
||||
#### Estimate $x_0$
|
||||
|
||||
$$x_0 \approx \hat{x}_0 = \frac{1}{\sqrt{\bar\alpha}}
|
||||
\Big( x_t - \sqrt{1 - \bar\alpha_t} \textcolor{lightgreen}{\epsilon_\theta}(x_t, t) \Big)$$
|
||||
"""
|
||||
# [gather](utils.html) $\bar\alpha_t$
|
||||
alpha_bar = gather(self.alpha_bar, t)
|
||||
|
||||
# $$x_0 \approx \hat{x}_0 = \frac{1}{\sqrt{\bar\alpha}}
|
||||
# \Big( x_t - \sqrt{1 - \bar\alpha_t} \textcolor{lightgreen}{\epsilon_\theta}(x_t, t) \Big)$$
|
||||
return (xt - (1 - alpha_bar) ** 0.5 * eps) / (alpha_bar ** 0.5)
|
||||
|
||||
|
||||
def main():
|
||||
"""Generate samples"""
|
||||
|
||||
# Training experiment run UUID
|
||||
run_uuid = "a44333ea251411ec8007d1a1762ed686"
|
||||
|
||||
# Start an evaluation
|
||||
experiment.evaluate()
|
||||
|
||||
# Create configs
|
||||
configs = Configs()
|
||||
# Load custom configuration of the training run
|
||||
configs_dict = experiment.load_configs(run_uuid)
|
||||
# Set configurations
|
||||
experiment.configs(configs, configs_dict)
|
||||
|
||||
# Initialize
|
||||
configs.init()
|
||||
|
||||
# Set PyTorch modules for saving and loading
|
||||
experiment.add_pytorch_models({'eps_model': configs.eps_model})
|
||||
|
||||
# Load training experiment
|
||||
experiment.load(run_uuid)
|
||||
|
||||
# Create sampler
|
||||
sampler = Sampler(diffusion=configs.diffusion,
|
||||
image_channels=configs.image_channels,
|
||||
image_size=configs.image_size,
|
||||
device=configs.device)
|
||||
|
||||
# Start evaluation
|
||||
with experiment.start():
|
||||
# No gradients
|
||||
with torch.no_grad():
|
||||
# Sample an image with an denoising animation
|
||||
sampler.sample_animation()
|
||||
|
||||
if False:
|
||||
# Get some images fro data
|
||||
data = next(iter(configs.data_loader)).to(configs.device)
|
||||
|
||||
# Create an interpolation animation
|
||||
sampler.interpolate_animate(data[0], data[1])
|
||||
|
||||
|
||||
#
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,295 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "AYV_dMVDxyc2",
|
||||
"pycharm": {
|
||||
"name": "#%% md\n"
|
||||
}
|
||||
},
|
||||
"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/diffusion/ddpm/experiment.ipynb)\n",
|
||||
"\n",
|
||||
"## [Denoising Diffusion Probabilistic Models (DDPM)](https://nn.labml.ai/diffusion/ddpm/index.html)\n",
|
||||
"\n",
|
||||
"This notebook trains a DDPM based model on MNIST digits dataset."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "AahG_i2y5tY9",
|
||||
"pycharm": {
|
||||
"name": "#%% md\n"
|
||||
}
|
||||
},
|
||||
"source": [
|
||||
"### Install the packages"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "ZCzmCrAIVg0L",
|
||||
"outputId": "cf107fb2-4d50-4c67-af34-367624553421",
|
||||
"pycharm": {
|
||||
"name": "#%%\n"
|
||||
}
|
||||
},
|
||||
"source": [
|
||||
"!pip install labml-nn --quiet"
|
||||
],
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "SE2VUQ6L5zxI",
|
||||
"pycharm": {
|
||||
"name": "#%% md\n"
|
||||
}
|
||||
},
|
||||
"source": [
|
||||
"### Imports"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"collapsed": false,
|
||||
"jupyter": {
|
||||
"outputs_hidden": false
|
||||
},
|
||||
"pycharm": {
|
||||
"name": "#%%\n"
|
||||
}
|
||||
},
|
||||
"source": [
|
||||
"from labml import experiment\n",
|
||||
"from labml_nn.diffusion.ddpm.experiment import Configs"
|
||||
],
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"pycharm": {
|
||||
"name": "#%% md\n"
|
||||
}
|
||||
},
|
||||
"source": [
|
||||
"### Create an experiment"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"collapsed": false,
|
||||
"jupyter": {
|
||||
"outputs_hidden": false
|
||||
},
|
||||
"pycharm": {
|
||||
"name": "#%%\n"
|
||||
}
|
||||
},
|
||||
"source": [
|
||||
"experiment.create(name=\"diffuse\", writers={'screen'})"
|
||||
],
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"pycharm": {
|
||||
"name": "#%% md\n"
|
||||
}
|
||||
},
|
||||
"source": [
|
||||
"### Configurations"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"collapsed": false,
|
||||
"jupyter": {
|
||||
"outputs_hidden": false
|
||||
},
|
||||
"pycharm": {
|
||||
"name": "#%%\n"
|
||||
}
|
||||
},
|
||||
"source": [
|
||||
"configs = Configs()"
|
||||
],
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"pycharm": {
|
||||
"name": "#%% md\n"
|
||||
}
|
||||
},
|
||||
"source": [
|
||||
"Set experiment configurations and assign a configurations dictionary to override configurations"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"collapsed": false,
|
||||
"jupyter": {
|
||||
"outputs_hidden": false
|
||||
},
|
||||
"pycharm": {
|
||||
"name": "#%%\n"
|
||||
}
|
||||
},
|
||||
"source": [
|
||||
"experiment.configs(configs, {\n",
|
||||
" 'dataset': 'MNIST',\n",
|
||||
" 'image_channels': 1,\n",
|
||||
" 'epochs': 5,\n",
|
||||
"})"
|
||||
],
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"pycharm": {
|
||||
"name": "#%% md\n"
|
||||
}
|
||||
},
|
||||
"source": [
|
||||
"Initializ"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"pycharm": {
|
||||
"name": "#%%\n"
|
||||
}
|
||||
},
|
||||
"source": [
|
||||
"configs.init()"
|
||||
],
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "EvI7MtgJ61w5",
|
||||
"pycharm": {
|
||||
"name": "#%% md\n"
|
||||
}
|
||||
},
|
||||
"source": [
|
||||
"Set PyTorch models for loading and saving"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 255
|
||||
},
|
||||
"id": "GDlt7dp-5ALt",
|
||||
"outputId": "e7548e8f-c541-4618-dc5a-1597cae42003",
|
||||
"pycharm": {
|
||||
"name": "#%%\n"
|
||||
}
|
||||
},
|
||||
"source": [
|
||||
"experiment.add_pytorch_models({'eps_model': configs.eps_model})"
|
||||
],
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "KJZRf8527GxL",
|
||||
"pycharm": {
|
||||
"name": "#%% md\n"
|
||||
}
|
||||
},
|
||||
"source": [
|
||||
"### Start the experiment and run the training loop."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 1000
|
||||
},
|
||||
"id": "aIAWo7Fw5DR8",
|
||||
"outputId": "db979785-bfe3-4eda-d3eb-8ccbe61053e5",
|
||||
"pycharm": {
|
||||
"name": "#%%\n"
|
||||
}
|
||||
},
|
||||
"source": [
|
||||
"# Start the experiment\n",
|
||||
"with experiment.start():\n",
|
||||
" configs.run()"
|
||||
],
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"pycharm": {
|
||||
"name": "#%%\n"
|
||||
}
|
||||
},
|
||||
"source": [],
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"accelerator": "GPU",
|
||||
"colab": {
|
||||
"collapsed_sections": [],
|
||||
"name": "Denoising Diffusion Probabilistic Models (DDPM)",
|
||||
"provenance": []
|
||||
},
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3 (ipykernel)",
|
||||
"language": "python",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"codemirror_mode": {
|
||||
"name": "ipython",
|
||||
"version": 3
|
||||
},
|
||||
"file_extension": ".py",
|
||||
"mimetype": "text/x-python",
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython3",
|
||||
"version": "3.8.12"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 4
|
||||
}
|
||||
@@ -0,0 +1,252 @@
|
||||
"""
|
||||
---
|
||||
title: Denoising Diffusion Probabilistic Models (DDPM) training
|
||||
summary: >
|
||||
Training code for
|
||||
Denoising Diffusion Probabilistic Model.
|
||||
---
|
||||
|
||||
# [Denoising Diffusion Probabilistic Models (DDPM)](index.html) training
|
||||
|
||||
[](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/diffusion/ddpm/experiment.ipynb)
|
||||
|
||||
This trains a DDPM based model on CelebA HQ dataset. 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/celebA` folder](#dataset_path).
|
||||
|
||||
The paper had used a exponential moving average of the model with a decay of $0.9999$. We have skipped this for
|
||||
simplicity.
|
||||
"""
|
||||
from typing import List
|
||||
|
||||
import torchvision
|
||||
from PIL import Image
|
||||
|
||||
import torch
|
||||
import torch.utils.data
|
||||
from labml import lab, tracker, experiment, monit
|
||||
from labml.configs import BaseConfigs, option
|
||||
from labml_nn.diffusion.ddpm import DenoiseDiffusion
|
||||
from labml_nn.diffusion.ddpm.unet import UNet
|
||||
from labml_nn.helpers.device import DeviceConfigs
|
||||
|
||||
|
||||
class Configs(BaseConfigs):
|
||||
"""
|
||||
## Configurations
|
||||
"""
|
||||
# Device to train the model on.
|
||||
# [`DeviceConfigs`](../../device.html)
|
||||
# picks up an available CUDA device or defaults to CPU.
|
||||
device: torch.device = DeviceConfigs()
|
||||
|
||||
# U-Net model for $\textcolor{lightgreen}{\epsilon_\theta}(x_t, t)$
|
||||
eps_model: UNet
|
||||
# [DDPM algorithm](index.html)
|
||||
diffusion: DenoiseDiffusion
|
||||
|
||||
# Number of channels in the image. $3$ for RGB.
|
||||
image_channels: int = 3
|
||||
# Image size
|
||||
image_size: int = 32
|
||||
# Number of channels in the initial feature map
|
||||
n_channels: int = 64
|
||||
# The list of channel numbers at each resolution.
|
||||
# The number of channels is `channel_multipliers[i] * n_channels`
|
||||
channel_multipliers: List[int] = [1, 2, 2, 4]
|
||||
# The list of booleans that indicate whether to use attention at each resolution
|
||||
is_attention: List[int] = [False, False, False, True]
|
||||
|
||||
# Number of time steps $T$
|
||||
n_steps: int = 1_000
|
||||
# Batch size
|
||||
batch_size: int = 64
|
||||
# Number of samples to generate
|
||||
n_samples: int = 16
|
||||
# Learning rate
|
||||
learning_rate: float = 2e-5
|
||||
|
||||
# Number of training epochs
|
||||
epochs: int = 1_000
|
||||
|
||||
# Dataset
|
||||
dataset: torch.utils.data.Dataset
|
||||
# Dataloader
|
||||
data_loader: torch.utils.data.DataLoader
|
||||
|
||||
# Adam optimizer
|
||||
optimizer: torch.optim.Adam
|
||||
|
||||
def init(self):
|
||||
# Create $\textcolor{lightgreen}{\epsilon_\theta}(x_t, t)$ model
|
||||
self.eps_model = UNet(
|
||||
image_channels=self.image_channels,
|
||||
n_channels=self.n_channels,
|
||||
ch_mults=self.channel_multipliers,
|
||||
is_attn=self.is_attention,
|
||||
).to(self.device)
|
||||
|
||||
# Create [DDPM class](index.html)
|
||||
self.diffusion = DenoiseDiffusion(
|
||||
eps_model=self.eps_model,
|
||||
n_steps=self.n_steps,
|
||||
device=self.device,
|
||||
)
|
||||
|
||||
# Create dataloader
|
||||
self.data_loader = torch.utils.data.DataLoader(self.dataset, self.batch_size, shuffle=True, pin_memory=True)
|
||||
# Create optimizer
|
||||
self.optimizer = torch.optim.Adam(self.eps_model.parameters(), lr=self.learning_rate)
|
||||
|
||||
# Image logging
|
||||
tracker.set_image("sample", True)
|
||||
|
||||
def sample(self):
|
||||
"""
|
||||
### Sample images
|
||||
"""
|
||||
with torch.no_grad():
|
||||
# $x_T \sim p(x_T) = \mathcal{N}(x_T; \mathbf{0}, \mathbf{I})$
|
||||
x = torch.randn([self.n_samples, self.image_channels, self.image_size, self.image_size],
|
||||
device=self.device)
|
||||
|
||||
# Remove noise for $T$ steps
|
||||
for t_ in monit.iterate('Sample', self.n_steps):
|
||||
# $t$
|
||||
t = self.n_steps - t_ - 1
|
||||
# Sample from $\textcolor{lightgreen}{p_\theta}(x_{t-1}|x_t)$
|
||||
x = self.diffusion.p_sample(x, x.new_full((self.n_samples,), t, dtype=torch.long))
|
||||
|
||||
# Log samples
|
||||
tracker.save('sample', x)
|
||||
|
||||
def train(self):
|
||||
"""
|
||||
### Train
|
||||
"""
|
||||
|
||||
# Iterate through the dataset
|
||||
for data in monit.iterate('Train', self.data_loader):
|
||||
# Increment global step
|
||||
tracker.add_global_step()
|
||||
# Move data to device
|
||||
data = data.to(self.device)
|
||||
|
||||
# Make the gradients zero
|
||||
self.optimizer.zero_grad()
|
||||
# Calculate loss
|
||||
loss = self.diffusion.loss(data)
|
||||
# Compute gradients
|
||||
loss.backward()
|
||||
# Take an optimization step
|
||||
self.optimizer.step()
|
||||
# Track the loss
|
||||
tracker.save('loss', loss)
|
||||
|
||||
def run(self):
|
||||
"""
|
||||
### Training loop
|
||||
"""
|
||||
for _ in monit.loop(self.epochs):
|
||||
# Train the model
|
||||
self.train()
|
||||
# Sample some images
|
||||
self.sample()
|
||||
# New line in the console
|
||||
tracker.new_line()
|
||||
|
||||
|
||||
class CelebADataset(torch.utils.data.Dataset):
|
||||
"""
|
||||
### CelebA HQ dataset
|
||||
"""
|
||||
|
||||
def __init__(self, image_size: int):
|
||||
super().__init__()
|
||||
|
||||
# CelebA images folder
|
||||
folder = lab.get_data_path() / 'celebA'
|
||||
# List of files
|
||||
self._files = [p for p in folder.glob(f'**/*.jpg')]
|
||||
|
||||
# Transformations to resize the image and convert to tensor
|
||||
self._transform = torchvision.transforms.Compose([
|
||||
torchvision.transforms.Resize(image_size),
|
||||
torchvision.transforms.ToTensor(),
|
||||
])
|
||||
|
||||
def __len__(self):
|
||||
"""
|
||||
Size of the dataset
|
||||
"""
|
||||
return len(self._files)
|
||||
|
||||
def __getitem__(self, index: int):
|
||||
"""
|
||||
Get an image
|
||||
"""
|
||||
img = Image.open(self._files[index])
|
||||
return self._transform(img)
|
||||
|
||||
|
||||
@option(Configs.dataset, 'CelebA')
|
||||
def celeb_dataset(c: Configs):
|
||||
"""
|
||||
Create CelebA dataset
|
||||
"""
|
||||
return CelebADataset(c.image_size)
|
||||
|
||||
|
||||
class MNISTDataset(torchvision.datasets.MNIST):
|
||||
"""
|
||||
### MNIST dataset
|
||||
"""
|
||||
|
||||
def __init__(self, image_size):
|
||||
transform = torchvision.transforms.Compose([
|
||||
torchvision.transforms.Resize(image_size),
|
||||
torchvision.transforms.ToTensor(),
|
||||
])
|
||||
|
||||
super().__init__(str(lab.get_data_path()), train=True, download=True, transform=transform)
|
||||
|
||||
def __getitem__(self, item):
|
||||
return super().__getitem__(item)[0]
|
||||
|
||||
|
||||
@option(Configs.dataset, 'MNIST')
|
||||
def mnist_dataset(c: Configs):
|
||||
"""
|
||||
Create MNIST dataset
|
||||
"""
|
||||
return MNISTDataset(c.image_size)
|
||||
|
||||
|
||||
def main():
|
||||
# Create experiment
|
||||
experiment.create(name='diffuse', writers={'screen', 'labml'})
|
||||
|
||||
# Create configurations
|
||||
configs = Configs()
|
||||
|
||||
# Set configurations. You can override the defaults by passing the values in the dictionary.
|
||||
experiment.configs(configs, {
|
||||
'dataset': 'CelebA', # 'MNIST'
|
||||
'image_channels': 3, # 1,
|
||||
'epochs': 100, # 5,
|
||||
})
|
||||
|
||||
# Initialize
|
||||
configs.init()
|
||||
|
||||
# Set models for saving and loading
|
||||
experiment.add_pytorch_models({'eps_model': configs.eps_model})
|
||||
|
||||
# Start and run the training loop
|
||||
with experiment.start():
|
||||
configs.run()
|
||||
|
||||
|
||||
#
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,15 @@
|
||||
# [Denoising Diffusion Probabilistic Models (DDPM)](https://nn.labml.ai/diffusion/ddpm/index.html)
|
||||
|
||||
[](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/diffusion/ddpm/experiment.ipynb)
|
||||
|
||||
This is a [PyTorch](https://pytorch.org) implementation/tutorial of the paper
|
||||
[Denoising Diffusion Probabilistic Models](https://arxiv.org/abs/2006.11239).
|
||||
|
||||
In simple terms, we get an image from data and add noise step by step.
|
||||
Then We train a model to predict that noise at each step and use the model to
|
||||
generate images.
|
||||
|
||||
Here is the [UNet model](https://nn.labml.ai/diffusion/ddpm/unet.html) that predicts the noise and
|
||||
[training code](https://nn.labml.ai/diffusion/ddpm/experiment.html).
|
||||
[This file](https://nn.labml.ai/diffusion/ddpm/evaluate.html) can generate samples and interpolations
|
||||
from a trained model.
|
||||
@@ -0,0 +1,415 @@
|
||||
"""
|
||||
---
|
||||
title: U-Net model for Denoising Diffusion Probabilistic Models (DDPM)
|
||||
summary: >
|
||||
UNet model for Denoising Diffusion Probabilistic Models (DDPM)
|
||||
---
|
||||
|
||||
# U-Net model for [Denoising Diffusion Probabilistic Models (DDPM)](index.html)
|
||||
|
||||
This is a [U-Net](../../unet/index.html) based model to predict noise
|
||||
$\textcolor{lightgreen}{\epsilon_\theta}(x_t, t)$.
|
||||
|
||||
U-Net is a gets it's name from the U shape in the model diagram.
|
||||
It processes a given image by progressively lowering (halving) the feature map resolution and then
|
||||
increasing the resolution.
|
||||
There are pass-through connection at each resolution.
|
||||
|
||||
|
||||
|
||||
This implementation contains a bunch of modifications to original U-Net (residual blocks, multi-head attention)
|
||||
and also adds time-step embeddings $t$.
|
||||
"""
|
||||
|
||||
import math
|
||||
from typing import Optional, Tuple, Union, List
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
|
||||
class Swish(nn.Module):
|
||||
"""
|
||||
### Swish activation function
|
||||
|
||||
$$x \cdot \sigma(x)$$
|
||||
"""
|
||||
|
||||
def forward(self, x):
|
||||
return x * torch.sigmoid(x)
|
||||
|
||||
|
||||
class TimeEmbedding(nn.Module):
|
||||
"""
|
||||
### Embeddings for $t$
|
||||
"""
|
||||
|
||||
def __init__(self, n_channels: int):
|
||||
"""
|
||||
* `n_channels` is the number of dimensions in the embedding
|
||||
"""
|
||||
super().__init__()
|
||||
self.n_channels = n_channels
|
||||
# First linear layer
|
||||
self.lin1 = nn.Linear(self.n_channels // 4, self.n_channels)
|
||||
# Activation
|
||||
self.act = Swish()
|
||||
# Second linear layer
|
||||
self.lin2 = nn.Linear(self.n_channels, self.n_channels)
|
||||
|
||||
def forward(self, t: torch.Tensor):
|
||||
# Create sinusoidal position embeddings
|
||||
# [same as those from the transformer](../../transformers/positional_encoding.html)
|
||||
#
|
||||
# \begin{align}
|
||||
# PE^{(1)}_{t,i} &= sin\Bigg(\frac{t}{10000^{\frac{i}{d - 1}}}\Bigg) \\
|
||||
# PE^{(2)}_{t,i} &= cos\Bigg(\frac{t}{10000^{\frac{i}{d - 1}}}\Bigg)
|
||||
# \end{align}
|
||||
#
|
||||
# where $d$ is `half_dim`
|
||||
half_dim = self.n_channels // 8
|
||||
emb = math.log(10_000) / (half_dim - 1)
|
||||
emb = torch.exp(torch.arange(half_dim, device=t.device) * -emb)
|
||||
emb = t[:, None] * emb[None, :]
|
||||
emb = torch.cat((emb.sin(), emb.cos()), dim=1)
|
||||
|
||||
# Transform with the MLP
|
||||
emb = self.act(self.lin1(emb))
|
||||
emb = self.lin2(emb)
|
||||
|
||||
#
|
||||
return emb
|
||||
|
||||
|
||||
class ResidualBlock(nn.Module):
|
||||
"""
|
||||
### Residual block
|
||||
|
||||
A residual block has two convolution layers with group normalization.
|
||||
Each resolution is processed with two residual blocks.
|
||||
"""
|
||||
|
||||
def __init__(self, in_channels: int, out_channels: int, time_channels: int,
|
||||
n_groups: int = 32, dropout: float = 0.1):
|
||||
"""
|
||||
* `in_channels` is the number of input channels
|
||||
* `out_channels` is the number of input channels
|
||||
* `time_channels` is the number channels in the time step ($t$) embeddings
|
||||
* `n_groups` is the number of groups for [group normalization](../../normalization/group_norm/index.html)
|
||||
* `dropout` is the dropout rate
|
||||
"""
|
||||
super().__init__()
|
||||
# Group normalization and the first convolution layer
|
||||
self.norm1 = nn.GroupNorm(n_groups, in_channels)
|
||||
self.act1 = Swish()
|
||||
self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=(3, 3), padding=(1, 1))
|
||||
|
||||
# Group normalization and the second convolution layer
|
||||
self.norm2 = nn.GroupNorm(n_groups, out_channels)
|
||||
self.act2 = Swish()
|
||||
self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=(3, 3), padding=(1, 1))
|
||||
|
||||
# If the number of input channels is not equal to the number of output channels we have to
|
||||
# project the shortcut connection
|
||||
if in_channels != out_channels:
|
||||
self.shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=(1, 1))
|
||||
else:
|
||||
self.shortcut = nn.Identity()
|
||||
|
||||
# Linear layer for time embeddings
|
||||
self.time_emb = nn.Linear(time_channels, out_channels)
|
||||
self.time_act = Swish()
|
||||
|
||||
self.dropout = nn.Dropout(dropout)
|
||||
|
||||
def forward(self, x: torch.Tensor, t: torch.Tensor):
|
||||
"""
|
||||
* `x` has shape `[batch_size, in_channels, height, width]`
|
||||
* `t` has shape `[batch_size, time_channels]`
|
||||
"""
|
||||
# First convolution layer
|
||||
h = self.conv1(self.act1(self.norm1(x)))
|
||||
# Add time embeddings
|
||||
h += self.time_emb(self.time_act(t))[:, :, None, None]
|
||||
# Second convolution layer
|
||||
h = self.conv2(self.dropout(self.act2(self.norm2(h))))
|
||||
|
||||
# Add the shortcut connection and return
|
||||
return h + self.shortcut(x)
|
||||
|
||||
|
||||
class AttentionBlock(nn.Module):
|
||||
"""
|
||||
### Attention block
|
||||
|
||||
This is similar to [transformer multi-head attention](../../transformers/mha.html).
|
||||
"""
|
||||
|
||||
def __init__(self, n_channels: int, n_heads: int = 1, d_k: int = None, n_groups: int = 32):
|
||||
"""
|
||||
* `n_channels` is the number of channels in the input
|
||||
* `n_heads` is the number of heads in multi-head attention
|
||||
* `d_k` is the number of dimensions in each head
|
||||
* `n_groups` is the number of groups for [group normalization](../../normalization/group_norm/index.html)
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
# Default `d_k`
|
||||
if d_k is None:
|
||||
d_k = n_channels
|
||||
# Normalization layer
|
||||
self.norm = nn.GroupNorm(n_groups, n_channels)
|
||||
# Projections for query, key and values
|
||||
self.projection = nn.Linear(n_channels, n_heads * d_k * 3)
|
||||
# Linear layer for final transformation
|
||||
self.output = nn.Linear(n_heads * d_k, n_channels)
|
||||
# Scale for dot-product attention
|
||||
self.scale = d_k ** -0.5
|
||||
#
|
||||
self.n_heads = n_heads
|
||||
self.d_k = d_k
|
||||
|
||||
def forward(self, x: torch.Tensor, t: Optional[torch.Tensor] = None):
|
||||
"""
|
||||
* `x` has shape `[batch_size, in_channels, height, width]`
|
||||
* `t` has shape `[batch_size, time_channels]`
|
||||
"""
|
||||
# `t` is not used, but it's kept in the arguments because for the attention layer function signature
|
||||
# to match with `ResidualBlock`.
|
||||
_ = t
|
||||
# Get shape
|
||||
batch_size, n_channels, height, width = x.shape
|
||||
# Change `x` to shape `[batch_size, seq, n_channels]`
|
||||
x = x.view(batch_size, n_channels, -1).permute(0, 2, 1)
|
||||
# Get query, key, and values (concatenated) and shape it to `[batch_size, seq, n_heads, 3 * d_k]`
|
||||
qkv = self.projection(x).view(batch_size, -1, self.n_heads, 3 * self.d_k)
|
||||
# Split query, key, and values. Each of them will have shape `[batch_size, seq, n_heads, d_k]`
|
||||
q, k, v = torch.chunk(qkv, 3, dim=-1)
|
||||
# Calculate scaled dot-product $\frac{Q K^\top}{\sqrt{d_k}}$
|
||||
attn = torch.einsum('bihd,bjhd->bijh', q, k) * self.scale
|
||||
# Softmax along the sequence dimension $\underset{seq}{softmax}\Bigg(\frac{Q K^\top}{\sqrt{d_k}}\Bigg)$
|
||||
attn = attn.softmax(dim=2)
|
||||
# Multiply by values
|
||||
res = torch.einsum('bijh,bjhd->bihd', attn, v)
|
||||
# Reshape to `[batch_size, seq, n_heads * d_k]`
|
||||
res = res.view(batch_size, -1, self.n_heads * self.d_k)
|
||||
# Transform to `[batch_size, seq, n_channels]`
|
||||
res = self.output(res)
|
||||
|
||||
# Add skip connection
|
||||
res += x
|
||||
|
||||
# Change to shape `[batch_size, in_channels, height, width]`
|
||||
res = res.permute(0, 2, 1).view(batch_size, n_channels, height, width)
|
||||
|
||||
#
|
||||
return res
|
||||
|
||||
|
||||
class DownBlock(nn.Module):
|
||||
"""
|
||||
### Down block
|
||||
|
||||
This combines `ResidualBlock` and `AttentionBlock`. These are used in the first half of U-Net at each resolution.
|
||||
"""
|
||||
|
||||
def __init__(self, in_channels: int, out_channels: int, time_channels: int, has_attn: bool):
|
||||
super().__init__()
|
||||
self.res = ResidualBlock(in_channels, out_channels, time_channels)
|
||||
if has_attn:
|
||||
self.attn = AttentionBlock(out_channels)
|
||||
else:
|
||||
self.attn = nn.Identity()
|
||||
|
||||
def forward(self, x: torch.Tensor, t: torch.Tensor):
|
||||
x = self.res(x, t)
|
||||
x = self.attn(x)
|
||||
return x
|
||||
|
||||
|
||||
class UpBlock(nn.Module):
|
||||
"""
|
||||
### Up block
|
||||
|
||||
This combines `ResidualBlock` and `AttentionBlock`. These are used in the second half of U-Net at each resolution.
|
||||
"""
|
||||
|
||||
def __init__(self, in_channels: int, out_channels: int, time_channels: int, has_attn: bool):
|
||||
super().__init__()
|
||||
# The input has `in_channels + out_channels` because we concatenate the output of the same resolution
|
||||
# from the first half of the U-Net
|
||||
self.res = ResidualBlock(in_channels + out_channels, out_channels, time_channels)
|
||||
if has_attn:
|
||||
self.attn = AttentionBlock(out_channels)
|
||||
else:
|
||||
self.attn = nn.Identity()
|
||||
|
||||
def forward(self, x: torch.Tensor, t: torch.Tensor):
|
||||
x = self.res(x, t)
|
||||
x = self.attn(x)
|
||||
return x
|
||||
|
||||
|
||||
class MiddleBlock(nn.Module):
|
||||
"""
|
||||
### Middle block
|
||||
|
||||
It combines a `ResidualBlock`, `AttentionBlock`, followed by another `ResidualBlock`.
|
||||
This block is applied at the lowest resolution of the U-Net.
|
||||
"""
|
||||
|
||||
def __init__(self, n_channels: int, time_channels: int):
|
||||
super().__init__()
|
||||
self.res1 = ResidualBlock(n_channels, n_channels, time_channels)
|
||||
self.attn = AttentionBlock(n_channels)
|
||||
self.res2 = ResidualBlock(n_channels, n_channels, time_channels)
|
||||
|
||||
def forward(self, x: torch.Tensor, t: torch.Tensor):
|
||||
x = self.res1(x, t)
|
||||
x = self.attn(x)
|
||||
x = self.res2(x, t)
|
||||
return x
|
||||
|
||||
|
||||
class Upsample(nn.Module):
|
||||
"""
|
||||
### Scale up the feature map by $2 \times$
|
||||
"""
|
||||
|
||||
def __init__(self, n_channels):
|
||||
super().__init__()
|
||||
self.conv = nn.ConvTranspose2d(n_channels, n_channels, (4, 4), (2, 2), (1, 1))
|
||||
|
||||
def forward(self, x: torch.Tensor, t: torch.Tensor):
|
||||
# `t` is not used, but it's kept in the arguments because for the attention layer function signature
|
||||
# to match with `ResidualBlock`.
|
||||
_ = t
|
||||
return self.conv(x)
|
||||
|
||||
|
||||
class Downsample(nn.Module):
|
||||
"""
|
||||
### Scale down the feature map by $\frac{1}{2} \times$
|
||||
"""
|
||||
|
||||
def __init__(self, n_channels):
|
||||
super().__init__()
|
||||
self.conv = nn.Conv2d(n_channels, n_channels, (3, 3), (2, 2), (1, 1))
|
||||
|
||||
def forward(self, x: torch.Tensor, t: torch.Tensor):
|
||||
# `t` is not used, but it's kept in the arguments because for the attention layer function signature
|
||||
# to match with `ResidualBlock`.
|
||||
_ = t
|
||||
return self.conv(x)
|
||||
|
||||
|
||||
class UNet(nn.Module):
|
||||
"""
|
||||
## U-Net
|
||||
"""
|
||||
|
||||
def __init__(self, image_channels: int = 3, n_channels: int = 64,
|
||||
ch_mults: Union[Tuple[int, ...], List[int]] = (1, 2, 2, 4),
|
||||
is_attn: Union[Tuple[bool, ...], List[bool]] = (False, False, True, True),
|
||||
n_blocks: int = 2):
|
||||
"""
|
||||
* `image_channels` is the number of channels in the image. $3$ for RGB.
|
||||
* `n_channels` is number of channels in the initial feature map that we transform the image into
|
||||
* `ch_mults` is the list of channel numbers at each resolution. The number of channels is `ch_mults[i] * n_channels`
|
||||
* `is_attn` is a list of booleans that indicate whether to use attention at each resolution
|
||||
* `n_blocks` is the number of `UpDownBlocks` at each resolution
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
# Number of resolutions
|
||||
n_resolutions = len(ch_mults)
|
||||
|
||||
# Project image into feature map
|
||||
self.image_proj = nn.Conv2d(image_channels, n_channels, kernel_size=(3, 3), padding=(1, 1))
|
||||
|
||||
# Time embedding layer. Time embedding has `n_channels * 4` channels
|
||||
self.time_emb = TimeEmbedding(n_channels * 4)
|
||||
|
||||
# #### First half of U-Net - decreasing resolution
|
||||
down = []
|
||||
# Number of channels
|
||||
out_channels = in_channels = n_channels
|
||||
# For each resolution
|
||||
for i in range(n_resolutions):
|
||||
# Number of output channels at this resolution
|
||||
out_channels = in_channels * ch_mults[i]
|
||||
# Add `n_blocks`
|
||||
for _ in range(n_blocks):
|
||||
down.append(DownBlock(in_channels, out_channels, n_channels * 4, is_attn[i]))
|
||||
in_channels = out_channels
|
||||
# Down sample at all resolutions except the last
|
||||
if i < n_resolutions - 1:
|
||||
down.append(Downsample(in_channels))
|
||||
|
||||
# Combine the set of modules
|
||||
self.down = nn.ModuleList(down)
|
||||
|
||||
# Middle block
|
||||
self.middle = MiddleBlock(out_channels, n_channels * 4, )
|
||||
|
||||
# #### Second half of U-Net - increasing resolution
|
||||
up = []
|
||||
# Number of channels
|
||||
in_channels = out_channels
|
||||
# For each resolution
|
||||
for i in reversed(range(n_resolutions)):
|
||||
# `n_blocks` at the same resolution
|
||||
out_channels = in_channels
|
||||
for _ in range(n_blocks):
|
||||
up.append(UpBlock(in_channels, out_channels, n_channels * 4, is_attn[i]))
|
||||
# Final block to reduce the number of channels
|
||||
out_channels = in_channels // ch_mults[i]
|
||||
up.append(UpBlock(in_channels, out_channels, n_channels * 4, is_attn[i]))
|
||||
in_channels = out_channels
|
||||
# Up sample at all resolutions except last
|
||||
if i > 0:
|
||||
up.append(Upsample(in_channels))
|
||||
|
||||
# Combine the set of modules
|
||||
self.up = nn.ModuleList(up)
|
||||
|
||||
# Final normalization and convolution layer
|
||||
self.norm = nn.GroupNorm(8, n_channels)
|
||||
self.act = Swish()
|
||||
self.final = nn.Conv2d(in_channels, image_channels, kernel_size=(3, 3), padding=(1, 1))
|
||||
|
||||
def forward(self, x: torch.Tensor, t: torch.Tensor):
|
||||
"""
|
||||
* `x` has shape `[batch_size, in_channels, height, width]`
|
||||
* `t` has shape `[batch_size]`
|
||||
"""
|
||||
|
||||
# Get time-step embeddings
|
||||
t = self.time_emb(t)
|
||||
|
||||
# Get image projection
|
||||
x = self.image_proj(x)
|
||||
|
||||
# `h` will store outputs at each resolution for skip connection
|
||||
h = [x]
|
||||
# First half of U-Net
|
||||
for m in self.down:
|
||||
x = m(x, t)
|
||||
h.append(x)
|
||||
|
||||
# Middle (bottom)
|
||||
x = self.middle(x, t)
|
||||
|
||||
# Second half of U-Net
|
||||
for m in self.up:
|
||||
if isinstance(m, Upsample):
|
||||
x = m(x, t)
|
||||
else:
|
||||
# Get the skip connection from first half of U-Net and concatenate
|
||||
s = h.pop()
|
||||
x = torch.cat((x, s), dim=1)
|
||||
#
|
||||
x = m(x, t)
|
||||
|
||||
# Final normalization and convolution
|
||||
return self.final(self.act(self.norm(x)))
|
||||
@@ -0,0 +1,16 @@
|
||||
"""
|
||||
---
|
||||
title: Utility functions for DDPM experiment
|
||||
summary: >
|
||||
Utility functions for DDPM experiment
|
||||
---
|
||||
|
||||
# Utility functions for [DDPM](index.html) experiemnt
|
||||
"""
|
||||
import torch.utils.data
|
||||
|
||||
|
||||
def gather(consts: torch.Tensor, t: torch.Tensor):
|
||||
"""Gather consts for $t$ and reshape to feature map shape"""
|
||||
c = consts.gather(-1, t)
|
||||
return c.reshape(-1, 1, 1, 1)
|
||||
@@ -0,0 +1,54 @@
|
||||
"""
|
||||
---
|
||||
title: Stable Diffusion
|
||||
summary: >
|
||||
Annotated PyTorch implementation/tutorial of stable diffusion.
|
||||
---
|
||||
|
||||
# Stable Diffusion
|
||||
|
||||
This is based on official stable diffusion repository
|
||||
[CompVis/stable-diffusion](https://github.com/CompVis/stable-diffusion).
|
||||
We have kept the model structure same so that open sourced weights could be directly loaded.
|
||||
Our implementation does not contain training code.
|
||||
|
||||
### [PromptArt](https://promptart.labml.ai)
|
||||
|
||||
|
||||
|
||||
We have deployed a stable diffusion based image generation service
|
||||
at [promptart.labml.ai](https://promptart.labml.ai)
|
||||
|
||||
### [Latent Diffusion Model](latent_diffusion.html)
|
||||
|
||||
The core is the [Latent Diffusion Model](latent_diffusion.html).
|
||||
It consists of:
|
||||
|
||||
* [AutoEncoder](model/autoencoder.html)
|
||||
* [U-Net](model/unet.html) with [attention](model/unet_attention.html)
|
||||
|
||||
We have also (optionally) integrated [Flash Attention](https://github.com/HazyResearch/flash-attention)
|
||||
into our [U-Net attention](model/unet_attention.html) which lets you speed up
|
||||
the performance by close to 50% on an RTX A6000 GPU.
|
||||
|
||||
The diffusion is conditioned based on [CLIP embeddings](model/clip_embedder.html).
|
||||
|
||||
### [Sampling Algorithms](sampler/index.html)
|
||||
|
||||
We have implemented the following [sampling algorithms](sampler/index.html):
|
||||
|
||||
* [Denoising Diffusion Probabilistic Models (DDPM) Sampling](sampler/ddpm.html)
|
||||
* [Denoising Diffusion Implicit Models (DDIM) Sampling](sampler/ddim.html)
|
||||
|
||||
### [Example Scripts](scripts/index.html)
|
||||
|
||||
Here are the image generation scripts:
|
||||
|
||||
* [Generate images from text prompts](scripts/text_to_image.html)
|
||||
* [Generate images based on a given image, guided by a prompt](scripts/image_to_image.html)
|
||||
* [Modify parts of a given image based on a text prompt](scripts/in_paint.html)
|
||||
|
||||
#### [Utilities](util.html)
|
||||
|
||||
[`util.py`](util.html) defines the utility functions.
|
||||
"""
|
||||
@@ -0,0 +1,145 @@
|
||||
"""
|
||||
---
|
||||
title: Latent Diffusion Models
|
||||
summary: >
|
||||
Annotated PyTorch implementation/tutorial of latent diffusion models from paper
|
||||
High-Resolution Image Synthesis with Latent Diffusion Models
|
||||
---
|
||||
|
||||
# Latent Diffusion Models
|
||||
|
||||
Latent diffusion models use an auto-encoder to map between image space and
|
||||
latent space. The diffusion model works on the latent space, which makes it
|
||||
a lot easier to train.
|
||||
It is based on paper
|
||||
[High-Resolution Image Synthesis with Latent Diffusion Models](https://arxiv.org/abs/2112.10752).
|
||||
|
||||
They use a pre-trained auto-encoder and train the diffusion U-Net on the latent
|
||||
space of the pre-trained auto-encoder.
|
||||
|
||||
For a simpler diffusion implementation refer to our [DDPM implementation](../ddpm/index.html).
|
||||
We use same notations for $\alpha_t$, $\beta_t$ schedules, etc.
|
||||
"""
|
||||
|
||||
from typing import List
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from labml_nn.diffusion.stable_diffusion.model.autoencoder import Autoencoder
|
||||
from labml_nn.diffusion.stable_diffusion.model.clip_embedder import CLIPTextEmbedder
|
||||
from labml_nn.diffusion.stable_diffusion.model.unet import UNetModel
|
||||
|
||||
|
||||
class DiffusionWrapper(nn.Module):
|
||||
"""
|
||||
*This is an empty wrapper class around the [U-Net](model/unet.html).
|
||||
We keep this to have the same model structure as
|
||||
[CompVis/stable-diffusion](https://github.com/CompVis/stable-diffusion)
|
||||
so that we do not have to map the checkpoint weights explicitly*.
|
||||
"""
|
||||
|
||||
def __init__(self, diffusion_model: UNetModel):
|
||||
super().__init__()
|
||||
self.diffusion_model = diffusion_model
|
||||
|
||||
def forward(self, x: torch.Tensor, time_steps: torch.Tensor, context: torch.Tensor):
|
||||
return self.diffusion_model(x, time_steps, context)
|
||||
|
||||
|
||||
class LatentDiffusion(nn.Module):
|
||||
"""
|
||||
## Latent diffusion model
|
||||
|
||||
This contains following components:
|
||||
|
||||
* [AutoEncoder](model/autoencoder.html)
|
||||
* [U-Net](model/unet.html) with [attention](model/unet_attention.html)
|
||||
* [CLIP embeddings generator](model/clip_embedder.html)
|
||||
"""
|
||||
model: DiffusionWrapper
|
||||
first_stage_model: Autoencoder
|
||||
cond_stage_model: CLIPTextEmbedder
|
||||
|
||||
def __init__(self,
|
||||
unet_model: UNetModel,
|
||||
autoencoder: Autoencoder,
|
||||
clip_embedder: CLIPTextEmbedder,
|
||||
latent_scaling_factor: float,
|
||||
n_steps: int,
|
||||
linear_start: float,
|
||||
linear_end: float,
|
||||
):
|
||||
"""
|
||||
:param unet_model: is the [U-Net](model/unet.html) that predicts noise
|
||||
$\epsilon_\text{cond}(x_t, c)$, in latent space
|
||||
:param autoencoder: is the [AutoEncoder](model/autoencoder.html)
|
||||
:param clip_embedder: is the [CLIP embeddings generator](model/clip_embedder.html)
|
||||
:param latent_scaling_factor: is the scaling factor for the latent space. The encodings of
|
||||
the autoencoder are scaled by this before feeding into the U-Net.
|
||||
:param n_steps: is the number of diffusion steps $T$.
|
||||
:param linear_start: is the start of the $\beta$ schedule.
|
||||
:param linear_end: is the end of the $\beta$ schedule.
|
||||
"""
|
||||
super().__init__()
|
||||
# Wrap the [U-Net](model/unet.html) to keep the same model structure as
|
||||
# [CompVis/stable-diffusion](https://github.com/CompVis/stable-diffusion).
|
||||
self.model = DiffusionWrapper(unet_model)
|
||||
# Auto-encoder and scaling factor
|
||||
self.first_stage_model = autoencoder
|
||||
self.latent_scaling_factor = latent_scaling_factor
|
||||
# [CLIP embeddings generator](model/clip_embedder.html)
|
||||
self.cond_stage_model = clip_embedder
|
||||
|
||||
# Number of steps $T$
|
||||
self.n_steps = n_steps
|
||||
|
||||
# $\beta$ schedule
|
||||
beta = torch.linspace(linear_start ** 0.5, linear_end ** 0.5, n_steps, dtype=torch.float64) ** 2
|
||||
self.beta = nn.Parameter(beta.to(torch.float32), requires_grad=False)
|
||||
# $\alpha_t = 1 - \beta_t$
|
||||
alpha = 1. - beta
|
||||
# $\bar\alpha_t = \prod_{s=1}^t \alpha_s$
|
||||
alpha_bar = torch.cumprod(alpha, dim=0)
|
||||
self.alpha_bar = nn.Parameter(alpha_bar.to(torch.float32), requires_grad=False)
|
||||
|
||||
@property
|
||||
def device(self):
|
||||
"""
|
||||
### Get model device
|
||||
"""
|
||||
return next(iter(self.model.parameters())).device
|
||||
|
||||
def get_text_conditioning(self, prompts: List[str]):
|
||||
"""
|
||||
### Get [CLIP embeddings](model/clip_embedder.html) for a list of text prompts
|
||||
"""
|
||||
return self.cond_stage_model(prompts)
|
||||
|
||||
def autoencoder_encode(self, image: torch.Tensor):
|
||||
"""
|
||||
### Get scaled latent space representation of the image
|
||||
|
||||
The encoder output is a distribution.
|
||||
We sample from that and multiply by the scaling factor.
|
||||
"""
|
||||
return self.latent_scaling_factor * self.first_stage_model.encode(image).sample()
|
||||
|
||||
def autoencoder_decode(self, z: torch.Tensor):
|
||||
"""
|
||||
### Get image from the latent representation
|
||||
|
||||
We scale down by the scaling factor and then decode.
|
||||
"""
|
||||
return self.first_stage_model.decode(z / self.latent_scaling_factor)
|
||||
|
||||
def forward(self, x: torch.Tensor, t: torch.Tensor, context: torch.Tensor):
|
||||
"""
|
||||
### Predict noise
|
||||
|
||||
Predict noise given the latent representation $x_t$, time step $t$, and the
|
||||
conditioning context $c$.
|
||||
|
||||
$$\epsilon_\text{cond}(x_t, c)$$
|
||||
"""
|
||||
return self.model(x, t, context)
|
||||
@@ -0,0 +1,13 @@
|
||||
"""
|
||||
---
|
||||
title: Modules used in stable diffusion
|
||||
summary: >
|
||||
Models and components for stable diffusion.
|
||||
---
|
||||
|
||||
# [Stable Diffusion](../index.html) Models
|
||||
|
||||
* [AutoEncoder](autoencoder.html)
|
||||
* [U-Net](unet.html) with [attention](unet_attention.html)
|
||||
* [CLIP embedder](clip_embedder.html).
|
||||
"""
|
||||
@@ -0,0 +1,433 @@
|
||||
"""
|
||||
---
|
||||
title: Autoencoder for Stable Diffusion
|
||||
summary: >
|
||||
Annotated PyTorch implementation/tutorial of the autoencoder
|
||||
for stable diffusion.
|
||||
---
|
||||
|
||||
# Autoencoder for [Stable Diffusion](../index.html)
|
||||
|
||||
This implements the auto-encoder model used to map between image space and latent space.
|
||||
|
||||
We have kept to the model definition and naming unchanged from
|
||||
[CompVis/stable-diffusion](https://github.com/CompVis/stable-diffusion)
|
||||
so that we can load the checkpoints directly.
|
||||
"""
|
||||
|
||||
from typing import List
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import nn
|
||||
|
||||
|
||||
class Autoencoder(nn.Module):
|
||||
"""
|
||||
## Autoencoder
|
||||
|
||||
This consists of the encoder and decoder modules.
|
||||
"""
|
||||
|
||||
def __init__(self, encoder: 'Encoder', decoder: 'Decoder', emb_channels: int, z_channels: int):
|
||||
"""
|
||||
:param encoder: is the encoder
|
||||
:param decoder: is the decoder
|
||||
:param emb_channels: is the number of dimensions in the quantized embedding space
|
||||
:param z_channels: is the number of channels in the embedding space
|
||||
"""
|
||||
super().__init__()
|
||||
self.encoder = encoder
|
||||
self.decoder = decoder
|
||||
# Convolution to map from embedding space to
|
||||
# quantized embedding space moments (mean and log variance)
|
||||
self.quant_conv = nn.Conv2d(2 * z_channels, 2 * emb_channels, 1)
|
||||
# Convolution to map from quantized embedding space back to
|
||||
# embedding space
|
||||
self.post_quant_conv = nn.Conv2d(emb_channels, z_channels, 1)
|
||||
|
||||
def encode(self, img: torch.Tensor) -> 'GaussianDistribution':
|
||||
"""
|
||||
### Encode images to latent representation
|
||||
|
||||
:param img: is the image tensor with shape `[batch_size, img_channels, img_height, img_width]`
|
||||
"""
|
||||
# Get embeddings with shape `[batch_size, z_channels * 2, z_height, z_height]`
|
||||
z = self.encoder(img)
|
||||
# Get the moments in the quantized embedding space
|
||||
moments = self.quant_conv(z)
|
||||
# Return the distribution
|
||||
return GaussianDistribution(moments)
|
||||
|
||||
def decode(self, z: torch.Tensor):
|
||||
"""
|
||||
### Decode images from latent representation
|
||||
|
||||
:param z: is the latent representation with shape `[batch_size, emb_channels, z_height, z_height]`
|
||||
"""
|
||||
# Map to embedding space from the quantized representation
|
||||
z = self.post_quant_conv(z)
|
||||
# Decode the image of shape `[batch_size, channels, height, width]`
|
||||
return self.decoder(z)
|
||||
|
||||
|
||||
class Encoder(nn.Module):
|
||||
"""
|
||||
## Encoder module
|
||||
"""
|
||||
|
||||
def __init__(self, *, channels: int, channel_multipliers: List[int], n_resnet_blocks: int,
|
||||
in_channels: int, z_channels: int):
|
||||
"""
|
||||
:param channels: is the number of channels in the first convolution layer
|
||||
:param channel_multipliers: are the multiplicative factors for the number of channels in the
|
||||
subsequent blocks
|
||||
:param n_resnet_blocks: is the number of resnet layers at each resolution
|
||||
:param in_channels: is the number of channels in the image
|
||||
:param z_channels: is the number of channels in the embedding space
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
# Number of blocks of different resolutions.
|
||||
# The resolution is halved at the end each top level block
|
||||
n_resolutions = len(channel_multipliers)
|
||||
|
||||
# Initial $3 \times 3$ convolution layer that maps the image to `channels`
|
||||
self.conv_in = nn.Conv2d(in_channels, channels, 3, stride=1, padding=1)
|
||||
|
||||
# Number of channels in each top level block
|
||||
channels_list = [m * channels for m in [1] + channel_multipliers]
|
||||
|
||||
# List of top-level blocks
|
||||
self.down = nn.ModuleList()
|
||||
# Create top-level blocks
|
||||
for i in range(n_resolutions):
|
||||
# Each top level block consists of multiple ResNet Blocks and down-sampling
|
||||
resnet_blocks = nn.ModuleList()
|
||||
# Add ResNet Blocks
|
||||
for _ in range(n_resnet_blocks):
|
||||
resnet_blocks.append(ResnetBlock(channels, channels_list[i + 1]))
|
||||
channels = channels_list[i + 1]
|
||||
# Top-level block
|
||||
down = nn.Module()
|
||||
down.block = resnet_blocks
|
||||
# Down-sampling at the end of each top level block except the last
|
||||
if i != n_resolutions - 1:
|
||||
down.downsample = DownSample(channels)
|
||||
else:
|
||||
down.downsample = nn.Identity()
|
||||
#
|
||||
self.down.append(down)
|
||||
|
||||
# Final ResNet blocks with attention
|
||||
self.mid = nn.Module()
|
||||
self.mid.block_1 = ResnetBlock(channels, channels)
|
||||
self.mid.attn_1 = AttnBlock(channels)
|
||||
self.mid.block_2 = ResnetBlock(channels, channels)
|
||||
|
||||
# Map to embedding space with a $3 \times 3$ convolution
|
||||
self.norm_out = normalization(channels)
|
||||
self.conv_out = nn.Conv2d(channels, 2 * z_channels, 3, stride=1, padding=1)
|
||||
|
||||
def forward(self, img: torch.Tensor):
|
||||
"""
|
||||
:param img: is the image tensor with shape `[batch_size, img_channels, img_height, img_width]`
|
||||
"""
|
||||
|
||||
# Map to `channels` with the initial convolution
|
||||
x = self.conv_in(img)
|
||||
|
||||
# Top-level blocks
|
||||
for down in self.down:
|
||||
# ResNet Blocks
|
||||
for block in down.block:
|
||||
x = block(x)
|
||||
# Down-sampling
|
||||
x = down.downsample(x)
|
||||
|
||||
# Final ResNet blocks with attention
|
||||
x = self.mid.block_1(x)
|
||||
x = self.mid.attn_1(x)
|
||||
x = self.mid.block_2(x)
|
||||
|
||||
# Normalize and map to embedding space
|
||||
x = self.norm_out(x)
|
||||
x = swish(x)
|
||||
x = self.conv_out(x)
|
||||
|
||||
#
|
||||
return x
|
||||
|
||||
|
||||
class Decoder(nn.Module):
|
||||
"""
|
||||
## Decoder module
|
||||
"""
|
||||
|
||||
def __init__(self, *, channels: int, channel_multipliers: List[int], n_resnet_blocks: int,
|
||||
out_channels: int, z_channels: int):
|
||||
"""
|
||||
:param channels: is the number of channels in the final convolution layer
|
||||
:param channel_multipliers: are the multiplicative factors for the number of channels in the
|
||||
previous blocks, in reverse order
|
||||
:param n_resnet_blocks: is the number of resnet layers at each resolution
|
||||
:param out_channels: is the number of channels in the image
|
||||
:param z_channels: is the number of channels in the embedding space
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
# Number of blocks of different resolutions.
|
||||
# The resolution is halved at the end each top level block
|
||||
num_resolutions = len(channel_multipliers)
|
||||
|
||||
# Number of channels in each top level block, in the reverse order
|
||||
channels_list = [m * channels for m in channel_multipliers]
|
||||
|
||||
# Number of channels in the top-level block
|
||||
channels = channels_list[-1]
|
||||
|
||||
# Initial $3 \times 3$ convolution layer that maps the embedding space to `channels`
|
||||
self.conv_in = nn.Conv2d(z_channels, channels, 3, stride=1, padding=1)
|
||||
|
||||
# ResNet blocks with attention
|
||||
self.mid = nn.Module()
|
||||
self.mid.block_1 = ResnetBlock(channels, channels)
|
||||
self.mid.attn_1 = AttnBlock(channels)
|
||||
self.mid.block_2 = ResnetBlock(channels, channels)
|
||||
|
||||
# List of top-level blocks
|
||||
self.up = nn.ModuleList()
|
||||
# Create top-level blocks
|
||||
for i in reversed(range(num_resolutions)):
|
||||
# Each top level block consists of multiple ResNet Blocks and up-sampling
|
||||
resnet_blocks = nn.ModuleList()
|
||||
# Add ResNet Blocks
|
||||
for _ in range(n_resnet_blocks + 1):
|
||||
resnet_blocks.append(ResnetBlock(channels, channels_list[i]))
|
||||
channels = channels_list[i]
|
||||
# Top-level block
|
||||
up = nn.Module()
|
||||
up.block = resnet_blocks
|
||||
# Up-sampling at the end of each top level block except the first
|
||||
if i != 0:
|
||||
up.upsample = UpSample(channels)
|
||||
else:
|
||||
up.upsample = nn.Identity()
|
||||
# Prepend to be consistent with the checkpoint
|
||||
self.up.insert(0, up)
|
||||
|
||||
# Map to image space with a $3 \times 3$ convolution
|
||||
self.norm_out = normalization(channels)
|
||||
self.conv_out = nn.Conv2d(channels, out_channels, 3, stride=1, padding=1)
|
||||
|
||||
def forward(self, z: torch.Tensor):
|
||||
"""
|
||||
:param z: is the embedding tensor with shape `[batch_size, z_channels, z_height, z_height]`
|
||||
"""
|
||||
|
||||
# Map to `channels` with the initial convolution
|
||||
h = self.conv_in(z)
|
||||
|
||||
# ResNet blocks with attention
|
||||
h = self.mid.block_1(h)
|
||||
h = self.mid.attn_1(h)
|
||||
h = self.mid.block_2(h)
|
||||
|
||||
# Top-level blocks
|
||||
for up in reversed(self.up):
|
||||
# ResNet Blocks
|
||||
for block in up.block:
|
||||
h = block(h)
|
||||
# Up-sampling
|
||||
h = up.upsample(h)
|
||||
|
||||
# Normalize and map to image space
|
||||
h = self.norm_out(h)
|
||||
h = swish(h)
|
||||
img = self.conv_out(h)
|
||||
|
||||
#
|
||||
return img
|
||||
|
||||
|
||||
class GaussianDistribution:
|
||||
"""
|
||||
## Gaussian Distribution
|
||||
"""
|
||||
|
||||
def __init__(self, parameters: torch.Tensor):
|
||||
"""
|
||||
:param parameters: are the means and log of variances of the embedding of shape
|
||||
`[batch_size, z_channels * 2, z_height, z_height]`
|
||||
"""
|
||||
# Split mean and log of variance
|
||||
self.mean, log_var = torch.chunk(parameters, 2, dim=1)
|
||||
# Clamp the log of variances
|
||||
self.log_var = torch.clamp(log_var, -30.0, 20.0)
|
||||
# Calculate standard deviation
|
||||
self.std = torch.exp(0.5 * self.log_var)
|
||||
|
||||
def sample(self):
|
||||
# Sample from the distribution
|
||||
return self.mean + self.std * torch.randn_like(self.std)
|
||||
|
||||
|
||||
class AttnBlock(nn.Module):
|
||||
"""
|
||||
## Attention block
|
||||
"""
|
||||
|
||||
def __init__(self, channels: int):
|
||||
"""
|
||||
:param channels: is the number of channels
|
||||
"""
|
||||
super().__init__()
|
||||
# Group normalization
|
||||
self.norm = normalization(channels)
|
||||
# Query, key and value mappings
|
||||
self.q = nn.Conv2d(channels, channels, 1)
|
||||
self.k = nn.Conv2d(channels, channels, 1)
|
||||
self.v = nn.Conv2d(channels, channels, 1)
|
||||
# Final $1 \times 1$ convolution layer
|
||||
self.proj_out = nn.Conv2d(channels, channels, 1)
|
||||
# Attention scaling factor
|
||||
self.scale = channels ** -0.5
|
||||
|
||||
def forward(self, x: torch.Tensor):
|
||||
"""
|
||||
:param x: is the tensor of shape `[batch_size, channels, height, width]`
|
||||
"""
|
||||
# Normalize `x`
|
||||
x_norm = self.norm(x)
|
||||
# Get query, key and vector embeddings
|
||||
q = self.q(x_norm)
|
||||
k = self.k(x_norm)
|
||||
v = self.v(x_norm)
|
||||
|
||||
# Reshape to query, key and vector embeedings from
|
||||
# `[batch_size, channels, height, width]` to
|
||||
# `[batch_size, channels, height * width]`
|
||||
b, c, h, w = q.shape
|
||||
q = q.view(b, c, h * w)
|
||||
k = k.view(b, c, h * w)
|
||||
v = v.view(b, c, h * w)
|
||||
|
||||
# Compute $\underset{seq}{softmax}\Bigg(\frac{Q K^\top}{\sqrt{d_{key}}}\Bigg)$
|
||||
attn = torch.einsum('bci,bcj->bij', q, k) * self.scale
|
||||
attn = F.softmax(attn, dim=2)
|
||||
|
||||
# Compute $\underset{seq}{softmax}\Bigg(\frac{Q K^\top}{\sqrt{d_{key}}}\Bigg)V$
|
||||
out = torch.einsum('bij,bcj->bci', attn, v)
|
||||
|
||||
# Reshape back to `[batch_size, channels, height, width]`
|
||||
out = out.view(b, c, h, w)
|
||||
# Final $1 \times 1$ convolution layer
|
||||
out = self.proj_out(out)
|
||||
|
||||
# Add residual connection
|
||||
return x + out
|
||||
|
||||
|
||||
class UpSample(nn.Module):
|
||||
"""
|
||||
## Up-sampling layer
|
||||
"""
|
||||
def __init__(self, channels: int):
|
||||
"""
|
||||
:param channels: is the number of channels
|
||||
"""
|
||||
super().__init__()
|
||||
# $3 \times 3$ convolution mapping
|
||||
self.conv = nn.Conv2d(channels, channels, 3, padding=1)
|
||||
|
||||
def forward(self, x: torch.Tensor):
|
||||
"""
|
||||
:param x: is the input feature map with shape `[batch_size, channels, height, width]`
|
||||
"""
|
||||
# Up-sample by a factor of $2$
|
||||
x = F.interpolate(x, scale_factor=2.0, mode="nearest")
|
||||
# Apply convolution
|
||||
return self.conv(x)
|
||||
|
||||
|
||||
class DownSample(nn.Module):
|
||||
"""
|
||||
## Down-sampling layer
|
||||
"""
|
||||
def __init__(self, channels: int):
|
||||
"""
|
||||
:param channels: is the number of channels
|
||||
"""
|
||||
super().__init__()
|
||||
# $3 \times 3$ convolution with stride length of $2$ to down-sample by a factor of $2$
|
||||
self.conv = nn.Conv2d(channels, channels, 3, stride=2, padding=0)
|
||||
|
||||
def forward(self, x: torch.Tensor):
|
||||
"""
|
||||
:param x: is the input feature map with shape `[batch_size, channels, height, width]`
|
||||
"""
|
||||
# Add padding
|
||||
x = F.pad(x, (0, 1, 0, 1), mode="constant", value=0)
|
||||
# Apply convolution
|
||||
return self.conv(x)
|
||||
|
||||
|
||||
class ResnetBlock(nn.Module):
|
||||
"""
|
||||
## ResNet Block
|
||||
"""
|
||||
def __init__(self, in_channels: int, out_channels: int):
|
||||
"""
|
||||
:param in_channels: is the number of channels in the input
|
||||
:param out_channels: is the number of channels in the output
|
||||
"""
|
||||
super().__init__()
|
||||
# First normalization and convolution layer
|
||||
self.norm1 = normalization(in_channels)
|
||||
self.conv1 = nn.Conv2d(in_channels, out_channels, 3, stride=1, padding=1)
|
||||
# Second normalization and convolution layer
|
||||
self.norm2 = normalization(out_channels)
|
||||
self.conv2 = nn.Conv2d(out_channels, out_channels, 3, stride=1, padding=1)
|
||||
# `in_channels` to `out_channels` mapping layer for residual connection
|
||||
if in_channels != out_channels:
|
||||
self.nin_shortcut = nn.Conv2d(in_channels, out_channels, 1, stride=1, padding=0)
|
||||
else:
|
||||
self.nin_shortcut = nn.Identity()
|
||||
|
||||
def forward(self, x: torch.Tensor):
|
||||
"""
|
||||
:param x: is the input feature map with shape `[batch_size, channels, height, width]`
|
||||
"""
|
||||
|
||||
h = x
|
||||
|
||||
# First normalization and convolution layer
|
||||
h = self.norm1(h)
|
||||
h = swish(h)
|
||||
h = self.conv1(h)
|
||||
|
||||
# Second normalization and convolution layer
|
||||
h = self.norm2(h)
|
||||
h = swish(h)
|
||||
h = self.conv2(h)
|
||||
|
||||
# Map and add residual
|
||||
return self.nin_shortcut(x) + h
|
||||
|
||||
|
||||
def swish(x: torch.Tensor):
|
||||
"""
|
||||
### Swish activation
|
||||
|
||||
$$x \cdot \sigma(x)$$
|
||||
"""
|
||||
return x * torch.sigmoid(x)
|
||||
|
||||
|
||||
def normalization(channels: int):
|
||||
"""
|
||||
### Group normalization
|
||||
|
||||
This is a helper function, with fixed number of groups and `eps`.
|
||||
"""
|
||||
return nn.GroupNorm(num_groups=32, num_channels=channels, eps=1e-6)
|
||||
@@ -0,0 +1,50 @@
|
||||
"""
|
||||
---
|
||||
title: CLIP Text Embedder
|
||||
summary: >
|
||||
CLIP embedder to get prompt embeddings for stable diffusion
|
||||
---
|
||||
|
||||
# CLIP Text Embedder
|
||||
|
||||
This is used to get prompt embeddings for [stable diffusion](../index.html).
|
||||
It uses HuggingFace Transformers CLIP model.
|
||||
"""
|
||||
|
||||
from typing import List
|
||||
|
||||
from torch import nn
|
||||
from transformers import CLIPTokenizer, CLIPTextModel
|
||||
|
||||
|
||||
class CLIPTextEmbedder(nn.Module):
|
||||
"""
|
||||
## CLIP Text Embedder
|
||||
"""
|
||||
|
||||
def __init__(self, version: str = "openai/clip-vit-large-patch14", device="cuda:0", max_length: int = 77):
|
||||
"""
|
||||
:param version: is the model version
|
||||
:param device: is the device
|
||||
:param max_length: is the max length of the tokenized prompt
|
||||
"""
|
||||
super().__init__()
|
||||
# Load the tokenizer
|
||||
self.tokenizer = CLIPTokenizer.from_pretrained(version)
|
||||
# Load the CLIP transformer
|
||||
self.transformer = CLIPTextModel.from_pretrained(version).eval()
|
||||
|
||||
self.device = device
|
||||
self.max_length = max_length
|
||||
|
||||
def forward(self, prompts: List[str]):
|
||||
"""
|
||||
:param prompts: are the list of prompts to embed
|
||||
"""
|
||||
# Tokenize the prompts
|
||||
batch_encoding = self.tokenizer(prompts, truncation=True, max_length=self.max_length, return_length=True,
|
||||
return_overflowing_tokens=False, padding="max_length", return_tensors="pt")
|
||||
# Get token ids
|
||||
tokens = batch_encoding["input_ids"].to(self.device)
|
||||
# Get CLIP embeddings
|
||||
return self.transformer(input_ids=tokens).last_hidden_state
|
||||
@@ -0,0 +1,344 @@
|
||||
"""
|
||||
---
|
||||
title: U-Net for Stable Diffusion
|
||||
summary: >
|
||||
Annotated PyTorch implementation/tutorial of the U-Net in stable diffusion.
|
||||
---
|
||||
|
||||
# U-Net for [Stable Diffusion](../index.html)
|
||||
|
||||
This implements the U-Net that
|
||||
gives $\epsilon_\text{cond}(x_t, c)$
|
||||
|
||||
We have kept to the model definition and naming unchanged from
|
||||
[CompVis/stable-diffusion](https://github.com/CompVis/stable-diffusion)
|
||||
so that we can load the checkpoints directly.
|
||||
"""
|
||||
|
||||
import math
|
||||
from typing import List
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from labml_nn.diffusion.stable_diffusion.model.unet_attention import SpatialTransformer
|
||||
|
||||
|
||||
class UNetModel(nn.Module):
|
||||
"""
|
||||
## U-Net model
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self, *,
|
||||
in_channels: int,
|
||||
out_channels: int,
|
||||
channels: int,
|
||||
n_res_blocks: int,
|
||||
attention_levels: List[int],
|
||||
channel_multipliers: List[int],
|
||||
n_heads: int,
|
||||
tf_layers: int = 1,
|
||||
d_cond: int = 768):
|
||||
"""
|
||||
:param in_channels: is the number of channels in the input feature map
|
||||
:param out_channels: is the number of channels in the output feature map
|
||||
:param channels: is the base channel count for the model
|
||||
:param n_res_blocks: number of residual blocks at each level
|
||||
:param attention_levels: are the levels at which attention should be performed
|
||||
:param channel_multipliers: are the multiplicative factors for number of channels for each level
|
||||
:param n_heads: is the number of attention heads in the transformers
|
||||
:param tf_layers: is the number of transformer layers in the transformers
|
||||
:param d_cond: is the size of the conditional embedding in the transformers
|
||||
"""
|
||||
super().__init__()
|
||||
self.channels = channels
|
||||
|
||||
# Number of levels
|
||||
levels = len(channel_multipliers)
|
||||
# Size time embeddings
|
||||
d_time_emb = channels * 4
|
||||
self.time_embed = nn.Sequential(
|
||||
nn.Linear(channels, d_time_emb),
|
||||
nn.SiLU(),
|
||||
nn.Linear(d_time_emb, d_time_emb),
|
||||
)
|
||||
|
||||
# Input half of the U-Net
|
||||
self.input_blocks = nn.ModuleList()
|
||||
# Initial $3 \times 3$ convolution that maps the input to `channels`.
|
||||
# The blocks are wrapped in `TimestepEmbedSequential` module because
|
||||
# different modules have different forward function signatures;
|
||||
# for example, convolution only accepts the feature map and
|
||||
# residual blocks accept the feature map and time embedding.
|
||||
# `TimestepEmbedSequential` calls them accordingly.
|
||||
self.input_blocks.append(TimestepEmbedSequential(
|
||||
nn.Conv2d(in_channels, channels, 3, padding=1)))
|
||||
# Number of channels at each block in the input half of U-Net
|
||||
input_block_channels = [channels]
|
||||
# Number of channels at each level
|
||||
channels_list = [channels * m for m in channel_multipliers]
|
||||
# Prepare levels
|
||||
for i in range(levels):
|
||||
# Add the residual blocks and attentions
|
||||
for _ in range(n_res_blocks):
|
||||
# Residual block maps from previous number of channels to the number of
|
||||
# channels in the current level
|
||||
layers = [ResBlock(channels, d_time_emb, out_channels=channels_list[i])]
|
||||
channels = channels_list[i]
|
||||
# Add transformer
|
||||
if i in attention_levels:
|
||||
layers.append(SpatialTransformer(channels, n_heads, tf_layers, d_cond))
|
||||
# Add them to the input half of the U-Net and keep track of the number of channels of
|
||||
# its output
|
||||
self.input_blocks.append(TimestepEmbedSequential(*layers))
|
||||
input_block_channels.append(channels)
|
||||
# Down sample at all levels except last
|
||||
if i != levels - 1:
|
||||
self.input_blocks.append(TimestepEmbedSequential(DownSample(channels)))
|
||||
input_block_channels.append(channels)
|
||||
|
||||
# The middle of the U-Net
|
||||
self.middle_block = TimestepEmbedSequential(
|
||||
ResBlock(channels, d_time_emb),
|
||||
SpatialTransformer(channels, n_heads, tf_layers, d_cond),
|
||||
ResBlock(channels, d_time_emb),
|
||||
)
|
||||
|
||||
# Second half of the U-Net
|
||||
self.output_blocks = nn.ModuleList([])
|
||||
# Prepare levels in reverse order
|
||||
for i in reversed(range(levels)):
|
||||
# Add the residual blocks and attentions
|
||||
for j in range(n_res_blocks + 1):
|
||||
# Residual block maps from previous number of channels plus the
|
||||
# skip connections from the input half of U-Net to the number of
|
||||
# channels in the current level.
|
||||
layers = [ResBlock(channels + input_block_channels.pop(), d_time_emb, out_channels=channels_list[i])]
|
||||
channels = channels_list[i]
|
||||
# Add transformer
|
||||
if i in attention_levels:
|
||||
layers.append(SpatialTransformer(channels, n_heads, tf_layers, d_cond))
|
||||
# Up-sample at every level after last residual block
|
||||
# except the last one.
|
||||
# Note that we are iterating in reverse; i.e. `i == 0` is the last.
|
||||
if i != 0 and j == n_res_blocks:
|
||||
layers.append(UpSample(channels))
|
||||
# Add to the output half of the U-Net
|
||||
self.output_blocks.append(TimestepEmbedSequential(*layers))
|
||||
|
||||
# Final normalization and $3 \times 3$ convolution
|
||||
self.out = nn.Sequential(
|
||||
normalization(channels),
|
||||
nn.SiLU(),
|
||||
nn.Conv2d(channels, out_channels, 3, padding=1),
|
||||
)
|
||||
|
||||
def time_step_embedding(self, time_steps: torch.Tensor, max_period: int = 10000):
|
||||
"""
|
||||
## Create sinusoidal time step embeddings
|
||||
|
||||
:param time_steps: are the time steps of shape `[batch_size]`
|
||||
:param max_period: controls the minimum frequency of the embeddings.
|
||||
"""
|
||||
# $\frac{c}{2}$; half the channels are sin and the other half is cos,
|
||||
half = self.channels // 2
|
||||
# $\frac{1}{10000^{\frac{2i}{c}}}$
|
||||
frequencies = torch.exp(
|
||||
-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
|
||||
).to(device=time_steps.device)
|
||||
# $\frac{t}{10000^{\frac{2i}{c}}}$
|
||||
args = time_steps[:, None].float() * frequencies[None]
|
||||
# $\cos\Bigg(\frac{t}{10000^{\frac{2i}{c}}}\Bigg)$ and $\sin\Bigg(\frac{t}{10000^{\frac{2i}{c}}}\Bigg)$
|
||||
return torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
|
||||
|
||||
def forward(self, x: torch.Tensor, time_steps: torch.Tensor, cond: torch.Tensor):
|
||||
"""
|
||||
:param x: is the input feature map of shape `[batch_size, channels, width, height]`
|
||||
:param time_steps: are the time steps of shape `[batch_size]`
|
||||
:param cond: conditioning of shape `[batch_size, n_cond, d_cond]`
|
||||
"""
|
||||
# To store the input half outputs for skip connections
|
||||
x_input_block = []
|
||||
|
||||
# Get time step embeddings
|
||||
t_emb = self.time_step_embedding(time_steps)
|
||||
t_emb = self.time_embed(t_emb)
|
||||
|
||||
# Input half of the U-Net
|
||||
for module in self.input_blocks:
|
||||
x = module(x, t_emb, cond)
|
||||
x_input_block.append(x)
|
||||
# Middle of the U-Net
|
||||
x = self.middle_block(x, t_emb, cond)
|
||||
# Output half of the U-Net
|
||||
for module in self.output_blocks:
|
||||
x = torch.cat([x, x_input_block.pop()], dim=1)
|
||||
x = module(x, t_emb, cond)
|
||||
|
||||
# Final normalization and $3 \times 3$ convolution
|
||||
return self.out(x)
|
||||
|
||||
|
||||
class TimestepEmbedSequential(nn.Sequential):
|
||||
"""
|
||||
### Sequential block for modules with different inputs
|
||||
|
||||
This sequential module can compose of different modules such as `ResBlock`,
|
||||
`nn.Conv` and `SpatialTransformer` and calls them with the matching signatures
|
||||
"""
|
||||
|
||||
def forward(self, x, t_emb, cond=None):
|
||||
for layer in self:
|
||||
if isinstance(layer, ResBlock):
|
||||
x = layer(x, t_emb)
|
||||
elif isinstance(layer, SpatialTransformer):
|
||||
x = layer(x, cond)
|
||||
else:
|
||||
x = layer(x)
|
||||
return x
|
||||
|
||||
|
||||
class UpSample(nn.Module):
|
||||
"""
|
||||
### Up-sampling layer
|
||||
"""
|
||||
|
||||
def __init__(self, channels: int):
|
||||
"""
|
||||
:param channels: is the number of channels
|
||||
"""
|
||||
super().__init__()
|
||||
# $3 \times 3$ convolution mapping
|
||||
self.conv = nn.Conv2d(channels, channels, 3, padding=1)
|
||||
|
||||
def forward(self, x: torch.Tensor):
|
||||
"""
|
||||
:param x: is the input feature map with shape `[batch_size, channels, height, width]`
|
||||
"""
|
||||
# Up-sample by a factor of $2$
|
||||
x = F.interpolate(x, scale_factor=2, mode="nearest")
|
||||
# Apply convolution
|
||||
return self.conv(x)
|
||||
|
||||
|
||||
class DownSample(nn.Module):
|
||||
"""
|
||||
## Down-sampling layer
|
||||
"""
|
||||
|
||||
def __init__(self, channels: int):
|
||||
"""
|
||||
:param channels: is the number of channels
|
||||
"""
|
||||
super().__init__()
|
||||
# $3 \times 3$ convolution with stride length of $2$ to down-sample by a factor of $2$
|
||||
self.op = nn.Conv2d(channels, channels, 3, stride=2, padding=1)
|
||||
|
||||
def forward(self, x: torch.Tensor):
|
||||
"""
|
||||
:param x: is the input feature map with shape `[batch_size, channels, height, width]`
|
||||
"""
|
||||
# Apply convolution
|
||||
return self.op(x)
|
||||
|
||||
|
||||
class ResBlock(nn.Module):
|
||||
"""
|
||||
## ResNet Block
|
||||
"""
|
||||
|
||||
def __init__(self, channels: int, d_t_emb: int, *, out_channels=None):
|
||||
"""
|
||||
:param channels: the number of input channels
|
||||
:param d_t_emb: the size of timestep embeddings
|
||||
:param out_channels: is the number of out channels. defaults to `channels.
|
||||
"""
|
||||
super().__init__()
|
||||
# `out_channels` not specified
|
||||
if out_channels is None:
|
||||
out_channels = channels
|
||||
|
||||
# First normalization and convolution
|
||||
self.in_layers = nn.Sequential(
|
||||
normalization(channels),
|
||||
nn.SiLU(),
|
||||
nn.Conv2d(channels, out_channels, 3, padding=1),
|
||||
)
|
||||
|
||||
# Time step embeddings
|
||||
self.emb_layers = nn.Sequential(
|
||||
nn.SiLU(),
|
||||
nn.Linear(d_t_emb, out_channels),
|
||||
)
|
||||
# Final convolution layer
|
||||
self.out_layers = nn.Sequential(
|
||||
normalization(out_channels),
|
||||
nn.SiLU(),
|
||||
nn.Dropout(0.),
|
||||
nn.Conv2d(out_channels, out_channels, 3, padding=1)
|
||||
)
|
||||
|
||||
# `channels` to `out_channels` mapping layer for residual connection
|
||||
if out_channels == channels:
|
||||
self.skip_connection = nn.Identity()
|
||||
else:
|
||||
self.skip_connection = nn.Conv2d(channels, out_channels, 1)
|
||||
|
||||
def forward(self, x: torch.Tensor, t_emb: torch.Tensor):
|
||||
"""
|
||||
:param x: is the input feature map with shape `[batch_size, channels, height, width]`
|
||||
:param t_emb: is the time step embeddings of shape `[batch_size, d_t_emb]`
|
||||
"""
|
||||
# Initial convolution
|
||||
h = self.in_layers(x)
|
||||
# Time step embeddings
|
||||
t_emb = self.emb_layers(t_emb).type(h.dtype)
|
||||
# Add time step embeddings
|
||||
h = h + t_emb[:, :, None, None]
|
||||
# Final convolution
|
||||
h = self.out_layers(h)
|
||||
# Add skip connection
|
||||
return self.skip_connection(x) + h
|
||||
|
||||
|
||||
class GroupNorm32(nn.GroupNorm):
|
||||
"""
|
||||
### Group normalization with float32 casting
|
||||
"""
|
||||
|
||||
def forward(self, x):
|
||||
return super().forward(x.float()).type(x.dtype)
|
||||
|
||||
|
||||
def normalization(channels):
|
||||
"""
|
||||
### Group normalization
|
||||
|
||||
This is a helper function, with fixed number of groups..
|
||||
"""
|
||||
return GroupNorm32(32, channels)
|
||||
|
||||
|
||||
def _test_time_embeddings():
|
||||
"""
|
||||
Test sinusoidal time step embeddings
|
||||
"""
|
||||
import matplotlib.pyplot as plt
|
||||
|
||||
plt.figure(figsize=(15, 5))
|
||||
m = UNetModel(in_channels=1, out_channels=1, channels=320, n_res_blocks=1, attention_levels=[],
|
||||
channel_multipliers=[],
|
||||
n_heads=1, tf_layers=1, d_cond=1)
|
||||
te = m.time_step_embedding(torch.arange(0, 1000))
|
||||
plt.plot(np.arange(1000), te[:, [50, 100, 190, 260]].numpy())
|
||||
plt.legend(["dim %d" % p for p in [50, 100, 190, 260]])
|
||||
plt.title("Time embeddings")
|
||||
plt.show()
|
||||
|
||||
|
||||
#
|
||||
if __name__ == '__main__':
|
||||
_test_time_embeddings()
|
||||
@@ -0,0 +1,309 @@
|
||||
"""
|
||||
---
|
||||
title: Transformer for Stable Diffusion U-Net
|
||||
summary: >
|
||||
Annotated PyTorch implementation/tutorial of the transformer
|
||||
for U-Net in stable diffusion.
|
||||
---
|
||||
|
||||
# Transformer for Stable Diffusion [U-Net](unet.html)
|
||||
|
||||
This implements the transformer module used in [U-Net](unet.html) that
|
||||
gives $\epsilon_\text{cond}(x_t, c)$
|
||||
|
||||
We have kept to the model definition and naming unchanged from
|
||||
[CompVis/stable-diffusion](https://github.com/CompVis/stable-diffusion)
|
||||
so that we can load the checkpoints directly.
|
||||
"""
|
||||
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import nn
|
||||
|
||||
|
||||
class SpatialTransformer(nn.Module):
|
||||
"""
|
||||
## Spatial Transformer
|
||||
"""
|
||||
|
||||
def __init__(self, channels: int, n_heads: int, n_layers: int, d_cond: int):
|
||||
"""
|
||||
:param channels: is the number of channels in the feature map
|
||||
:param n_heads: is the number of attention heads
|
||||
:param n_layers: is the number of transformer layers
|
||||
:param d_cond: is the size of the conditional embedding
|
||||
"""
|
||||
super().__init__()
|
||||
# Initial group normalization
|
||||
self.norm = torch.nn.GroupNorm(num_groups=32, num_channels=channels, eps=1e-6, affine=True)
|
||||
# Initial $1 \times 1$ convolution
|
||||
self.proj_in = nn.Conv2d(channels, channels, kernel_size=1, stride=1, padding=0)
|
||||
|
||||
# Transformer layers
|
||||
self.transformer_blocks = nn.ModuleList(
|
||||
[BasicTransformerBlock(channels, n_heads, channels // n_heads, d_cond=d_cond) for _ in range(n_layers)]
|
||||
)
|
||||
|
||||
# Final $1 \times 1$ convolution
|
||||
self.proj_out = nn.Conv2d(channels, channels, kernel_size=1, stride=1, padding=0)
|
||||
|
||||
def forward(self, x: torch.Tensor, cond: torch.Tensor):
|
||||
"""
|
||||
:param x: is the feature map of shape `[batch_size, channels, height, width]`
|
||||
:param cond: is the conditional embeddings of shape `[batch_size, n_cond, d_cond]`
|
||||
"""
|
||||
# Get shape `[batch_size, channels, height, width]`
|
||||
b, c, h, w = x.shape
|
||||
# For residual connection
|
||||
x_in = x
|
||||
# Normalize
|
||||
x = self.norm(x)
|
||||
# Initial $1 \times 1$ convolution
|
||||
x = self.proj_in(x)
|
||||
# Transpose and reshape from `[batch_size, channels, height, width]`
|
||||
# to `[batch_size, height * width, channels]`
|
||||
x = x.permute(0, 2, 3, 1).view(b, h * w, c)
|
||||
# Apply the transformer layers
|
||||
for block in self.transformer_blocks:
|
||||
x = block(x, cond)
|
||||
# Reshape and transpose from `[batch_size, height * width, channels]`
|
||||
# to `[batch_size, channels, height, width]`
|
||||
x = x.view(b, h, w, c).permute(0, 3, 1, 2)
|
||||
# Final $1 \times 1$ convolution
|
||||
x = self.proj_out(x)
|
||||
# Add residual
|
||||
return x + x_in
|
||||
|
||||
|
||||
class BasicTransformerBlock(nn.Module):
|
||||
"""
|
||||
### Transformer Layer
|
||||
"""
|
||||
|
||||
def __init__(self, d_model: int, n_heads: int, d_head: int, d_cond: int):
|
||||
"""
|
||||
:param d_model: is the input embedding size
|
||||
:param n_heads: is the number of attention heads
|
||||
:param d_head: is the size of a attention head
|
||||
:param d_cond: is the size of the conditional embeddings
|
||||
"""
|
||||
super().__init__()
|
||||
# Self-attention layer and pre-norm layer
|
||||
self.attn1 = CrossAttention(d_model, d_model, n_heads, d_head)
|
||||
self.norm1 = nn.LayerNorm(d_model)
|
||||
# Cross attention layer and pre-norm layer
|
||||
self.attn2 = CrossAttention(d_model, d_cond, n_heads, d_head)
|
||||
self.norm2 = nn.LayerNorm(d_model)
|
||||
# Feed-forward network and pre-norm layer
|
||||
self.ff = FeedForward(d_model)
|
||||
self.norm3 = nn.LayerNorm(d_model)
|
||||
|
||||
def forward(self, x: torch.Tensor, cond: torch.Tensor):
|
||||
"""
|
||||
:param x: are the input embeddings of shape `[batch_size, height * width, d_model]`
|
||||
:param cond: is the conditional embeddings of shape `[batch_size, n_cond, d_cond]`
|
||||
"""
|
||||
# Self attention
|
||||
x = self.attn1(self.norm1(x)) + x
|
||||
# Cross-attention with conditioning
|
||||
x = self.attn2(self.norm2(x), cond=cond) + x
|
||||
# Feed-forward network
|
||||
x = self.ff(self.norm3(x)) + x
|
||||
#
|
||||
return x
|
||||
|
||||
|
||||
class CrossAttention(nn.Module):
|
||||
"""
|
||||
### Cross Attention Layer
|
||||
|
||||
This falls-back to self-attention when conditional embeddings are not specified.
|
||||
"""
|
||||
|
||||
use_flash_attention: bool = False
|
||||
|
||||
def __init__(self, d_model: int, d_cond: int, n_heads: int, d_head: int, is_inplace: bool = True):
|
||||
"""
|
||||
:param d_model: is the input embedding size
|
||||
:param n_heads: is the number of attention heads
|
||||
:param d_head: is the size of a attention head
|
||||
:param d_cond: is the size of the conditional embeddings
|
||||
:param is_inplace: specifies whether to perform the attention softmax computation inplace to
|
||||
save memory
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
self.is_inplace = is_inplace
|
||||
self.n_heads = n_heads
|
||||
self.d_head = d_head
|
||||
|
||||
# Attention scaling factor
|
||||
self.scale = d_head ** -0.5
|
||||
|
||||
# Query, key and value mappings
|
||||
d_attn = d_head * n_heads
|
||||
self.to_q = nn.Linear(d_model, d_attn, bias=False)
|
||||
self.to_k = nn.Linear(d_cond, d_attn, bias=False)
|
||||
self.to_v = nn.Linear(d_cond, d_attn, bias=False)
|
||||
|
||||
# Final linear layer
|
||||
self.to_out = nn.Sequential(nn.Linear(d_attn, d_model))
|
||||
|
||||
# Setup [flash attention](https://github.com/HazyResearch/flash-attention).
|
||||
# Flash attention is only used if it's installed
|
||||
# and `CrossAttention.use_flash_attention` is set to `True`.
|
||||
try:
|
||||
# You can install flash attention by cloning their Github repo,
|
||||
# [https://github.com/HazyResearch/flash-attention](https://github.com/HazyResearch/flash-attention)
|
||||
# and then running `python setup.py install`
|
||||
from flash_attn.flash_attention import FlashAttention
|
||||
self.flash = FlashAttention()
|
||||
# Set the scale for scaled dot-product attention.
|
||||
self.flash.softmax_scale = self.scale
|
||||
# Set to `None` if it's not installed
|
||||
except ImportError:
|
||||
self.flash = None
|
||||
|
||||
def forward(self, x: torch.Tensor, cond: Optional[torch.Tensor] = None):
|
||||
"""
|
||||
:param x: are the input embeddings of shape `[batch_size, height * width, d_model]`
|
||||
:param cond: is the conditional embeddings of shape `[batch_size, n_cond, d_cond]`
|
||||
"""
|
||||
|
||||
# If `cond` is `None` we perform self attention
|
||||
has_cond = cond is not None
|
||||
if not has_cond:
|
||||
cond = x
|
||||
|
||||
# Get query, key and value vectors
|
||||
q = self.to_q(x)
|
||||
k = self.to_k(cond)
|
||||
v = self.to_v(cond)
|
||||
|
||||
# Use flash attention if it's available and the head size is less than or equal to `128`
|
||||
if CrossAttention.use_flash_attention and self.flash is not None and not has_cond and self.d_head <= 128:
|
||||
return self.flash_attention(q, k, v)
|
||||
# Otherwise, fallback to normal attention
|
||||
else:
|
||||
return self.normal_attention(q, k, v)
|
||||
|
||||
def flash_attention(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor):
|
||||
"""
|
||||
#### Flash Attention
|
||||
|
||||
:param q: are the query vectors before splitting heads, of shape `[batch_size, seq, d_attn]`
|
||||
:param k: are the query vectors before splitting heads, of shape `[batch_size, seq, d_attn]`
|
||||
:param v: are the query vectors before splitting heads, of shape `[batch_size, seq, d_attn]`
|
||||
"""
|
||||
|
||||
# Get batch size and number of elements along sequence axis (`width * height`)
|
||||
batch_size, seq_len, _ = q.shape
|
||||
|
||||
# Stack `q`, `k`, `v` vectors for flash attention, to get a single tensor of
|
||||
# shape `[batch_size, seq_len, 3, n_heads * d_head]`
|
||||
qkv = torch.stack((q, k, v), dim=2)
|
||||
# Split the heads
|
||||
qkv = qkv.view(batch_size, seq_len, 3, self.n_heads, self.d_head)
|
||||
|
||||
# Flash attention works for head sizes `32`, `64` and `128`, so we have to pad the heads to
|
||||
# fit this size.
|
||||
if self.d_head <= 32:
|
||||
pad = 32 - self.d_head
|
||||
elif self.d_head <= 64:
|
||||
pad = 64 - self.d_head
|
||||
elif self.d_head <= 128:
|
||||
pad = 128 - self.d_head
|
||||
else:
|
||||
raise ValueError(f'Head size ${self.d_head} too large for Flash Attention')
|
||||
|
||||
# Pad the heads
|
||||
if pad:
|
||||
qkv = torch.cat((qkv, qkv.new_zeros(batch_size, seq_len, 3, self.n_heads, pad)), dim=-1)
|
||||
|
||||
# Compute attention
|
||||
# $$\underset{seq}{softmax}\Bigg(\frac{Q K^\top}{\sqrt{d_{key}}}\Bigg)V$$
|
||||
# This gives a tensor of shape `[batch_size, seq_len, n_heads, d_padded]`
|
||||
out, _ = self.flash(qkv)
|
||||
# Truncate the extra head size
|
||||
out = out[:, :, :, :self.d_head]
|
||||
# Reshape to `[batch_size, seq_len, n_heads * d_head]`
|
||||
out = out.reshape(batch_size, seq_len, self.n_heads * self.d_head)
|
||||
|
||||
# Map to `[batch_size, height * width, d_model]` with a linear layer
|
||||
return self.to_out(out)
|
||||
|
||||
def normal_attention(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor):
|
||||
"""
|
||||
#### Normal Attention
|
||||
|
||||
:param q: are the query vectors before splitting heads, of shape `[batch_size, seq, d_attn]`
|
||||
:param k: are the query vectors before splitting heads, of shape `[batch_size, seq, d_attn]`
|
||||
:param v: are the query vectors before splitting heads, of shape `[batch_size, seq, d_attn]`
|
||||
"""
|
||||
|
||||
# Split them to heads of shape `[batch_size, seq_len, n_heads, d_head]`
|
||||
q = q.view(*q.shape[:2], self.n_heads, -1)
|
||||
k = k.view(*k.shape[:2], self.n_heads, -1)
|
||||
v = v.view(*v.shape[:2], self.n_heads, -1)
|
||||
|
||||
# Calculate attention $\frac{Q K^\top}{\sqrt{d_{key}}}$
|
||||
attn = torch.einsum('bihd,bjhd->bhij', q, k) * self.scale
|
||||
|
||||
# Compute softmax
|
||||
# $$\underset{seq}{softmax}\Bigg(\frac{Q K^\top}{\sqrt{d_{key}}}\Bigg)$$
|
||||
if self.is_inplace:
|
||||
half = attn.shape[0] // 2
|
||||
attn[half:] = attn[half:].softmax(dim=-1)
|
||||
attn[:half] = attn[:half].softmax(dim=-1)
|
||||
else:
|
||||
attn = attn.softmax(dim=-1)
|
||||
|
||||
# Compute attention output
|
||||
# $$\underset{seq}{softmax}\Bigg(\frac{Q K^\top}{\sqrt{d_{key}}}\Bigg)V$$
|
||||
out = torch.einsum('bhij,bjhd->bihd', attn, v)
|
||||
# Reshape to `[batch_size, height * width, n_heads * d_head]`
|
||||
out = out.reshape(*out.shape[:2], -1)
|
||||
# Map to `[batch_size, height * width, d_model]` with a linear layer
|
||||
return self.to_out(out)
|
||||
|
||||
|
||||
class FeedForward(nn.Module):
|
||||
"""
|
||||
### Feed-Forward Network
|
||||
"""
|
||||
|
||||
def __init__(self, d_model: int, d_mult: int = 4):
|
||||
"""
|
||||
:param d_model: is the input embedding size
|
||||
:param d_mult: is multiplicative factor for the hidden layer size
|
||||
"""
|
||||
super().__init__()
|
||||
self.net = nn.Sequential(
|
||||
GeGLU(d_model, d_model * d_mult),
|
||||
nn.Dropout(0.),
|
||||
nn.Linear(d_model * d_mult, d_model)
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor):
|
||||
return self.net(x)
|
||||
|
||||
|
||||
class GeGLU(nn.Module):
|
||||
"""
|
||||
### GeGLU Activation
|
||||
|
||||
$$\text{GeGLU}(x) = (xW + b) * \text{GELU}(xV + c)$$
|
||||
"""
|
||||
|
||||
def __init__(self, d_in: int, d_out: int):
|
||||
super().__init__()
|
||||
# Combined linear projections $xW + b$ and $xV + c$
|
||||
self.proj = nn.Linear(d_in, d_out * 2)
|
||||
|
||||
def forward(self, x: torch.Tensor):
|
||||
# Get $xW + b$ and $xV + c$
|
||||
x, gate = self.proj(x).chunk(2, dim=-1)
|
||||
# $\text{GeGLU}(x) = (xW + b) * \text{GELU}(xV + c)$
|
||||
return x * F.gelu(gate)
|
||||
@@ -0,0 +1,126 @@
|
||||
"""
|
||||
---
|
||||
title: Sampling algorithms for stable diffusion
|
||||
summary: >
|
||||
Annotated PyTorch implementation/tutorial of
|
||||
sampling algorithms
|
||||
for stable diffusion model.
|
||||
---
|
||||
|
||||
# Sampling algorithms for [stable diffusion](../index.html)
|
||||
|
||||
We have implemented the following [sampling algorithms](sampler/index.html):
|
||||
|
||||
* [Denoising Diffusion Probabilistic Models (DDPM) Sampling](ddpm.html)
|
||||
* [Denoising Diffusion Implicit Models (DDIM) Sampling](ddim.html)
|
||||
"""
|
||||
|
||||
from typing import Optional, List
|
||||
|
||||
import torch
|
||||
|
||||
from labml_nn.diffusion.stable_diffusion.latent_diffusion import LatentDiffusion
|
||||
|
||||
|
||||
class DiffusionSampler:
|
||||
"""
|
||||
## Base class for sampling algorithms
|
||||
"""
|
||||
model: LatentDiffusion
|
||||
|
||||
def __init__(self, model: LatentDiffusion):
|
||||
"""
|
||||
:param model: is the model to predict noise $\epsilon_\text{cond}(x_t, c)$
|
||||
"""
|
||||
super().__init__()
|
||||
# Set the model $\epsilon_\text{cond}(x_t, c)$
|
||||
self.model = model
|
||||
# Get number of steps the model was trained with $T$
|
||||
self.n_steps = model.n_steps
|
||||
|
||||
def get_eps(self, x: torch.Tensor, t: torch.Tensor, c: torch.Tensor, *,
|
||||
uncond_scale: float, uncond_cond: Optional[torch.Tensor]):
|
||||
"""
|
||||
## Get $\epsilon(x_t, c)$
|
||||
|
||||
:param x: is $x_t$ of shape `[batch_size, channels, height, width]`
|
||||
:param t: is $t$ of shape `[batch_size]`
|
||||
:param c: is the conditional embeddings $c$ of shape `[batch_size, emb_size]`
|
||||
:param uncond_scale: is the unconditional guidance scale $s$. This is used for
|
||||
$\epsilon_\theta(x_t, c) = s\epsilon_\text{cond}(x_t, c) + (s - 1)\epsilon_\text{cond}(x_t, c_u)$
|
||||
:param uncond_cond: is the conditional embedding for empty prompt $c_u$
|
||||
"""
|
||||
# When the scale $s = 1$
|
||||
# $$\epsilon_\theta(x_t, c) = \epsilon_\text{cond}(x_t, c)$$
|
||||
if uncond_cond is None or uncond_scale == 1.:
|
||||
return self.model(x, t, c)
|
||||
|
||||
# Duplicate $x_t$ and $t$
|
||||
x_in = torch.cat([x] * 2)
|
||||
t_in = torch.cat([t] * 2)
|
||||
# Concatenated $c$ and $c_u$
|
||||
c_in = torch.cat([uncond_cond, c])
|
||||
# Get $\epsilon_\text{cond}(x_t, c)$ and $\epsilon_\text{cond}(x_t, c_u)$
|
||||
e_t_uncond, e_t_cond = self.model(x_in, t_in, c_in).chunk(2)
|
||||
# Calculate
|
||||
# $$\epsilon_\theta(x_t, c) = s\epsilon_\text{cond}(x_t, c) + (s - 1)\epsilon_\text{cond}(x_t, c_u)$$
|
||||
e_t = e_t_uncond + uncond_scale * (e_t_cond - e_t_uncond)
|
||||
|
||||
#
|
||||
return e_t
|
||||
|
||||
def sample(self,
|
||||
shape: List[int],
|
||||
cond: torch.Tensor,
|
||||
repeat_noise: bool = False,
|
||||
temperature: float = 1.,
|
||||
x_last: Optional[torch.Tensor] = None,
|
||||
uncond_scale: float = 1.,
|
||||
uncond_cond: Optional[torch.Tensor] = None,
|
||||
skip_steps: int = 0,
|
||||
):
|
||||
"""
|
||||
### Sampling Loop
|
||||
|
||||
:param shape: is the shape of the generated images in the
|
||||
form `[batch_size, channels, height, width]`
|
||||
:param cond: is the conditional embeddings $c$
|
||||
:param temperature: is the noise temperature (random noise gets multiplied by this)
|
||||
:param x_last: is $x_T$. If not provided random noise will be used.
|
||||
:param uncond_scale: is the unconditional guidance scale $s$. This is used for
|
||||
$\epsilon_\theta(x_t, c) = s\epsilon_\text{cond}(x_t, c) + (s - 1)\epsilon_\text{cond}(x_t, c_u)$
|
||||
:param uncond_cond: is the conditional embedding for empty prompt $c_u$
|
||||
:param skip_steps: is the number of time steps to skip.
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
def paint(self, x: torch.Tensor, cond: torch.Tensor, t_start: int, *,
|
||||
orig: Optional[torch.Tensor] = None,
|
||||
mask: Optional[torch.Tensor] = None, orig_noise: Optional[torch.Tensor] = None,
|
||||
uncond_scale: float = 1.,
|
||||
uncond_cond: Optional[torch.Tensor] = None,
|
||||
):
|
||||
"""
|
||||
### Painting Loop
|
||||
|
||||
:param x: is $x_{T'}$ of shape `[batch_size, channels, height, width]`
|
||||
:param cond: is the conditional embeddings $c$
|
||||
:param t_start: is the sampling step to start from, $T'$
|
||||
:param orig: is the original image in latent page which we are in paining.
|
||||
:param mask: is the mask to keep the original image.
|
||||
:param orig_noise: is fixed noise to be added to the original image.
|
||||
:param uncond_scale: is the unconditional guidance scale $s$. This is used for
|
||||
$\epsilon_\theta(x_t, c) = s\epsilon_\text{cond}(x_t, c) + (s - 1)\epsilon_\text{cond}(x_t, c_u)$
|
||||
:param uncond_cond: is the conditional embedding for empty prompt $c_u$
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
|
||||
def q_sample(self, x0: torch.Tensor, index: int, noise: Optional[torch.Tensor] = None):
|
||||
"""
|
||||
### Sample from $q(x_t|x_0)$
|
||||
|
||||
:param x0: is $x_0$ of shape `[batch_size, channels, height, width]`
|
||||
:param index: is the time step $t$ index
|
||||
:param noise: is the noise, $\epsilon$
|
||||
"""
|
||||
raise NotImplementedError()
|
||||
@@ -0,0 +1,300 @@
|
||||
"""
|
||||
---
|
||||
title: Denoising Diffusion Implicit Models (DDIM) Sampling
|
||||
summary: >
|
||||
Annotated PyTorch implementation/tutorial of
|
||||
Denoising Diffusion Implicit Models (DDIM) Sampling
|
||||
for stable diffusion model.
|
||||
---
|
||||
|
||||
# Denoising Diffusion Implicit Models (DDIM) Sampling
|
||||
|
||||
This implements DDIM sampling from the paper
|
||||
[Denoising Diffusion Implicit Models](https://arxiv.org/abs/2010.02502)
|
||||
"""
|
||||
|
||||
from typing import Optional, List
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from labml import monit
|
||||
from labml_nn.diffusion.stable_diffusion.latent_diffusion import LatentDiffusion
|
||||
from labml_nn.diffusion.stable_diffusion.sampler import DiffusionSampler
|
||||
|
||||
|
||||
class DDIMSampler(DiffusionSampler):
|
||||
"""
|
||||
## DDIM Sampler
|
||||
|
||||
This extends the [`DiffusionSampler` base class](index.html).
|
||||
|
||||
DDIM samples images by repeatedly removing noise by sampling step by step using,
|
||||
|
||||
\begin{align}
|
||||
x_{\tau_{i-1}} &= \sqrt{\alpha_{\tau_{i-1}}}\Bigg(
|
||||
\frac{x_{\tau_i} - \sqrt{1 - \alpha_{\tau_i}}\epsilon_\theta(x_{\tau_i})}{\sqrt{\alpha_{\tau_i}}}
|
||||
\Bigg) \\
|
||||
&+ \sqrt{1 - \alpha_{\tau_{i- 1}} - \sigma_{\tau_i}^2} \cdot \epsilon_\theta(x_{\tau_i}) \\
|
||||
&+ \sigma_{\tau_i} \epsilon_{\tau_i}
|
||||
\end{align}
|
||||
|
||||
where $\epsilon_{\tau_i}$ is random noise,
|
||||
$\tau$ is a subsequence of $[1,2,\dots,T]$ of length $S$,
|
||||
and
|
||||
$$\sigma_{\tau_i} =
|
||||
\eta \sqrt{\frac{1 - \alpha_{\tau_{i-1}}}{1 - \alpha_{\tau_i}}}
|
||||
\sqrt{1 - \frac{\alpha_{\tau_i}}{\alpha_{\tau_{i-1}}}}$$
|
||||
|
||||
Note that, $\alpha_t$ in DDIM paper refers to ${\color{lightgreen}\bar\alpha_t}$ from [DDPM](ddpm.html).
|
||||
"""
|
||||
|
||||
model: LatentDiffusion
|
||||
|
||||
def __init__(self, model: LatentDiffusion, n_steps: int, ddim_discretize: str = "uniform", ddim_eta: float = 0.):
|
||||
"""
|
||||
:param model: is the model to predict noise $\epsilon_\text{cond}(x_t, c)$
|
||||
:param n_steps: is the number of DDIM sampling steps, $S$
|
||||
:param ddim_discretize: specifies how to extract $\tau$ from $[1,2,\dots,T]$.
|
||||
It can be either `uniform` or `quad`.
|
||||
:param ddim_eta: is $\eta$ used to calculate $\sigma_{\tau_i}$. $\eta = 0$ makes the
|
||||
sampling process deterministic.
|
||||
"""
|
||||
super().__init__(model)
|
||||
# Number of steps, $T$
|
||||
self.n_steps = model.n_steps
|
||||
|
||||
# Calculate $\tau$ to be uniformly distributed across $[1,2,\dots,T]$
|
||||
if ddim_discretize == 'uniform':
|
||||
c = self.n_steps // n_steps
|
||||
self.time_steps = np.asarray(list(range(0, self.n_steps, c))) + 1
|
||||
# Calculate $\tau$ to be quadratically distributed across $[1,2,\dots,T]$
|
||||
elif ddim_discretize == 'quad':
|
||||
self.time_steps = ((np.linspace(0, np.sqrt(self.n_steps * .8), n_steps)) ** 2).astype(int) + 1
|
||||
else:
|
||||
raise NotImplementedError(ddim_discretize)
|
||||
|
||||
with torch.no_grad():
|
||||
# Get ${\color{lightgreen}\bar\alpha_t}$
|
||||
alpha_bar = self.model.alpha_bar
|
||||
|
||||
# $\alpha_{\tau_i}$
|
||||
self.ddim_alpha = alpha_bar[self.time_steps].clone().to(torch.float32)
|
||||
# $\sqrt{\alpha_{\tau_i}}$
|
||||
self.ddim_alpha_sqrt = torch.sqrt(self.ddim_alpha)
|
||||
# $\alpha_{\tau_{i-1}}$
|
||||
self.ddim_alpha_prev = torch.cat([alpha_bar[0:1], alpha_bar[self.time_steps[:-1]]])
|
||||
|
||||
# $$\sigma_{\tau_i} =
|
||||
# \eta \sqrt{\frac{1 - \alpha_{\tau_{i-1}}}{1 - \alpha_{\tau_i}}}
|
||||
# \sqrt{1 - \frac{\alpha_{\tau_i}}{\alpha_{\tau_{i-1}}}}$$
|
||||
self.ddim_sigma = (ddim_eta *
|
||||
((1 - self.ddim_alpha_prev) / (1 - self.ddim_alpha) *
|
||||
(1 - self.ddim_alpha / self.ddim_alpha_prev)) ** .5)
|
||||
|
||||
# $\sqrt{1 - \alpha_{\tau_i}}$
|
||||
self.ddim_sqrt_one_minus_alpha = (1. - self.ddim_alpha) ** .5
|
||||
|
||||
@torch.no_grad()
|
||||
def sample(self,
|
||||
shape: List[int],
|
||||
cond: torch.Tensor,
|
||||
repeat_noise: bool = False,
|
||||
temperature: float = 1.,
|
||||
x_last: Optional[torch.Tensor] = None,
|
||||
uncond_scale: float = 1.,
|
||||
uncond_cond: Optional[torch.Tensor] = None,
|
||||
skip_steps: int = 0,
|
||||
):
|
||||
"""
|
||||
### Sampling Loop
|
||||
|
||||
:param shape: is the shape of the generated images in the
|
||||
form `[batch_size, channels, height, width]`
|
||||
:param cond: is the conditional embeddings $c$
|
||||
:param temperature: is the noise temperature (random noise gets multiplied by this)
|
||||
:param x_last: is $x_{\tau_S}$. If not provided random noise will be used.
|
||||
:param uncond_scale: is the unconditional guidance scale $s$. This is used for
|
||||
$\epsilon_\theta(x_t, c) = s\epsilon_\text{cond}(x_t, c) + (s - 1)\epsilon_\text{cond}(x_t, c_u)$
|
||||
:param uncond_cond: is the conditional embedding for empty prompt $c_u$
|
||||
:param skip_steps: is the number of time steps to skip $i'$. We start sampling from $S - i'$.
|
||||
And `x_last` is then $x_{\tau_{S - i'}}$.
|
||||
"""
|
||||
|
||||
# Get device and batch size
|
||||
device = self.model.device
|
||||
bs = shape[0]
|
||||
|
||||
# Get $x_{\tau_S}$
|
||||
x = x_last if x_last is not None else torch.randn(shape, device=device)
|
||||
|
||||
# Time steps to sample at $\tau_{S - i'}, \tau_{S - i' - 1}, \dots, \tau_1$
|
||||
time_steps = np.flip(self.time_steps)[skip_steps:]
|
||||
|
||||
for i, step in monit.enum('Sample', time_steps):
|
||||
# Index $i$ in the list $[\tau_1, \tau_2, \dots, \tau_S]$
|
||||
index = len(time_steps) - i - 1
|
||||
# Time step $\tau_i$
|
||||
ts = x.new_full((bs,), step, dtype=torch.long)
|
||||
|
||||
# Sample $x_{\tau_{i-1}}$
|
||||
x, pred_x0, e_t = self.p_sample(x, cond, ts, step, index=index,
|
||||
repeat_noise=repeat_noise,
|
||||
temperature=temperature,
|
||||
uncond_scale=uncond_scale,
|
||||
uncond_cond=uncond_cond)
|
||||
|
||||
# Return $x_0$
|
||||
return x
|
||||
|
||||
@torch.no_grad()
|
||||
def p_sample(self, x: torch.Tensor, c: torch.Tensor, t: torch.Tensor, step: int, index: int, *,
|
||||
repeat_noise: bool = False,
|
||||
temperature: float = 1.,
|
||||
uncond_scale: float = 1.,
|
||||
uncond_cond: Optional[torch.Tensor] = None):
|
||||
"""
|
||||
### Sample $x_{\tau_{i-1}}$
|
||||
|
||||
:param x: is $x_{\tau_i}$ of shape `[batch_size, channels, height, width]`
|
||||
:param c: is the conditional embeddings $c$ of shape `[batch_size, emb_size]`
|
||||
:param t: is $\tau_i$ of shape `[batch_size]`
|
||||
:param step: is the step $\tau_i$ as an integer
|
||||
:param index: is index $i$ in the list $[\tau_1, \tau_2, \dots, \tau_S]$
|
||||
:param repeat_noise: specified whether the noise should be same for all samples in the batch
|
||||
:param temperature: is the noise temperature (random noise gets multiplied by this)
|
||||
:param uncond_scale: is the unconditional guidance scale $s$. This is used for
|
||||
$\epsilon_\theta(x_t, c) = s\epsilon_\text{cond}(x_t, c) + (s - 1)\epsilon_\text{cond}(x_t, c_u)$
|
||||
:param uncond_cond: is the conditional embedding for empty prompt $c_u$
|
||||
"""
|
||||
|
||||
# Get $\epsilon_\theta(x_{\tau_i})$
|
||||
e_t = self.get_eps(x, t, c,
|
||||
uncond_scale=uncond_scale,
|
||||
uncond_cond=uncond_cond)
|
||||
|
||||
# Calculate $x_{\tau_{i - 1}}$ and predicted $x_0$
|
||||
x_prev, pred_x0 = self.get_x_prev_and_pred_x0(e_t, index, x,
|
||||
temperature=temperature,
|
||||
repeat_noise=repeat_noise)
|
||||
|
||||
#
|
||||
return x_prev, pred_x0, e_t
|
||||
|
||||
def get_x_prev_and_pred_x0(self, e_t: torch.Tensor, index: int, x: torch.Tensor, *,
|
||||
temperature: float,
|
||||
repeat_noise: bool):
|
||||
"""
|
||||
### Sample $x_{\tau_{i-1}}$ given $\epsilon_\theta(x_{\tau_i})$
|
||||
"""
|
||||
|
||||
# $\alpha_{\tau_i}$
|
||||
alpha = self.ddim_alpha[index]
|
||||
# $\alpha_{\tau_{i-1}}$
|
||||
alpha_prev = self.ddim_alpha_prev[index]
|
||||
# $\sigma_{\tau_i}$
|
||||
sigma = self.ddim_sigma[index]
|
||||
# $\sqrt{1 - \alpha_{\tau_i}}$
|
||||
sqrt_one_minus_alpha = self.ddim_sqrt_one_minus_alpha[index]
|
||||
|
||||
# Current prediction for $x_0$,
|
||||
# $$\frac{x_{\tau_i} - \sqrt{1 - \alpha_{\tau_i}}\epsilon_\theta(x_{\tau_i})}{\sqrt{\alpha_{\tau_i}}}$$
|
||||
pred_x0 = (x - sqrt_one_minus_alpha * e_t) / (alpha ** 0.5)
|
||||
# Direction pointing to $x_t$
|
||||
# $$\sqrt{1 - \alpha_{\tau_{i- 1}} - \sigma_{\tau_i}^2} \cdot \epsilon_\theta(x_{\tau_i})$$
|
||||
dir_xt = (1. - alpha_prev - sigma ** 2).sqrt() * e_t
|
||||
|
||||
# No noise is added, when $\eta = 0$
|
||||
if sigma == 0.:
|
||||
noise = 0.
|
||||
# If same noise is used for all samples in the batch
|
||||
elif repeat_noise:
|
||||
noise = torch.randn((1, *x.shape[1:]), device=x.device)
|
||||
# Different noise for each sample
|
||||
else:
|
||||
noise = torch.randn(x.shape, device=x.device)
|
||||
|
||||
# Multiply noise by the temperature
|
||||
noise = noise * temperature
|
||||
|
||||
# \begin{align}
|
||||
# x_{\tau_{i-1}} &= \sqrt{\alpha_{\tau_{i-1}}}\Bigg(
|
||||
# \frac{x_{\tau_i} - \sqrt{1 - \alpha_{\tau_i}}\epsilon_\theta(x_{\tau_i})}{\sqrt{\alpha_{\tau_i}}}
|
||||
# \Bigg) \\
|
||||
# &+ \sqrt{1 - \alpha_{\tau_{i- 1}} - \sigma_{\tau_i}^2} \cdot \epsilon_\theta(x_{\tau_i}) \\
|
||||
# &+ \sigma_{\tau_i} \epsilon_{\tau_i}
|
||||
# \end{align}
|
||||
x_prev = (alpha_prev ** 0.5) * pred_x0 + dir_xt + sigma * noise
|
||||
|
||||
#
|
||||
return x_prev, pred_x0
|
||||
|
||||
@torch.no_grad()
|
||||
def q_sample(self, x0: torch.Tensor, index: int, noise: Optional[torch.Tensor] = None):
|
||||
"""
|
||||
### Sample from $q_{\sigma,\tau}(x_{\tau_i}|x_0)$
|
||||
|
||||
$$q_{\sigma,\tau}(x_t|x_0) =
|
||||
\mathcal{N} \Big(x_t; \sqrt{\alpha_{\tau_i}} x_0, (1-\alpha_{\tau_i}) \mathbf{I} \Big)$$
|
||||
|
||||
:param x0: is $x_0$ of shape `[batch_size, channels, height, width]`
|
||||
:param index: is the time step $\tau_i$ index $i$
|
||||
:param noise: is the noise, $\epsilon$
|
||||
"""
|
||||
|
||||
# Random noise, if noise is not specified
|
||||
if noise is None:
|
||||
noise = torch.randn_like(x0)
|
||||
|
||||
# Sample from
|
||||
# $$q_{\sigma,\tau}(x_t|x_0) =
|
||||
# \mathcal{N} \Big(x_t; \sqrt{\alpha_{\tau_i}} x_0, (1-\alpha_{\tau_i}) \mathbf{I} \Big)$$
|
||||
return self.ddim_alpha_sqrt[index] * x0 + self.ddim_sqrt_one_minus_alpha[index] * noise
|
||||
|
||||
@torch.no_grad()
|
||||
def paint(self, x: torch.Tensor, cond: torch.Tensor, t_start: int, *,
|
||||
orig: Optional[torch.Tensor] = None,
|
||||
mask: Optional[torch.Tensor] = None, orig_noise: Optional[torch.Tensor] = None,
|
||||
uncond_scale: float = 1.,
|
||||
uncond_cond: Optional[torch.Tensor] = None,
|
||||
):
|
||||
"""
|
||||
### Painting Loop
|
||||
|
||||
:param x: is $x_{S'}$ of shape `[batch_size, channels, height, width]`
|
||||
:param cond: is the conditional embeddings $c$
|
||||
:param t_start: is the sampling step to start from, $S'$
|
||||
:param orig: is the original image in latent page which we are in paining.
|
||||
If this is not provided, it'll be an image to image transformation.
|
||||
:param mask: is the mask to keep the original image.
|
||||
:param orig_noise: is fixed noise to be added to the original image.
|
||||
:param uncond_scale: is the unconditional guidance scale $s$. This is used for
|
||||
$\epsilon_\theta(x_t, c) = s\epsilon_\text{cond}(x_t, c) + (s - 1)\epsilon_\text{cond}(x_t, c_u)$
|
||||
:param uncond_cond: is the conditional embedding for empty prompt $c_u$
|
||||
"""
|
||||
# Get batch size
|
||||
bs = x.shape[0]
|
||||
|
||||
# Time steps to sample at $\tau_{S`}, \tau_{S' - 1}, \dots, \tau_1$
|
||||
time_steps = np.flip(self.time_steps[:t_start])
|
||||
|
||||
for i, step in monit.enum('Paint', time_steps):
|
||||
# Index $i$ in the list $[\tau_1, \tau_2, \dots, \tau_S]$
|
||||
index = len(time_steps) - i - 1
|
||||
# Time step $\tau_i$
|
||||
ts = x.new_full((bs,), step, dtype=torch.long)
|
||||
|
||||
# Sample $x_{\tau_{i-1}}$
|
||||
x, _, _ = self.p_sample(x, cond, ts, step, index=index,
|
||||
uncond_scale=uncond_scale,
|
||||
uncond_cond=uncond_cond)
|
||||
|
||||
# Replace the masked area with original image
|
||||
if orig is not None:
|
||||
# Get the $q_{\sigma,\tau}(x_{\tau_i}|x_0)$ for original image in latent space
|
||||
orig_t = self.q_sample(orig, index, noise=orig_noise)
|
||||
# Replace the masked area
|
||||
x = orig_t * mask + x * (1 - mask)
|
||||
|
||||
#
|
||||
return x
|
||||
@@ -0,0 +1,226 @@
|
||||
"""
|
||||
---
|
||||
title: Denoising Diffusion Probabilistic Models (DDPM) Sampling
|
||||
summary: >
|
||||
Annotated PyTorch implementation/tutorial of
|
||||
Denoising Diffusion Probabilistic Models (DDPM) Sampling
|
||||
for stable diffusion model.
|
||||
---
|
||||
|
||||
# Denoising Diffusion Probabilistic Models (DDPM) Sampling
|
||||
|
||||
For a simpler DDPM implementation refer to our [DDPM implementation](../../ddpm/index.html).
|
||||
We use same notations for $\alpha_t$, $\beta_t$ schedules, etc.
|
||||
"""
|
||||
|
||||
from typing import Optional, List
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from labml import monit
|
||||
from labml_nn.diffusion.stable_diffusion.latent_diffusion import LatentDiffusion
|
||||
from labml_nn.diffusion.stable_diffusion.sampler import DiffusionSampler
|
||||
|
||||
|
||||
class DDPMSampler(DiffusionSampler):
|
||||
"""
|
||||
## DDPM Sampler
|
||||
|
||||
This extends the [`DiffusionSampler` base class](index.html).
|
||||
|
||||
DDPM samples images by repeatedly removing noise by sampling step by step from
|
||||
$p_\theta(x_{t-1} | x_t)$,
|
||||
|
||||
\begin{align}
|
||||
|
||||
p_\theta(x_{t-1} | x_t) &= \mathcal{N}\big(x_{t-1}; \mu_\theta(x_t, t), \tilde\beta_t \mathbf{I} \big) \\
|
||||
|
||||
\mu_t(x_t, t) &= \frac{\sqrt{\bar\alpha_{t-1}}\beta_t}{1 - \bar\alpha_t}x_0
|
||||
+ \frac{\sqrt{\alpha_t}(1 - \bar\alpha_{t-1})}{1-\bar\alpha_t}x_t \\
|
||||
|
||||
\tilde\beta_t &= \frac{1 - \bar\alpha_{t-1}}{1 - \bar\alpha_t} \beta_t \\
|
||||
|
||||
x_0 &= \frac{1}{\sqrt{\bar\alpha_t}} x_t - \Big(\sqrt{\frac{1}{\bar\alpha_t} - 1}\Big)\epsilon_\theta \\
|
||||
|
||||
\end{align}
|
||||
"""
|
||||
|
||||
model: LatentDiffusion
|
||||
|
||||
def __init__(self, model: LatentDiffusion):
|
||||
"""
|
||||
:param model: is the model to predict noise $\epsilon_\text{cond}(x_t, c)$
|
||||
"""
|
||||
super().__init__(model)
|
||||
|
||||
# Sampling steps $1, 2, \dots, T$
|
||||
self.time_steps = np.asarray(list(range(self.n_steps)))
|
||||
|
||||
with torch.no_grad():
|
||||
# $\bar\alpha_t$
|
||||
alpha_bar = self.model.alpha_bar
|
||||
# $\beta_t$ schedule
|
||||
beta = self.model.beta
|
||||
# $\bar\alpha_{t-1}$
|
||||
alpha_bar_prev = torch.cat([alpha_bar.new_tensor([1.]), alpha_bar[:-1]])
|
||||
|
||||
# $\sqrt{\bar\alpha}$
|
||||
self.sqrt_alpha_bar = alpha_bar ** .5
|
||||
# $\sqrt{1 - \bar\alpha}$
|
||||
self.sqrt_1m_alpha_bar = (1. - alpha_bar) ** .5
|
||||
# $\frac{1}{\sqrt{\bar\alpha_t}}$
|
||||
self.sqrt_recip_alpha_bar = alpha_bar ** -.5
|
||||
# $\sqrt{\frac{1}{\bar\alpha_t} - 1}$
|
||||
self.sqrt_recip_m1_alpha_bar = (1 / alpha_bar - 1) ** .5
|
||||
|
||||
# $\frac{1 - \bar\alpha_{t-1}}{1 - \bar\alpha_t} \beta_t$
|
||||
variance = beta * (1. - alpha_bar_prev) / (1. - alpha_bar)
|
||||
# Clamped log of $\tilde\beta_t$
|
||||
self.log_var = torch.log(torch.clamp(variance, min=1e-20))
|
||||
# $\frac{\sqrt{\bar\alpha_{t-1}}\beta_t}{1 - \bar\alpha_t}$
|
||||
self.mean_x0_coef = beta * (alpha_bar_prev ** .5) / (1. - alpha_bar)
|
||||
# $\frac{\sqrt{\alpha_t}(1 - \bar\alpha_{t-1})}{1-\bar\alpha_t}$
|
||||
self.mean_xt_coef = (1. - alpha_bar_prev) * ((1 - beta) ** 0.5) / (1. - alpha_bar)
|
||||
|
||||
@torch.no_grad()
|
||||
def sample(self,
|
||||
shape: List[int],
|
||||
cond: torch.Tensor,
|
||||
repeat_noise: bool = False,
|
||||
temperature: float = 1.,
|
||||
x_last: Optional[torch.Tensor] = None,
|
||||
uncond_scale: float = 1.,
|
||||
uncond_cond: Optional[torch.Tensor] = None,
|
||||
skip_steps: int = 0,
|
||||
):
|
||||
"""
|
||||
### Sampling Loop
|
||||
|
||||
:param shape: is the shape of the generated images in the
|
||||
form `[batch_size, channels, height, width]`
|
||||
:param cond: is the conditional embeddings $c$
|
||||
:param temperature: is the noise temperature (random noise gets multiplied by this)
|
||||
:param x_last: is $x_T$. If not provided random noise will be used.
|
||||
:param uncond_scale: is the unconditional guidance scale $s$. This is used for
|
||||
$\epsilon_\theta(x_t, c) = s\epsilon_\text{cond}(x_t, c) + (s - 1)\epsilon_\text{cond}(x_t, c_u)$
|
||||
:param uncond_cond: is the conditional embedding for empty prompt $c_u$
|
||||
:param skip_steps: is the number of time steps to skip $t'$. We start sampling from $T - t'$.
|
||||
And `x_last` is then $x_{T - t'}$.
|
||||
"""
|
||||
|
||||
# Get device and batch size
|
||||
device = self.model.device
|
||||
bs = shape[0]
|
||||
|
||||
# Get $x_T$
|
||||
x = x_last if x_last is not None else torch.randn(shape, device=device)
|
||||
|
||||
# Time steps to sample at $T - t', T - t' - 1, \dots, 1$
|
||||
time_steps = np.flip(self.time_steps)[skip_steps:]
|
||||
|
||||
# Sampling loop
|
||||
for step in monit.iterate('Sample', time_steps):
|
||||
# Time step $t$
|
||||
ts = x.new_full((bs,), step, dtype=torch.long)
|
||||
|
||||
# Sample $x_{t-1}$
|
||||
x, pred_x0, e_t = self.p_sample(x, cond, ts, step,
|
||||
repeat_noise=repeat_noise,
|
||||
temperature=temperature,
|
||||
uncond_scale=uncond_scale,
|
||||
uncond_cond=uncond_cond)
|
||||
|
||||
# Return $x_0$
|
||||
return x
|
||||
|
||||
@torch.no_grad()
|
||||
def p_sample(self, x: torch.Tensor, c: torch.Tensor, t: torch.Tensor, step: int,
|
||||
repeat_noise: bool = False,
|
||||
temperature: float = 1.,
|
||||
uncond_scale: float = 1., uncond_cond: Optional[torch.Tensor] = None):
|
||||
"""
|
||||
### Sample $x_{t-1}$ from $p_\theta(x_{t-1} | x_t)$
|
||||
|
||||
:param x: is $x_t$ of shape `[batch_size, channels, height, width]`
|
||||
:param c: is the conditional embeddings $c$ of shape `[batch_size, emb_size]`
|
||||
:param t: is $t$ of shape `[batch_size]`
|
||||
:param step: is the step $t$ as an integer
|
||||
:repeat_noise: specified whether the noise should be same for all samples in the batch
|
||||
:param temperature: is the noise temperature (random noise gets multiplied by this)
|
||||
:param uncond_scale: is the unconditional guidance scale $s$. This is used for
|
||||
$\epsilon_\theta(x_t, c) = s\epsilon_\text{cond}(x_t, c) + (s - 1)\epsilon_\text{cond}(x_t, c_u)$
|
||||
:param uncond_cond: is the conditional embedding for empty prompt $c_u$
|
||||
"""
|
||||
|
||||
# Get $\epsilon_\theta$
|
||||
e_t = self.get_eps(x, t, c,
|
||||
uncond_scale=uncond_scale,
|
||||
uncond_cond=uncond_cond)
|
||||
|
||||
# Get batch size
|
||||
bs = x.shape[0]
|
||||
|
||||
# $\frac{1}{\sqrt{\bar\alpha_t}}$
|
||||
sqrt_recip_alpha_bar = x.new_full((bs, 1, 1, 1), self.sqrt_recip_alpha_bar[step])
|
||||
# $\sqrt{\frac{1}{\bar\alpha_t} - 1}$
|
||||
sqrt_recip_m1_alpha_bar = x.new_full((bs, 1, 1, 1), self.sqrt_recip_m1_alpha_bar[step])
|
||||
|
||||
# Calculate $x_0$ with current $\epsilon_\theta$
|
||||
#
|
||||
# $$x_0 = \frac{1}{\sqrt{\bar\alpha_t}} x_t - \Big(\sqrt{\frac{1}{\bar\alpha_t} - 1}\Big)\epsilon_\theta$$
|
||||
x0 = sqrt_recip_alpha_bar * x - sqrt_recip_m1_alpha_bar * e_t
|
||||
|
||||
# $\frac{\sqrt{\bar\alpha_{t-1}}\beta_t}{1 - \bar\alpha_t}$
|
||||
mean_x0_coef = x.new_full((bs, 1, 1, 1), self.mean_x0_coef[step])
|
||||
# $\frac{\sqrt{\alpha_t}(1 - \bar\alpha_{t-1})}{1-\bar\alpha_t}$
|
||||
mean_xt_coef = x.new_full((bs, 1, 1, 1), self.mean_xt_coef[step])
|
||||
|
||||
# Calculate $\mu_t(x_t, t)$
|
||||
#
|
||||
# $$\mu_t(x_t, t) = \frac{\sqrt{\bar\alpha_{t-1}}\beta_t}{1 - \bar\alpha_t}x_0
|
||||
# + \frac{\sqrt{\alpha_t}(1 - \bar\alpha_{t-1})}{1-\bar\alpha_t}x_t$$
|
||||
mean = mean_x0_coef * x0 + mean_xt_coef * x
|
||||
# $\log \tilde\beta_t$
|
||||
log_var = x.new_full((bs, 1, 1, 1), self.log_var[step])
|
||||
|
||||
# Do not add noise when $t = 1$ (final step sampling process).
|
||||
# Note that `step` is `0` when $t = 1$)
|
||||
if step == 0:
|
||||
noise = 0
|
||||
# If same noise is used for all samples in the batch
|
||||
elif repeat_noise:
|
||||
noise = torch.randn((1, *x.shape[1:]))
|
||||
# Different noise for each sample
|
||||
else:
|
||||
noise = torch.randn(x.shape)
|
||||
|
||||
# Multiply noise by the temperature
|
||||
noise = noise * temperature
|
||||
|
||||
# Sample from,
|
||||
#
|
||||
# $$p_\theta(x_{t-1} | x_t) = \mathcal{N}\big(x_{t-1}; \mu_\theta(x_t, t), \tilde\beta_t \mathbf{I} \big)$$
|
||||
x_prev = mean + (0.5 * log_var).exp() * noise
|
||||
|
||||
#
|
||||
return x_prev, x0, e_t
|
||||
|
||||
@torch.no_grad()
|
||||
def q_sample(self, x0: torch.Tensor, index: int, noise: Optional[torch.Tensor] = None):
|
||||
"""
|
||||
### Sample from $q(x_t|x_0)$
|
||||
|
||||
$$q(x_t|x_0) = \mathcal{N} \Big(x_t; \sqrt{\bar\alpha_t} x_0, (1-\bar\alpha_t) \mathbf{I} \Big)$$
|
||||
|
||||
:param x0: is $x_0$ of shape `[batch_size, channels, height, width]`
|
||||
:param index: is the time step $t$ index
|
||||
:param noise: is the noise, $\epsilon$
|
||||
"""
|
||||
|
||||
# Random noise, if noise is not specified
|
||||
if noise is None:
|
||||
noise = torch.randn_like(x0)
|
||||
|
||||
# Sample from $\mathcal{N} \Big(x_t; \sqrt{\bar\alpha_t} x_0, (1-\bar\alpha_t) \mathbf{I} \Big)$
|
||||
return self.sqrt_alpha_bar[index] * x0 + self.sqrt_1m_alpha_bar[index] * noise
|
||||
@@ -0,0 +1,13 @@
|
||||
"""
|
||||
---
|
||||
title: Scripts to show example usages stable diffusion
|
||||
summary: >
|
||||
Annotated PyTorch implementation/tutorial of example usages of stable diffusion
|
||||
---
|
||||
|
||||
# Scripts to show example usages [stable diffusion](../index.html)
|
||||
|
||||
* [Prompt to image diffusion](text_to_image.html)
|
||||
* [Image to image diffusion](image_to_image.html)
|
||||
* [In-painting](in_paint.html)
|
||||
"""
|
||||
@@ -0,0 +1,149 @@
|
||||
"""
|
||||
---
|
||||
title: Generate images using stable diffusion with a prompt from a given image
|
||||
summary: >
|
||||
Generate images using stable diffusion with a prompt from a given image
|
||||
---
|
||||
|
||||
# Generate images using [stable diffusion](../index.html) with a prompt from a given image
|
||||
"""
|
||||
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
|
||||
from labml import lab, monit
|
||||
from labml_nn.diffusion.stable_diffusion.sampler.ddim import DDIMSampler
|
||||
from labml_nn.diffusion.stable_diffusion.util import load_model, load_img, save_images, set_seed
|
||||
|
||||
|
||||
class Img2Img:
|
||||
"""
|
||||
### Image to image class
|
||||
"""
|
||||
|
||||
def __init__(self, *, checkpoint_path: Path,
|
||||
ddim_steps: int = 50,
|
||||
ddim_eta: float = 0.0):
|
||||
"""
|
||||
:param checkpoint_path: is the path of the checkpoint
|
||||
:param ddim_steps: is the number of sampling steps
|
||||
:param ddim_eta: is the [DDIM sampling](../sampler/ddim.html) $\eta$ constant
|
||||
"""
|
||||
self.ddim_steps = ddim_steps
|
||||
|
||||
# Load [latent diffusion model](../latent_diffusion.html)
|
||||
self.model = load_model(checkpoint_path)
|
||||
# Get device
|
||||
self.device = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")
|
||||
# Move the model to device
|
||||
self.model.to(self.device)
|
||||
|
||||
# Initialize [DDIM sampler](../sampler/ddim.html)
|
||||
self.sampler = DDIMSampler(self.model,
|
||||
n_steps=ddim_steps,
|
||||
ddim_eta=ddim_eta)
|
||||
|
||||
@torch.no_grad()
|
||||
def __call__(self, *,
|
||||
dest_path: str,
|
||||
orig_img: str,
|
||||
strength: float,
|
||||
batch_size: int = 3,
|
||||
prompt: str,
|
||||
uncond_scale: float = 5.0,
|
||||
):
|
||||
"""
|
||||
:param dest_path: is the path to store the generated images
|
||||
:param orig_img: is the image to transform
|
||||
:param strength: specifies how much of the original image should not be preserved
|
||||
:param batch_size: is the number of images to generate in a batch
|
||||
:param prompt: is the prompt to generate images with
|
||||
:param uncond_scale: is the unconditional guidance scale $s$. This is used for
|
||||
$\epsilon_\theta(x_t, c) = s\epsilon_\text{cond}(x_t, c) + (s - 1)\epsilon_\text{cond}(x_t, c_u)$
|
||||
"""
|
||||
# Make a batch of prompts
|
||||
prompts = batch_size * [prompt]
|
||||
# Load image
|
||||
orig_image = load_img(orig_img).to(self.device)
|
||||
# Encode the image in the latent space and make `batch_size` copies of it
|
||||
orig = self.model.autoencoder_encode(orig_image).repeat(batch_size, 1, 1, 1)
|
||||
|
||||
# Get the number of steps to diffuse the original
|
||||
assert 0. <= strength <= 1., 'can only work with strength in [0.0, 1.0]'
|
||||
t_index = int(strength * self.ddim_steps)
|
||||
|
||||
# AMP auto casting
|
||||
with torch.cuda.amp.autocast():
|
||||
# In unconditional scaling is not $1$ get the embeddings for empty prompts (no conditioning).
|
||||
if uncond_scale != 1.0:
|
||||
un_cond = self.model.get_text_conditioning(batch_size * [""])
|
||||
else:
|
||||
un_cond = None
|
||||
# Get the prompt embeddings
|
||||
cond = self.model.get_text_conditioning(prompts)
|
||||
# Add noise to the original image
|
||||
x = self.sampler.q_sample(orig, t_index)
|
||||
# Reconstruct from the noisy image
|
||||
x = self.sampler.paint(x, cond, t_index,
|
||||
uncond_scale=uncond_scale,
|
||||
uncond_cond=un_cond)
|
||||
# Decode the image from the [autoencoder](../model/autoencoder.html)
|
||||
images = self.model.autoencoder_decode(x)
|
||||
|
||||
# Save images
|
||||
save_images(images, dest_path, 'img_')
|
||||
|
||||
|
||||
def main():
|
||||
"""
|
||||
### CLI
|
||||
"""
|
||||
parser = argparse.ArgumentParser()
|
||||
|
||||
parser.add_argument(
|
||||
"--prompt",
|
||||
type=str,
|
||||
nargs="?",
|
||||
default="a painting of a cute monkey playing guitar",
|
||||
help="the prompt to render"
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--orig-img",
|
||||
type=str,
|
||||
nargs="?",
|
||||
help="path to the input image"
|
||||
)
|
||||
|
||||
parser.add_argument("--batch_size", type=int, default=4, help="batch size", )
|
||||
parser.add_argument("--steps", type=int, default=50, help="number of ddim sampling steps")
|
||||
|
||||
parser.add_argument("--scale", type=float, default=5.0,
|
||||
help="unconditional guidance scale: "
|
||||
"eps = eps(x, empty) + scale * (eps(x, cond) - eps(x, empty))")
|
||||
|
||||
parser.add_argument("--strength", type=float, default=0.75,
|
||||
help="strength for noise: "
|
||||
" 1.0 corresponds to full destruction of information in init image")
|
||||
|
||||
opt = parser.parse_args()
|
||||
set_seed(42)
|
||||
|
||||
img2img = Img2Img(checkpoint_path=lab.get_data_path() / 'stable-diffusion' / 'sd-v1-4.ckpt',
|
||||
ddim_steps=opt.steps)
|
||||
|
||||
with monit.section('Generate'):
|
||||
img2img(
|
||||
dest_path='outputs',
|
||||
orig_img=opt.orig_img,
|
||||
strength=opt.strength,
|
||||
batch_size=opt.batch_size,
|
||||
prompt=opt.prompt,
|
||||
uncond_scale=opt.scale)
|
||||
|
||||
|
||||
#
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,166 @@
|
||||
"""
|
||||
---
|
||||
title: In-paint images using stable diffusion with a prompt
|
||||
summary: >
|
||||
In-paint images using stable diffusion with a prompt
|
||||
---
|
||||
|
||||
# In-paint images using [stable diffusion](../index.html) with a prompt
|
||||
"""
|
||||
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
|
||||
from labml import lab, monit
|
||||
from labml_nn.diffusion.stable_diffusion.latent_diffusion import LatentDiffusion
|
||||
from labml_nn.diffusion.stable_diffusion.sampler import DiffusionSampler
|
||||
from labml_nn.diffusion.stable_diffusion.sampler.ddim import DDIMSampler
|
||||
from labml_nn.diffusion.stable_diffusion.util import load_model, save_images, load_img, set_seed
|
||||
|
||||
|
||||
class InPaint:
|
||||
"""
|
||||
### Image in-painting class
|
||||
"""
|
||||
model: LatentDiffusion
|
||||
sampler: DiffusionSampler
|
||||
|
||||
def __init__(self, *, checkpoint_path: Path,
|
||||
ddim_steps: int = 50,
|
||||
ddim_eta: float = 0.0):
|
||||
"""
|
||||
:param checkpoint_path: is the path of the checkpoint
|
||||
:param ddim_steps: is the number of sampling steps
|
||||
:param ddim_eta: is the [DDIM sampling](../sampler/ddim.html) $\eta$ constant
|
||||
"""
|
||||
self.ddim_steps = ddim_steps
|
||||
|
||||
# Load [latent diffusion model](../latent_diffusion.html)
|
||||
self.model = load_model(checkpoint_path)
|
||||
# Get device
|
||||
self.device = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")
|
||||
# Move the model to device
|
||||
self.model.to(self.device)
|
||||
|
||||
# Initialize [DDIM sampler](../sampler/ddim.html)
|
||||
self.sampler = DDIMSampler(self.model,
|
||||
n_steps=ddim_steps,
|
||||
ddim_eta=ddim_eta)
|
||||
|
||||
@torch.no_grad()
|
||||
def __call__(self, *,
|
||||
dest_path: str,
|
||||
orig_img: str,
|
||||
strength: float,
|
||||
batch_size: int = 3,
|
||||
prompt: str,
|
||||
uncond_scale: float = 5.0,
|
||||
mask: Optional[torch.Tensor] = None,
|
||||
):
|
||||
"""
|
||||
:param dest_path: is the path to store the generated images
|
||||
:param orig_img: is the image to transform
|
||||
:param strength: specifies how much of the original image should not be preserved
|
||||
:param batch_size: is the number of images to generate in a batch
|
||||
:param prompt: is the prompt to generate images with
|
||||
:param uncond_scale: is the unconditional guidance scale $s$. This is used for
|
||||
$\epsilon_\theta(x_t, c) = s\epsilon_\text{cond}(x_t, c) + (s - 1)\epsilon_\text{cond}(x_t, c_u)$
|
||||
"""
|
||||
# Make a batch of prompts
|
||||
prompts = batch_size * [prompt]
|
||||
# Load image
|
||||
orig_image = load_img(orig_img).to(self.device)
|
||||
# Encode the image in the latent space and make `batch_size` copies of it
|
||||
orig = self.model.autoencoder_encode(orig_image).repeat(batch_size, 1, 1, 1)
|
||||
# If `mask` is not provided,
|
||||
# we set a sample mask to preserve the bottom half of the image
|
||||
if mask is None:
|
||||
mask = torch.zeros_like(orig, device=self.device)
|
||||
mask[:, :, mask.shape[2] // 2:, :] = 1.
|
||||
else:
|
||||
mask = mask.to(self.device)
|
||||
# Noise diffuse the original image
|
||||
orig_noise = torch.randn(orig.shape, device=self.device)
|
||||
|
||||
# Get the number of steps to diffuse the original
|
||||
assert 0. <= strength <= 1., 'can only work with strength in [0.0, 1.0]'
|
||||
t_index = int(strength * self.ddim_steps)
|
||||
|
||||
# AMP auto casting
|
||||
with torch.cuda.amp.autocast():
|
||||
# In unconditional scaling is not $1$ get the embeddings for empty prompts (no conditioning).
|
||||
if uncond_scale != 1.0:
|
||||
un_cond = self.model.get_text_conditioning(batch_size * [""])
|
||||
else:
|
||||
un_cond = None
|
||||
# Get the prompt embeddings
|
||||
cond = self.model.get_text_conditioning(prompts)
|
||||
# Add noise to the original image
|
||||
x = self.sampler.q_sample(orig, t_index, noise=orig_noise)
|
||||
# Reconstruct from the noisy image, while preserving the masked area
|
||||
x = self.sampler.paint(x, cond, t_index,
|
||||
orig=orig,
|
||||
mask=mask,
|
||||
orig_noise=orig_noise,
|
||||
uncond_scale=uncond_scale,
|
||||
uncond_cond=un_cond)
|
||||
# Decode the image from the [autoencoder](../model/autoencoder.html)
|
||||
images = self.model.autoencoder_decode(x)
|
||||
|
||||
# Save images
|
||||
save_images(images, dest_path, 'paint_')
|
||||
|
||||
|
||||
def main():
|
||||
"""
|
||||
### CLI
|
||||
"""
|
||||
parser = argparse.ArgumentParser()
|
||||
|
||||
parser.add_argument(
|
||||
"--prompt",
|
||||
type=str,
|
||||
nargs="?",
|
||||
default="a painting of a cute monkey playing guitar",
|
||||
help="the prompt to render"
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--orig-img",
|
||||
type=str,
|
||||
nargs="?",
|
||||
help="path to the input image"
|
||||
)
|
||||
|
||||
parser.add_argument("--batch_size", type=int, default=4, help="batch size", )
|
||||
parser.add_argument("--steps", type=int, default=50, help="number of sampling steps")
|
||||
|
||||
parser.add_argument("--scale", type=float, default=5.0,
|
||||
help="unconditional guidance scale: "
|
||||
"eps = eps(x, empty) + scale * (eps(x, cond) - eps(x, empty))")
|
||||
|
||||
parser.add_argument("--strength", type=float, default=0.75,
|
||||
help="strength for noise: "
|
||||
" 1.0 corresponds to full destruction of information in init image")
|
||||
|
||||
opt = parser.parse_args()
|
||||
set_seed(42)
|
||||
|
||||
in_paint = InPaint(checkpoint_path=lab.get_data_path() / 'stable-diffusion' / 'sd-v1-4.ckpt',
|
||||
ddim_steps=opt.steps)
|
||||
|
||||
with monit.section('Generate'):
|
||||
in_paint(dest_path='outputs',
|
||||
orig_img=opt.orig_img,
|
||||
strength=opt.strength,
|
||||
batch_size=opt.batch_size,
|
||||
prompt=opt.prompt,
|
||||
uncond_scale=opt.scale)
|
||||
|
||||
|
||||
#
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,158 @@
|
||||
"""
|
||||
---
|
||||
title: Generate images using stable diffusion with a prompt
|
||||
summary: >
|
||||
Generate images using stable diffusion with a prompt
|
||||
---
|
||||
|
||||
# Generate images using [stable diffusion](../index.html) with a prompt
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
|
||||
from labml import lab, monit
|
||||
from labml_nn.diffusion.stable_diffusion.latent_diffusion import LatentDiffusion
|
||||
from labml_nn.diffusion.stable_diffusion.sampler.ddim import DDIMSampler
|
||||
from labml_nn.diffusion.stable_diffusion.sampler.ddpm import DDPMSampler
|
||||
from labml_nn.diffusion.stable_diffusion.util import load_model, save_images, set_seed
|
||||
|
||||
|
||||
class Txt2Img:
|
||||
"""
|
||||
### Text to image class
|
||||
"""
|
||||
model: LatentDiffusion
|
||||
|
||||
def __init__(self, *,
|
||||
checkpoint_path: Path,
|
||||
sampler_name: str,
|
||||
n_steps: int = 50,
|
||||
ddim_eta: float = 0.0,
|
||||
):
|
||||
"""
|
||||
:param checkpoint_path: is the path of the checkpoint
|
||||
:param sampler_name: is the name of the [sampler](../sampler/index.html)
|
||||
:param n_steps: is the number of sampling steps
|
||||
:param ddim_eta: is the [DDIM sampling](../sampler/ddim.html) $\eta$ constant
|
||||
"""
|
||||
# Load [latent diffusion model](../latent_diffusion.html)
|
||||
self.model = load_model(checkpoint_path)
|
||||
# Get device
|
||||
self.device = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")
|
||||
# Move the model to device
|
||||
self.model.to(self.device)
|
||||
|
||||
# Initialize [sampler](../sampler/index.html)
|
||||
if sampler_name == 'ddim':
|
||||
self.sampler = DDIMSampler(self.model,
|
||||
n_steps=n_steps,
|
||||
ddim_eta=ddim_eta)
|
||||
elif sampler_name == 'ddpm':
|
||||
self.sampler = DDPMSampler(self.model)
|
||||
|
||||
@torch.no_grad()
|
||||
def __call__(self, *,
|
||||
dest_path: str,
|
||||
batch_size: int = 3,
|
||||
prompt: str,
|
||||
h: int = 512, w: int = 512,
|
||||
uncond_scale: float = 7.5,
|
||||
):
|
||||
"""
|
||||
:param dest_path: is the path to store the generated images
|
||||
:param batch_size: is the number of images to generate in a batch
|
||||
:param prompt: is the prompt to generate images with
|
||||
:param h: is the height of the image
|
||||
:param w: is the width of the image
|
||||
:param uncond_scale: is the unconditional guidance scale $s$. This is used for
|
||||
$\epsilon_\theta(x_t, c) = s\epsilon_\text{cond}(x_t, c) + (s - 1)\epsilon_\text{cond}(x_t, c_u)$
|
||||
"""
|
||||
# Number of channels in the image
|
||||
c = 4
|
||||
# Image to latent space resolution reduction
|
||||
f = 8
|
||||
|
||||
# Make a batch of prompts
|
||||
prompts = batch_size * [prompt]
|
||||
|
||||
# AMP auto casting
|
||||
with torch.cuda.amp.autocast():
|
||||
# In unconditional scaling is not $1$ get the embeddings for empty prompts (no conditioning).
|
||||
if uncond_scale != 1.0:
|
||||
un_cond = self.model.get_text_conditioning(batch_size * [""])
|
||||
else:
|
||||
un_cond = None
|
||||
# Get the prompt embeddings
|
||||
cond = self.model.get_text_conditioning(prompts)
|
||||
# [Sample in the latent space](../sampler/index.html).
|
||||
# `x` will be of shape `[batch_size, c, h / f, w / f]`
|
||||
x = self.sampler.sample(cond=cond,
|
||||
shape=[batch_size, c, h // f, w // f],
|
||||
uncond_scale=uncond_scale,
|
||||
uncond_cond=un_cond)
|
||||
# Decode the image from the [autoencoder](../model/autoencoder.html)
|
||||
images = self.model.autoencoder_decode(x)
|
||||
|
||||
# Save images
|
||||
save_images(images, dest_path, 'txt_')
|
||||
|
||||
|
||||
def main():
|
||||
"""
|
||||
### CLI
|
||||
"""
|
||||
parser = argparse.ArgumentParser()
|
||||
|
||||
parser.add_argument(
|
||||
"--prompt",
|
||||
type=str,
|
||||
nargs="?",
|
||||
default="a painting of a virus monster playing guitar",
|
||||
help="the prompt to render"
|
||||
)
|
||||
|
||||
parser.add_argument("--batch_size", type=int, default=4, help="batch size")
|
||||
|
||||
parser.add_argument(
|
||||
'--sampler',
|
||||
dest='sampler_name',
|
||||
choices=['ddim', 'ddpm'],
|
||||
default='ddim',
|
||||
help=f'Set the sampler.',
|
||||
)
|
||||
|
||||
parser.add_argument("--flash", action='store_true', help="whether to use flash attention")
|
||||
|
||||
parser.add_argument("--steps", type=int, default=50, help="number of sampling steps")
|
||||
|
||||
parser.add_argument("--scale", type=float, default=7.5,
|
||||
help="unconditional guidance scale: "
|
||||
"eps = eps(x, empty) + scale * (eps(x, cond) - eps(x, empty))")
|
||||
|
||||
opt = parser.parse_args()
|
||||
|
||||
set_seed(42)
|
||||
|
||||
# Set flash attention
|
||||
from labml_nn.diffusion.stable_diffusion.model.unet_attention import CrossAttention
|
||||
CrossAttention.use_flash_attention = opt.flash
|
||||
|
||||
#
|
||||
txt2img = Txt2Img(checkpoint_path=lab.get_data_path() / 'stable-diffusion' / 'sd-v1-4.ckpt',
|
||||
sampler_name=opt.sampler_name,
|
||||
n_steps=opt.steps)
|
||||
|
||||
with monit.section('Generate'):
|
||||
txt2img(dest_path='outputs',
|
||||
batch_size=opt.batch_size,
|
||||
prompt=opt.prompt,
|
||||
uncond_scale=opt.scale)
|
||||
|
||||
|
||||
#
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,151 @@
|
||||
"""
|
||||
---
|
||||
title: Utility functions for stable diffusion
|
||||
summary: >
|
||||
Utility functions for stable diffusion
|
||||
---
|
||||
|
||||
# Utility functions for [stable diffusion](index.html)
|
||||
"""
|
||||
|
||||
import os
|
||||
import random
|
||||
from pathlib import Path
|
||||
|
||||
import PIL
|
||||
import numpy as np
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
from labml import monit
|
||||
from labml.logger import inspect
|
||||
from labml_nn.diffusion.stable_diffusion.latent_diffusion import LatentDiffusion
|
||||
from labml_nn.diffusion.stable_diffusion.model.autoencoder import Encoder, Decoder, Autoencoder
|
||||
from labml_nn.diffusion.stable_diffusion.model.clip_embedder import CLIPTextEmbedder
|
||||
from labml_nn.diffusion.stable_diffusion.model.unet import UNetModel
|
||||
|
||||
|
||||
def set_seed(seed: int):
|
||||
"""
|
||||
### Set random seeds
|
||||
"""
|
||||
random.seed(seed)
|
||||
np.random.seed(seed)
|
||||
torch.manual_seed(seed)
|
||||
torch.cuda.manual_seed_all(seed)
|
||||
|
||||
|
||||
def load_model(path: Path = None) -> LatentDiffusion:
|
||||
"""
|
||||
### Load [`LatentDiffusion` model](latent_diffusion.html)
|
||||
"""
|
||||
|
||||
# Initialize the autoencoder
|
||||
with monit.section('Initialize autoencoder'):
|
||||
encoder = Encoder(z_channels=4,
|
||||
in_channels=3,
|
||||
channels=128,
|
||||
channel_multipliers=[1, 2, 4, 4],
|
||||
n_resnet_blocks=2)
|
||||
|
||||
decoder = Decoder(out_channels=3,
|
||||
z_channels=4,
|
||||
channels=128,
|
||||
channel_multipliers=[1, 2, 4, 4],
|
||||
n_resnet_blocks=2)
|
||||
|
||||
autoencoder = Autoencoder(emb_channels=4,
|
||||
encoder=encoder,
|
||||
decoder=decoder,
|
||||
z_channels=4)
|
||||
|
||||
# Initialize the CLIP text embedder
|
||||
with monit.section('Initialize CLIP Embedder'):
|
||||
clip_text_embedder = CLIPTextEmbedder()
|
||||
|
||||
# Initialize the U-Net
|
||||
with monit.section('Initialize U-Net'):
|
||||
unet_model = UNetModel(in_channels=4,
|
||||
out_channels=4,
|
||||
channels=320,
|
||||
attention_levels=[0, 1, 2],
|
||||
n_res_blocks=2,
|
||||
channel_multipliers=[1, 2, 4, 4],
|
||||
n_heads=8,
|
||||
tf_layers=1,
|
||||
d_cond=768)
|
||||
|
||||
# Initialize the Latent Diffusion model
|
||||
with monit.section('Initialize Latent Diffusion model'):
|
||||
model = LatentDiffusion(linear_start=0.00085,
|
||||
linear_end=0.0120,
|
||||
n_steps=1000,
|
||||
latent_scaling_factor=0.18215,
|
||||
|
||||
autoencoder=autoencoder,
|
||||
clip_embedder=clip_text_embedder,
|
||||
unet_model=unet_model)
|
||||
|
||||
# Load the checkpoint
|
||||
with monit.section(f"Loading model from {path}"):
|
||||
checkpoint = torch.load(path, map_location="cpu")
|
||||
|
||||
# Set model state
|
||||
with monit.section('Load state'):
|
||||
missing_keys, extra_keys = model.load_state_dict(checkpoint["state_dict"], strict=False)
|
||||
|
||||
# Debugging output
|
||||
inspect(global_step=checkpoint.get('global_step', -1), missing_keys=missing_keys, extra_keys=extra_keys,
|
||||
_expand=True)
|
||||
|
||||
#
|
||||
model.eval()
|
||||
return model
|
||||
|
||||
|
||||
def load_img(path: str):
|
||||
"""
|
||||
### Load an image
|
||||
|
||||
This loads an image from a file and returns a PyTorch tensor.
|
||||
|
||||
:param path: is the path of the image
|
||||
"""
|
||||
# Open Image
|
||||
image = Image.open(path).convert("RGB")
|
||||
# Get image size
|
||||
w, h = image.size
|
||||
# Resize to a multiple of 32
|
||||
w = w - w % 32
|
||||
h = h - h % 32
|
||||
image = image.resize((w, h), resample=PIL.Image.LANCZOS)
|
||||
# Convert to numpy and map to `[-1, 1]` for `[0, 255]`
|
||||
image = np.array(image).astype(np.float32) * (2. / 255.0) - 1
|
||||
# Transpose to shape `[batch_size, channels, height, width]`
|
||||
image = image[None].transpose(0, 3, 1, 2)
|
||||
# Convert to torch
|
||||
return torch.from_numpy(image)
|
||||
|
||||
|
||||
def save_images(images: torch.Tensor, dest_path: str, prefix: str = '', img_format: str = 'jpeg'):
|
||||
"""
|
||||
### Save a images
|
||||
|
||||
:param images: is the tensor with images of shape `[batch_size, channels, height, width]`
|
||||
:param dest_path: is the folder to save images in
|
||||
:param prefix: is the prefix to add to file names
|
||||
:param img_format: is the image format
|
||||
"""
|
||||
|
||||
# Create the destination folder
|
||||
os.makedirs(dest_path, exist_ok=True)
|
||||
|
||||
# Map images to `[0, 1]` space and clip
|
||||
images = torch.clamp((images + 1.0) / 2.0, min=0.0, max=1.0)
|
||||
# Transpose to `[batch_size, height, width, channels]` and convert to numpy
|
||||
images = images.cpu().permute(0, 2, 3, 1).numpy()
|
||||
|
||||
# Save images
|
||||
for i, img in enumerate(images):
|
||||
img = Image.fromarray((255. * img).astype(np.uint8))
|
||||
img.save(os.path.join(dest_path, f"{prefix}{i:05}.{img_format}"), format=img_format)
|
||||
@@ -0,0 +1,243 @@
|
||||
"""
|
||||
---
|
||||
title: Distilling the Knowledge in a Neural Network
|
||||
summary: >
|
||||
PyTorch implementation and tutorial of the paper
|
||||
Distilling the Knowledge in a Neural Network.
|
||||
---
|
||||
|
||||
# Distilling the Knowledge in a Neural Network
|
||||
|
||||
This is a [PyTorch](https://pytorch.org) implementation/tutorial of the paper
|
||||
[Distilling the Knowledge in a Neural Network](https://arxiv.org/abs/1503.02531).
|
||||
|
||||
It's a way of training a small network using the knowledge in a trained larger network;
|
||||
i.e. distilling the knowledge from the large network.
|
||||
|
||||
A large model with regularization or an ensemble of models (using dropout) generalizes
|
||||
better than a small model when trained directly on the data and labels.
|
||||
However, a small model can be trained to generalize better with help of a large model.
|
||||
Smaller models are better in production: faster, less compute, less memory.
|
||||
|
||||
The output probabilities of a trained model give more information than the labels
|
||||
because it assigns non-zero probabilities to incorrect classes as well.
|
||||
These probabilities tell us that a sample has a chance of belonging to certain classes.
|
||||
For instance, when classifying digits, when given an image of digit *7*,
|
||||
a generalized model will give a high probability to 7 and a small but non-zero
|
||||
probability to 2, while assigning almost zero probability to other digits.
|
||||
Distillation uses this information to train a small model better.
|
||||
|
||||
## Soft Targets
|
||||
|
||||
The probabilities are usually computed with a softmax operation,
|
||||
|
||||
$$q_i = \frac{\exp (z_i)}{\sum_j \exp (z_j)}$$
|
||||
|
||||
where $q_i$ is the probability for class $i$ and $z_i$ is the logit.
|
||||
|
||||
We train the small model to minimize the Cross entropy or KL Divergence between its output
|
||||
probability distribution and the large network's output probability distribution
|
||||
(soft targets).
|
||||
|
||||
One of the problems here is that the probabilities assigned to incorrect classes by the
|
||||
large network are often very small and don't contribute to the loss.
|
||||
So they soften the probabilities by applying a temperature $T$,
|
||||
|
||||
$$q_i = \frac{\exp (\frac{z_i}{T})}{\sum_j \exp (\frac{z_j}{T})}$$
|
||||
|
||||
where higher values for $T$ will produce softer probabilities.
|
||||
|
||||
## Training
|
||||
|
||||
Paper suggests adding a second loss term for predicting the actual labels
|
||||
when training the small model.
|
||||
We calculate the composite loss as the weighted sum of the two loss terms:
|
||||
soft targets and actual labels.
|
||||
|
||||
The dataset for distillation is called *the transfer set*, and the paper
|
||||
suggests using the same training data.
|
||||
|
||||
## Our experiment
|
||||
|
||||
We train on CIFAR-10 dataset.
|
||||
We [train a large model](large.html) that has $14,728,266$ parameters
|
||||
with dropout and it gives an accuracy of 85% on the validation set.
|
||||
A [small model](small.html) with $437,034$ parameters
|
||||
gives an accuracy of 80%.
|
||||
|
||||
We then train the small model with distillation from the large model,
|
||||
and it gives an accuracy of 82%; a 2% increase in the accuracy.
|
||||
"""
|
||||
|
||||
import torch
|
||||
import torch.nn.functional
|
||||
from torch import nn
|
||||
|
||||
from labml import experiment, tracker
|
||||
from labml.configs import option
|
||||
from labml_nn.helpers.trainer import BatchIndex
|
||||
from labml_nn.distillation.large import LargeModel
|
||||
from labml_nn.distillation.small import SmallModel
|
||||
from labml_nn.experiments.cifar10 import CIFAR10Configs
|
||||
|
||||
|
||||
class Configs(CIFAR10Configs):
|
||||
"""
|
||||
## Configurations
|
||||
|
||||
This extends from [`CIFAR10Configs`](../experiments/cifar10.html) which defines all the
|
||||
dataset related configurations, optimizer, and a training loop.
|
||||
"""
|
||||
# The small model
|
||||
model: SmallModel
|
||||
# The large model
|
||||
large: LargeModel
|
||||
# KL Divergence loss for soft targets
|
||||
kl_div_loss = nn.KLDivLoss(log_target=True)
|
||||
# Cross entropy loss for true label loss
|
||||
loss_func = nn.CrossEntropyLoss()
|
||||
# Temperature, $T$
|
||||
temperature: float = 5.
|
||||
# Weight for soft targets loss.
|
||||
#
|
||||
# The gradients produced by soft targets get scaled by $\frac{1}{T^2}$.
|
||||
# To compensate for this the paper suggests scaling the soft targets loss
|
||||
# by a factor of $T^2$
|
||||
soft_targets_weight: float = 100.
|
||||
# Weight for true label cross entropy loss
|
||||
label_loss_weight: float = 0.5
|
||||
|
||||
def step(self, batch: any, batch_idx: BatchIndex):
|
||||
"""
|
||||
### Training/validation step
|
||||
|
||||
We define a custom training/validation step to include the distillation
|
||||
"""
|
||||
|
||||
# Training/Evaluation mode for the small model
|
||||
self.model.train(self.mode.is_train)
|
||||
# Large model in evaluation mode
|
||||
self.large.eval()
|
||||
|
||||
# Move data to the device
|
||||
data, target = batch[0].to(self.device), batch[1].to(self.device)
|
||||
|
||||
# Update global step (number of samples processed) when in training mode
|
||||
if self.mode.is_train:
|
||||
tracker.add_global_step(len(data))
|
||||
|
||||
# Get the output logits, $v_i$, from the large model
|
||||
with torch.no_grad():
|
||||
large_logits = self.large(data)
|
||||
|
||||
# Get the output logits, $z_i$, from the small model
|
||||
output = self.model(data)
|
||||
|
||||
# Soft targets
|
||||
# $$p_i = \frac{\exp (\frac{v_i}{T})}{\sum_j \exp (\frac{v_j}{T})}$$
|
||||
soft_targets = nn.functional.log_softmax(large_logits / self.temperature, dim=-1)
|
||||
# Temperature adjusted probabilities of the small model
|
||||
# $$q_i = \frac{\exp (\frac{z_i}{T})}{\sum_j \exp (\frac{z_j}{T})}$$
|
||||
soft_prob = nn.functional.log_softmax(output / self.temperature, dim=-1)
|
||||
|
||||
# Calculate the soft targets loss
|
||||
soft_targets_loss = self.kl_div_loss(soft_prob, soft_targets)
|
||||
# Calculate the true label loss
|
||||
label_loss = self.loss_func(output, target)
|
||||
# Weighted sum of the two losses
|
||||
loss = self.soft_targets_weight * soft_targets_loss + self.label_loss_weight * label_loss
|
||||
# Log the losses
|
||||
tracker.add({"loss.kl_div.": soft_targets_loss,
|
||||
"loss.nll": label_loss,
|
||||
"loss.": loss})
|
||||
|
||||
# Calculate and log accuracy
|
||||
self.accuracy(output, target)
|
||||
self.accuracy.track()
|
||||
|
||||
# Train the model
|
||||
if self.mode.is_train:
|
||||
# Calculate gradients
|
||||
loss.backward()
|
||||
# Take optimizer step
|
||||
self.optimizer.step()
|
||||
# Log the model parameters and gradients on last batch of every epoch
|
||||
if batch_idx.is_last:
|
||||
tracker.add('model', self.model)
|
||||
# Clear the gradients
|
||||
self.optimizer.zero_grad()
|
||||
|
||||
# Save the tracked metrics
|
||||
tracker.save()
|
||||
|
||||
|
||||
@option(Configs.large)
|
||||
def _large_model(c: Configs):
|
||||
"""
|
||||
### Create large model
|
||||
"""
|
||||
return LargeModel().to(c.device)
|
||||
|
||||
|
||||
@option(Configs.model)
|
||||
def _small_student_model(c: Configs):
|
||||
"""
|
||||
### Create small model
|
||||
"""
|
||||
return SmallModel().to(c.device)
|
||||
|
||||
|
||||
def get_saved_model(run_uuid: str, checkpoint: int):
|
||||
"""
|
||||
### Load [trained large model](large.html)
|
||||
"""
|
||||
|
||||
from labml_nn.distillation.large import Configs as LargeConfigs
|
||||
|
||||
# In evaluation mode (no recording)
|
||||
experiment.evaluate()
|
||||
# Initialize configs of the large model training experiment
|
||||
conf = LargeConfigs()
|
||||
# Load saved configs
|
||||
experiment.configs(conf, experiment.load_configs(run_uuid))
|
||||
# Set models for saving/loading
|
||||
experiment.add_pytorch_models({'model': conf.model})
|
||||
# Set which run and checkpoint to load
|
||||
experiment.load(run_uuid, checkpoint)
|
||||
# Start the experiment - this will load the model, and prepare everything
|
||||
experiment.start()
|
||||
|
||||
# Return the model
|
||||
return conf.model
|
||||
|
||||
|
||||
def main(run_uuid: str, checkpoint: int):
|
||||
"""
|
||||
Train a small model with distillation
|
||||
"""
|
||||
# Load saved model
|
||||
large_model = get_saved_model(run_uuid, checkpoint)
|
||||
# Create experiment
|
||||
experiment.create(name='distillation', comment='cifar10')
|
||||
# Create configurations
|
||||
conf = Configs()
|
||||
# Set the loaded large model
|
||||
conf.large = large_model
|
||||
# Load configurations
|
||||
experiment.configs(conf, {
|
||||
'optimizer.optimizer': 'Adam',
|
||||
'optimizer.learning_rate': 2.5e-4,
|
||||
'model': '_small_student_model',
|
||||
})
|
||||
# Set model for saving/loading
|
||||
experiment.add_pytorch_models({'model': conf.model})
|
||||
# Start experiment from scratch
|
||||
experiment.load(None, None)
|
||||
# Start the experiment and run the training loop
|
||||
with experiment.start():
|
||||
conf.run()
|
||||
|
||||
|
||||
#
|
||||
if __name__ == '__main__':
|
||||
main('d46cd53edaec11eb93c38d6538aee7d6', 1_000_000)
|
||||
@@ -0,0 +1,89 @@
|
||||
"""
|
||||
---
|
||||
title: Train a large model on CIFAR 10
|
||||
summary: >
|
||||
Train a large model on CIFAR 10 for distillation.
|
||||
---
|
||||
|
||||
# Train a large model on CIFAR 10
|
||||
|
||||
This trains a large model on CIFAR 10 for [distillation](index.html).
|
||||
"""
|
||||
|
||||
import torch.nn as nn
|
||||
|
||||
from labml import experiment, logger
|
||||
from labml.configs import option
|
||||
from labml_nn.experiments.cifar10 import CIFAR10Configs, CIFAR10VGGModel
|
||||
from labml_nn.normalization.batch_norm import BatchNorm
|
||||
|
||||
|
||||
class Configs(CIFAR10Configs):
|
||||
"""
|
||||
## Configurations
|
||||
|
||||
We use [`CIFAR10Configs`](../experiments/cifar10.html) which defines all the
|
||||
dataset related configurations, optimizer, and a training loop.
|
||||
"""
|
||||
pass
|
||||
|
||||
|
||||
class LargeModel(CIFAR10VGGModel):
|
||||
"""
|
||||
### VGG style model for CIFAR-10 classification
|
||||
|
||||
This derives from the [generic VGG style architecture](../experiments/cifar10.html).
|
||||
"""
|
||||
|
||||
def conv_block(self, in_channels, out_channels) -> nn.Module:
|
||||
"""
|
||||
Create a convolution layer and the activations
|
||||
"""
|
||||
return nn.Sequential(
|
||||
# Dropout
|
||||
nn.Dropout(0.1),
|
||||
# Convolution layer
|
||||
nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
|
||||
# Batch normalization
|
||||
BatchNorm(out_channels, track_running_stats=False),
|
||||
# ReLU activation
|
||||
nn.ReLU(inplace=True),
|
||||
)
|
||||
|
||||
def __init__(self):
|
||||
# Create a model with given convolution sizes (channels)
|
||||
super().__init__([[64, 64], [128, 128], [256, 256, 256], [512, 512, 512], [512, 512, 512]])
|
||||
|
||||
|
||||
@option(Configs.model)
|
||||
def _large_model(c: Configs):
|
||||
"""
|
||||
### Create model
|
||||
"""
|
||||
return LargeModel().to(c.device)
|
||||
|
||||
|
||||
def main():
|
||||
# Create experiment
|
||||
experiment.create(name='cifar10', comment='large model')
|
||||
# Create configurations
|
||||
conf = Configs()
|
||||
# Load configurations
|
||||
experiment.configs(conf, {
|
||||
'optimizer.optimizer': 'Adam',
|
||||
'optimizer.learning_rate': 2.5e-4,
|
||||
'is_save_models': True,
|
||||
'epochs': 20,
|
||||
})
|
||||
# Set model for saving/loading
|
||||
experiment.add_pytorch_models({'model': conf.model})
|
||||
# Print number of parameters in the model
|
||||
logger.inspect(params=(sum(p.numel() for p in conf.model.parameters() if p.requires_grad)))
|
||||
# Start the experiment and run the training loop
|
||||
with experiment.start():
|
||||
conf.run()
|
||||
|
||||
|
||||
#
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,20 @@
|
||||
# [Distilling the Knowledge in a Neural Network](https://nn.labml.ai/distillation/index.html)
|
||||
|
||||
This is a [PyTorch](https://pytorch.org) implementation/tutorial of the paper
|
||||
[Distilling the Knowledge in a Neural Network](https://arxiv.org/abs/1503.02531).
|
||||
|
||||
It's a way of training a small network using the knowledge in a trained larger network;
|
||||
i.e. distilling the knowledge from the large network.
|
||||
|
||||
A large model with regularization or an ensemble of models (using dropout) generalizes
|
||||
better than a small model when trained directly on the data and labels.
|
||||
However, a small model can be trained to generalize better with help of a large model.
|
||||
Smaller models are better in production: faster, less compute, less memory.
|
||||
|
||||
The output probabilities of a trained model give more information than the labels
|
||||
because it assigns non-zero probabilities to incorrect classes as well.
|
||||
These probabilities tell us that a sample has a chance of belonging to certain classes.
|
||||
For instance, when classifying digits, when given an image of digit *7*,
|
||||
a generalized model will give a high probability to 7 and a small but non-zero
|
||||
probability to 2, while assigning almost zero probability to other digits.
|
||||
Distillation uses this information to train a small model better.
|
||||
@@ -0,0 +1,85 @@
|
||||
"""
|
||||
---
|
||||
title: Train a small model on CIFAR 10
|
||||
summary: >
|
||||
Train a small model on CIFAR 10 to test how much distillation benefits.
|
||||
---
|
||||
|
||||
# Train a small model on CIFAR 10
|
||||
|
||||
This trains a small model on CIFAR 10 to test how much [distillation](index.html) benefits.
|
||||
"""
|
||||
|
||||
import torch.nn as nn
|
||||
|
||||
from labml import experiment, logger
|
||||
from labml.configs import option
|
||||
from labml_nn.experiments.cifar10 import CIFAR10Configs, CIFAR10VGGModel
|
||||
from labml_nn.normalization.batch_norm import BatchNorm
|
||||
|
||||
|
||||
class Configs(CIFAR10Configs):
|
||||
"""
|
||||
## Configurations
|
||||
|
||||
We use [`CIFAR10Configs`](../experiments/cifar10.html) which defines all the
|
||||
dataset related configurations, optimizer, and a training loop.
|
||||
"""
|
||||
pass
|
||||
|
||||
|
||||
class SmallModel(CIFAR10VGGModel):
|
||||
"""
|
||||
### VGG style model for CIFAR-10 classification
|
||||
|
||||
This derives from the [generic VGG style architecture](../experiments/cifar10.html).
|
||||
"""
|
||||
|
||||
def conv_block(self, in_channels, out_channels) -> nn.Module:
|
||||
"""
|
||||
Create a convolution layer and the activations
|
||||
"""
|
||||
return nn.Sequential(
|
||||
# Convolution layer
|
||||
nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
|
||||
# Batch normalization
|
||||
BatchNorm(out_channels, track_running_stats=False),
|
||||
# ReLU activation
|
||||
nn.ReLU(inplace=True),
|
||||
)
|
||||
|
||||
def __init__(self):
|
||||
# Create a model with given convolution sizes (channels)
|
||||
super().__init__([[32, 32], [64, 64], [128], [128], [128]])
|
||||
|
||||
|
||||
@option(Configs.model)
|
||||
def _small_model(c: Configs):
|
||||
"""
|
||||
### Create model
|
||||
"""
|
||||
return SmallModel().to(c.device)
|
||||
|
||||
|
||||
def main():
|
||||
# Create experiment
|
||||
experiment.create(name='cifar10', comment='small model')
|
||||
# Create configurations
|
||||
conf = Configs()
|
||||
# Load configurations
|
||||
experiment.configs(conf, {
|
||||
'optimizer.optimizer': 'Adam',
|
||||
'optimizer.learning_rate': 2.5e-4,
|
||||
})
|
||||
# Set model for saving/loading
|
||||
experiment.add_pytorch_models({'model': conf.model})
|
||||
# Print number of parameters in the model
|
||||
logger.inspect(params=(sum(p.numel() for p in conf.model.parameters() if p.requires_grad)))
|
||||
# Start the experiment and run the training loop
|
||||
with experiment.start():
|
||||
conf.run()
|
||||
|
||||
|
||||
#
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,260 @@
|
||||
"""
|
||||
---
|
||||
title: Arithmetic Dataset
|
||||
summary: >
|
||||
This creates arithmetic problems.
|
||||
---
|
||||
|
||||
*This is based on code by [Georges Harik (@gharik)](https://twitter.com/gharik).*
|
||||
"""
|
||||
|
||||
import random
|
||||
import string
|
||||
from typing import List
|
||||
|
||||
import torch
|
||||
from labml.logger import Text
|
||||
from torch.utils.data import DataLoader, Dataset
|
||||
|
||||
from labml import monit, logger, tracker
|
||||
from labml.configs import option
|
||||
from labml_nn.experiments.nlp_autoregression import NLPAutoRegressionConfigs, transpose_batch
|
||||
|
||||
|
||||
class ArithmeticDataset(Dataset):
|
||||
"""
|
||||
## Arithmetic Dataset
|
||||
|
||||
This creates arithmetic addition problems and solutions with workings.
|
||||
We've only implemented addition so far.
|
||||
|
||||
It's based on a character level tokenization.
|
||||
"""
|
||||
|
||||
def __init__(self, seq_len: int, max_digits: int, n_sequences: int):
|
||||
"""
|
||||
:param seq_len: is the sequence length of generated math problems.
|
||||
We fill as many problems as possible upto this length
|
||||
:max_digits: is the maximum number of digits in the operand integers
|
||||
:n_sequences: is the number of sequences per epoch
|
||||
"""
|
||||
self.n_sequences = n_sequences
|
||||
self.max_digits = max_digits
|
||||
self.seq_len = seq_len
|
||||
# Token id to string
|
||||
self.itos = list(string.digits + 'xe =\n?+;')
|
||||
# Character to token id
|
||||
self.stoi = {c: i for i, c in enumerate(self.itos)}
|
||||
|
||||
@staticmethod
|
||||
def make_int(n_digits: int):
|
||||
"""
|
||||
Generates an integer with `n_digit` number of digits
|
||||
"""
|
||||
res = 0
|
||||
for i in range(n_digits):
|
||||
d = random.randrange(1, 11) if i == 0 else random.randrange(0, 11)
|
||||
res = res * 10 + d
|
||||
|
||||
return res
|
||||
|
||||
@staticmethod
|
||||
def get_add_explanation(x: int, y: int):
|
||||
"""
|
||||
Generates the workings for `x + y`.
|
||||
For example for `11+29` it generates
|
||||
`1e0+9e0+0e0=10e0 1e0+2e0+1e0=4e0`.
|
||||
"""
|
||||
|
||||
carry = 0
|
||||
e = 0
|
||||
explanation = []
|
||||
while x > 0 or y > 0 or carry > 0:
|
||||
rx, ry = x % 10, y % 10
|
||||
total = rx + ry + carry
|
||||
explanation.append(f"{rx}e{e}+{ry}e{e}+{carry}e{e}=={total}e{e}")
|
||||
x, y, carry = x // 10, y // 10, total // 10
|
||||
e += 1
|
||||
|
||||
return ' '.join(explanation)
|
||||
|
||||
# Make a problem with a pre_explanation or not
|
||||
def make_add_problem(self):
|
||||
"""
|
||||
Creates an arithmetic addition problem with workings and answer.
|
||||
"""
|
||||
x = self.make_int(n_digits=random.randrange(1, self.max_digits + 1))
|
||||
y = self.make_int(n_digits=random.randrange(1, self.max_digits + 1))
|
||||
|
||||
explanation = self.get_add_explanation(x, y)
|
||||
return f"x={x}+{y}; {explanation} x=={x + y}\n"
|
||||
|
||||
def get_qa(self):
|
||||
"""
|
||||
Get arithmetic problem and answer. This is used for evaluation.
|
||||
"""
|
||||
x = self.make_int(n_digits=random.randrange(1, self.max_digits + 1))
|
||||
y = self.make_int(n_digits=random.randrange(1, self.max_digits + 1))
|
||||
|
||||
return f'x={x}+{y};', f'{x + y}'
|
||||
|
||||
def get_packed_math_input(self):
|
||||
"""
|
||||
Generate multiple problems and pack them into a sequence.
|
||||
"""
|
||||
s_enc = []
|
||||
while len(s_enc) <= self.seq_len:
|
||||
s_part = self.make_add_problem()
|
||||
s_part_enc = self.encode('?' + s_part)
|
||||
s_enc = s_enc + s_part_enc
|
||||
return s_enc
|
||||
|
||||
def encode(self, s: str):
|
||||
"""
|
||||
Encode a given string
|
||||
"""
|
||||
return [self.stoi[c] for c in s]
|
||||
|
||||
def decode(self, arr: List[int]):
|
||||
"""
|
||||
Decode a list of token ids
|
||||
"""
|
||||
return ''.join([self.itos[c] for c in arr])
|
||||
|
||||
def __getitem__(self, idx: int):
|
||||
"""
|
||||
Get a input and target pair for auto-regressive modelling
|
||||
"""
|
||||
s = torch.tensor(self.get_packed_math_input())
|
||||
return s[:self.seq_len], s[1:self.seq_len + 1]
|
||||
|
||||
def __len__(self):
|
||||
"""
|
||||
Number of sequences per epoch
|
||||
"""
|
||||
return self.n_sequences
|
||||
|
||||
|
||||
class ArithmeticAutoregression(NLPAutoRegressionConfigs):
|
||||
"""
|
||||
## Arithmetic Task Experiment Configurations
|
||||
"""
|
||||
# Maximum number of digits per operand integer
|
||||
max_digits: int = 4
|
||||
# Number of training sequences per epoch
|
||||
train_sequences_per_epoch: int = 2 ** 12
|
||||
# Training data loader
|
||||
train_loader: DataLoader = 'arithmetic_train_loader'
|
||||
# Number of problems in evaluation
|
||||
n_tests: int = 64
|
||||
# No need of a validation dataset
|
||||
validator = None
|
||||
# Number of times to run evaluations per epoch
|
||||
inner_iterations = 4
|
||||
# Number of tokens in the vocabulary
|
||||
n_tokens = len(ArithmeticDataset(1, 1, 1).itos)
|
||||
|
||||
@torch.no_grad()
|
||||
def sample(self):
|
||||
"""
|
||||
### Evaluation
|
||||
|
||||
We use the sampling function to evaluate the model on a set of problems
|
||||
"""
|
||||
|
||||
# Skip in the first epoch
|
||||
if self.training_loop.idx < 1:
|
||||
return
|
||||
|
||||
# Create a dataset to generate problems
|
||||
dataset = ArithmeticDataset(self.seq_len, self.max_digits, 1)
|
||||
# Get a set of problems and answers
|
||||
qa = [dataset.get_qa() for _ in range(self.n_tests)]
|
||||
# Collect the problems only
|
||||
questions = [p[0] for p in qa]
|
||||
|
||||
# Create a tensor with only the initial token
|
||||
data = torch.tensor([[dataset.stoi[p[0]] for p in questions]])
|
||||
# Move to device
|
||||
data = data.to(self.device)
|
||||
|
||||
# Number of sequences that have completed
|
||||
finished = torch.zeros((len(questions),)).bool().to(self.device)
|
||||
# Token id of the new line character - this marks end of the answer
|
||||
new_line = dataset.stoi['\n']
|
||||
|
||||
# Sampled results
|
||||
results = [p[0] for p in questions]
|
||||
|
||||
# Sample upto sequence length
|
||||
for i in monit.iterate('Sample', self.seq_len - 1):
|
||||
# If all the sequences have completed we skip this
|
||||
if finished.sum() == len(finished):
|
||||
continue
|
||||
|
||||
# Get the model output
|
||||
output, *_ = self.model(data)
|
||||
# Get the model prediction (greedy)
|
||||
output = output[-1].argmax(dim=-1)
|
||||
|
||||
# Find which sequences have finished
|
||||
finished = finished | (output == new_line)
|
||||
# Skip if all have finished
|
||||
if finished.sum() == len(finished):
|
||||
continue
|
||||
|
||||
# Override with the question
|
||||
for j, p in enumerate(questions):
|
||||
if len(p) > i + 1:
|
||||
output[j] = dataset.stoi[p[i + 1]]
|
||||
|
||||
# Add the next token to the input
|
||||
data = torch.cat([data, output[None, :]], dim=0)
|
||||
|
||||
# Get the sampled results
|
||||
for j, c in enumerate(output):
|
||||
results[j] += dataset.itos[c]
|
||||
|
||||
# Discard everything after the answer in the results
|
||||
results = [r.split('\n')[0] for r in results]
|
||||
|
||||
# Log a sample
|
||||
res_sample = results[0].split(';')
|
||||
logger.log([(res_sample[0], Text.key), (';', Text.subtle), (';'.join(res_sample[1:]), Text.none)])
|
||||
|
||||
# Get the answers
|
||||
results = [r.split('x==')[-1] for r in results]
|
||||
|
||||
# Count the number of correct answers
|
||||
correct = 0
|
||||
for r, _qa in zip(results, qa):
|
||||
if r == _qa[1]:
|
||||
correct += 1
|
||||
|
||||
# Log the score
|
||||
tracker.save('score', correct / len(results))
|
||||
|
||||
|
||||
@option(ArithmeticAutoregression.train_loader)
|
||||
def arithmetic_train_loader(c: ArithmeticAutoregression):
|
||||
"""
|
||||
Training data loader
|
||||
"""
|
||||
return DataLoader(ArithmeticDataset(c.seq_len, c.max_digits, c.train_sequences_per_epoch),
|
||||
batch_size=c.batch_size,
|
||||
collate_fn=transpose_batch,
|
||||
num_workers=4)
|
||||
|
||||
|
||||
def _test():
|
||||
"""
|
||||
Code to test generated problems
|
||||
"""
|
||||
dataset = ArithmeticDataset(256, 8, 10)
|
||||
|
||||
print(dataset.decode(dataset.get_packed_math_input()))
|
||||
|
||||
|
||||
#
|
||||
if __name__ == '__main__':
|
||||
_test()
|
||||
@@ -0,0 +1,111 @@
|
||||
"""
|
||||
---
|
||||
title: CIFAR10 Experiment
|
||||
summary: >
|
||||
This is a reusable trainer for CIFAR10 dataset
|
||||
---
|
||||
|
||||
# CIFAR10 Experiment
|
||||
"""
|
||||
from typing import List
|
||||
|
||||
import torch.nn as nn
|
||||
|
||||
from labml import lab
|
||||
from labml.configs import option
|
||||
from labml_nn.helpers.datasets import CIFAR10Configs as CIFAR10DatasetConfigs
|
||||
from labml_nn.experiments.mnist import MNISTConfigs
|
||||
|
||||
|
||||
class CIFAR10Configs(CIFAR10DatasetConfigs, MNISTConfigs):
|
||||
"""
|
||||
## Configurations
|
||||
|
||||
This extends from [CIFAR 10 dataset configurations](../helpers/datasets.html)
|
||||
and [`MNISTConfigs`](mnist.html).
|
||||
"""
|
||||
# Use CIFAR10 dataset by default
|
||||
dataset_name: str = 'CIFAR10'
|
||||
|
||||
|
||||
@option(CIFAR10Configs.train_dataset)
|
||||
def cifar10_train_augmented():
|
||||
"""
|
||||
### Augmented CIFAR 10 train dataset
|
||||
"""
|
||||
from torchvision.datasets import CIFAR10
|
||||
from torchvision.transforms import transforms
|
||||
return CIFAR10(str(lab.get_data_path()),
|
||||
train=True,
|
||||
download=True,
|
||||
transform=transforms.Compose([
|
||||
# Pad and crop
|
||||
transforms.RandomCrop(32, padding=4),
|
||||
# Random horizontal flip
|
||||
transforms.RandomHorizontalFlip(),
|
||||
#
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
|
||||
]))
|
||||
|
||||
|
||||
@option(CIFAR10Configs.valid_dataset)
|
||||
def cifar10_valid_no_augment():
|
||||
"""
|
||||
### Non-augmented CIFAR 10 validation dataset
|
||||
"""
|
||||
from torchvision.datasets import CIFAR10
|
||||
from torchvision.transforms import transforms
|
||||
return CIFAR10(str(lab.get_data_path()),
|
||||
train=False,
|
||||
download=True,
|
||||
transform=transforms.Compose([
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
|
||||
]))
|
||||
|
||||
|
||||
class CIFAR10VGGModel(nn.Module):
|
||||
"""
|
||||
### VGG model for CIFAR-10 classification
|
||||
"""
|
||||
|
||||
def conv_block(self, in_channels, out_channels) -> nn.Module:
|
||||
"""
|
||||
Convolution and activation combined
|
||||
"""
|
||||
return nn.Sequential(
|
||||
nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
|
||||
nn.ReLU(inplace=True),
|
||||
)
|
||||
|
||||
def __init__(self, blocks: List[List[int]]):
|
||||
super().__init__()
|
||||
|
||||
# 5 $2 \times 2$ pooling layers will produce a output of size $1 \ times 1$.
|
||||
# CIFAR 10 image size is $32 \times 32$
|
||||
assert len(blocks) == 5
|
||||
layers = []
|
||||
# RGB channels
|
||||
in_channels = 3
|
||||
# Number of channels in each layer in each block
|
||||
for block in blocks:
|
||||
# Convolution, Normalization and Activation layers
|
||||
for channels in block:
|
||||
layers += self.conv_block(in_channels, channels)
|
||||
in_channels = channels
|
||||
# Max pooling at end of each block
|
||||
layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
|
||||
|
||||
# Create a sequential model with the layers
|
||||
self.layers = nn.Sequential(*layers)
|
||||
# Final logits layer
|
||||
self.fc = nn.Linear(in_channels, 10)
|
||||
|
||||
def forward(self, x):
|
||||
# The VGG layers
|
||||
x = self.layers(x)
|
||||
# Reshape for classification layer
|
||||
x = x.view(x.shape[0], -1)
|
||||
# Final linear layer
|
||||
return self.fc(x)
|
||||
@@ -0,0 +1,111 @@
|
||||
"""
|
||||
---
|
||||
title: MNIST Experiment
|
||||
summary: >
|
||||
This is a reusable trainer for MNIST dataset
|
||||
---
|
||||
|
||||
# MNIST Experiment
|
||||
"""
|
||||
|
||||
import torch.nn as nn
|
||||
import torch.utils.data
|
||||
|
||||
from labml import tracker
|
||||
from labml.configs import option
|
||||
from labml_nn.helpers.datasets import MNISTConfigs as MNISTDatasetConfigs
|
||||
from labml_nn.helpers.device import DeviceConfigs
|
||||
from labml_nn.helpers.metrics import Accuracy
|
||||
from labml_nn.helpers.trainer import TrainValidConfigs, BatchIndex
|
||||
from labml_nn.optimizers.configs import OptimizerConfigs
|
||||
|
||||
|
||||
class MNISTConfigs(MNISTDatasetConfigs, TrainValidConfigs):
|
||||
"""
|
||||
<a id="MNISTConfigs"></a>
|
||||
|
||||
## Trainer configurations
|
||||
"""
|
||||
|
||||
# Optimizer
|
||||
optimizer: torch.optim.Adam
|
||||
# Training device
|
||||
device: torch.device = DeviceConfigs()
|
||||
|
||||
# Classification model
|
||||
model: nn.Module
|
||||
# Number of epochs to train for
|
||||
epochs: int = 10
|
||||
|
||||
# Number of times to switch between training and validation within an epoch
|
||||
inner_iterations = 10
|
||||
|
||||
# Accuracy function
|
||||
accuracy = Accuracy()
|
||||
# Loss function
|
||||
loss_func = nn.CrossEntropyLoss()
|
||||
|
||||
def init(self):
|
||||
"""
|
||||
### Initialization
|
||||
"""
|
||||
# Set tracker configurations
|
||||
tracker.set_scalar("loss.*", True)
|
||||
tracker.set_scalar("accuracy.*", True)
|
||||
# Add accuracy as a state module.
|
||||
# The name is probably confusing, since it's meant to store
|
||||
# states between training and validation for RNNs.
|
||||
# This will keep the accuracy metric stats separate for training and validation.
|
||||
self.state_modules = [self.accuracy]
|
||||
|
||||
def step(self, batch: any, batch_idx: BatchIndex):
|
||||
"""
|
||||
### Training or validation step
|
||||
"""
|
||||
|
||||
# Training/Evaluation mode
|
||||
self.model.train(self.mode.is_train)
|
||||
|
||||
# Move data to the device
|
||||
data, target = batch[0].to(self.device), batch[1].to(self.device)
|
||||
|
||||
# Update global step (number of samples processed) when in training mode
|
||||
if self.mode.is_train:
|
||||
tracker.add_global_step(len(data))
|
||||
|
||||
# Get model outputs.
|
||||
output = self.model(data)
|
||||
|
||||
# Calculate and log loss
|
||||
loss = self.loss_func(output, target)
|
||||
tracker.add("loss.", loss)
|
||||
|
||||
# Calculate and log accuracy
|
||||
self.accuracy(output, target)
|
||||
self.accuracy.track()
|
||||
|
||||
# Train the model
|
||||
if self.mode.is_train:
|
||||
# Calculate gradients
|
||||
loss.backward()
|
||||
# Take optimizer step
|
||||
self.optimizer.step()
|
||||
# Log the model parameters and gradients on last batch of every epoch
|
||||
if batch_idx.is_last:
|
||||
tracker.add('model', self.model)
|
||||
# Clear the gradients
|
||||
self.optimizer.zero_grad()
|
||||
|
||||
# Save the tracked metrics
|
||||
tracker.save()
|
||||
|
||||
|
||||
@option(MNISTConfigs.optimizer)
|
||||
def _optimizer(c: MNISTConfigs):
|
||||
"""
|
||||
### Default optimizer configurations
|
||||
"""
|
||||
opt_conf = OptimizerConfigs()
|
||||
opt_conf.parameters = c.model.parameters()
|
||||
opt_conf.optimizer = 'Adam'
|
||||
return opt_conf
|
||||
@@ -0,0 +1,331 @@
|
||||
"""
|
||||
---
|
||||
title: NLP auto-regression trainer
|
||||
summary: >
|
||||
This is a reusable trainer for auto-regressive tasks
|
||||
---
|
||||
|
||||
# Auto-regressive NLP model trainer
|
||||
"""
|
||||
|
||||
from typing import Callable
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from labml import lab, monit, logger, tracker
|
||||
from labml.configs import option
|
||||
from labml.logger import Text
|
||||
from labml_nn.helpers.datasets import TextDataset, SequentialDataLoader, SequentialUnBatchedDataset, TextFileDataset
|
||||
from labml_nn.helpers.device import DeviceConfigs
|
||||
from labml_nn.helpers.metrics import Accuracy
|
||||
from labml_nn.helpers.trainer import TrainValidConfigs, BatchIndex
|
||||
from labml_nn.optimizers.configs import OptimizerConfigs
|
||||
from torch.utils.data import DataLoader, RandomSampler
|
||||
|
||||
|
||||
class CrossEntropyLoss(nn.Module):
|
||||
"""
|
||||
### Cross entropy loss
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.loss = nn.CrossEntropyLoss()
|
||||
|
||||
def forward(self, outputs, targets):
|
||||
return self.loss(outputs.view(-1, outputs.shape[-1]), targets.view(-1))
|
||||
|
||||
|
||||
class NLPAutoRegressionConfigs(TrainValidConfigs):
|
||||
"""
|
||||
<a id="NLPAutoRegressionConfigs"></a>
|
||||
|
||||
## Trainer configurations
|
||||
|
||||
This has the basic configurations for NLP auto-regressive task training.
|
||||
All the properties are configurable.
|
||||
"""
|
||||
|
||||
# Optimizer
|
||||
optimizer: torch.optim.Adam
|
||||
# Training device
|
||||
device: torch.device = DeviceConfigs()
|
||||
|
||||
# Autoregressive model
|
||||
model: nn.Module
|
||||
# Text dataset
|
||||
text: TextDataset
|
||||
# Batch size
|
||||
batch_size: int = 16
|
||||
# Length of the sequence, or context size
|
||||
seq_len: int = 512
|
||||
# Number of token in vocabulary
|
||||
n_tokens: int
|
||||
# Tokenizer
|
||||
tokenizer: Callable = 'character'
|
||||
|
||||
# Text prompt to start sampling (for illustration)
|
||||
prompt: str
|
||||
# The token separator when sampling (blank for character level tokenization)
|
||||
prompt_separator: str
|
||||
|
||||
# Whether to periodically save models
|
||||
is_save_models = True
|
||||
|
||||
# Loss function
|
||||
loss_func = CrossEntropyLoss()
|
||||
# Accuracy function
|
||||
accuracy = Accuracy()
|
||||
# Model embedding size
|
||||
d_model: int = 512
|
||||
# Gradient clipping
|
||||
grad_norm_clip: float = 1.0
|
||||
|
||||
# Training data loader
|
||||
train_loader: DataLoader = 'shuffled_train_loader'
|
||||
# Validation data loader
|
||||
valid_loader: DataLoader = 'shuffled_valid_loader'
|
||||
|
||||
# Data loaders shuffle with replacement
|
||||
dataloader_shuffle_with_replacement: bool = False
|
||||
|
||||
# Whether to log model parameters and gradients (once per epoch).
|
||||
# These are summarized stats per layer, but it could still lead
|
||||
# to many indicators for very deep networks.
|
||||
is_log_model_params_grads: bool = False
|
||||
|
||||
# Whether to log model activations (once per epoch).
|
||||
# These are summarized stats per layer, but it could still lead
|
||||
# to many indicators for very deep networks.
|
||||
is_log_model_activations: bool = False
|
||||
|
||||
def init(self):
|
||||
"""
|
||||
### Initialization
|
||||
"""
|
||||
# Set tracker configurations
|
||||
tracker.set_scalar("accuracy.*", True)
|
||||
tracker.set_scalar("loss.*", True)
|
||||
tracker.set_text("sampled", False)
|
||||
# Add accuracy as a state module.
|
||||
# The name is probably confusing, since it's meant to store
|
||||
# states between training and validation for RNNs.
|
||||
# This will keep the accuracy metric stats separate for training and validation.
|
||||
self.state_modules = [self.accuracy]
|
||||
|
||||
def other_metrics(self, output: torch.Tensor, target: torch.Tensor):
|
||||
"""Override to calculate and log other metrics"""
|
||||
pass
|
||||
|
||||
def step(self, batch: any, batch_idx: BatchIndex):
|
||||
"""
|
||||
### Training or validation step
|
||||
"""
|
||||
|
||||
# Set training/eval mode
|
||||
self.model.train(self.mode.is_train)
|
||||
|
||||
# Move data to the device
|
||||
data, target = batch[0].to(self.device), batch[1].to(self.device)
|
||||
|
||||
# Update global step (number of tokens processed) when in training mode
|
||||
if self.mode.is_train:
|
||||
tracker.add_global_step(data.shape[0] * data.shape[1])
|
||||
|
||||
# Get model outputs.
|
||||
# It's returning a tuple for states when using RNNs.
|
||||
# This is not implemented yet. 😜
|
||||
output, *_ = self.model(data)
|
||||
|
||||
# Calculate and log loss
|
||||
loss = self.loss_func(output, target)
|
||||
tracker.add("loss.", loss)
|
||||
|
||||
# Calculate and log accuracy
|
||||
self.accuracy(output, target)
|
||||
self.accuracy.track()
|
||||
|
||||
self.other_metrics(output, target)
|
||||
|
||||
# Train the model
|
||||
if self.mode.is_train:
|
||||
# Calculate gradients
|
||||
loss.backward()
|
||||
# Clip gradients
|
||||
torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=self.grad_norm_clip)
|
||||
# Take optimizer step
|
||||
self.optimizer.step()
|
||||
# Log the model parameters and gradients on last batch of every epoch
|
||||
if batch_idx.is_last and self.is_log_model_params_grads:
|
||||
tracker.add('model', self.model)
|
||||
# Clear the gradients
|
||||
self.optimizer.zero_grad()
|
||||
|
||||
# Save the tracked metrics
|
||||
tracker.save()
|
||||
|
||||
def sample(self):
|
||||
"""
|
||||
### Sampling function to generate samples periodically while training
|
||||
"""
|
||||
|
||||
# Starting prompt
|
||||
prompt = self.prompt
|
||||
# Collect output for printing
|
||||
log = [(prompt, Text.subtle)]
|
||||
# Sample 25 tokens
|
||||
for i in monit.iterate('Sample', 25):
|
||||
# Tokenize the prompt
|
||||
data = self.text.text_to_i(prompt).unsqueeze(-1)
|
||||
data = data.to(self.device)
|
||||
# Get the model output
|
||||
output, *_ = self.model(data)
|
||||
# Get the model prediction (greedy)
|
||||
output = output.argmax(dim=-1).squeeze()
|
||||
# Add the prediction to prompt
|
||||
prompt += self.prompt_separator + self.text.itos[output[-1]]
|
||||
# Add the prediction for logging
|
||||
log += [(self.prompt_separator + self.text.itos[output[-1]], Text.value)]
|
||||
|
||||
tracker.add({'sampled': prompt})
|
||||
# Print the sampled output
|
||||
logger.log(log)
|
||||
|
||||
|
||||
@option(NLPAutoRegressionConfigs.optimizer)
|
||||
def _optimizer(c: NLPAutoRegressionConfigs):
|
||||
"""
|
||||
### Default [optimizer configurations](../optimizers/configs.html)
|
||||
"""
|
||||
|
||||
optimizer = OptimizerConfigs()
|
||||
optimizer.parameters = c.model.parameters()
|
||||
optimizer.optimizer = 'Adam'
|
||||
optimizer.d_model = c.d_model
|
||||
|
||||
return optimizer
|
||||
|
||||
|
||||
@option(NLPAutoRegressionConfigs.n_tokens)
|
||||
def _n_tokens(c: NLPAutoRegressionConfigs):
|
||||
"""
|
||||
Get number of tokens
|
||||
"""
|
||||
return c.text.n_tokens
|
||||
|
||||
|
||||
@option(NLPAutoRegressionConfigs.tokenizer)
|
||||
def basic_english():
|
||||
"""
|
||||
### Basic english tokenizer
|
||||
|
||||
We use character level tokenizer in this experiment.
|
||||
You can switch by setting,
|
||||
|
||||
```
|
||||
'tokenizer': 'basic_english',
|
||||
```
|
||||
|
||||
in the configurations dictionary when starting the experiment.
|
||||
"""
|
||||
|
||||
from torchtext.data import get_tokenizer
|
||||
return get_tokenizer('basic_english')
|
||||
|
||||
|
||||
def character_tokenizer(x: str):
|
||||
"""
|
||||
### Character level tokenizer
|
||||
"""
|
||||
return list(x)
|
||||
|
||||
|
||||
@option(NLPAutoRegressionConfigs.tokenizer)
|
||||
def character():
|
||||
"""
|
||||
### Character level tokenizer configuration
|
||||
"""
|
||||
return character_tokenizer
|
||||
|
||||
|
||||
@option(NLPAutoRegressionConfigs.text)
|
||||
def tiny_shakespeare(c: NLPAutoRegressionConfigs):
|
||||
"""
|
||||
### Tiny Shakespeare dataset
|
||||
|
||||
It will download from the url if not present
|
||||
"""
|
||||
return TextFileDataset(
|
||||
lab.get_data_path() / 'tiny_shakespeare.txt',
|
||||
c.tokenizer,
|
||||
url='https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt')
|
||||
|
||||
|
||||
@option(NLPAutoRegressionConfigs.train_loader)
|
||||
def sequential_train_loader(c: NLPAutoRegressionConfigs):
|
||||
"""
|
||||
### Sequential training data loader
|
||||
"""
|
||||
return SequentialDataLoader(text=c.text.train,
|
||||
dataset=c.text,
|
||||
batch_size=c.batch_size,
|
||||
seq_len=c.seq_len)
|
||||
|
||||
|
||||
@option(NLPAutoRegressionConfigs.valid_loader)
|
||||
def sequential_valid_loader(c: NLPAutoRegressionConfigs):
|
||||
"""
|
||||
### Sequential validation data loader
|
||||
"""
|
||||
return SequentialDataLoader(text=c.text.valid,
|
||||
dataset=c.text,
|
||||
batch_size=c.batch_size,
|
||||
seq_len=c.seq_len)
|
||||
|
||||
|
||||
def transpose_batch(batch):
|
||||
"""
|
||||
### Transpose batch
|
||||
|
||||
`DataLoader` collects the batches on the first dimension.
|
||||
We need to transpose it to be sequence first.
|
||||
"""
|
||||
|
||||
transposed_data = list(zip(*batch))
|
||||
# Stack the batch along the second dimension `dim=1`
|
||||
src = torch.stack(transposed_data[0], dim=1)
|
||||
tgt = torch.stack(transposed_data[1], dim=1)
|
||||
|
||||
return src, tgt
|
||||
|
||||
|
||||
@option(NLPAutoRegressionConfigs.train_loader)
|
||||
def shuffled_train_loader(c: NLPAutoRegressionConfigs):
|
||||
"""
|
||||
### Shuffled training data loader
|
||||
"""
|
||||
dataset = SequentialUnBatchedDataset(text=c.text.train,
|
||||
dataset=c.text,
|
||||
seq_len=c.seq_len)
|
||||
sampler = RandomSampler(dataset, replacement=c.dataloader_shuffle_with_replacement)
|
||||
|
||||
return DataLoader(dataset,
|
||||
batch_size=c.batch_size,
|
||||
collate_fn=transpose_batch,
|
||||
sampler=sampler)
|
||||
|
||||
|
||||
@option(NLPAutoRegressionConfigs.valid_loader)
|
||||
def shuffled_valid_loader(c: NLPAutoRegressionConfigs):
|
||||
"""
|
||||
### Shuffled validation data loader
|
||||
"""
|
||||
dataset = SequentialUnBatchedDataset(text=c.text.valid,
|
||||
dataset=c.text,
|
||||
seq_len=c.seq_len)
|
||||
sampler = RandomSampler(dataset, replacement=c.dataloader_shuffle_with_replacement)
|
||||
|
||||
return DataLoader(dataset,
|
||||
batch_size=c.batch_size,
|
||||
collate_fn=transpose_batch,
|
||||
sampler=sampler)
|
||||
@@ -0,0 +1,289 @@
|
||||
"""
|
||||
---
|
||||
title: NLP classification trainer
|
||||
summary: >
|
||||
This is a reusable trainer for classification tasks
|
||||
---
|
||||
|
||||
# NLP model trainer for classification
|
||||
"""
|
||||
|
||||
from collections import Counter
|
||||
from typing import Callable
|
||||
|
||||
import torchtext
|
||||
import torchtext.vocab
|
||||
from torchtext.vocab import Vocab
|
||||
|
||||
import torch
|
||||
from labml import lab, tracker, monit
|
||||
from labml.configs import option
|
||||
from labml_nn.helpers.device import DeviceConfigs
|
||||
from labml_nn.helpers.metrics import Accuracy
|
||||
from labml_nn.helpers.trainer import TrainValidConfigs, BatchIndex
|
||||
from labml_nn.optimizers.configs import OptimizerConfigs
|
||||
from torch import nn
|
||||
from torch.utils.data import DataLoader
|
||||
|
||||
|
||||
class NLPClassificationConfigs(TrainValidConfigs):
|
||||
"""
|
||||
<a id="NLPClassificationConfigs"></a>
|
||||
|
||||
## Trainer configurations
|
||||
|
||||
This has the basic configurations for NLP classification task training.
|
||||
All the properties are configurable.
|
||||
"""
|
||||
|
||||
# Optimizer
|
||||
optimizer: torch.optim.Adam
|
||||
# Training device
|
||||
device: torch.device = DeviceConfigs()
|
||||
|
||||
# Autoregressive model
|
||||
model: nn.Module
|
||||
# Batch size
|
||||
batch_size: int = 16
|
||||
# Length of the sequence, or context size
|
||||
seq_len: int = 512
|
||||
# Vocabulary
|
||||
vocab: Vocab = 'ag_news'
|
||||
# Number of token in vocabulary
|
||||
n_tokens: int
|
||||
# Number of classes
|
||||
n_classes: int = 'ag_news'
|
||||
# Tokenizer
|
||||
tokenizer: Callable = 'character'
|
||||
|
||||
# Whether to periodically save models
|
||||
is_save_models = True
|
||||
|
||||
# Loss function
|
||||
loss_func = nn.CrossEntropyLoss()
|
||||
# Accuracy function
|
||||
accuracy = Accuracy()
|
||||
# Model embedding size
|
||||
d_model: int = 512
|
||||
# Gradient clipping
|
||||
grad_norm_clip: float = 1.0
|
||||
|
||||
# Training data loader
|
||||
train_loader: DataLoader = 'ag_news'
|
||||
# Validation data loader
|
||||
valid_loader: DataLoader = 'ag_news'
|
||||
|
||||
# Whether to log model parameters and gradients (once per epoch).
|
||||
# These are summarized stats per layer, but it could still lead
|
||||
# to many indicators for very deep networks.
|
||||
is_log_model_params_grads: bool = False
|
||||
|
||||
# Whether to log model activations (once per epoch).
|
||||
# These are summarized stats per layer, but it could still lead
|
||||
# to many indicators for very deep networks.
|
||||
is_log_model_activations: bool = False
|
||||
|
||||
def init(self):
|
||||
"""
|
||||
### Initialization
|
||||
"""
|
||||
# Set tracker configurations
|
||||
tracker.set_scalar("accuracy.*", True)
|
||||
tracker.set_scalar("loss.*", True)
|
||||
# Add accuracy as a state module.
|
||||
# The name is probably confusing, since it's meant to store
|
||||
# states between training and validation for RNNs.
|
||||
# This will keep the accuracy metric stats separate for training and validation.
|
||||
self.state_modules = [self.accuracy]
|
||||
|
||||
def step(self, batch: any, batch_idx: BatchIndex):
|
||||
"""
|
||||
### Training or validation step
|
||||
"""
|
||||
|
||||
# Move data to the device
|
||||
data, target = batch[0].to(self.device), batch[1].to(self.device)
|
||||
|
||||
# Update global step (number of tokens processed) when in training mode
|
||||
if self.mode.is_train:
|
||||
tracker.add_global_step(data.shape[1])
|
||||
|
||||
# Get model outputs.
|
||||
# It's returning a tuple for states when using RNNs.
|
||||
# This is not implemented yet. 😜
|
||||
output, *_ = self.model(data)
|
||||
|
||||
# Calculate and log loss
|
||||
loss = self.loss_func(output, target)
|
||||
tracker.add("loss.", loss)
|
||||
|
||||
# Calculate and log accuracy
|
||||
self.accuracy(output, target)
|
||||
self.accuracy.track()
|
||||
|
||||
# Train the model
|
||||
if self.mode.is_train:
|
||||
# Calculate gradients
|
||||
loss.backward()
|
||||
# Clip gradients
|
||||
torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=self.grad_norm_clip)
|
||||
# Take optimizer step
|
||||
self.optimizer.step()
|
||||
# Log the model parameters and gradients on last batch of every epoch
|
||||
if batch_idx.is_last and self.is_log_model_params_grads:
|
||||
tracker.add('model', self.model)
|
||||
# Clear the gradients
|
||||
self.optimizer.zero_grad()
|
||||
|
||||
# Save the tracked metrics
|
||||
tracker.save()
|
||||
|
||||
|
||||
@option(NLPClassificationConfigs.optimizer)
|
||||
def _optimizer(c: NLPClassificationConfigs):
|
||||
"""
|
||||
### Default [optimizer configurations](../optimizers/configs.html)
|
||||
"""
|
||||
|
||||
optimizer = OptimizerConfigs()
|
||||
optimizer.parameters = c.model.parameters()
|
||||
optimizer.optimizer = 'Adam'
|
||||
optimizer.d_model = c.d_model
|
||||
|
||||
return optimizer
|
||||
|
||||
|
||||
@option(NLPClassificationConfigs.tokenizer)
|
||||
def basic_english():
|
||||
"""
|
||||
### Basic english tokenizer
|
||||
|
||||
We use character level tokenizer in this experiment.
|
||||
You can switch by setting,
|
||||
|
||||
```
|
||||
'tokenizer': 'basic_english',
|
||||
```
|
||||
|
||||
in the configurations dictionary when starting the experiment.
|
||||
|
||||
"""
|
||||
from torchtext.data import get_tokenizer
|
||||
return get_tokenizer('basic_english')
|
||||
|
||||
|
||||
def character_tokenizer(x: str):
|
||||
"""
|
||||
### Character level tokenizer
|
||||
"""
|
||||
return list(x)
|
||||
|
||||
|
||||
@option(NLPClassificationConfigs.tokenizer)
|
||||
def character():
|
||||
"""
|
||||
Character level tokenizer configuration
|
||||
"""
|
||||
return character_tokenizer
|
||||
|
||||
|
||||
@option(NLPClassificationConfigs.n_tokens)
|
||||
def _n_tokens(c: NLPClassificationConfigs):
|
||||
"""
|
||||
Get number of tokens
|
||||
"""
|
||||
return len(c.vocab) + 2
|
||||
|
||||
|
||||
class CollateFunc:
|
||||
"""
|
||||
## Function to load data into batches
|
||||
"""
|
||||
|
||||
def __init__(self, tokenizer, vocab: Vocab, seq_len: int, padding_token: int, classifier_token: int):
|
||||
"""
|
||||
* `tokenizer` is the tokenizer function
|
||||
* `vocab` is the vocabulary
|
||||
* `seq_len` is the length of the sequence
|
||||
* `padding_token` is the token used for padding when the `seq_len` is larger than the text length
|
||||
* `classifier_token` is the `[CLS]` token which we set at end of the input
|
||||
"""
|
||||
self.classifier_token = classifier_token
|
||||
self.padding_token = padding_token
|
||||
self.seq_len = seq_len
|
||||
self.vocab = vocab
|
||||
self.tokenizer = tokenizer
|
||||
|
||||
def __call__(self, batch):
|
||||
"""
|
||||
* `batch` is the batch of data collected by the `DataLoader`
|
||||
"""
|
||||
|
||||
# Input data tensor, initialized with `padding_token`
|
||||
data = torch.full((self.seq_len, len(batch)), self.padding_token, dtype=torch.long)
|
||||
# Empty labels tensor
|
||||
labels = torch.zeros(len(batch), dtype=torch.long)
|
||||
|
||||
# Loop through the samples
|
||||
for (i, (_label, _text)) in enumerate(batch):
|
||||
# Set the label
|
||||
labels[i] = int(_label) - 1
|
||||
# Tokenize the input text
|
||||
_text = [self.vocab[token] for token in self.tokenizer(_text)]
|
||||
# Truncate upto `seq_len`
|
||||
_text = _text[:self.seq_len]
|
||||
# Transpose and add to data
|
||||
data[:len(_text), i] = data.new_tensor(_text)
|
||||
|
||||
# Set the final token in the sequence to `[CLS]`
|
||||
data[-1, :] = self.classifier_token
|
||||
|
||||
#
|
||||
return data, labels
|
||||
|
||||
|
||||
@option([NLPClassificationConfigs.n_classes,
|
||||
NLPClassificationConfigs.vocab,
|
||||
NLPClassificationConfigs.train_loader,
|
||||
NLPClassificationConfigs.valid_loader])
|
||||
def ag_news(c: NLPClassificationConfigs):
|
||||
"""
|
||||
### AG News dataset
|
||||
|
||||
This loads the AG News dataset and the set the values for
|
||||
`n_classes`, `vocab`, `train_loader`, and `valid_loader`.
|
||||
"""
|
||||
|
||||
# Get training and validation datasets
|
||||
train, valid = torchtext.datasets.AG_NEWS(root=str(lab.get_data_path() / 'ag_news'), split=('train', 'test'))
|
||||
|
||||
# Load data to memory
|
||||
with monit.section('Load data'):
|
||||
from labml_nn.utils import MapStyleDataset
|
||||
|
||||
# Create [map-style datasets](../utils.html#map_style_dataset)
|
||||
train, valid = MapStyleDataset(train), MapStyleDataset(valid)
|
||||
|
||||
# Get tokenizer
|
||||
tokenizer = c.tokenizer
|
||||
|
||||
# Create a counter
|
||||
counter = Counter()
|
||||
# Collect tokens from training dataset
|
||||
for (label, line) in train:
|
||||
counter.update(tokenizer(line))
|
||||
# Collect tokens from validation dataset
|
||||
for (label, line) in valid:
|
||||
counter.update(tokenizer(line))
|
||||
# Create vocabulary
|
||||
vocab = torchtext.vocab.vocab(counter, min_freq=1)
|
||||
|
||||
# Create training data loader
|
||||
train_loader = DataLoader(train, batch_size=c.batch_size, shuffle=True,
|
||||
collate_fn=CollateFunc(tokenizer, vocab, c.seq_len, len(vocab), len(vocab) + 1))
|
||||
# Create validation data loader
|
||||
valid_loader = DataLoader(valid, batch_size=c.batch_size, shuffle=True,
|
||||
collate_fn=CollateFunc(tokenizer, vocab, c.seq_len, len(vocab), len(vocab) + 1))
|
||||
|
||||
# Return `n_classes`, `vocab`, `train_loader`, and `valid_loader`
|
||||
return 4, vocab, train_loader, valid_loader
|
||||
@@ -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).
|
||||
@@ -0,0 +1,12 @@
|
||||
"""
|
||||
---
|
||||
title: Graph Neural Networks
|
||||
summary: >
|
||||
A set of PyTorch implementations/tutorials related to graph neural networks
|
||||
---
|
||||
|
||||
# Graph Neural Networks
|
||||
|
||||
* [Graph Attention Networks (GAT)](gat/index.html)
|
||||
* [Graph Attention Networks v2 (GATv2)](gatv2/index.html)
|
||||
"""
|
||||
@@ -0,0 +1,196 @@
|
||||
"""
|
||||
---
|
||||
title: Graph Attention Networks (GAT)
|
||||
summary: >
|
||||
A PyTorch implementation/tutorial of Graph Attention Networks.
|
||||
---
|
||||
|
||||
# Graph Attention Networks (GAT)
|
||||
|
||||
This is a [PyTorch](https://pytorch.org) implementation of the paper
|
||||
[Graph Attention Networks](https://arxiv.org/abs/1710.10903).
|
||||
|
||||
GATs work on graph data.
|
||||
A graph consists of nodes and edges connecting nodes.
|
||||
For example, in Cora dataset the nodes are research papers and the edges are citations that
|
||||
connect the papers.
|
||||
|
||||
GAT uses masked self-attention, kind of similar to [transformers](../../transformers/mha.html).
|
||||
GAT consists of graph attention layers stacked on top of each other.
|
||||
Each graph attention layer gets node embeddings as inputs and outputs transformed embeddings.
|
||||
The node embeddings pay attention to the embeddings of other nodes it's connected to.
|
||||
The details of graph attention layers are included alongside the implementation.
|
||||
|
||||
Here is [the training code](experiment.html) for training
|
||||
a two-layer GAT on Cora dataset.
|
||||
"""
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
|
||||
class GraphAttentionLayer(nn.Module):
|
||||
"""
|
||||
## Graph attention layer
|
||||
|
||||
This is a single graph attention layer.
|
||||
A GAT is made up of multiple such layers.
|
||||
|
||||
It takes
|
||||
$$\mathbf{h} = \{ \overrightarrow{h_1}, \overrightarrow{h_2}, \dots, \overrightarrow{h_N} \}$$,
|
||||
where $\overrightarrow{h_i} \in \mathbb{R}^F$ as input
|
||||
and outputs
|
||||
$$\mathbf{h'} = \{ \overrightarrow{h'_1}, \overrightarrow{h'_2}, \dots, \overrightarrow{h'_N} \}$$,
|
||||
where $\overrightarrow{h'_i} \in \mathbb{R}^{F'}$.
|
||||
"""
|
||||
def __init__(self, in_features: int, out_features: int, n_heads: int,
|
||||
is_concat: bool = True,
|
||||
dropout: float = 0.6,
|
||||
leaky_relu_negative_slope: float = 0.2):
|
||||
"""
|
||||
* `in_features`, $F$, is the number of input features per node
|
||||
* `out_features`, $F'$, is the number of output features per node
|
||||
* `n_heads`, $K$, is the number of attention heads
|
||||
* `is_concat` whether the multi-head results should be concatenated or averaged
|
||||
* `dropout` is the dropout probability
|
||||
* `leaky_relu_negative_slope` is the negative slope for leaky relu activation
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
self.is_concat = is_concat
|
||||
self.n_heads = n_heads
|
||||
|
||||
# Calculate the number of dimensions per head
|
||||
if is_concat:
|
||||
assert out_features % n_heads == 0
|
||||
# If we are concatenating the multiple heads
|
||||
self.n_hidden = out_features // n_heads
|
||||
else:
|
||||
# If we are averaging the multiple heads
|
||||
self.n_hidden = out_features
|
||||
|
||||
# Linear layer for initial transformation;
|
||||
# i.e. to transform the node embeddings before self-attention
|
||||
self.linear = nn.Linear(in_features, self.n_hidden * n_heads, bias=False)
|
||||
# Linear layer to compute attention score $e_{ij}$
|
||||
self.attn = nn.Linear(self.n_hidden * 2, 1, bias=False)
|
||||
# The activation for attention score $e_{ij}$
|
||||
self.activation = nn.LeakyReLU(negative_slope=leaky_relu_negative_slope)
|
||||
# Softmax to compute attention $\alpha_{ij}$
|
||||
self.softmax = nn.Softmax(dim=1)
|
||||
# Dropout layer to be applied for attention
|
||||
self.dropout = nn.Dropout(dropout)
|
||||
|
||||
def forward(self, h: torch.Tensor, adj_mat: torch.Tensor):
|
||||
"""
|
||||
* `h`, $\mathbf{h}$ is the input node embeddings of shape `[n_nodes, in_features]`.
|
||||
* `adj_mat` is the adjacency matrix of shape `[n_nodes, n_nodes, n_heads]`.
|
||||
We use shape `[n_nodes, n_nodes, 1]` since the adjacency is the same for each head.
|
||||
|
||||
Adjacency matrix represent the edges (or connections) among nodes.
|
||||
`adj_mat[i][j]` is `True` if there is an edge from node `i` to node `j`.
|
||||
"""
|
||||
|
||||
# Number of nodes
|
||||
n_nodes = h.shape[0]
|
||||
# The initial transformation,
|
||||
# $$\overrightarrow{g^k_i} = \mathbf{W}^k \overrightarrow{h_i}$$
|
||||
# for each head.
|
||||
# We do single linear transformation and then split it up for each head.
|
||||
g = self.linear(h).view(n_nodes, self.n_heads, self.n_hidden)
|
||||
|
||||
# #### Calculate attention score
|
||||
#
|
||||
# We calculate these for each head $k$. *We have omitted $\cdot^k$ for simplicity*.
|
||||
#
|
||||
# $$e_{ij} = a(\mathbf{W} \overrightarrow{h_i}, \mathbf{W} \overrightarrow{h_j}) =
|
||||
# a(\overrightarrow{g_i}, \overrightarrow{g_j})$$
|
||||
#
|
||||
# $e_{ij}$ is the attention score (importance) from node $j$ to node $i$.
|
||||
# We calculate this for each head.
|
||||
#
|
||||
# $a$ is the attention mechanism, that calculates the attention score.
|
||||
# The paper concatenates
|
||||
# $\overrightarrow{g_i}$, $\overrightarrow{g_j}$
|
||||
# and does a linear transformation with a weight vector $\mathbf{a} \in \mathbb{R}^{2 F'}$
|
||||
# followed by a $\text{LeakyReLU}$.
|
||||
#
|
||||
# $$e_{ij} = \text{LeakyReLU} \Big(
|
||||
# \mathbf{a}^\top \Big[
|
||||
# \overrightarrow{g_i} \Vert \overrightarrow{g_j}
|
||||
# \Big] \Big)$$
|
||||
|
||||
# First we calculate
|
||||
# $\Big[\overrightarrow{g_i} \Vert \overrightarrow{g_j} \Big]$
|
||||
# for all pairs of $i, j$.
|
||||
#
|
||||
# `g_repeat` gets
|
||||
# $$\{\overrightarrow{g_1}, \overrightarrow{g_2}, \dots, \overrightarrow{g_N},
|
||||
# \overrightarrow{g_1}, \overrightarrow{g_2}, \dots, \overrightarrow{g_N}, ...\}$$
|
||||
# where each node embedding is repeated `n_nodes` times.
|
||||
g_repeat = g.repeat(n_nodes, 1, 1)
|
||||
# `g_repeat_interleave` gets
|
||||
# $$\{\overrightarrow{g_1}, \overrightarrow{g_1}, \dots, \overrightarrow{g_1},
|
||||
# \overrightarrow{g_2}, \overrightarrow{g_2}, \dots, \overrightarrow{g_2}, ...\}$$
|
||||
# where each node embedding is repeated `n_nodes` times.
|
||||
g_repeat_interleave = g.repeat_interleave(n_nodes, dim=0)
|
||||
# Now we concatenate to get
|
||||
# $$\{\overrightarrow{g_1} \Vert \overrightarrow{g_1},
|
||||
# \overrightarrow{g_1} \Vert \overrightarrow{g_2},
|
||||
# \dots, \overrightarrow{g_1} \Vert \overrightarrow{g_N},
|
||||
# \overrightarrow{g_2} \Vert \overrightarrow{g_1},
|
||||
# \overrightarrow{g_2} \Vert \overrightarrow{g_2},
|
||||
# \dots, \overrightarrow{g_2} \Vert \overrightarrow{g_N}, ...\}$$
|
||||
g_concat = torch.cat([g_repeat_interleave, g_repeat], dim=-1)
|
||||
# Reshape so that `g_concat[i, j]` is $\overrightarrow{g_i} \Vert \overrightarrow{g_j}$
|
||||
g_concat = g_concat.view(n_nodes, n_nodes, self.n_heads, 2 * self.n_hidden)
|
||||
|
||||
# Calculate
|
||||
# $$e_{ij} = \text{LeakyReLU} \Big(
|
||||
# \mathbf{a}^\top \Big[
|
||||
# \overrightarrow{g_i} \Vert \overrightarrow{g_j}
|
||||
# \Big] \Big)$$
|
||||
# `e` is of shape `[n_nodes, n_nodes, n_heads, 1]`
|
||||
e = self.activation(self.attn(g_concat))
|
||||
# Remove the last dimension of size `1`
|
||||
e = e.squeeze(-1)
|
||||
|
||||
# The adjacency matrix should have shape
|
||||
# `[n_nodes, n_nodes, n_heads]` or`[n_nodes, n_nodes, 1]`
|
||||
assert adj_mat.shape[0] == 1 or adj_mat.shape[0] == n_nodes
|
||||
assert adj_mat.shape[1] == 1 or adj_mat.shape[1] == n_nodes
|
||||
assert adj_mat.shape[2] == 1 or adj_mat.shape[2] == self.n_heads
|
||||
# Mask $e_{ij}$ based on adjacency matrix.
|
||||
# $e_{ij}$ is set to $- \infty$ if there is no edge from $i$ to $j$.
|
||||
e = e.masked_fill(adj_mat == 0, float('-inf'))
|
||||
|
||||
# We then normalize attention scores (or coefficients)
|
||||
# $$\alpha_{ij} = \text{softmax}_j(e_{ij}) =
|
||||
# \frac{\exp(e_{ij})}{\sum_{k \in \mathcal{N}_i} \exp(e_{ik})}$$
|
||||
#
|
||||
# where $\mathcal{N}_i$ is the set of nodes connected to $i$.
|
||||
#
|
||||
# We do this by setting unconnected $e_{ij}$ to $- \infty$ which
|
||||
# makes $\exp(e_{ij}) \sim 0$ for unconnected pairs.
|
||||
a = self.softmax(e)
|
||||
|
||||
# Apply dropout regularization
|
||||
a = self.dropout(a)
|
||||
|
||||
# Calculate final output for each head
|
||||
# $$\overrightarrow{h'^k_i} = \sum_{j \in \mathcal{N}_i} \alpha^k_{ij} \overrightarrow{g^k_j}$$
|
||||
#
|
||||
# *Note:* The paper includes the final activation $\sigma$ in $\overrightarrow{h_i}$
|
||||
# We have omitted this from the Graph Attention Layer implementation
|
||||
# and use it on the GAT model to match with how other PyTorch modules are defined -
|
||||
# activation as a separate layer.
|
||||
attn_res = torch.einsum('ijh,jhf->ihf', a, g)
|
||||
|
||||
# Concatenate the heads
|
||||
if self.is_concat:
|
||||
# $$\overrightarrow{h'_i} = \Bigg\Vert_{k=1}^{K} \overrightarrow{h'^k_i}$$
|
||||
return attn_res.reshape(n_nodes, self.n_heads * self.n_hidden)
|
||||
# Take the mean of the heads
|
||||
else:
|
||||
# $$\overrightarrow{h'_i} = \frac{1}{K} \sum_{k=1}^{K} \overrightarrow{h'^k_i}$$
|
||||
return attn_res.mean(dim=1)
|
||||
@@ -0,0 +1,309 @@
|
||||
"""
|
||||
---
|
||||
title: Train a Graph Attention Network (GAT) on Cora dataset
|
||||
summary: >
|
||||
This trains is a Graph Attention Network (GAT) on Cora dataset
|
||||
---
|
||||
|
||||
# Train a Graph Attention Network (GAT) on Cora dataset
|
||||
"""
|
||||
|
||||
from typing import Dict
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from labml import lab, monit, tracker, experiment
|
||||
from labml.configs import BaseConfigs, option, calculate
|
||||
from labml.utils import download
|
||||
from labml_nn.helpers.device import DeviceConfigs
|
||||
from labml_nn.graphs.gat import GraphAttentionLayer
|
||||
from labml_nn.optimizers.configs import OptimizerConfigs
|
||||
|
||||
|
||||
class CoraDataset:
|
||||
"""
|
||||
## [Cora Dataset](https://linqs.soe.ucsc.edu/data)
|
||||
|
||||
Cora dataset is a dataset of research papers.
|
||||
For each paper we are given a binary feature vector that indicates the presence of words.
|
||||
Each paper is classified into one of 7 classes.
|
||||
The dataset also has the citation network.
|
||||
|
||||
The papers are the nodes of the graph and the edges are the citations.
|
||||
|
||||
The task is to classify the nodes to the 7 classes with feature vectors and
|
||||
citation network as input.
|
||||
"""
|
||||
# Labels for each node
|
||||
labels: torch.Tensor
|
||||
# Set of class names and an unique integer index
|
||||
classes: Dict[str, int]
|
||||
# Feature vectors for all nodes
|
||||
features: torch.Tensor
|
||||
# Adjacency matrix with the edge information.
|
||||
# `adj_mat[i][j]` is `True` if there is an edge from `i` to `j`.
|
||||
adj_mat: torch.Tensor
|
||||
|
||||
@staticmethod
|
||||
def _download():
|
||||
"""
|
||||
Download the dataset
|
||||
"""
|
||||
if not (lab.get_data_path() / 'cora').exists():
|
||||
download.download_file('https://linqs-data.soe.ucsc.edu/public/lbc/cora.tgz',
|
||||
lab.get_data_path() / 'cora.tgz')
|
||||
download.extract_tar(lab.get_data_path() / 'cora.tgz', lab.get_data_path())
|
||||
|
||||
def __init__(self, include_edges: bool = True):
|
||||
"""
|
||||
Load the dataset
|
||||
"""
|
||||
|
||||
# Whether to include edges.
|
||||
# This is test how much accuracy is lost if we ignore the citation network.
|
||||
self.include_edges = include_edges
|
||||
|
||||
# Download dataset
|
||||
self._download()
|
||||
|
||||
# Read the paper ids, feature vectors, and labels
|
||||
with monit.section('Read content file'):
|
||||
content = np.genfromtxt(str(lab.get_data_path() / 'cora/cora.content'), dtype=np.dtype(str))
|
||||
# Load the citations, it's a list of pairs of integers.
|
||||
with monit.section('Read citations file'):
|
||||
citations = np.genfromtxt(str(lab.get_data_path() / 'cora/cora.cites'), dtype=np.int32)
|
||||
|
||||
# Get the feature vectors
|
||||
features = torch.tensor(np.array(content[:, 1:-1], dtype=np.float32))
|
||||
# Normalize the feature vectors
|
||||
self.features = features / features.sum(dim=1, keepdim=True)
|
||||
|
||||
# Get the class names and assign an unique integer to each of them
|
||||
self.classes = {s: i for i, s in enumerate(set(content[:, -1]))}
|
||||
# Get the labels as those integers
|
||||
self.labels = torch.tensor([self.classes[i] for i in content[:, -1]], dtype=torch.long)
|
||||
|
||||
# Get the paper ids
|
||||
paper_ids = np.array(content[:, 0], dtype=np.int32)
|
||||
# Map of paper id to index
|
||||
ids_to_idx = {id_: i for i, id_ in enumerate(paper_ids)}
|
||||
|
||||
# Empty adjacency matrix - an identity matrix
|
||||
self.adj_mat = torch.eye(len(self.labels), dtype=torch.bool)
|
||||
|
||||
# Mark the citations in the adjacency matrix
|
||||
if self.include_edges:
|
||||
for e in citations:
|
||||
# The pair of paper indexes
|
||||
e1, e2 = ids_to_idx[e[0]], ids_to_idx[e[1]]
|
||||
# We build a symmetrical graph, where if paper $i$ referenced
|
||||
# paper $j$ we place an adge from $i$ to $j$ as well as an edge
|
||||
# from $j$ to $i$.
|
||||
self.adj_mat[e1][e2] = True
|
||||
self.adj_mat[e2][e1] = True
|
||||
|
||||
|
||||
class GAT(nn.Module):
|
||||
"""
|
||||
## Graph Attention Network (GAT)
|
||||
|
||||
This graph attention network has two [graph attention layers](index.html).
|
||||
"""
|
||||
|
||||
def __init__(self, in_features: int, n_hidden: int, n_classes: int, n_heads: int, dropout: float):
|
||||
"""
|
||||
* `in_features` is the number of features per node
|
||||
* `n_hidden` is the number of features in the first graph attention layer
|
||||
* `n_classes` is the number of classes
|
||||
* `n_heads` is the number of heads in the graph attention layers
|
||||
* `dropout` is the dropout probability
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
# First graph attention layer where we concatenate the heads
|
||||
self.layer1 = GraphAttentionLayer(in_features, n_hidden, n_heads, is_concat=True, dropout=dropout)
|
||||
# Activation function after first graph attention layer
|
||||
self.activation = nn.ELU()
|
||||
# Final graph attention layer where we average the heads
|
||||
self.output = GraphAttentionLayer(n_hidden, n_classes, 1, is_concat=False, dropout=dropout)
|
||||
# Dropout
|
||||
self.dropout = nn.Dropout(dropout)
|
||||
|
||||
def forward(self, x: torch.Tensor, adj_mat: torch.Tensor):
|
||||
"""
|
||||
* `x` is the features vectors of shape `[n_nodes, in_features]`
|
||||
* `adj_mat` is the adjacency matrix of the form
|
||||
`[n_nodes, n_nodes, n_heads]` or `[n_nodes, n_nodes, 1]`
|
||||
"""
|
||||
# Apply dropout to the input
|
||||
x = self.dropout(x)
|
||||
# First graph attention layer
|
||||
x = self.layer1(x, adj_mat)
|
||||
# Activation function
|
||||
x = self.activation(x)
|
||||
# Dropout
|
||||
x = self.dropout(x)
|
||||
# Output layer (without activation) for logits
|
||||
return self.output(x, adj_mat)
|
||||
|
||||
|
||||
def accuracy(output: torch.Tensor, labels: torch.Tensor):
|
||||
"""
|
||||
A simple function to calculate the accuracy
|
||||
"""
|
||||
return output.argmax(dim=-1).eq(labels).sum().item() / len(labels)
|
||||
|
||||
|
||||
class Configs(BaseConfigs):
|
||||
"""
|
||||
## Configurations
|
||||
"""
|
||||
|
||||
# Model
|
||||
model: GAT
|
||||
# Number of nodes to train on
|
||||
training_samples: int = 500
|
||||
# Number of features per node in the input
|
||||
in_features: int
|
||||
# Number of features in the first graph attention layer
|
||||
n_hidden: int = 64
|
||||
# Number of heads
|
||||
n_heads: int = 8
|
||||
# Number of classes for classification
|
||||
n_classes: int
|
||||
# Dropout probability
|
||||
dropout: float = 0.6
|
||||
# Whether to include the citation network
|
||||
include_edges: bool = True
|
||||
# Dataset
|
||||
dataset: CoraDataset
|
||||
# Number of training iterations
|
||||
epochs: int = 1_000
|
||||
# Loss function
|
||||
loss_func = nn.CrossEntropyLoss()
|
||||
# Device to train on
|
||||
#
|
||||
# This creates configs for device, so that
|
||||
# we can change the device by passing a config value
|
||||
device: torch.device = DeviceConfigs()
|
||||
# Optimizer
|
||||
optimizer: torch.optim.Adam
|
||||
|
||||
def run(self):
|
||||
"""
|
||||
### Training loop
|
||||
|
||||
We do full batch training since the dataset is small.
|
||||
If we were to sample and train we will have to sample a set of
|
||||
nodes for each training step along with the edges that span
|
||||
across those selected nodes.
|
||||
"""
|
||||
# Move the feature vectors to the device
|
||||
features = self.dataset.features.to(self.device)
|
||||
# Move the labels to the device
|
||||
labels = self.dataset.labels.to(self.device)
|
||||
# Move the adjacency matrix to the device
|
||||
edges_adj = self.dataset.adj_mat.to(self.device)
|
||||
# Add an empty third dimension for the heads
|
||||
edges_adj = edges_adj.unsqueeze(-1)
|
||||
|
||||
# Random indexes
|
||||
idx_rand = torch.randperm(len(labels))
|
||||
# Nodes for training
|
||||
idx_train = idx_rand[:self.training_samples]
|
||||
# Nodes for validation
|
||||
idx_valid = idx_rand[self.training_samples:]
|
||||
|
||||
# Training loop
|
||||
for epoch in monit.loop(self.epochs):
|
||||
# Set the model to training mode
|
||||
self.model.train()
|
||||
# Make all the gradients zero
|
||||
self.optimizer.zero_grad()
|
||||
# Evaluate the model
|
||||
output = self.model(features, edges_adj)
|
||||
# Get the loss for training nodes
|
||||
loss = self.loss_func(output[idx_train], labels[idx_train])
|
||||
# Calculate gradients
|
||||
loss.backward()
|
||||
# Take optimization step
|
||||
self.optimizer.step()
|
||||
# Log the loss
|
||||
tracker.add('loss.train', loss)
|
||||
# Log the accuracy
|
||||
tracker.add('accuracy.train', accuracy(output[idx_train], labels[idx_train]))
|
||||
|
||||
# Set mode to evaluation mode for validation
|
||||
self.model.eval()
|
||||
|
||||
# No need to compute gradients
|
||||
with torch.no_grad():
|
||||
# Evaluate the model again
|
||||
output = self.model(features, edges_adj)
|
||||
# Calculate the loss for validation nodes
|
||||
loss = self.loss_func(output[idx_valid], labels[idx_valid])
|
||||
# Log the loss
|
||||
tracker.add('loss.valid', loss)
|
||||
# Log the accuracy
|
||||
tracker.add('accuracy.valid', accuracy(output[idx_valid], labels[idx_valid]))
|
||||
|
||||
# Save logs
|
||||
tracker.save()
|
||||
|
||||
|
||||
@option(Configs.dataset)
|
||||
def cora_dataset(c: Configs):
|
||||
"""
|
||||
Create Cora dataset
|
||||
"""
|
||||
return CoraDataset(c.include_edges)
|
||||
|
||||
|
||||
# Get the number of classes
|
||||
calculate(Configs.n_classes, lambda c: len(c.dataset.classes))
|
||||
# Number of features in the input
|
||||
calculate(Configs.in_features, lambda c: c.dataset.features.shape[1])
|
||||
|
||||
|
||||
@option(Configs.model)
|
||||
def gat_model(c: Configs):
|
||||
"""
|
||||
Create GAT model
|
||||
"""
|
||||
return GAT(c.in_features, c.n_hidden, c.n_classes, c.n_heads, c.dropout).to(c.device)
|
||||
|
||||
|
||||
@option(Configs.optimizer)
|
||||
def _optimizer(c: Configs):
|
||||
"""
|
||||
Create configurable optimizer
|
||||
"""
|
||||
opt_conf = OptimizerConfigs()
|
||||
opt_conf.parameters = c.model.parameters()
|
||||
return opt_conf
|
||||
|
||||
|
||||
def main():
|
||||
# Create configurations
|
||||
conf = Configs()
|
||||
# Create an experiment
|
||||
experiment.create(name='gat')
|
||||
# Calculate configurations.
|
||||
experiment.configs(conf, {
|
||||
# Adam optimizer
|
||||
'optimizer.optimizer': 'Adam',
|
||||
'optimizer.learning_rate': 5e-3,
|
||||
'optimizer.weight_decay': 5e-4,
|
||||
})
|
||||
|
||||
# Start and watch the experiment
|
||||
with experiment.start():
|
||||
# Run the training
|
||||
conf.run()
|
||||
|
||||
|
||||
#
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,18 @@
|
||||
# [Graph Attention Networks (GAT)](https://nn.labml.ai/graphs/gat/index.html)
|
||||
|
||||
This is a [PyTorch](https://pytorch.org) implementation of the paper
|
||||
[Graph Attention Networks](https://arxiv.org/abs/1710.10903).
|
||||
|
||||
GATs work on graph data.
|
||||
A graph consists of nodes and edges connecting nodes.
|
||||
For example, in Cora dataset the nodes are research papers and the edges are citations that
|
||||
connect the papers.
|
||||
|
||||
GAT uses masked self-attention, kind of similar to [transformers](https://nn.labml.ai/transformers/mha.html).
|
||||
GAT consists of graph attention layers stacked on top of each other.
|
||||
Each graph attention layer gets node embeddings as inputs and outputs transformed embeddings.
|
||||
The node embeddings pay attention to the embeddings of other nodes it's connected to.
|
||||
The details of graph attention layers are included alongside the implementation.
|
||||
|
||||
Here is [the training code](https://nn.labml.ai/graphs/gat/experiment.html) for training
|
||||
a two-layer GAT on Cora dataset.
|
||||
@@ -0,0 +1,236 @@
|
||||
"""
|
||||
---
|
||||
title: Graph Attention Networks v2 (GATv2)
|
||||
summary: >
|
||||
A PyTorch implementation/tutorial of Graph Attention Networks v2.
|
||||
---
|
||||
# Graph Attention Networks v2 (GATv2)
|
||||
This is a [PyTorch](https://pytorch.org) implementation of the GATv2 operator from the paper
|
||||
[How Attentive are Graph Attention Networks?](https://arxiv.org/abs/2105.14491).
|
||||
|
||||
GATv2s work on graph data similar to [GAT](../gat/index.html).
|
||||
A graph consists of nodes and edges connecting nodes.
|
||||
For example, in Cora dataset the nodes are research papers and the edges are citations that
|
||||
connect the papers.
|
||||
|
||||
The GATv2 operator fixes the static attention problem of the standard [GAT](../gat/index.html).
|
||||
Static attention is when the attention to the key nodes has the same rank (order) for any query node.
|
||||
[GAT](../gat/index.html) computes attention from query node $i$ to key node $j$ as,
|
||||
|
||||
\begin{align}
|
||||
e_{ij} &= \text{LeakyReLU} \Big(\mathbf{a}^\top \Big[
|
||||
\mathbf{W} \overrightarrow{h_i} \Vert \mathbf{W} \overrightarrow{h_j}
|
||||
\Big] \Big) \\
|
||||
&=
|
||||
\text{LeakyReLU} \Big(\mathbf{a}_1^\top \mathbf{W} \overrightarrow{h_i} +
|
||||
\mathbf{a}_2^\top \mathbf{W} \overrightarrow{h_j}
|
||||
\Big)
|
||||
\end{align}
|
||||
|
||||
Note that for any query node $i$, the attention rank ($argsort$) of keys depends only
|
||||
on $\mathbf{a}_2^\top \mathbf{W} \overrightarrow{h_j}$.
|
||||
Therefore the attention rank of keys remains the same (*static*) for all queries.
|
||||
|
||||
GATv2 allows dynamic attention by changing the attention mechanism,
|
||||
|
||||
\begin{align}
|
||||
e_{ij} &= \mathbf{a}^\top \text{LeakyReLU} \Big( \mathbf{W} \Big[
|
||||
\overrightarrow{h_i} \Vert \overrightarrow{h_j}
|
||||
\Big] \Big) \\
|
||||
&= \mathbf{a}^\top \text{LeakyReLU} \Big(
|
||||
\mathbf{W}_l \overrightarrow{h_i} + \mathbf{W}_r \overrightarrow{h_j}
|
||||
\Big)
|
||||
\end{align}
|
||||
|
||||
The paper shows that GATs static attention mechanism fails on some graph problems
|
||||
with a synthetic dictionary lookup dataset.
|
||||
It's a fully connected bipartite graph where one set of nodes (query nodes)
|
||||
have a key associated with it
|
||||
and the other set of nodes have both a key and a value associated with it.
|
||||
The goal is to predict the values of query nodes.
|
||||
GAT fails on this task because of its limited static attention.
|
||||
|
||||
Here is [the training code](experiment.html) for training
|
||||
a two-layer GATv2 on Cora dataset.
|
||||
"""
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
|
||||
|
||||
class GraphAttentionV2Layer(nn.Module):
|
||||
"""
|
||||
## Graph attention v2 layer
|
||||
This is a single graph attention v2 layer.
|
||||
A GATv2 is made up of multiple such layers.
|
||||
It takes
|
||||
$$\mathbf{h} = \{ \overrightarrow{h_1}, \overrightarrow{h_2}, \dots, \overrightarrow{h_N} \}$$,
|
||||
where $\overrightarrow{h_i} \in \mathbb{R}^F$ as input
|
||||
and outputs
|
||||
$$\mathbf{h'} = \{ \overrightarrow{h'_1}, \overrightarrow{h'_2}, \dots, \overrightarrow{h'_N} \}$$,
|
||||
where $\overrightarrow{h'_i} \in \mathbb{R}^{F'}$.
|
||||
"""
|
||||
|
||||
def __init__(self, in_features: int, out_features: int, n_heads: int,
|
||||
is_concat: bool = True,
|
||||
dropout: float = 0.6,
|
||||
leaky_relu_negative_slope: float = 0.2,
|
||||
share_weights: bool = False):
|
||||
"""
|
||||
* `in_features`, $F$, is the number of input features per node
|
||||
* `out_features`, $F'$, is the number of output features per node
|
||||
* `n_heads`, $K$, is the number of attention heads
|
||||
* `is_concat` whether the multi-head results should be concatenated or averaged
|
||||
* `dropout` is the dropout probability
|
||||
* `leaky_relu_negative_slope` is the negative slope for leaky relu activation
|
||||
* `share_weights` if set to `True`, the same matrix will be applied to the source and the target node of every edge
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
self.is_concat = is_concat
|
||||
self.n_heads = n_heads
|
||||
self.share_weights = share_weights
|
||||
|
||||
# Calculate the number of dimensions per head
|
||||
if is_concat:
|
||||
assert out_features % n_heads == 0
|
||||
# If we are concatenating the multiple heads
|
||||
self.n_hidden = out_features // n_heads
|
||||
else:
|
||||
# If we are averaging the multiple heads
|
||||
self.n_hidden = out_features
|
||||
|
||||
# Linear layer for initial source transformation;
|
||||
# i.e. to transform the source node embeddings before self-attention
|
||||
self.linear_l = nn.Linear(in_features, self.n_hidden * n_heads, bias=False)
|
||||
# If `share_weights` is `True` the same linear layer is used for the target nodes
|
||||
if share_weights:
|
||||
self.linear_r = self.linear_l
|
||||
else:
|
||||
self.linear_r = nn.Linear(in_features, self.n_hidden * n_heads, bias=False)
|
||||
# Linear layer to compute attention score $e_{ij}$
|
||||
self.attn = nn.Linear(self.n_hidden, 1, bias=False)
|
||||
# The activation for attention score $e_{ij}$
|
||||
self.activation = nn.LeakyReLU(negative_slope=leaky_relu_negative_slope)
|
||||
# Softmax to compute attention $\alpha_{ij}$
|
||||
self.softmax = nn.Softmax(dim=1)
|
||||
# Dropout layer to be applied for attention
|
||||
self.dropout = nn.Dropout(dropout)
|
||||
|
||||
def forward(self, h: torch.Tensor, adj_mat: torch.Tensor):
|
||||
"""
|
||||
* `h`, $\mathbf{h}$ is the input node embeddings of shape `[n_nodes, in_features]`.
|
||||
* `adj_mat` is the adjacency matrix of shape `[n_nodes, n_nodes, n_heads]`.
|
||||
We use shape `[n_nodes, n_nodes, 1]` since the adjacency is the same for each head.
|
||||
Adjacency matrix represent the edges (or connections) among nodes.
|
||||
`adj_mat[i][j]` is `True` if there is an edge from node `i` to node `j`.
|
||||
"""
|
||||
|
||||
# Number of nodes
|
||||
n_nodes = h.shape[0]
|
||||
# The initial transformations,
|
||||
# $$\overrightarrow{{g_l}^k_i} = \mathbf{W_l}^k \overrightarrow{h_i}$$
|
||||
# $$\overrightarrow{{g_r}^k_i} = \mathbf{W_r}^k \overrightarrow{h_i}$$
|
||||
# for each head.
|
||||
# We do two linear transformations and then split it up for each head.
|
||||
g_l = self.linear_l(h).view(n_nodes, self.n_heads, self.n_hidden)
|
||||
g_r = self.linear_r(h).view(n_nodes, self.n_heads, self.n_hidden)
|
||||
|
||||
# #### Calculate attention score
|
||||
#
|
||||
# We calculate these for each head $k$. *We have omitted $\cdot^k$ for simplicity*.
|
||||
#
|
||||
# $$e_{ij} = a(\mathbf{W_l} \overrightarrow{h_i}, \mathbf{W_r} \overrightarrow{h_j}) =
|
||||
# a(\overrightarrow{{g_l}_i}, \overrightarrow{{g_r}_j})$$
|
||||
#
|
||||
# $e_{ij}$ is the attention score (importance) from node $j$ to node $i$.
|
||||
# We calculate this for each head.
|
||||
#
|
||||
# $a$ is the attention mechanism, that calculates the attention score.
|
||||
# The paper sums
|
||||
# $\overrightarrow{{g_l}_i}$, $\overrightarrow{{g_r}_j}$
|
||||
# followed by a $\text{LeakyReLU}$
|
||||
# and does a linear transformation with a weight vector $\mathbf{a} \in \mathbb{R}^{F'}$
|
||||
#
|
||||
#
|
||||
# $$e_{ij} = \mathbf{a}^\top \text{LeakyReLU} \Big(
|
||||
# \Big[
|
||||
# \overrightarrow{{g_l}_i} + \overrightarrow{{g_r}_j}
|
||||
# \Big] \Big)$$
|
||||
# Note: The paper desrcibes $e_{ij}$ as
|
||||
# $$e_{ij} = \mathbf{a}^\top \text{LeakyReLU} \Big( \mathbf{W}
|
||||
# \Big[
|
||||
# \overrightarrow{h_i} \Vert \overrightarrow{h_j}
|
||||
# \Big] \Big)$$
|
||||
# which is equivalent to the definition we use here.
|
||||
|
||||
# First we calculate
|
||||
# $\Big[\overrightarrow{{g_l}_i} + \overrightarrow{{g_r}_j} \Big]$
|
||||
# for all pairs of $i, j$.
|
||||
#
|
||||
# `g_l_repeat` gets
|
||||
# $$\{\overrightarrow{{g_l}_1}, \overrightarrow{{g_l}_2}, \dots, \overrightarrow{{g_l}_N},
|
||||
# \overrightarrow{{g_l}_1}, \overrightarrow{{g_l}_2}, \dots, \overrightarrow{{g_l}_N}, ...\}$$
|
||||
# where each node embedding is repeated `n_nodes` times.
|
||||
g_l_repeat = g_l.repeat(n_nodes, 1, 1)
|
||||
# `g_r_repeat_interleave` gets
|
||||
# $$\{\overrightarrow{{g_r}_1}, \overrightarrow{{g_r}_1}, \dots, \overrightarrow{{g_r}_1},
|
||||
# \overrightarrow{{g_r}_2}, \overrightarrow{{g_r}_2}, \dots, \overrightarrow{{g_r}_2}, ...\}$$
|
||||
# where each node embedding is repeated `n_nodes` times.
|
||||
g_r_repeat_interleave = g_r.repeat_interleave(n_nodes, dim=0)
|
||||
# Now we add the two tensors to get
|
||||
# $$\{\overrightarrow{{g_l}_1} + \overrightarrow{{g_r}_1},
|
||||
# \overrightarrow{{g_l}_1} + \overrightarrow{{g_r}_2},
|
||||
# \dots, \overrightarrow{{g_l}_1} +\overrightarrow{{g_r}_N},
|
||||
# \overrightarrow{{g_l}_2} + \overrightarrow{{g_r}_1},
|
||||
# \overrightarrow{{g_l}_2} + \overrightarrow{{g_r}_2},
|
||||
# \dots, \overrightarrow{{g_l}_2} + \overrightarrow{{g_r}_N}, ...\}$$
|
||||
g_sum = g_l_repeat + g_r_repeat_interleave
|
||||
# Reshape so that `g_sum[i, j]` is $\overrightarrow{{g_l}_i} + \overrightarrow{{g_r}_j}$
|
||||
g_sum = g_sum.view(n_nodes, n_nodes, self.n_heads, self.n_hidden)
|
||||
|
||||
# Calculate
|
||||
# $$e_{ij} = \mathbf{a}^\top \text{LeakyReLU} \Big(
|
||||
# \Big[
|
||||
# \overrightarrow{{g_l}_i} + \overrightarrow{{g_r}_j}
|
||||
# \Big] \Big)$$
|
||||
# `e` is of shape `[n_nodes, n_nodes, n_heads, 1]`
|
||||
e = self.attn(self.activation(g_sum))
|
||||
# Remove the last dimension of size `1`
|
||||
e = e.squeeze(-1)
|
||||
|
||||
# The adjacency matrix should have shape
|
||||
# `[n_nodes, n_nodes, n_heads]` or`[n_nodes, n_nodes, 1]`
|
||||
assert adj_mat.shape[0] == 1 or adj_mat.shape[0] == n_nodes
|
||||
assert adj_mat.shape[1] == 1 or adj_mat.shape[1] == n_nodes
|
||||
assert adj_mat.shape[2] == 1 or adj_mat.shape[2] == self.n_heads
|
||||
# Mask $e_{ij}$ based on adjacency matrix.
|
||||
# $e_{ij}$ is set to $- \infty$ if there is no edge from $i$ to $j$.
|
||||
e = e.masked_fill(adj_mat == 0, float('-inf'))
|
||||
|
||||
# We then normalize attention scores (or coefficients)
|
||||
# $$\alpha_{ij} = \text{softmax}_j(e_{ij}) =
|
||||
# \frac{\exp(e_{ij})}{\sum_{j' \in \mathcal{N}_i} \exp(e_{ij'})}$$
|
||||
#
|
||||
# where $\mathcal{N}_i$ is the set of nodes connected to $i$.
|
||||
#
|
||||
# We do this by setting unconnected $e_{ij}$ to $- \infty$ which
|
||||
# makes $\exp(e_{ij}) \sim 0$ for unconnected pairs.
|
||||
a = self.softmax(e)
|
||||
|
||||
# Apply dropout regularization
|
||||
a = self.dropout(a)
|
||||
|
||||
# Calculate final output for each head
|
||||
# $$\overrightarrow{h'^k_i} = \sum_{j \in \mathcal{N}_i} \alpha^k_{ij} \overrightarrow{{g_r}_{j,k}}$$
|
||||
attn_res = torch.einsum('ijh,jhf->ihf', a, g_r)
|
||||
|
||||
# Concatenate the heads
|
||||
if self.is_concat:
|
||||
# $$\overrightarrow{h'_i} = \Bigg\Vert_{k=1}^{K} \overrightarrow{h'^k_i}$$
|
||||
return attn_res.reshape(n_nodes, self.n_heads * self.n_hidden)
|
||||
# Take the mean of the heads
|
||||
else:
|
||||
# $$\overrightarrow{h'_i} = \frac{1}{K} \sum_{k=1}^{K} \overrightarrow{h'^k_i}$$
|
||||
return attn_res.mean(dim=1)
|
||||
@@ -0,0 +1,113 @@
|
||||
"""
|
||||
---
|
||||
title: Train a Graph Attention Network v2 (GATv2) on Cora dataset
|
||||
summary: >
|
||||
This trains is a Graph Attention Network v2 (GATv2) on Cora dataset
|
||||
---
|
||||
|
||||
# Train a Graph Attention Network v2 (GATv2) on Cora dataset
|
||||
"""
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from labml import experiment
|
||||
from labml.configs import option
|
||||
from labml_nn.graphs.gat.experiment import Configs as GATConfigs
|
||||
from labml_nn.graphs.gatv2 import GraphAttentionV2Layer
|
||||
|
||||
|
||||
class GATv2(nn.Module):
|
||||
"""
|
||||
## Graph Attention Network v2 (GATv2)
|
||||
|
||||
This graph attention network has two [graph attention layers](index.html).
|
||||
"""
|
||||
|
||||
def __init__(self, in_features: int, n_hidden: int, n_classes: int, n_heads: int, dropout: float,
|
||||
share_weights: bool = True):
|
||||
"""
|
||||
* `in_features` is the number of features per node
|
||||
* `n_hidden` is the number of features in the first graph attention layer
|
||||
* `n_classes` is the number of classes
|
||||
* `n_heads` is the number of heads in the graph attention layers
|
||||
* `dropout` is the dropout probability
|
||||
* `share_weights` if set to True, the same matrix will be applied to the source and the target node of every edge
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
# First graph attention layer where we concatenate the heads
|
||||
self.layer1 = GraphAttentionV2Layer(in_features, n_hidden, n_heads,
|
||||
is_concat=True, dropout=dropout, share_weights=share_weights)
|
||||
# Activation function after first graph attention layer
|
||||
self.activation = nn.ELU()
|
||||
# Final graph attention layer where we average the heads
|
||||
self.output = GraphAttentionV2Layer(n_hidden, n_classes, 1,
|
||||
is_concat=False, dropout=dropout, share_weights=share_weights)
|
||||
# Dropout
|
||||
self.dropout = nn.Dropout(dropout)
|
||||
|
||||
def forward(self, x: torch.Tensor, adj_mat: torch.Tensor):
|
||||
"""
|
||||
* `x` is the features vectors of shape `[n_nodes, in_features]`
|
||||
* `adj_mat` is the adjacency matrix of the form
|
||||
`[n_nodes, n_nodes, n_heads]` or `[n_nodes, n_nodes, 1]`
|
||||
"""
|
||||
# Apply dropout to the input
|
||||
x = self.dropout(x)
|
||||
# First graph attention layer
|
||||
x = self.layer1(x, adj_mat)
|
||||
# Activation function
|
||||
x = self.activation(x)
|
||||
# Dropout
|
||||
x = self.dropout(x)
|
||||
# Output layer (without activation) for logits
|
||||
return self.output(x, adj_mat)
|
||||
|
||||
|
||||
class Configs(GATConfigs):
|
||||
"""
|
||||
## Configurations
|
||||
|
||||
Since the experiment is same as [GAT experiment](../gat/experiment.html) but with
|
||||
[GATv2 model](index.html) we extend the same configs and change the model.
|
||||
"""
|
||||
|
||||
# Whether to share weights for source and target nodes of edges
|
||||
share_weights: bool = False
|
||||
# Set the model
|
||||
model: GATv2 = 'gat_v2_model'
|
||||
|
||||
|
||||
@option(Configs.model)
|
||||
def gat_v2_model(c: Configs):
|
||||
"""
|
||||
Create GATv2 model
|
||||
"""
|
||||
return GATv2(c.in_features, c.n_hidden, c.n_classes, c.n_heads, c.dropout, c.share_weights).to(c.device)
|
||||
|
||||
|
||||
def main():
|
||||
# Create configurations
|
||||
conf = Configs()
|
||||
# Create an experiment
|
||||
experiment.create(name='gatv2')
|
||||
# Calculate configurations.
|
||||
experiment.configs(conf, {
|
||||
# Adam optimizer
|
||||
'optimizer.optimizer': 'Adam',
|
||||
'optimizer.learning_rate': 5e-3,
|
||||
'optimizer.weight_decay': 5e-4,
|
||||
|
||||
'dropout': 0.7,
|
||||
})
|
||||
|
||||
# Start and watch the experiment
|
||||
with experiment.start():
|
||||
# Run the training
|
||||
conf.run()
|
||||
|
||||
|
||||
#
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,17 @@
|
||||
# [Graph Attention Networks v2 (GATv2)](https://nn.labml.ai/graphs/gatv2/index.html)
|
||||
|
||||
This is a [PyTorch](https://pytorch.org) implementation of the GATv2 operator from the paper
|
||||
[How Attentive are Graph Attention Networks?](https://arxiv.org/abs/2105.14491).
|
||||
|
||||
GATv2s work on graph data.
|
||||
A graph consists of nodes and edges connecting nodes.
|
||||
For example, in Cora dataset the nodes are research papers and the edges are citations that
|
||||
connect the papers.
|
||||
|
||||
The GATv2 operator fixes the static attention problem of the standard GAT:
|
||||
since the linear layers in the standard GAT are applied right after each other, the ranking
|
||||
of attended nodes is unconditioned on the query node.
|
||||
In contrast, in GATv2, every node can attend to any other node.
|
||||
|
||||
Here is [the training code](https://nn.labml.ai/graphs/gatv2/experiment.html) for training
|
||||
a two-layer GATv2 on Cora dataset.
|
||||
@@ -0,0 +1,322 @@
|
||||
import random
|
||||
from pathlib import PurePath, Path
|
||||
from typing import List, Callable, Dict, Optional
|
||||
|
||||
from torchvision import datasets, transforms
|
||||
|
||||
import torch
|
||||
from labml import lab
|
||||
from labml import monit
|
||||
from labml.configs import BaseConfigs
|
||||
from labml.configs import aggregate, option
|
||||
from labml.utils.download import download_file
|
||||
from torch.utils.data import DataLoader
|
||||
from torch.utils.data import IterableDataset, Dataset
|
||||
|
||||
|
||||
def _mnist_dataset(is_train, transform):
|
||||
return datasets.MNIST(str(lab.get_data_path()),
|
||||
train=is_train,
|
||||
download=True,
|
||||
transform=transform)
|
||||
|
||||
|
||||
class MNISTConfigs(BaseConfigs):
|
||||
"""
|
||||
Configurable MNIST data set.
|
||||
|
||||
Arguments:
|
||||
dataset_name (str): name of the data set, ``MNIST``
|
||||
dataset_transforms (torchvision.transforms.Compose): image transformations
|
||||
train_dataset (torchvision.datasets.MNIST): training dataset
|
||||
valid_dataset (torchvision.datasets.MNIST): validation dataset
|
||||
|
||||
train_loader (torch.utils.data.DataLoader): training data loader
|
||||
valid_loader (torch.utils.data.DataLoader): validation data loader
|
||||
|
||||
train_batch_size (int): training batch size
|
||||
valid_batch_size (int): validation batch size
|
||||
|
||||
train_loader_shuffle (bool): whether to shuffle training data
|
||||
valid_loader_shuffle (bool): whether to shuffle validation data
|
||||
"""
|
||||
|
||||
dataset_name: str = 'MNIST'
|
||||
dataset_transforms: transforms.Compose
|
||||
train_dataset: datasets.MNIST
|
||||
valid_dataset: datasets.MNIST
|
||||
|
||||
train_loader: DataLoader
|
||||
valid_loader: DataLoader
|
||||
|
||||
train_batch_size: int = 64
|
||||
valid_batch_size: int = 1024
|
||||
|
||||
train_loader_shuffle: bool = True
|
||||
valid_loader_shuffle: bool = False
|
||||
|
||||
|
||||
@option(MNISTConfigs.dataset_transforms)
|
||||
def mnist_transforms():
|
||||
return transforms.Compose([
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize((0.1307,), (0.3081,))
|
||||
])
|
||||
|
||||
|
||||
@option(MNISTConfigs.train_dataset)
|
||||
def mnist_train_dataset(c: MNISTConfigs):
|
||||
return _mnist_dataset(True, c.dataset_transforms)
|
||||
|
||||
|
||||
@option(MNISTConfigs.valid_dataset)
|
||||
def mnist_valid_dataset(c: MNISTConfigs):
|
||||
return _mnist_dataset(False, c.dataset_transforms)
|
||||
|
||||
|
||||
@option(MNISTConfigs.train_loader)
|
||||
def mnist_train_loader(c: MNISTConfigs):
|
||||
return DataLoader(c.train_dataset,
|
||||
batch_size=c.train_batch_size,
|
||||
shuffle=c.train_loader_shuffle)
|
||||
|
||||
|
||||
@option(MNISTConfigs.valid_loader)
|
||||
def mnist_valid_loader(c: MNISTConfigs):
|
||||
return DataLoader(c.valid_dataset,
|
||||
batch_size=c.valid_batch_size,
|
||||
shuffle=c.valid_loader_shuffle)
|
||||
|
||||
|
||||
aggregate(MNISTConfigs.dataset_name, 'MNIST',
|
||||
(MNISTConfigs.dataset_transforms, 'mnist_transforms'),
|
||||
(MNISTConfigs.train_dataset, 'mnist_train_dataset'),
|
||||
(MNISTConfigs.valid_dataset, 'mnist_valid_dataset'),
|
||||
(MNISTConfigs.train_loader, 'mnist_train_loader'),
|
||||
(MNISTConfigs.valid_loader, 'mnist_valid_loader'))
|
||||
|
||||
|
||||
def _cifar_dataset(is_train, transform):
|
||||
return datasets.CIFAR10(str(lab.get_data_path()),
|
||||
train=is_train,
|
||||
download=True,
|
||||
transform=transform)
|
||||
|
||||
|
||||
class CIFAR10Configs(BaseConfigs):
|
||||
"""
|
||||
Configurable CIFAR 10 data set.
|
||||
|
||||
Arguments:
|
||||
dataset_name (str): name of the data set, ``CIFAR10``
|
||||
dataset_transforms (torchvision.transforms.Compose): image transformations
|
||||
train_dataset (torchvision.datasets.CIFAR10): training dataset
|
||||
valid_dataset (torchvision.datasets.CIFAR10): validation dataset
|
||||
|
||||
train_loader (torch.utils.data.DataLoader): training data loader
|
||||
valid_loader (torch.utils.data.DataLoader): validation data loader
|
||||
|
||||
train_batch_size (int): training batch size
|
||||
valid_batch_size (int): validation batch size
|
||||
|
||||
train_loader_shuffle (bool): whether to shuffle training data
|
||||
valid_loader_shuffle (bool): whether to shuffle validation data
|
||||
"""
|
||||
dataset_name: str = 'CIFAR10'
|
||||
dataset_transforms: transforms.Compose
|
||||
train_dataset: datasets.CIFAR10
|
||||
valid_dataset: datasets.CIFAR10
|
||||
|
||||
train_loader: DataLoader
|
||||
valid_loader: DataLoader
|
||||
|
||||
train_batch_size: int = 64
|
||||
valid_batch_size: int = 1024
|
||||
|
||||
train_loader_shuffle: bool = True
|
||||
valid_loader_shuffle: bool = False
|
||||
|
||||
|
||||
@CIFAR10Configs.calc(CIFAR10Configs.dataset_transforms)
|
||||
def cifar10_transforms():
|
||||
return transforms.Compose([
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
|
||||
])
|
||||
|
||||
|
||||
@CIFAR10Configs.calc(CIFAR10Configs.train_dataset)
|
||||
def cifar10_train_dataset(c: CIFAR10Configs):
|
||||
return _cifar_dataset(True, c.dataset_transforms)
|
||||
|
||||
|
||||
@CIFAR10Configs.calc(CIFAR10Configs.valid_dataset)
|
||||
def cifar10_valid_dataset(c: CIFAR10Configs):
|
||||
return _cifar_dataset(False, c.dataset_transforms)
|
||||
|
||||
|
||||
@CIFAR10Configs.calc(CIFAR10Configs.train_loader)
|
||||
def cifar10_train_loader(c: CIFAR10Configs):
|
||||
return DataLoader(c.train_dataset,
|
||||
batch_size=c.train_batch_size,
|
||||
shuffle=c.train_loader_shuffle)
|
||||
|
||||
|
||||
@CIFAR10Configs.calc(CIFAR10Configs.valid_loader)
|
||||
def cifar10_valid_loader(c: CIFAR10Configs):
|
||||
return DataLoader(c.valid_dataset,
|
||||
batch_size=c.valid_batch_size,
|
||||
shuffle=c.valid_loader_shuffle)
|
||||
|
||||
|
||||
CIFAR10Configs.aggregate(CIFAR10Configs.dataset_name, 'CIFAR10',
|
||||
(CIFAR10Configs.dataset_transforms, 'cifar10_transforms'),
|
||||
(CIFAR10Configs.train_dataset, 'cifar10_train_dataset'),
|
||||
(CIFAR10Configs.valid_dataset, 'cifar10_valid_dataset'),
|
||||
(CIFAR10Configs.train_loader, 'cifar10_train_loader'),
|
||||
(CIFAR10Configs.valid_loader, 'cifar10_valid_loader'))
|
||||
|
||||
|
||||
class TextDataset:
|
||||
itos: List[str]
|
||||
stoi: Dict[str, int]
|
||||
n_tokens: int
|
||||
train: str
|
||||
valid: str
|
||||
standard_tokens: List[str] = []
|
||||
|
||||
@staticmethod
|
||||
def load(path: PurePath):
|
||||
with open(str(path), 'r') as f:
|
||||
return f.read()
|
||||
|
||||
def __init__(self, path: PurePath, tokenizer: Callable, train: str, valid: str, test: str, *,
|
||||
n_tokens: Optional[int] = None,
|
||||
stoi: Optional[Dict[str, int]] = None,
|
||||
itos: Optional[List[str]] = None):
|
||||
self.test = test
|
||||
self.valid = valid
|
||||
self.train = train
|
||||
self.tokenizer = tokenizer
|
||||
self.path = path
|
||||
|
||||
if n_tokens or stoi or itos:
|
||||
assert stoi and itos and n_tokens
|
||||
self.n_tokens = n_tokens
|
||||
self.stoi = stoi
|
||||
self.itos = itos
|
||||
else:
|
||||
self.n_tokens = len(self.standard_tokens)
|
||||
self.stoi = {t: i for i, t in enumerate(self.standard_tokens)}
|
||||
|
||||
with monit.section("Tokenize"):
|
||||
tokens = self.tokenizer(self.train) + self.tokenizer(self.valid)
|
||||
tokens = sorted(list(set(tokens)))
|
||||
|
||||
for t in monit.iterate("Build vocabulary", tokens):
|
||||
self.stoi[t] = self.n_tokens
|
||||
self.n_tokens += 1
|
||||
|
||||
self.itos = [''] * self.n_tokens
|
||||
for t, n in self.stoi.items():
|
||||
self.itos[n] = t
|
||||
|
||||
def text_to_i(self, text: str) -> torch.Tensor:
|
||||
tokens = self.tokenizer(text)
|
||||
return torch.tensor([self.stoi[s] for s in tokens if s in self.stoi], dtype=torch.long)
|
||||
|
||||
def __repr__(self):
|
||||
return f'{len(self.train) / 1_000_000 :,.2f}M, {len(self.valid) / 1_000_000 :,.2f}M - {str(self.path)}'
|
||||
|
||||
|
||||
class SequentialDataLoader(IterableDataset):
|
||||
def __init__(self, *, text: str, dataset: TextDataset,
|
||||
batch_size: int, seq_len: int):
|
||||
self.seq_len = seq_len
|
||||
data = dataset.text_to_i(text)
|
||||
n_batch = data.shape[0] // batch_size
|
||||
data = data.narrow(0, 0, n_batch * batch_size)
|
||||
data = data.view(batch_size, -1).t().contiguous()
|
||||
self.data = data
|
||||
|
||||
def __len__(self):
|
||||
return self.data.shape[0] // self.seq_len
|
||||
|
||||
def __iter__(self):
|
||||
self.idx = 0
|
||||
return self
|
||||
|
||||
def __next__(self):
|
||||
if self.idx >= self.data.shape[0] - 1:
|
||||
raise StopIteration()
|
||||
|
||||
seq_len = min(self.seq_len, self.data.shape[0] - 1 - self.idx)
|
||||
i = self.idx + seq_len
|
||||
data = self.data[self.idx: i]
|
||||
target = self.data[self.idx + 1: i + 1]
|
||||
self.idx = i
|
||||
return data, target
|
||||
|
||||
def __getitem__(self, idx):
|
||||
seq_len = min(self.seq_len, self.data.shape[0] - 1 - idx)
|
||||
i = idx + seq_len
|
||||
data = self.data[idx: i]
|
||||
target = self.data[idx + 1: i + 1]
|
||||
return data, target
|
||||
|
||||
|
||||
class SequentialUnBatchedDataset(Dataset):
|
||||
def __init__(self, *, text: str, dataset: TextDataset,
|
||||
seq_len: int,
|
||||
is_random_offset: bool = True):
|
||||
self.is_random_offset = is_random_offset
|
||||
self.seq_len = seq_len
|
||||
self.data = dataset.text_to_i(text)
|
||||
|
||||
def __len__(self):
|
||||
return (self.data.shape[0] - 1) // self.seq_len
|
||||
|
||||
def __getitem__(self, idx):
|
||||
start = idx * self.seq_len
|
||||
assert start + self.seq_len + 1 <= self.data.shape[0]
|
||||
if self.is_random_offset:
|
||||
start += random.randint(0, min(self.seq_len - 1, self.data.shape[0] - (start + self.seq_len + 1)))
|
||||
|
||||
end = start + self.seq_len
|
||||
data = self.data[start: end]
|
||||
target = self.data[start + 1: end + 1]
|
||||
return data, target
|
||||
|
||||
|
||||
class TextFileDataset(TextDataset):
|
||||
standard_tokens = []
|
||||
|
||||
def __init__(self, path: PurePath, tokenizer: Callable, *,
|
||||
url: Optional[str] = None,
|
||||
filter_subset: Optional[int] = None):
|
||||
path = Path(path)
|
||||
if not path.exists():
|
||||
if not url:
|
||||
raise FileNotFoundError(str(path))
|
||||
else:
|
||||
download_file(url, path)
|
||||
|
||||
with monit.section("Load data"):
|
||||
text = self.load(path)
|
||||
if filter_subset:
|
||||
text = text[:filter_subset]
|
||||
split = int(len(text) * .9)
|
||||
train = text[:split]
|
||||
valid = text[split:]
|
||||
|
||||
super().__init__(path, tokenizer, train, valid, '')
|
||||
|
||||
|
||||
def _test_tiny_shakespeare():
|
||||
from labml import lab
|
||||
_ = TextFileDataset(lab.get_data_path() / 'tiny_shakespeare.txt', lambda x: list(x),
|
||||
url='https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt')
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
_test_tiny_shakespeare()
|
||||
@@ -0,0 +1,70 @@
|
||||
import torch
|
||||
|
||||
from labml.configs import BaseConfigs, hyperparams, option
|
||||
|
||||
|
||||
class DeviceInfo:
|
||||
def __init__(self, *,
|
||||
use_cuda: bool,
|
||||
cuda_device: int):
|
||||
self.use_cuda = use_cuda
|
||||
self.cuda_device = cuda_device
|
||||
self.cuda_count = torch.cuda.device_count()
|
||||
|
||||
self.is_cuda = self.use_cuda and torch.cuda.is_available()
|
||||
if not self.is_cuda:
|
||||
self.device = torch.device('cpu')
|
||||
else:
|
||||
if self.cuda_device < self.cuda_count:
|
||||
self.device = torch.device('cuda', self.cuda_device)
|
||||
else:
|
||||
self.device = torch.device('cuda', self.cuda_count - 1)
|
||||
|
||||
def __str__(self):
|
||||
if not self.is_cuda:
|
||||
return "CPU"
|
||||
|
||||
if self.cuda_device < self.cuda_count:
|
||||
return f"GPU:{self.cuda_device} - {torch.cuda.get_device_name(self.cuda_device)}"
|
||||
else:
|
||||
return (f"GPU:{self.cuda_count - 1}({self.cuda_device}) "
|
||||
f"- {torch.cuda.get_device_name(self.cuda_count - 1)}")
|
||||
|
||||
|
||||
class DeviceConfigs(BaseConfigs):
|
||||
r"""
|
||||
This is a configurable module to get a single device to train model on.
|
||||
It can pick up CUDA devices and it will fall back to CPU if they are not available.
|
||||
|
||||
It has other small advantages such as being able to view the
|
||||
actual device name on configurations view of
|
||||
`labml app <https://github.com/labmlai/labml/tree/master/app>`_
|
||||
|
||||
Arguments:
|
||||
cuda_device (int): The CUDA device number. Defaults to ``0``.
|
||||
use_cuda (bool): Whether to use CUDA devices. Defaults to ``True``.
|
||||
"""
|
||||
cuda_device: int = 0
|
||||
use_cuda: bool = True
|
||||
|
||||
device_info: DeviceInfo
|
||||
|
||||
device: torch.device
|
||||
|
||||
def __init__(self):
|
||||
super().__init__(_primary='device')
|
||||
|
||||
|
||||
@option(DeviceConfigs.device)
|
||||
def _device(c: DeviceConfigs):
|
||||
return c.device_info.device
|
||||
|
||||
|
||||
hyperparams(DeviceConfigs.cuda_device, DeviceConfigs.use_cuda,
|
||||
is_hyperparam=False)
|
||||
|
||||
|
||||
@option(DeviceConfigs.device_info)
|
||||
def _device_info(c: DeviceConfigs):
|
||||
return DeviceInfo(use_cuda=c.use_cuda,
|
||||
cuda_device=c.cuda_device)
|
||||
@@ -0,0 +1,85 @@
|
||||
import dataclasses
|
||||
from abc import ABC
|
||||
|
||||
import torch
|
||||
from labml import tracker
|
||||
|
||||
|
||||
class StateModule:
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
# def __call__(self):
|
||||
# raise NotImplementedError
|
||||
|
||||
def create_state(self) -> any:
|
||||
raise NotImplementedError
|
||||
|
||||
def set_state(self, data: any):
|
||||
raise NotImplementedError
|
||||
|
||||
def on_epoch_start(self):
|
||||
raise NotImplementedError
|
||||
|
||||
def on_epoch_end(self):
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
class Metric(StateModule, ABC):
|
||||
def track(self):
|
||||
pass
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class AccuracyState:
|
||||
samples: int = 0
|
||||
correct: int = 0
|
||||
|
||||
def reset(self):
|
||||
self.samples = 0
|
||||
self.correct = 0
|
||||
|
||||
|
||||
class Accuracy(Metric):
|
||||
data: AccuracyState
|
||||
|
||||
def __init__(self, ignore_index: int = -1):
|
||||
super().__init__()
|
||||
self.ignore_index = ignore_index
|
||||
|
||||
def __call__(self, output: torch.Tensor, target: torch.Tensor):
|
||||
output = output.view(-1, output.shape[-1])
|
||||
target = target.view(-1)
|
||||
pred = output.argmax(dim=-1)
|
||||
mask = target == self.ignore_index
|
||||
pred.masked_fill_(mask, self.ignore_index)
|
||||
n_masked = mask.sum().item()
|
||||
self.data.correct += pred.eq(target).sum().item() - n_masked
|
||||
self.data.samples += len(target) - n_masked
|
||||
|
||||
def create_state(self):
|
||||
return AccuracyState()
|
||||
|
||||
def set_state(self, data: any):
|
||||
self.data = data
|
||||
|
||||
def on_epoch_start(self):
|
||||
self.data.reset()
|
||||
|
||||
def on_epoch_end(self):
|
||||
self.track()
|
||||
|
||||
def track(self):
|
||||
if self.data.samples == 0:
|
||||
return
|
||||
tracker.add("accuracy.", self.data.correct / self.data.samples)
|
||||
|
||||
|
||||
class AccuracyDirect(Accuracy):
|
||||
data: AccuracyState
|
||||
|
||||
def __call__(self, output: torch.Tensor, target: torch.Tensor):
|
||||
output = output.view(-1)
|
||||
target = target.view(-1)
|
||||
self.data.correct += output.eq(target).sum().item()
|
||||
self.data.samples += len(target)
|
||||
@@ -0,0 +1,97 @@
|
||||
from typing import Tuple
|
||||
|
||||
import torch
|
||||
from labml import tracker
|
||||
|
||||
from labml.configs import BaseConfigs, option, meta_config
|
||||
|
||||
|
||||
class OptimizerConfigs(BaseConfigs):
|
||||
r"""
|
||||
This creates a configurable optimizer.
|
||||
|
||||
Arguments:
|
||||
learning_rate (float): Learning rate of the optimizer. Defaults to ``0.01``.
|
||||
momentum (float): Momentum of the optimizer. Defaults to ``0.5``.
|
||||
parameters: Model parameters to optimize.
|
||||
d_model (int): Embedding size of the model (for Noam optimizer).
|
||||
betas (Tuple[float, float]): Betas for Adam optimizer. Defaults to ``(0.9, 0.999)``.
|
||||
eps (float): Epsilon for Adam/RMSProp optimizers. Defaults to ``1e-8``.
|
||||
step_factor (int): Step factor for Noam optimizer. Defaults to ``1024``.
|
||||
|
||||
Also there is a better (more options) implementation in ``labml_nn``.
|
||||
`We recommend using that <https://nn.labml.ai/optimizers/configs.html>`_.
|
||||
"""
|
||||
|
||||
optimizer: torch.optim.Adam
|
||||
learning_rate: float = 0.01
|
||||
momentum: float = 0.5
|
||||
parameters: any
|
||||
d_model: int
|
||||
betas: Tuple[float, float] = (0.9, 0.999)
|
||||
eps: float = 1e-8
|
||||
step_factor: int = 1024
|
||||
|
||||
def __init__(self):
|
||||
super().__init__(_primary='optimizer')
|
||||
|
||||
|
||||
meta_config(OptimizerConfigs.parameters)
|
||||
|
||||
|
||||
@option(OptimizerConfigs.optimizer, 'SGD')
|
||||
def sgd_optimizer(c: OptimizerConfigs):
|
||||
return torch.optim.SGD(c.parameters, c.learning_rate, c.momentum)
|
||||
|
||||
|
||||
@option(OptimizerConfigs.optimizer, 'Adam')
|
||||
def adam_optimizer(c: OptimizerConfigs):
|
||||
return torch.optim.Adam(c.parameters, lr=c.learning_rate,
|
||||
betas=c.betas, eps=c.eps)
|
||||
|
||||
|
||||
class NoamOpt:
|
||||
def __init__(self, model_size: int, learning_rate: float, warmup: int, step_factor: int, optimizer):
|
||||
self.step_factor = step_factor
|
||||
self.optimizer = optimizer
|
||||
self.warmup = warmup
|
||||
self.learning_rate = learning_rate
|
||||
self.model_size = model_size
|
||||
self._rate = 0
|
||||
|
||||
def step(self):
|
||||
rate = self.rate(tracker.get_global_step() / self.step_factor)
|
||||
for p in self.optimizer.param_groups:
|
||||
p['lr'] = rate
|
||||
self._rate = rate
|
||||
self.optimizer.step()
|
||||
|
||||
def rate(self, step):
|
||||
factor = self.model_size ** (-0.5) * min(step ** (-0.5), step * self.warmup ** (-1.5))
|
||||
return self.learning_rate * factor
|
||||
|
||||
def zero_grad(self):
|
||||
self.optimizer.zero_grad()
|
||||
|
||||
|
||||
@option(OptimizerConfigs.optimizer, 'Noam')
|
||||
def noam_optimizer(c: OptimizerConfigs):
|
||||
optimizer = torch.optim.Adam(c.parameters, lr=0.0, betas=c.betas, eps=c.eps)
|
||||
return NoamOpt(c.d_model, 1, 2000, c.step_factor, optimizer)
|
||||
|
||||
|
||||
def _test_noam_optimizer():
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
|
||||
opts = [NoamOpt(512, 1, 4000, None),
|
||||
NoamOpt(512, 1, 8000, None),
|
||||
NoamOpt(2048, 1, 2000, None)]
|
||||
plt.plot(np.arange(1, 20000), [[opt.rate(i) for opt in opts] for i in range(1, 20000)])
|
||||
plt.legend(["512:4000", "512:8000", "256:4000"])
|
||||
plt.title("Optimizer")
|
||||
plt.show()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
_test_noam_optimizer()
|
||||
@@ -0,0 +1,84 @@
|
||||
from typing import Tuple, List
|
||||
|
||||
|
||||
class Schedule:
|
||||
def __call__(self, x):
|
||||
raise NotImplementedError()
|
||||
|
||||
|
||||
class Flat(Schedule):
|
||||
def __init__(self, value):
|
||||
self.__value = value
|
||||
|
||||
def __call__(self, x):
|
||||
return self.__value
|
||||
|
||||
def __str__(self):
|
||||
return f"Schedule({self.__value})"
|
||||
|
||||
|
||||
class Dynamic(Schedule):
|
||||
def __init__(self, value):
|
||||
self.__value = value
|
||||
|
||||
def __call__(self, x):
|
||||
return self.__value
|
||||
|
||||
def update(self, value):
|
||||
self.__value = value
|
||||
|
||||
def __str__(self):
|
||||
return "Dynamic"
|
||||
|
||||
|
||||
class Piecewise(Schedule):
|
||||
"""
|
||||
## Piecewise schedule
|
||||
"""
|
||||
|
||||
def __init__(self, endpoints: List[Tuple[float, float]], outside_value: float = None):
|
||||
"""
|
||||
### Initialize
|
||||
|
||||
`endpoints` is list of pairs `(x, y)`.
|
||||
The values between endpoints are linearly interpolated.
|
||||
`y` values outside the range covered by `x` are
|
||||
`outside_value`.
|
||||
"""
|
||||
|
||||
# `(x, y)` pairs should be sorted
|
||||
indexes = [e[0] for e in endpoints]
|
||||
assert indexes == sorted(indexes)
|
||||
|
||||
self._outside_value = outside_value
|
||||
self._endpoints = endpoints
|
||||
|
||||
def __call__(self, x):
|
||||
"""
|
||||
### Find `y` for given `x`
|
||||
"""
|
||||
|
||||
# iterate through each segment
|
||||
for (x1, y1), (x2, y2) in zip(self._endpoints[:-1], self._endpoints[1:]):
|
||||
# interpolate if `x` is within the segment
|
||||
if x1 <= x < x2:
|
||||
dx = float(x - x1) / (x2 - x1)
|
||||
return y1 + dx * (y2 - y1)
|
||||
|
||||
# return outside value otherwise
|
||||
return self._outside_value
|
||||
|
||||
def __str__(self):
|
||||
endpoints = ", ".join([f"({e[0]}, {e[1]})" for e in self._endpoints])
|
||||
return f"Schedule[{endpoints}, {self._outside_value}]"
|
||||
|
||||
|
||||
class RelativePiecewise(Piecewise):
|
||||
def __init__(self, relative_endpoits: List[Tuple[float, float]], total_steps: int):
|
||||
endpoints = []
|
||||
for e in relative_endpoits:
|
||||
index = int(total_steps * e[0])
|
||||
assert index >= 0
|
||||
endpoints.append((index, e[1]))
|
||||
|
||||
super().__init__(endpoints, outside_value=relative_endpoits[-1][1])
|
||||
@@ -0,0 +1,511 @@
|
||||
import signal
|
||||
import typing
|
||||
from typing import Dict, List, Callable
|
||||
from typing import Optional, Tuple, Any, Collection
|
||||
|
||||
import torch.optim
|
||||
import torch.optim
|
||||
import torch.utils.data
|
||||
import torch.utils.data
|
||||
from labml import tracker, logger, monit
|
||||
from labml.configs import BaseConfigs, meta_config, option
|
||||
from labml.internal.monitor import Loop
|
||||
from labml.logger import Text
|
||||
from torch import nn
|
||||
from .device import DeviceConfigs
|
||||
from .metrics import StateModule
|
||||
|
||||
|
||||
class TrainingLoopIterator(Collection):
|
||||
def __init__(self, start: int, total: int, step: Optional[int]):
|
||||
self.step = step
|
||||
self.total = total
|
||||
self.start = start
|
||||
self.i = None
|
||||
|
||||
def __iter__(self):
|
||||
self.i = None
|
||||
return self
|
||||
|
||||
def __next__(self):
|
||||
if self.step is not None:
|
||||
if self.i is None:
|
||||
self.i = self.start
|
||||
else:
|
||||
self.i += self.step
|
||||
else:
|
||||
if self.i is None:
|
||||
self.i = 0
|
||||
else:
|
||||
self.i += 1
|
||||
|
||||
if self.i >= self.total:
|
||||
raise StopIteration()
|
||||
|
||||
if self.step is None:
|
||||
return tracker.get_global_step()
|
||||
else:
|
||||
return self.i
|
||||
|
||||
def __len__(self) -> int:
|
||||
if self.step is not None:
|
||||
return (self.total - self.start) // self.step
|
||||
else:
|
||||
return self.total
|
||||
|
||||
def __contains__(self, x: object) -> bool:
|
||||
return False
|
||||
|
||||
|
||||
class TrainingLoop:
|
||||
_iter: Optional[TrainingLoopIterator]
|
||||
__loop: Loop
|
||||
__signal_received: Optional[Tuple[Any, Any]]
|
||||
|
||||
def __init__(self, *,
|
||||
loop_count: int,
|
||||
loop_step: Optional[int],
|
||||
log_new_line_interval: int,
|
||||
log_write_interval: int,
|
||||
is_loop_on_interrupt: bool):
|
||||
self.__loop_count = loop_count
|
||||
self.__loop_step = loop_step
|
||||
self.__log_new_line_interval = log_new_line_interval
|
||||
self.__log_write_interval = log_write_interval
|
||||
self.__last_write_step = 0
|
||||
self.__last_new_line_step = 0
|
||||
self.__last_save_step = 0
|
||||
self.__signal_received = None
|
||||
self.__is_loop_on_interrupt = is_loop_on_interrupt
|
||||
self._iter = None
|
||||
|
||||
def __iter__(self):
|
||||
self._iter = TrainingLoopIterator(tracker.get_global_step(),
|
||||
self.__loop_count,
|
||||
self.__loop_step)
|
||||
|
||||
self.__loop = monit.loop(typing.cast(Collection, self._iter))
|
||||
|
||||
iter(self.__loop)
|
||||
try:
|
||||
self.old_handler = signal.signal(signal.SIGINT, self.__handler)
|
||||
except ValueError:
|
||||
pass
|
||||
return self
|
||||
|
||||
@property
|
||||
def idx(self):
|
||||
if not self._iter:
|
||||
return 0
|
||||
if not self._iter.i:
|
||||
return 0
|
||||
if self.__loop_step is None:
|
||||
return self._iter.i
|
||||
return self._iter.i / self.__loop_step
|
||||
|
||||
def __finish(self):
|
||||
try:
|
||||
signal.signal(signal.SIGINT, self.old_handler)
|
||||
except ValueError:
|
||||
pass
|
||||
tracker.save()
|
||||
tracker.new_line()
|
||||
|
||||
def __next__(self):
|
||||
if self.__signal_received is not None:
|
||||
logger.log('\nKilling Loop.', Text.danger)
|
||||
monit.finish_loop()
|
||||
self.__finish()
|
||||
raise StopIteration("SIGINT")
|
||||
|
||||
try:
|
||||
global_step = next(self.__loop)
|
||||
except StopIteration as e:
|
||||
self.__finish()
|
||||
raise e
|
||||
|
||||
tracker.set_global_step(global_step)
|
||||
|
||||
if global_step - self.__last_write_step >= self.__log_write_interval:
|
||||
tracker.save()
|
||||
self.__last_write_step = global_step
|
||||
if global_step - self.__last_new_line_step >= self.__log_new_line_interval:
|
||||
tracker.new_line()
|
||||
self.__last_new_line_step = global_step
|
||||
|
||||
return global_step
|
||||
|
||||
def __handler(self, sig, frame):
|
||||
# Pass second interrupt without delaying
|
||||
if self.__signal_received is not None:
|
||||
logger.log('\nSIGINT received twice. Stopping...', Text.danger)
|
||||
self.old_handler(*self.__signal_received)
|
||||
return
|
||||
|
||||
if self.__is_loop_on_interrupt:
|
||||
# Store the interrupt signal for later
|
||||
self.__signal_received = (sig, frame)
|
||||
logger.log('\nSIGINT received. Delaying KeyboardInterrupt.', Text.danger)
|
||||
else:
|
||||
self.__finish()
|
||||
logger.log('Killing loop...', Text.danger)
|
||||
self.old_handler(sig, frame)
|
||||
|
||||
def __str__(self):
|
||||
return "LabTrainingLoop"
|
||||
|
||||
|
||||
class TrainingLoopConfigs(BaseConfigs):
|
||||
r"""
|
||||
This is a configurable training loop. You can extend this class for your configurations
|
||||
if it involves a training loop.
|
||||
|
||||
>>> for step in conf.training_loop:
|
||||
>>> ...
|
||||
|
||||
Arguments:
|
||||
loop_count (int): Total number of steps. Defaults to ``10``.
|
||||
loop_step (int): Number of steps to increment per iteration. Defaults to ``1``.
|
||||
log_new_line_interval (int): The interval (in steps) to print a new line to the screen.
|
||||
Defaults to ``1``.
|
||||
log_write_interval (int): The interval (in steps) to call :func:`labml.tracker.save`.
|
||||
Defaults to ``1``.
|
||||
is_loop_on_interrupt (bool): Whether to handle keyboard interrupts and wait until a iteration is complete.
|
||||
Defaults to ``False``.
|
||||
"""
|
||||
loop_count: int = 10
|
||||
loop_step: int = 1
|
||||
log_new_line_interval: int = 1
|
||||
log_write_interval: int = 1
|
||||
is_loop_on_interrupt: bool = False
|
||||
|
||||
training_loop: TrainingLoop
|
||||
|
||||
|
||||
@option(TrainingLoopConfigs.training_loop)
|
||||
def _loop_configs(c: TrainingLoopConfigs):
|
||||
return TrainingLoop(loop_count=c.loop_count,
|
||||
loop_step=c.loop_step,
|
||||
log_new_line_interval=c.log_new_line_interval,
|
||||
log_write_interval=c.log_write_interval,
|
||||
is_loop_on_interrupt=c.is_loop_on_interrupt)
|
||||
|
||||
|
||||
meta_config(TrainingLoopConfigs.loop_step,
|
||||
TrainingLoopConfigs.loop_count,
|
||||
TrainingLoopConfigs.log_new_line_interval,
|
||||
TrainingLoopConfigs.log_write_interval,
|
||||
TrainingLoopConfigs.is_loop_on_interrupt)
|
||||
|
||||
|
||||
class ModeState:
|
||||
def __init__(self):
|
||||
self._rollback_stack = []
|
||||
|
||||
self.is_train = False
|
||||
self.is_optimize = False
|
||||
|
||||
def _enter(self, mode: Dict[str, any]):
|
||||
rollback = {}
|
||||
for k, v in mode.items():
|
||||
if v is None:
|
||||
continue
|
||||
rollback[k] = getattr(self, k)
|
||||
setattr(self, k, v)
|
||||
|
||||
self._rollback_stack.append(rollback)
|
||||
|
||||
return len(self._rollback_stack)
|
||||
|
||||
def _exit(self, n: int):
|
||||
assert n == len(self._rollback_stack)
|
||||
|
||||
rollback = self._rollback_stack[-1]
|
||||
self._rollback_stack.pop(-1)
|
||||
|
||||
for k, v in rollback.items():
|
||||
setattr(self, k, v)
|
||||
|
||||
def update(self, *,
|
||||
is_train: Optional[bool] = None,
|
||||
is_optimize: Optional[bool] = None):
|
||||
return Mode(self,
|
||||
is_train=is_train,
|
||||
is_optimize=is_optimize)
|
||||
|
||||
|
||||
class Mode:
|
||||
def __init__(self, mode: ModeState, **kwargs: any):
|
||||
self.mode = mode
|
||||
self.update = {}
|
||||
for k, v in kwargs.items():
|
||||
if v is not None:
|
||||
self.update[k] = v
|
||||
|
||||
self.idx = -1
|
||||
|
||||
def __enter__(self):
|
||||
self.idx = self.mode._enter(self.update)
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
self.mode._exit(self.idx)
|
||||
|
||||
|
||||
class Trainer:
|
||||
def __init__(self, *,
|
||||
name: str,
|
||||
mode: ModeState,
|
||||
data_loader: torch.utils.data.DataLoader,
|
||||
inner_iterations: int,
|
||||
state_modules: List[StateModule],
|
||||
is_track_time: bool,
|
||||
step: Callable[[any, 'BatchIndex'], None]):
|
||||
self.is_track_time = is_track_time
|
||||
self.mode = mode
|
||||
self.name = name
|
||||
self.step = step
|
||||
self.state_modules = state_modules
|
||||
self.__iterable = None
|
||||
self.__states = [sm.create_state() for sm in self.state_modules]
|
||||
self.inner_iterations = inner_iterations
|
||||
self.data_loader = data_loader
|
||||
self._batch_index = BatchIndex(len(self.data_loader), self.inner_iterations)
|
||||
|
||||
def set_data_loader(self, data_loader: torch.utils.data.DataLoader):
|
||||
self.data_loader = data_loader
|
||||
self._batch_index = BatchIndex(len(data_loader), self.inner_iterations)
|
||||
self.__iterable = None
|
||||
|
||||
def __call__(self):
|
||||
for sm, s in zip(self.state_modules, self.__states):
|
||||
sm.set_state(s)
|
||||
|
||||
if self.__iterable is None or self._batch_index.completed:
|
||||
self.__iterable = iter(self.data_loader)
|
||||
self._batch_index.reset(len(self.data_loader), self.inner_iterations)
|
||||
for sm in self.state_modules:
|
||||
sm.on_epoch_start()
|
||||
with torch.set_grad_enabled(self.mode.is_train):
|
||||
self.__iterate()
|
||||
|
||||
if self._batch_index.completed:
|
||||
for sm in self.state_modules:
|
||||
sm.on_epoch_end()
|
||||
|
||||
def __iterate(self):
|
||||
with monit.section(self.name, is_partial=True, is_track=self.is_track_time):
|
||||
if self._batch_index.idx == 0:
|
||||
monit.progress(0)
|
||||
while not self._batch_index.iteration_completed:
|
||||
batch = next(self.__iterable)
|
||||
|
||||
self.step(batch, self._batch_index)
|
||||
|
||||
self._batch_index.step()
|
||||
monit.progress(self._batch_index.epoch_progress)
|
||||
|
||||
self._batch_index.step_inner()
|
||||
|
||||
|
||||
class BatchIndex:
|
||||
idx: int
|
||||
total: int
|
||||
iteration: int
|
||||
total_iterations: int
|
||||
|
||||
def __init__(self, total: int, total_iterations: int):
|
||||
self.total_iterations = total_iterations
|
||||
self.total = total
|
||||
|
||||
def is_interval(self, interval: int):
|
||||
if interval <= 0:
|
||||
return False
|
||||
if self.idx + 1 == self.total:
|
||||
return True
|
||||
else:
|
||||
return (self.idx + 1) % interval == 0
|
||||
|
||||
@property
|
||||
def is_last(self):
|
||||
return self.idx + 1 == self.total
|
||||
|
||||
@property
|
||||
def completed(self):
|
||||
return self.iteration >= self.total_iterations
|
||||
|
||||
@property
|
||||
def iteration_completed(self):
|
||||
# // is important so that the last step happens on the last iteration
|
||||
return self.idx >= (self.iteration + 1) * self.total // self.total_iterations
|
||||
|
||||
@property
|
||||
def epoch_progress(self):
|
||||
return self.idx / self.total
|
||||
|
||||
def step(self):
|
||||
self.idx += 1
|
||||
|
||||
def step_inner(self):
|
||||
self.iteration += 1
|
||||
|
||||
def reset(self, total: int, total_iterations: int):
|
||||
self.total = total
|
||||
self.total_iterations = total_iterations
|
||||
self.idx = 0
|
||||
self.iteration = 0
|
||||
|
||||
|
||||
class TrainValidConfigs(TrainingLoopConfigs):
|
||||
r"""
|
||||
This is a configurable module that you can extend for experiments that involve a
|
||||
training and validation datasets (i.e. most DL experiments).
|
||||
|
||||
Arguments:
|
||||
epochs (int): Number of epochs to train on. Defaults to ``10``.
|
||||
train_loader (torch.utils.data.DataLoader): Training data loader.
|
||||
valid_loader (torch.utils.data.DataLoader): Training data loader.
|
||||
inner_iterations (int): Number of times to switch between training and validation
|
||||
within an epoch. Defaults to ``1``.
|
||||
|
||||
You can override ``init``, ``step`` functions. There is also a ``sample`` function
|
||||
that you can override to generate samples ever time it switches between training and validation.
|
||||
"""
|
||||
state_modules: List[StateModule]
|
||||
|
||||
mode: ModeState
|
||||
|
||||
epochs: int = 10
|
||||
|
||||
trainer: Trainer
|
||||
validator: Trainer
|
||||
train_loader: torch.utils.data.DataLoader
|
||||
valid_loader: torch.utils.data.DataLoader
|
||||
|
||||
loop_count = '_data_loop_count'
|
||||
loop_step = None
|
||||
|
||||
inner_iterations: int = 1
|
||||
|
||||
is_track_time: bool = False
|
||||
|
||||
def init(self):
|
||||
pass
|
||||
|
||||
def step(self, batch: Any, batch_idx: BatchIndex):
|
||||
raise NotImplementedError
|
||||
|
||||
def run_step(self):
|
||||
for i in range(self.inner_iterations):
|
||||
with tracker.namespace('sample'):
|
||||
self.sample()
|
||||
with self.mode.update(is_train=True):
|
||||
with tracker.namespace('train'):
|
||||
self.trainer()
|
||||
if self.validator:
|
||||
with tracker.namespace('valid'):
|
||||
self.validator()
|
||||
tracker.save()
|
||||
|
||||
def run(self):
|
||||
with monit.section("Initialize"):
|
||||
self.init()
|
||||
_ = self.validator
|
||||
_ = self.trainer
|
||||
for _ in self.training_loop:
|
||||
self.run_step()
|
||||
|
||||
def sample(self):
|
||||
pass
|
||||
|
||||
|
||||
@option(TrainValidConfigs.trainer)
|
||||
def _default_trainer(c: TrainValidConfigs):
|
||||
return Trainer(name='Train',
|
||||
mode=c.mode,
|
||||
data_loader=c.train_loader,
|
||||
inner_iterations=c.inner_iterations,
|
||||
state_modules=c.state_modules,
|
||||
is_track_time=c.is_track_time,
|
||||
step=c.step)
|
||||
|
||||
|
||||
@option(TrainValidConfigs.validator)
|
||||
def _default_validator(c: TrainValidConfigs):
|
||||
return Trainer(name='Valid',
|
||||
mode=c.mode,
|
||||
data_loader=c.valid_loader,
|
||||
inner_iterations=c.inner_iterations,
|
||||
state_modules=c.state_modules,
|
||||
is_track_time=c.is_track_time,
|
||||
step=c.step)
|
||||
|
||||
|
||||
@option(TrainValidConfigs.loop_count)
|
||||
def _data_loop_count(c: TrainValidConfigs):
|
||||
return c.epochs
|
||||
|
||||
|
||||
class SimpleTrainValidConfigs(TrainValidConfigs):
|
||||
r"""
|
||||
This is a configurable module that works for many standard DL experiments.
|
||||
|
||||
Arguments:
|
||||
model: A PyTorch model.
|
||||
optimizer: A PyTorch optimizer to update model.
|
||||
device: The device to train the model on. This defaults to a configurable device
|
||||
loss_function: A function to calculate the loss. This should accept ``model_output, target`` as
|
||||
arguments.
|
||||
update_batches (int): Number of batches to accumulate before taking an optimizer step.
|
||||
Defaults to ``1``.
|
||||
log_save_batches (int): How often to call :func:`labml.tracker.save`.
|
||||
"""
|
||||
optimizer: torch.optim.Adam
|
||||
model: nn.Module
|
||||
device: torch.device = DeviceConfigs()
|
||||
|
||||
loss_func: nn.Module
|
||||
|
||||
update_batches: int = 1
|
||||
log_save_batches: int = 1
|
||||
|
||||
state_modules: List[StateModule] = []
|
||||
|
||||
def init(self):
|
||||
pass
|
||||
|
||||
def step(self, batch: Any, batch_idx: BatchIndex):
|
||||
self.model.train(self.mode.is_train)
|
||||
data, target = batch[0].to(self.device), batch[1].to(self.device)
|
||||
|
||||
if self.mode.is_train:
|
||||
tracker.add_global_step(len(data))
|
||||
|
||||
with monit.section("model"):
|
||||
output = self.model(data)
|
||||
|
||||
loss = self.loss_func(output, target)
|
||||
tracker.add("loss.", loss)
|
||||
|
||||
if self.mode.is_train:
|
||||
with monit.section('backward'):
|
||||
loss.backward()
|
||||
|
||||
if batch_idx.is_interval(self.update_batches):
|
||||
with monit.section('optimize'):
|
||||
self.optimizer.step()
|
||||
self.optimizer.zero_grad()
|
||||
|
||||
if batch_idx.is_interval(self.log_save_batches):
|
||||
tracker.save()
|
||||
|
||||
|
||||
meta_config(SimpleTrainValidConfigs.update_batches,
|
||||
)
|
||||
|
||||
|
||||
@option(SimpleTrainValidConfigs.optimizer)
|
||||
def _default_optimizer(c: SimpleTrainValidConfigs):
|
||||
from .optimizer import OptimizerConfigs
|
||||
opt_conf = OptimizerConfigs()
|
||||
opt_conf.parameters = c.model.parameters()
|
||||
return opt_conf
|
||||
@@ -0,0 +1,8 @@
|
||||
"""
|
||||
---
|
||||
title: HyperNetworks
|
||||
summary: A PyTorch implementation/tutorial of HyperLSTM introduced in paper HyperNetworks.
|
||||
---
|
||||
|
||||
## [HyperLSTM](hyper_lstm.html)
|
||||
"""
|
||||
@@ -0,0 +1,129 @@
|
||||
{
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"name": "HyperLSTM",
|
||||
"provenance": [],
|
||||
"collapsed_sections": []
|
||||
},
|
||||
"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/hypernetworks/experiment.ipynb) \n",
|
||||
"\n",
|
||||
"## HyperLSTM\n",
|
||||
"\n",
|
||||
"This is an experiment training Shakespear dataset with HyperLSTM from paper HyperNetworks."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "ZCzmCrAIVg0L"
|
||||
},
|
||||
"source": [
|
||||
"!pip install labml-nn"
|
||||
],
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "0hJXx_g0wS2C"
|
||||
},
|
||||
"source": [
|
||||
"from labml import experiment\n",
|
||||
"from labml_nn.hypernetworks.experiment import Configs"
|
||||
],
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 255
|
||||
},
|
||||
"id": "WQ8VGpMGwZuj",
|
||||
"outputId": "5833cc50-26a8-496e-e729-88f42b3f4651"
|
||||
},
|
||||
"source": [
|
||||
"# Create experiment\n",
|
||||
"experiment.create(name=\"hyper_lstm\", comment='')\n",
|
||||
"# Create configs\n",
|
||||
"conf = Configs()\n",
|
||||
"# Load configurations\n",
|
||||
"experiment.configs(conf,\n",
|
||||
" # A dictionary of configurations to override\n",
|
||||
" {'tokenizer': 'character',\n",
|
||||
" 'text': 'tiny_shakespeare',\n",
|
||||
" 'optimizer.learning_rate': 2.5e-4,\n",
|
||||
" 'optimizer.optimizer': 'Adam',\n",
|
||||
" 'prompt': 'It is',\n",
|
||||
" 'prompt_separator': '',\n",
|
||||
"\n",
|
||||
" 'rnn_model': 'hyper_lstm',\n",
|
||||
"\n",
|
||||
" 'train_loader': 'shuffled_train_loader',\n",
|
||||
" 'valid_loader': 'shuffled_valid_loader',\n",
|
||||
"\n",
|
||||
" 'seq_len': 512,\n",
|
||||
" 'epochs': 128,\n",
|
||||
" 'batch_size': 2,\n",
|
||||
" 'inner_iterations': 25})\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"# Set models for saving and loading\n",
|
||||
"experiment.add_pytorch_models({'model': conf.model})\n",
|
||||
"\n",
|
||||
"conf.init()"
|
||||
],
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/",
|
||||
"height": 425
|
||||
},
|
||||
"id": "f07vAOaHwumr",
|
||||
"outputId": "6b51205e-3852-4dce-f7a7-f3ba4066ba21"
|
||||
},
|
||||
"source": [
|
||||
"# Start the experiment\n",
|
||||
"with experiment.start():\n",
|
||||
" # `TrainValidConfigs.run`\n",
|
||||
" conf.run()"
|
||||
],
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "crH6MzKmw-SY"
|
||||
},
|
||||
"source": [
|
||||
""
|
||||
],
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,101 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from labml import experiment
|
||||
from labml.configs import option
|
||||
from labml.utils.pytorch import get_modules
|
||||
|
||||
from labml_nn.experiments.nlp_autoregression import NLPAutoRegressionConfigs
|
||||
from labml_nn.hypernetworks.hyper_lstm import HyperLSTM
|
||||
from labml_nn.lstm import LSTM
|
||||
|
||||
|
||||
class AutoregressiveModel(nn.Module):
|
||||
"""
|
||||
## Auto regressive model
|
||||
"""
|
||||
|
||||
def __init__(self, n_vocab: int, d_model: int, rnn_model: nn.Module):
|
||||
super().__init__()
|
||||
# Token embedding module
|
||||
self.src_embed = nn.Embedding(n_vocab, d_model)
|
||||
self.lstm = rnn_model
|
||||
self.generator = nn.Linear(d_model, n_vocab)
|
||||
|
||||
def forward(self, x: torch.Tensor):
|
||||
x = self.src_embed(x)
|
||||
# Embed the tokens (`src`) and run it through the the transformer
|
||||
res, state = self.lstm(x)
|
||||
# Generate logits of the next token
|
||||
return self.generator(res), state
|
||||
|
||||
|
||||
class Configs(NLPAutoRegressionConfigs):
|
||||
"""
|
||||
## Configurations
|
||||
|
||||
The default configs can and will be over-ridden when we start the experiment
|
||||
"""
|
||||
|
||||
model: AutoregressiveModel
|
||||
rnn_model: nn.Module
|
||||
|
||||
d_model: int = 512
|
||||
n_rhn: int = 16
|
||||
n_z: int = 16
|
||||
|
||||
|
||||
@option(Configs.model)
|
||||
def autoregressive_model(c: Configs):
|
||||
"""
|
||||
Initialize the auto-regressive model
|
||||
"""
|
||||
m = AutoregressiveModel(c.n_tokens, c.d_model, c.rnn_model)
|
||||
return m.to(c.device)
|
||||
|
||||
|
||||
@option(Configs.rnn_model)
|
||||
def hyper_lstm(c: Configs):
|
||||
return HyperLSTM(c.d_model, c.d_model, c.n_rhn, c.n_z, 1)
|
||||
|
||||
|
||||
@option(Configs.rnn_model)
|
||||
def lstm(c: Configs):
|
||||
return LSTM(c.d_model, c.d_model, 1)
|
||||
|
||||
|
||||
def main():
|
||||
# Create experiment
|
||||
experiment.create(name="hyper_lstm", comment='')
|
||||
# Create configs
|
||||
conf = Configs()
|
||||
# Load configurations
|
||||
experiment.configs(conf,
|
||||
# A dictionary of configurations to override
|
||||
{'tokenizer': 'character',
|
||||
'text': 'tiny_shakespeare',
|
||||
'optimizer.learning_rate': 2.5e-4,
|
||||
'optimizer.optimizer': 'Adam',
|
||||
'prompt': 'It is',
|
||||
'prompt_separator': '',
|
||||
|
||||
'rnn_model': 'hyper_lstm',
|
||||
|
||||
'train_loader': 'shuffled_train_loader',
|
||||
'valid_loader': 'shuffled_valid_loader',
|
||||
|
||||
'seq_len': 512,
|
||||
'epochs': 128,
|
||||
'batch_size': 2,
|
||||
'inner_iterations': 25})
|
||||
|
||||
# Set models for saving and loading
|
||||
experiment.add_pytorch_models(get_modules(conf))
|
||||
|
||||
# Start the experiment
|
||||
with experiment.start():
|
||||
# `TrainValidConfigs.run`
|
||||
conf.run()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,272 @@
|
||||
"""
|
||||
---
|
||||
title: HyperNetworks - HyperLSTM
|
||||
summary: A PyTorch implementation/tutorial of HyperLSTM introduced in paper HyperNetworks.
|
||||
---
|
||||
|
||||
# HyperNetworks - HyperLSTM
|
||||
|
||||
We have implemented HyperLSTM introduced in paper
|
||||
[HyperNetworks](https://arxiv.org/abs/1609.09106), with annotations
|
||||
using [PyTorch](https://pytorch.org).
|
||||
[This blog post](https://blog.otoro.net/2016/09/28/hyper-networks/)
|
||||
by David Ha gives a good explanation of HyperNetworks.
|
||||
|
||||
We have an experiment that trains a HyperLSTM to predict text on Shakespeare dataset.
|
||||
Here's the link to code: [`experiment.py`](experiment.html)
|
||||
|
||||
[](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/hypernetworks/experiment.ipynb)
|
||||
|
||||
HyperNetworks use a smaller network to generate weights of a larger network.
|
||||
There are two variants: static hyper-networks and dynamic hyper-networks.
|
||||
Static HyperNetworks have smaller networks that generate weights (kernels)
|
||||
of a convolutional network. Dynamic HyperNetworks generate parameters of a
|
||||
recurrent neural network
|
||||
for each step. This is an implementation of the latter.
|
||||
|
||||
## Dynamic HyperNetworks
|
||||
In a RNN the parameters stay constant for each step.
|
||||
Dynamic HyperNetworks generate different parameters for each step.
|
||||
HyperLSTM has the structure of a LSTM but the parameters of
|
||||
each step are changed by a smaller LSTM network.
|
||||
|
||||
In the basic form, a Dynamic HyperNetwork has a smaller recurrent network that generates
|
||||
a feature vector corresponding to each parameter tensor of the larger recurrent network.
|
||||
Let's say the larger network has some parameter $\textcolor{cyan}{W_h}$ the smaller network generates a feature
|
||||
vector $z_h$ and we dynamically compute $\textcolor{cyan}{W_h}$ as a linear transformation of $z_h$.
|
||||
For instance $\textcolor{cyan}{W_h} = \langle W_{hz}, z_h \rangle$ where
|
||||
$W_{hz}$ is a 3-d tensor parameter and $\langle . \rangle$ is a tensor-vector multiplication.
|
||||
$z_h$ is usually a linear transformation of the output of the smaller recurrent network.
|
||||
|
||||
### Weight scaling instead of computing
|
||||
|
||||
Large recurrent networks have large dynamically computed parameters.
|
||||
These are calculated using linear transformation of feature vector $z$.
|
||||
And this transformation requires an even larger weight tensor.
|
||||
That is, when $\textcolor{cyan}{W_h}$ has shape $N_h \times N_h$,
|
||||
$W_{hz}$ will be $N_h \times N_h \times N_z$.
|
||||
|
||||
To overcome this, we compute the weight parameters of the recurrent network by
|
||||
dynamically scaling each row of a matrix of same size.
|
||||
|
||||
\begin{align}
|
||||
d(z) = W_{hz} z_h \\
|
||||
\\
|
||||
\textcolor{cyan}{W_h} =
|
||||
\begin{pmatrix}
|
||||
d_0(z) W_{hd_0} \\
|
||||
d_1(z) W_{hd_1} \\
|
||||
... \\
|
||||
d_{N_h}(z) W_{hd_{N_h}} \\
|
||||
\end{pmatrix}
|
||||
\end{align}
|
||||
|
||||
where $W_{hd}$ is a $N_h \times N_h$ parameter matrix.
|
||||
|
||||
We can further optimize this when we compute $\textcolor{cyan}{W_h} h$,
|
||||
as
|
||||
$$\textcolor{lightgreen}{d(z) \odot (W_{hd} h)}$$
|
||||
where $\odot$ stands for element-wise multiplication.
|
||||
"""
|
||||
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from labml_nn.lstm import LSTMCell
|
||||
|
||||
|
||||
class HyperLSTMCell(nn.Module):
|
||||
"""
|
||||
## HyperLSTM Cell
|
||||
|
||||
For HyperLSTM the smaller network and the larger network both have the LSTM structure.
|
||||
This is defined in Appendix A.2.2 in the paper.
|
||||
"""
|
||||
|
||||
def __init__(self, input_size: int, hidden_size: int, hyper_size: int, n_z: int):
|
||||
"""
|
||||
`input_size` is the size of the input $x_t$,
|
||||
`hidden_size` is the size of the LSTM, and
|
||||
`hyper_size` is the size of the smaller LSTM that alters the weights of the larger outer LSTM.
|
||||
`n_z` is the size of the feature vectors used to alter the LSTM weights.
|
||||
|
||||
We use the output of the smaller LSTM to compute $z_h^{i,f,g,o}$, $z_x^{i,f,g,o}$ and
|
||||
$z_b^{i,f,g,o}$ using linear transformations.
|
||||
We calculate $d_h^{i,f,g,o}(z_h^{i,f,g,o})$, $d_x^{i,f,g,o}(z_x^{i,f,g,o})$, and
|
||||
$d_b^{i,f,g,o}(z_b^{i,f,g,o})$ from these, using linear transformations again.
|
||||
These are then used to scale the rows of weight and bias tensors of the main LSTM.
|
||||
|
||||
📝 Since the computation of $z$ and $d$ are two sequential linear transformations
|
||||
these can be combined into a single linear transformation.
|
||||
However we've implemented this separately so that it matches with the description
|
||||
in the paper.
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
# The input to the hyperLSTM is
|
||||
# $$
|
||||
# \hat{x}_t = \begin{pmatrix}
|
||||
# h_{t-1} \\
|
||||
# x_t
|
||||
# \end{pmatrix}
|
||||
# $$
|
||||
# where $x_t$ is the input and $h_{t-1}$ is the output of the outer LSTM at previous step.
|
||||
# So the input size is `hidden_size + input_size`.
|
||||
#
|
||||
# The output of hyperLSTM is $\hat{h}_t$ and $\hat{c}_t$.
|
||||
self.hyper = LSTMCell(hidden_size + input_size, hyper_size, layer_norm=True)
|
||||
|
||||
# $$z_h^{i,f,g,o} = lin_{h}^{i,f,g,o}(\hat{h}_t)$$
|
||||
# 🤔 In the paper it was specified as
|
||||
# $$z_h^{i,f,g,o} = lin_{h}^{i,f,g,o}(\hat{h}_{\textcolor{red}{t-1}})$$
|
||||
# I feel that it's a typo.
|
||||
self.z_h = nn.Linear(hyper_size, 4 * n_z)
|
||||
# $$z_x^{i,f,g,o} = lin_x^{i,f,g,o}(\hat{h}_t)$$
|
||||
self.z_x = nn.Linear(hyper_size, 4 * n_z)
|
||||
# $$z_b^{i,f,g,o} = lin_b^{i,f,g,o}(\hat{h}_t)$$
|
||||
self.z_b = nn.Linear(hyper_size, 4 * n_z, bias=False)
|
||||
|
||||
# $$d_h^{i,f,g,o}(z_h^{i,f,g,o}) = lin_{dh}^{i,f,g,o}(z_h^{i,f,g,o})$$
|
||||
d_h = [nn.Linear(n_z, hidden_size, bias=False) for _ in range(4)]
|
||||
self.d_h = nn.ModuleList(d_h)
|
||||
# $$d_x^{i,f,g,o}(z_x^{i,f,g,o}) = lin_{dx}^{i,f,g,o}(z_x^{i,f,g,o})$$
|
||||
d_x = [nn.Linear(n_z, hidden_size, bias=False) for _ in range(4)]
|
||||
self.d_x = nn.ModuleList(d_x)
|
||||
# $$d_b^{i,f,g,o}(z_b^{i,f,g,o}) = lin_{db}^{i,f,g,o}(z_b^{i,f,g,o})$$
|
||||
d_b = [nn.Linear(n_z, hidden_size) for _ in range(4)]
|
||||
self.d_b = nn.ModuleList(d_b)
|
||||
|
||||
# The weight matrices $W_h^{i,f,g,o}$
|
||||
self.w_h = nn.ParameterList([nn.Parameter(torch.zeros(hidden_size, hidden_size)) for _ in range(4)])
|
||||
# The weight matrices $W_x^{i,f,g,o}$
|
||||
self.w_x = nn.ParameterList([nn.Parameter(torch.zeros(hidden_size, input_size)) for _ in range(4)])
|
||||
|
||||
# Layer normalization
|
||||
self.layer_norm = nn.ModuleList([nn.LayerNorm(hidden_size) for _ in range(4)])
|
||||
self.layer_norm_c = nn.LayerNorm(hidden_size)
|
||||
|
||||
def forward(self, x: torch.Tensor,
|
||||
h: torch.Tensor, c: torch.Tensor,
|
||||
h_hat: torch.Tensor, c_hat: torch.Tensor):
|
||||
# $$
|
||||
# \hat{x}_t = \begin{pmatrix}
|
||||
# h_{t-1} \\
|
||||
# x_t
|
||||
# \end{pmatrix}
|
||||
# $$
|
||||
x_hat = torch.cat((h, x), dim=-1)
|
||||
# $$\hat{h}_t, \hat{c}_t = lstm(\hat{x}_t, \hat{h}_{t-1}, \hat{c}_{t-1})$$
|
||||
h_hat, c_hat = self.hyper(x_hat, h_hat, c_hat)
|
||||
|
||||
# $$z_h^{i,f,g,o} = lin_{h}^{i,f,g,o}(\hat{h}_t)$$
|
||||
z_h = self.z_h(h_hat).chunk(4, dim=-1)
|
||||
# $$z_x^{i,f,g,o} = lin_x^{i,f,g,o}(\hat{h}_t)$$
|
||||
z_x = self.z_x(h_hat).chunk(4, dim=-1)
|
||||
# $$z_b^{i,f,g,o} = lin_b^{i,f,g,o}(\hat{h}_t)$$
|
||||
z_b = self.z_b(h_hat).chunk(4, dim=-1)
|
||||
|
||||
# We calculate $i$, $f$, $g$ and $o$ in a loop
|
||||
ifgo = []
|
||||
for i in range(4):
|
||||
# $$d_h^{i,f,g,o}(z_h^{i,f,g,o}) = lin_{dh}^{i,f,g,o}(z_h^{i,f,g,o})$$
|
||||
d_h = self.d_h[i](z_h[i])
|
||||
# $$d_x^{i,f,g,o}(z_x^{i,f,g,o}) = lin_{dx}^{i,f,g,o}(z_x^{i,f,g,o})$$
|
||||
d_x = self.d_x[i](z_x[i])
|
||||
|
||||
# \begin{align}
|
||||
# {i,f,g,o} = LN(&\textcolor{lightgreen}{d_h^{i,f,g,o}(z_h) \odot (W_h^{i,f,g,o} h_{t-1})} \\
|
||||
# + &\textcolor{lightgreen}{d_x^{i,f,g,o}(z_x) \odot (W_h^{i,f,g,o} x_t)} \\
|
||||
# + &d_b^{i,f,g,o}(z_b))
|
||||
# \end{align}
|
||||
y = d_h * torch.einsum('ij,bj->bi', self.w_h[i], h) + \
|
||||
d_x * torch.einsum('ij,bj->bi', self.w_x[i], x) + \
|
||||
self.d_b[i](z_b[i])
|
||||
|
||||
ifgo.append(self.layer_norm[i](y))
|
||||
|
||||
# $$i_t, f_t, g_t, o_t$$
|
||||
i, f, g, o = ifgo
|
||||
|
||||
# $$c_t = \sigma(f_t) \odot c_{t-1} + \sigma(i_t) \odot \tanh(g_t) $$
|
||||
c_next = torch.sigmoid(f) * c + torch.sigmoid(i) * torch.tanh(g)
|
||||
|
||||
# $$h_t = \sigma(o_t) \odot \tanh(LN(c_t))$$
|
||||
h_next = torch.sigmoid(o) * torch.tanh(self.layer_norm_c(c_next))
|
||||
|
||||
return h_next, c_next, h_hat, c_hat
|
||||
|
||||
|
||||
class HyperLSTM(nn.Module):
|
||||
"""
|
||||
# HyperLSTM module
|
||||
"""
|
||||
|
||||
def __init__(self, input_size: int, hidden_size: int, hyper_size: int, n_z: int, n_layers: int):
|
||||
"""
|
||||
Create a network of `n_layers` of HyperLSTM.
|
||||
"""
|
||||
|
||||
super().__init__()
|
||||
|
||||
# Store sizes to initialize state
|
||||
self.n_layers = n_layers
|
||||
self.hidden_size = hidden_size
|
||||
self.hyper_size = hyper_size
|
||||
|
||||
# Create cells for each layer. Note that only the first layer gets the input directly.
|
||||
# Rest of the layers get the input from the layer below
|
||||
self.cells = nn.ModuleList([HyperLSTMCell(input_size, hidden_size, hyper_size, n_z)] +
|
||||
[HyperLSTMCell(hidden_size, hidden_size, hyper_size, n_z) for _ in
|
||||
range(n_layers - 1)])
|
||||
|
||||
def forward(self, x: torch.Tensor,
|
||||
state: Optional[Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]] = None):
|
||||
"""
|
||||
* `x` has shape `[n_steps, batch_size, input_size]` and
|
||||
* `state` is a tuple of $h, c, \hat{h}, \hat{c}$.
|
||||
$h, c$ have shape `[batch_size, hidden_size]` and
|
||||
$\hat{h}, \hat{c}$ have shape `[batch_size, hyper_size]`.
|
||||
"""
|
||||
n_steps, batch_size = x.shape[:2]
|
||||
|
||||
# Initialize the state with zeros if `None`
|
||||
if state is None:
|
||||
h = [x.new_zeros(batch_size, self.hidden_size) for _ in range(self.n_layers)]
|
||||
c = [x.new_zeros(batch_size, self.hidden_size) for _ in range(self.n_layers)]
|
||||
h_hat = [x.new_zeros(batch_size, self.hyper_size) for _ in range(self.n_layers)]
|
||||
c_hat = [x.new_zeros(batch_size, self.hyper_size) for _ in range(self.n_layers)]
|
||||
#
|
||||
else:
|
||||
(h, c, h_hat, c_hat) = state
|
||||
# Reverse stack the tensors to get the states of each layer
|
||||
#
|
||||
# 📝 You can just work with the tensor itself but this is easier to debug
|
||||
h, c = list(torch.unbind(h)), list(torch.unbind(c))
|
||||
h_hat, c_hat = list(torch.unbind(h_hat)), list(torch.unbind(c_hat))
|
||||
|
||||
# Collect the outputs of the final layer at each step
|
||||
out = []
|
||||
for t in range(n_steps):
|
||||
# Input to the first layer is the input itself
|
||||
inp = x[t]
|
||||
# Loop through the layers
|
||||
for layer in range(self.n_layers):
|
||||
# Get the state of the layer
|
||||
h[layer], c[layer], h_hat[layer], c_hat[layer] = \
|
||||
self.cells[layer](inp, h[layer], c[layer], h_hat[layer], c_hat[layer])
|
||||
# Input to the next layer is the state of this layer
|
||||
inp = h[layer]
|
||||
# Collect the output $h$ of the final layer
|
||||
out.append(h[-1])
|
||||
|
||||
# Stack the outputs and states
|
||||
out = torch.stack(out)
|
||||
h = torch.stack(h)
|
||||
c = torch.stack(c)
|
||||
h_hat = torch.stack(h_hat)
|
||||
c_hat = torch.stack(c_hat)
|
||||
|
||||
#
|
||||
return out, (h, c, h_hat, c_hat)
|
||||
@@ -0,0 +1,151 @@
|
||||
"""
|
||||
---
|
||||
title: Low-Rank Adaptation (LoRA)
|
||||
summary: >
|
||||
Annotated implementation of RoRA from paper
|
||||
LoRA: Low-Rank Adaptation of Large Language Models
|
||||
---
|
||||
|
||||
# Low-Rank Adaptation (LoRA)
|
||||
|
||||
This is an implementation of
|
||||
[Low-Rank Adaptation (LoRA)](https://arxiv.org/abs/2106.09685)
|
||||
in [PyTorch](https://pytorch.org).
|
||||
|
||||
Low-Rank Adaptation (LoRA) freezes pre-trained model weights and injects
|
||||
trainable rank decomposition matrices into each layer of the transformer.
|
||||
This makes it possible to efficiently fine-tune large language models by
|
||||
reducing trainable parameters by a large factor.
|
||||
|
||||
Here's [the training code](experiment.html) for training a GPT2 model with LoRA
|
||||
on Tiny Shakespeare dataset.
|
||||
"""
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class Linear(nn.Module):
|
||||
"""
|
||||
## LoRA Linear Layer
|
||||
|
||||
LoRA linear layer adds a low-rank decomposition to the pre-trained
|
||||
weight matrix ($W_0 \in \mathbb{R}^{d \times k}$)
|
||||
of the linear layer.
|
||||
|
||||
$$W_0 + \Delta W = W_0 + BA$$
|
||||
|
||||
, where $B \in \mathbb{R}^{d \times r}$, $A \in \mathbb{R}^{r \times k}$,
|
||||
and the rank $r \ll min(d, k)$.
|
||||
|
||||
All parameters are frozen except $A$ and $B$.
|
||||
|
||||
$\Delta W$ is initialized to be zero at the beginning of the training.
|
||||
|
||||
They multiple $x \Delta W^T$ by $\frac{\alpha}{r}$ where $\alpha$ is a hyper-parameter.
|
||||
Once $\alpha$ is tuned it can be kept the same when varying $r$.
|
||||
"""
|
||||
|
||||
def __init__(self, in_features: int, out_features: int, bias: bool,
|
||||
r: int, alpha: int = None):
|
||||
"""
|
||||
:param in_features: is the number of input features of the linear layer
|
||||
:param out_features: is the number of output features of the linear layer
|
||||
:param bias: is a flag indicating if there is a bias parameter
|
||||
:param r: is the rank of the decomposition $r$
|
||||
:param alpha: is the scaling factor $\alpha$
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
# Set $\alpha = r$ is not provided. i.e. make the scaling factor $\frac{\alpha}{r} = 1$.
|
||||
if alpha is None:
|
||||
alpha = r
|
||||
|
||||
# The pre-trained weight $W_0$
|
||||
self.weight = nn.Parameter(torch.empty((out_features, in_features)))
|
||||
# Freeze it
|
||||
self.weight.requires_grad = False
|
||||
|
||||
if bias:
|
||||
# Bias parameter $b_0$ (also frozen)
|
||||
self.bias = nn.Parameter(torch.empty(out_features))
|
||||
self.bias.requires_grad = False
|
||||
else:
|
||||
# No bias parameter
|
||||
self.bias = None
|
||||
|
||||
# scaling factor $\frac{\alpha}{r}$
|
||||
self.scaling = alpha / r
|
||||
# Matrix $A \in \mathbb{R}^{r \times k}$
|
||||
self.lora_a = nn.Parameter(torch.empty((r, in_features)))
|
||||
# Matrix $B \in \mathbb{R}^{d \times r}$, we keep $A$ and $B$ transposed
|
||||
self.lora_b = nn.Parameter(torch.empty((out_features, r)))
|
||||
|
||||
with torch.no_grad():
|
||||
# Initialize $A$ similar to a weight matrix in a normal linear layer
|
||||
nn.init.kaiming_uniform_(self.lora_a, a=5 ** 0.5)
|
||||
# Initialize $B$ to $0$ so that $\Delta W = BA$ is $0$ at initialization
|
||||
nn.init.zeros_(self.lora_b)
|
||||
|
||||
def forward(self, x: torch.Tensor):
|
||||
# Compute $x W_0^T + b_0$
|
||||
result = nn.functional.linear(x, self.weight, bias=self.bias)
|
||||
|
||||
# Add $\frac{\alpha}{r} x \Delta W^T = \frac{\alpha}{r} x {(BA)}^T = \frac{\alpha}{r} x A^T B^T$
|
||||
result += (x @ self.lora_a.T @ self.lora_b.T) * self.scaling
|
||||
|
||||
#
|
||||
return result
|
||||
|
||||
|
||||
class Embedding(nn.Module):
|
||||
"""
|
||||
## LoRA Embedding Layer
|
||||
|
||||
Similar to LoRA linear layer this adds a low-rank decomposition to the pre-trained
|
||||
embedding weights matrix ($W_0 \in \mathbb{R}^{d \times k}$).
|
||||
|
||||
$$W_0 + \Delta W = W_0 + BA$$
|
||||
"""
|
||||
|
||||
def __init__(self, num_embeddings: int, embedding_dim: int,
|
||||
r: int, alpha: int = None):
|
||||
"""
|
||||
|
||||
:param num_embeddings: is the number of embeddings
|
||||
:param embedding_dim: is the number embedding dimensions
|
||||
:param r: is the rank of the decomposition $r$
|
||||
:param alpha: is the scaling factor $\alpha$
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
# Set $\alpha = r$ is not provided. i.e. make the scaling factor $\frac{\alpha}{r} = 1$.
|
||||
if alpha is None:
|
||||
alpha = r
|
||||
|
||||
# The pre-trained embedding weights $W_0^T$ (frozen)
|
||||
self.weight = nn.Parameter(torch.empty((num_embeddings, embedding_dim)))
|
||||
self.weight.requires_grad = False
|
||||
|
||||
# scaling factor $\frac{\alpha}{r}$
|
||||
self.scaling = alpha / r
|
||||
# Matrix $A \in \mathbb{R}^{r \times k}$
|
||||
self.lora_a = nn.Parameter(torch.empty((r, num_embeddings)))
|
||||
# Matrix $B \in \mathbb{R}^{d \times r}$
|
||||
self.lora_b = nn.Parameter(torch.empty((embedding_dim, r)))
|
||||
|
||||
with torch.no_grad():
|
||||
# Initialize $A$ with a normal distribution
|
||||
nn.init.normal_(self.lora_a)
|
||||
# Initialize $B$ to $0$ so that $\Delta W = BA$ is $0$ at initialization
|
||||
nn.init.zeros_(self.lora_b)
|
||||
|
||||
def forward(self, x: torch.Tensor):
|
||||
# Compute the embeddings $\text{onehot}(x) W_0$
|
||||
result = nn.functional.embedding(x, self.weight)
|
||||
|
||||
# Add $\frac{\alpha}{r} \text{onehot}(x) \Delta W^T = \frac{\alpha}{r} \text{onehot}(x) A^T B^T$
|
||||
result += (nn.functional.embedding(x, self.lora_a.T) @ self.lora_b.T) * self.scaling
|
||||
|
||||
#
|
||||
return result
|
||||
@@ -0,0 +1,97 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"metadata": {},
|
||||
"cell_type": "code",
|
||||
"outputs": [],
|
||||
"execution_count": null,
|
||||
"source": "!pip install labml-nn",
|
||||
"id": "c5ed37230628ee76"
|
||||
},
|
||||
{
|
||||
"metadata": {},
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"from labml_nn.lora.experiment import Trainer\n",
|
||||
"from labml import experiment"
|
||||
],
|
||||
"id": "1b9da2e59ffce5d5",
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"id": "initial_id",
|
||||
"metadata": {
|
||||
"collapsed": true
|
||||
},
|
||||
"source": "experiment.create(name=\"lora_gpt2\")",
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
},
|
||||
{
|
||||
"metadata": {},
|
||||
"cell_type": "code",
|
||||
"source": "trainer = Trainer()",
|
||||
"id": "31c9bc08eca2592",
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
},
|
||||
{
|
||||
"metadata": {},
|
||||
"cell_type": "code",
|
||||
"source": "experiment.configs(trainer)",
|
||||
"id": "fb6ce74326558948",
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
},
|
||||
{
|
||||
"metadata": {},
|
||||
"cell_type": "code",
|
||||
"source": "trainer.initialize()",
|
||||
"id": "1456cfab47dee3b",
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
},
|
||||
{
|
||||
"metadata": {},
|
||||
"cell_type": "code",
|
||||
"source": [
|
||||
"with experiment.start():\n",
|
||||
" trainer.run()"
|
||||
],
|
||||
"id": "3fe4068fd2df9094",
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
},
|
||||
{
|
||||
"metadata": {},
|
||||
"cell_type": "code",
|
||||
"source": "",
|
||||
"id": "d3c3c723ebbe854a",
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {
|
||||
"display_name": "Python (ml)",
|
||||
"language": "python",
|
||||
"name": "ml"
|
||||
},
|
||||
"language_info": {
|
||||
"codemirror_mode": {
|
||||
"name": "ipython",
|
||||
"version": 2
|
||||
},
|
||||
"file_extension": ".py",
|
||||
"mimetype": "text/x-python",
|
||||
"name": "python",
|
||||
"nbconvert_exporter": "python",
|
||||
"pygments_lexer": "ipython2",
|
||||
"version": "2.7.6"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 5
|
||||
}
|
||||
@@ -0,0 +1,189 @@
|
||||
"""
|
||||
---
|
||||
title: Finetune GPT-2 with LoRA
|
||||
summary: This is training code with notes for fine-tuning pre-trained GPT-2 model with LoRA.
|
||||
---
|
||||
|
||||
# Finetune [GPT-2](gpt2.html) with [LoRA](index.html)
|
||||
|
||||
Here's a Colab notebook for training a feedback transformer on Tiny Shakespeare dataset.
|
||||
|
||||
[](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/lora/experiment.ipynb)
|
||||
"""
|
||||
|
||||
import torch
|
||||
from torch.optim import Adam
|
||||
from torch.utils.data import DataLoader, TensorDataset
|
||||
from transformers import AutoTokenizer, AutoModelForCausalLM
|
||||
|
||||
from labml import lab, monit, tracker
|
||||
from labml.configs import BaseConfigs, option
|
||||
from labml.utils.download import download_file
|
||||
from labml_nn.helpers.device import DeviceConfigs
|
||||
from labml_nn.lora.gpt2 import GPTModel
|
||||
|
||||
|
||||
class Trainer(BaseConfigs):
|
||||
"""
|
||||
## Trainer configurations and the training loop
|
||||
|
||||
The default configs can and will be over-ridden when we start the experiment
|
||||
"""
|
||||
device: torch.device = DeviceConfigs()
|
||||
|
||||
# GPT-2 configs
|
||||
layer_norm_epsilon: float = 1e-05
|
||||
d_model: int = 768
|
||||
n_layers: int = 12
|
||||
n_heads: int = 12
|
||||
n_positions: int = 1024
|
||||
vocab_size: int = 50257
|
||||
|
||||
# Training configs
|
||||
epochs: int = 10
|
||||
batch_size: int = 32
|
||||
learning_rate: float = 1e-4
|
||||
context_len: int = 512
|
||||
|
||||
# LoRA rank
|
||||
lora_r: int = 32
|
||||
|
||||
# Dataset
|
||||
text: TensorDataset = "tiny_shakespeare"
|
||||
# Huggingface tokenizer
|
||||
tokenizer = AutoTokenizer.from_pretrained("gpt2")
|
||||
# [GPT2 model](gpt2.html)
|
||||
model: GPTModel
|
||||
# Optimizer
|
||||
optimizer: torch.optim.Adam
|
||||
# Cross entropy loss
|
||||
loss_func = torch.nn.CrossEntropyLoss()
|
||||
# Dataloader
|
||||
data_loader: DataLoader
|
||||
|
||||
def _load_pretrained_weights(self):
|
||||
"""
|
||||
### Load pre-trained [GPT-2 from huggingface](https://huggingface.co/openai-community/gpt2)
|
||||
"""
|
||||
|
||||
# Load the huggingface model and get the parameters
|
||||
hf_model = AutoModelForCausalLM.from_pretrained("gpt2")
|
||||
state_dict = hf_model.state_dict()
|
||||
|
||||
# Transformer embedding and prediction layer parameter mapping (`hf: ours`)
|
||||
mapping = {
|
||||
'transformer.wte.weight': 'token_embedding.weight',
|
||||
'transformer.wpe.weight': 'position_embedding.weight',
|
||||
'transformer.ln_f.weight': 'final_norm.weight',
|
||||
'transformer.ln_f.bias': 'final_norm.bias',
|
||||
'lm_head.weight': 'lm_head.weight'
|
||||
}
|
||||
|
||||
# Mapping (`hf: ours`) of decoder layers
|
||||
for i in range(12):
|
||||
mapping[f'transformer.h.{i}.ln_1.weight'] = f'blocks.{i}.attn_norm.weight'
|
||||
mapping[f'transformer.h.{i}.ln_1.bias'] = f'blocks.{i}.attn_norm.bias'
|
||||
mapping[f'transformer.h.{i}.attn.c_attn.weight'] = f'blocks.{i}.attn.qkv_projection.weight'
|
||||
mapping[f'transformer.h.{i}.attn.c_attn.bias'] = f'blocks.{i}.attn.qkv_projection.bias'
|
||||
mapping[f'transformer.h.{i}.attn.c_proj.weight'] = f'blocks.{i}.attn.output_projection.weight'
|
||||
mapping[f'transformer.h.{i}.attn.c_proj.bias'] = f'blocks.{i}.attn.output_projection.bias'
|
||||
mapping[f'transformer.h.{i}.ln_2.weight'] = f'blocks.{i}.ffn_norm.weight'
|
||||
mapping[f'transformer.h.{i}.ln_2.bias'] = f'blocks.{i}.ffn_norm.bias'
|
||||
mapping[f'transformer.h.{i}.mlp.c_fc.weight'] = f'blocks.{i}.ffn.linear_in.weight'
|
||||
mapping[f'transformer.h.{i}.mlp.c_fc.bias'] = f'blocks.{i}.ffn.linear_in.bias'
|
||||
mapping[f'transformer.h.{i}.mlp.c_proj.weight'] = f'blocks.{i}.ffn.linear_out.weight'
|
||||
mapping[f'transformer.h.{i}.mlp.c_proj.bias'] = f'blocks.{i}.ffn.linear_out.bias'
|
||||
|
||||
# Move the parameters based on mapping
|
||||
new_state_dict = {}
|
||||
for old_key, new_key in mapping.items():
|
||||
if old_key in state_dict:
|
||||
new_state_dict[new_key] = state_dict[old_key]
|
||||
|
||||
# GPT-2 hugging face uses 1D Convolution layers. We need to transpose those weights since we use linear layers
|
||||
convo_layers = ([f'blocks.{i}.ffn.linear_in.weight' for i in range(12)] +
|
||||
[f'blocks.{i}.ffn.linear_out.weight' for i in range(12)] +
|
||||
[f'blocks.{i}.attn.qkv_projection.weight' for i in range(12)] +
|
||||
[f'blocks.{i}.attn.output_projection.weight' for i in range(12)])
|
||||
|
||||
for layer in convo_layers:
|
||||
new_state_dict[layer] = torch.transpose(new_state_dict[layer], 0, 1)
|
||||
|
||||
# Load out model. We use `strict = False` because the state does not have LoRA weights
|
||||
missing_keys, unexpected_keys = self.model.load_state_dict(new_state_dict, strict=False)
|
||||
|
||||
# make sure that only lora weights are not loaded
|
||||
assert all('lora' in key for key in missing_keys)
|
||||
assert not unexpected_keys
|
||||
|
||||
def initialize(self):
|
||||
"""
|
||||
### Initialize the model, optimizer and dataloader
|
||||
"""
|
||||
# Initialize the [GPT2 model](gpt2.html)
|
||||
self.model = GPTModel(
|
||||
layer_norm_epsilon=self.layer_norm_epsilon,
|
||||
d_model=self.d_model,
|
||||
n_layers=self.n_layers,
|
||||
n_heads=self.n_heads,
|
||||
n_positions=self.n_positions,
|
||||
vocab_size=self.vocab_size,
|
||||
r=self.lora_r,
|
||||
)
|
||||
self.model.to(self.device)
|
||||
# Load pre-trained model weights
|
||||
self._load_pretrained_weights()
|
||||
|
||||
# Initialize the optimizer
|
||||
self.optimizer = Adam(self.model.parameters(), lr=self.learning_rate)
|
||||
|
||||
# Initialize the data loader
|
||||
self.data_loader = DataLoader(self.text, batch_size=self.batch_size, shuffle=True)
|
||||
|
||||
def run(self):
|
||||
"""
|
||||
### Training loop
|
||||
"""
|
||||
|
||||
for _ in monit.loop(self.epochs):
|
||||
# `inputs` has shape `[batch_size, seq_len]`
|
||||
for (inputs,) in monit.iterate('Train', self.data_loader):
|
||||
# Move `inputs` to device
|
||||
inputs = inputs.to(self.device)
|
||||
# Call the model, with the all but the last token
|
||||
logits = self.model(inputs[:, :-1])
|
||||
# Get cross entropy loss
|
||||
loss = self.loss_func(logits.reshape(-1, logits.shape[-1]), inputs[:, 1:].reshape(-1))
|
||||
|
||||
# Make gradients 0
|
||||
self.optimizer.zero_grad()
|
||||
# Compute gradients
|
||||
loss.backward()
|
||||
# Optimize
|
||||
self.optimizer.step()
|
||||
|
||||
# Log the loss
|
||||
tracker.save({'loss': loss})
|
||||
tracker.add_global_step()
|
||||
#
|
||||
tracker.new_line()
|
||||
|
||||
|
||||
@option(Trainer.text)
|
||||
def tiny_shakespeare(c: Trainer):
|
||||
"""
|
||||
### Tiny Shakespeare dataset
|
||||
|
||||
It will download from the url if not present
|
||||
"""
|
||||
path = lab.get_data_path() / 'tiny_shakespeare.txt'
|
||||
if not path.exists():
|
||||
download_file("https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt", path)
|
||||
with open(path, 'r', encoding='utf-8') as f:
|
||||
text = f.read()
|
||||
|
||||
tokens = c.tokenizer.encode(text)
|
||||
num_batches = len(tokens) // (c.batch_size * c.context_len)
|
||||
tokens = tokens[:num_batches * c.batch_size * c.context_len]
|
||||
input_ids = torch.tensor(tokens).view(-1, c.context_len)
|
||||
return TensorDataset(input_ids)
|
||||
@@ -0,0 +1,192 @@
|
||||
"""
|
||||
---
|
||||
title: GPT-2 with LoRA
|
||||
summary: GPT-2 implementation with LoRA modules
|
||||
---
|
||||
|
||||
# GPT-2 with [LoRA modules](index.html)
|
||||
|
||||
Here's [the training code](experiment.html) for training a GPT2 model with LoRA
|
||||
on Tiny Shakespeare dataset.
|
||||
"""
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from labml_nn.lora import Linear, Embedding
|
||||
|
||||
|
||||
class FFN(nn.Module):
|
||||
"""
|
||||
### Feedforward Network
|
||||
"""
|
||||
|
||||
def __init__(self, d_model: int, d_ff: int, r: int):
|
||||
"""
|
||||
:param d_model: is the number of dimensions
|
||||
:param d_ff: is the size of the hidden dimension
|
||||
:param r: is the lora rank
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
# The linear layers and the activation
|
||||
self.linear_in = Linear(d_model, d_ff, r=r, bias=True)
|
||||
self.linear_out = Linear(d_ff, d_model, r=r, bias=True)
|
||||
self.act = nn.GELU()
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
:param x: is the embeddings tensor with shape `[batch_size, seq_len, d_model]`
|
||||
"""
|
||||
x = self.linear_in(x)
|
||||
x = self.act(x)
|
||||
x = self.linear_out(x)
|
||||
return x
|
||||
|
||||
|
||||
class MultiHeadAttention(nn.Module):
|
||||
"""
|
||||
### Multi-Head Attention
|
||||
"""
|
||||
|
||||
def __init__(self, d_model: int, n_heads: int, r: int):
|
||||
"""
|
||||
:param d_model: is the number of dimensions in the embeddings
|
||||
:param n_heads: is the number of heads
|
||||
:param r: is the lora rank
|
||||
"""
|
||||
super().__init__()
|
||||
self.d_model = d_model
|
||||
self.n_heads = n_heads
|
||||
self.d_head = d_model // n_heads
|
||||
|
||||
# Linear transformation for QKV
|
||||
self.qkv_projection = Linear(d_model, d_model * 3, r=r, bias=True)
|
||||
# Output projection
|
||||
self.output_projection = Linear(d_model, d_model, r=r, bias=True)
|
||||
|
||||
def _split_heads(self, x: torch.Tensor):
|
||||
"""
|
||||
:param x: is the tensor with shape `[batch_size, seq_len, d_model]`
|
||||
"""
|
||||
# Split last dimension to `[n_heads, d_head]`
|
||||
x = x.view(x.shape[:-1] + (self.n_heads, self.d_head))
|
||||
# Reorder to `[batch_size, head, seq_length, d_head]`
|
||||
return x.permute(0, 2, 1, 3)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
:param x: is the embeddings tensor with shape `[batch_size, seq_len, d_model]`
|
||||
"""
|
||||
batch_size, seq_length, _ = x.shape
|
||||
|
||||
# Get query, key and value
|
||||
q, k, v = self.qkv_projection(x).split(self.d_model, dim=-1)
|
||||
|
||||
# Transform them from shape `[batch_size, seq_len, d_model]` to `[batch_size, head, seq_length, d_head]`
|
||||
q = self._split_heads(q)
|
||||
k = self._split_heads(k)
|
||||
v = self._split_heads(v)
|
||||
|
||||
# Apply causal attention
|
||||
attn_output = torch.nn.functional.scaled_dot_product_attention(q, k, v, is_causal=True)
|
||||
|
||||
# Transform them from shape `[batch_size, head, seq_length, d_head]` to `[batch_size, seq_len, d_model]`
|
||||
attn_output = attn_output.permute(0, 2, 1, 3).reshape(batch_size, seq_length, self.d_model)
|
||||
|
||||
# Final project
|
||||
return self.output_projection(attn_output)
|
||||
|
||||
|
||||
class Block(nn.Module):
|
||||
"""
|
||||
### Decoder block
|
||||
"""
|
||||
|
||||
def __init__(self, d_model: int, n_heads: int, layer_norm_epsilon: float, r: int):
|
||||
"""
|
||||
:param d_model: is the number of dimensions in the embeddings
|
||||
:param n_heads: is the number of heads
|
||||
:param layer_norm_epsilon: is the layer norm epsilon
|
||||
:param r: is the lora rank
|
||||
"""
|
||||
super().__init__()
|
||||
# Attention pre-normalization layer
|
||||
self.attn_norm = nn.LayerNorm(d_model, eps=layer_norm_epsilon)
|
||||
# Attention layer
|
||||
self.attn = MultiHeadAttention(d_model, n_heads, r)
|
||||
# FFN pre-normalization layer
|
||||
self.ffn_norm = nn.LayerNorm(d_model, eps=layer_norm_epsilon)
|
||||
# Feed-forward network
|
||||
self.ffn = FFN(d_model, d_model * 4, r)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
:param x: is the embeddings tensor with shape `[batch_size, seq_len, d_model]`
|
||||
"""
|
||||
# Attention
|
||||
x = x + self.attn(self.attn_norm(x))
|
||||
# FFN
|
||||
x = x + self.ffn(self.ffn_norm(x))
|
||||
|
||||
return x
|
||||
|
||||
|
||||
class GPTModel(nn.Module):
|
||||
"""
|
||||
## GPT2 Model
|
||||
"""
|
||||
|
||||
def __init__(self, *, d_model: int,
|
||||
n_heads: int, n_layers: int,
|
||||
n_positions: int,
|
||||
layer_norm_epsilon: float,
|
||||
vocab_size: int, r: int):
|
||||
"""
|
||||
:param d_model: is the number of dimensions in the embeddings
|
||||
:param n_heads: is the number of attention heads
|
||||
:param n_layers: is the number of decoder layers
|
||||
:param n_positions: is the number of positional embeddings
|
||||
:param layer_norm_epsilon: is the layer norm epsilon
|
||||
:param vocab_size: is the vocabulary size
|
||||
:param r: is the lora rank
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
# Token and absolute positional embeddings
|
||||
self.token_embedding = Embedding(vocab_size, d_model, r=r)
|
||||
self.position_embedding = Embedding(n_positions, d_model, r=r)
|
||||
|
||||
# Decoder blocks
|
||||
self.blocks = nn.ModuleList([Block(d_model, n_heads, layer_norm_epsilon, r=r)
|
||||
for _ in range(n_layers)])
|
||||
|
||||
# Final layer norm
|
||||
self.final_norm = nn.LayerNorm(d_model, eps=layer_norm_epsilon)
|
||||
# Projection layer to logit space
|
||||
self.lm_head = Linear(d_model, vocab_size, r=r, bias=False)
|
||||
|
||||
def forward(self, input_ids: torch.Tensor):
|
||||
"""
|
||||
:param input_ids: has shape `[batch_size, seq_len]`
|
||||
"""
|
||||
batch_size, seq_len = input_ids.shape
|
||||
|
||||
# Get token embeddings
|
||||
token_embeddings = self.token_embedding(input_ids)
|
||||
# Get position ids
|
||||
position_ids = torch.arange(seq_len, device=input_ids.device)[None, :]
|
||||
# Get position embeddings
|
||||
position_embeddings = self.position_embedding(position_ids)
|
||||
|
||||
# Add position embeddings
|
||||
x = token_embeddings + position_embeddings
|
||||
|
||||
# Run through transformer blocks
|
||||
for block in self.blocks:
|
||||
x = block(x)
|
||||
|
||||
# Final normalization
|
||||
x = self.final_norm(x)
|
||||
# Get logits from projection layer
|
||||
return self.lm_head(x)
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user