chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,17 @@
|
||||
"""
|
||||
---
|
||||
title: Normalization Layers
|
||||
summary: >
|
||||
A set of PyTorch implementations/tutorials of normalization layers.
|
||||
---
|
||||
|
||||
# Normalization Layers
|
||||
|
||||
* [Batch Normalization](batch_norm/index.html)
|
||||
* [Layer Normalization](layer_norm/index.html)
|
||||
* [Instance Normalization](instance_norm/index.html)
|
||||
* [Group Normalization](group_norm/index.html)
|
||||
* [Weight Standardization](weight_standardization/index.html)
|
||||
* [Batch-Channel Normalization](batch_channel_norm/index.html)
|
||||
* [DeepNorm](deep_norm/index.html)
|
||||
"""
|
||||
@@ -0,0 +1,235 @@
|
||||
"""
|
||||
---
|
||||
title: Batch-Channel Normalization
|
||||
summary: >
|
||||
A PyTorch implementation/tutorial of Batch-Channel Normalization.
|
||||
---
|
||||
|
||||
# Batch-Channel Normalization
|
||||
|
||||
This is a [PyTorch](https://pytorch.org) implementation of Batch-Channel Normalization from the paper
|
||||
[Micro-Batch Training with Batch-Channel Normalization and Weight Standardization](https://arxiv.org/abs/1903.10520).
|
||||
We also have an [annotated implementation of Weight Standardization](../weight_standardization/index.html).
|
||||
|
||||
Batch-Channel Normalization performs batch normalization followed
|
||||
by a channel normalization (similar to a [Group Normalization](../group_norm/index.html).
|
||||
When the batch size is small a running mean and variance is used for
|
||||
batch normalization.
|
||||
|
||||
Here is [the training code](../weight_standardization/experiment.html) for training
|
||||
a VGG network that uses weight standardization to classify CIFAR-10 data.
|
||||
|
||||
[](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/normalization/weight_standardization/experiment.ipynb)
|
||||
"""
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from labml_nn.normalization.batch_norm import BatchNorm
|
||||
|
||||
|
||||
class BatchChannelNorm(nn.Module):
|
||||
"""
|
||||
## Batch-Channel Normalization
|
||||
|
||||
This first performs a batch normalization - either [normal batch norm](../batch_norm/index.html)
|
||||
or a batch norm with
|
||||
estimated mean and variance (exponential mean/variance over multiple batches).
|
||||
Then a channel normalization performed.
|
||||
"""
|
||||
|
||||
def __init__(self, channels: int, groups: int,
|
||||
eps: float = 1e-5, momentum: float = 0.1, estimate: bool = True):
|
||||
"""
|
||||
* `channels` is the number of features in the input
|
||||
* `groups` is the number of groups the features are divided into
|
||||
* `eps` is $\epsilon$, used in $\sqrt{Var[x^{(k)}] + \epsilon}$ for numerical stability
|
||||
* `momentum` is the momentum in taking the exponential moving average
|
||||
* `estimate` is whether to use running mean and variance for batch norm
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
# Use estimated batch norm or normal batch norm.
|
||||
if estimate:
|
||||
self.batch_norm = EstimatedBatchNorm(channels,
|
||||
eps=eps, momentum=momentum)
|
||||
else:
|
||||
self.batch_norm = BatchNorm(channels,
|
||||
eps=eps, momentum=momentum)
|
||||
|
||||
# Channel normalization
|
||||
self.channel_norm = ChannelNorm(channels, groups, eps)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.batch_norm(x)
|
||||
return self.channel_norm(x)
|
||||
|
||||
|
||||
class EstimatedBatchNorm(nn.Module):
|
||||
"""
|
||||
## Estimated Batch Normalization
|
||||
|
||||
When input $X \in \mathbb{R}^{B \times C \times H \times W}$ is a batch of image representations,
|
||||
where $B$ is the batch size, $C$ is the number of channels, $H$ is the height and $W$ is the width.
|
||||
$\gamma \in \mathbb{R}^{C}$ and $\beta \in \mathbb{R}^{C}$.
|
||||
|
||||
$$\dot{X}_{\cdot, C, \cdot, \cdot} = \gamma_C
|
||||
\frac{X_{\cdot, C, \cdot, \cdot} - \hat{\mu}_C}{\hat{\sigma}_C}
|
||||
+ \beta_C$$
|
||||
|
||||
where,
|
||||
|
||||
\begin{align}
|
||||
\hat{\mu}_C &\longleftarrow (1 - r)\hat{\mu}_C + r \frac{1}{B H W} \sum_{b,h,w} X_{b,c,h,w} \\
|
||||
\hat{\sigma}^2_C &\longleftarrow (1 - r)\hat{\sigma}^2_C + r \frac{1}{B H W} \sum_{b,h,w} \big(X_{b,c,h,w} - \hat{\mu}_C \big)^2
|
||||
\end{align}
|
||||
|
||||
are the running mean and variances. $r$ is the momentum for calculating the exponential mean.
|
||||
"""
|
||||
def __init__(self, channels: int,
|
||||
eps: float = 1e-5, momentum: float = 0.1, affine: bool = True):
|
||||
"""
|
||||
* `channels` is the number of features in the input
|
||||
* `eps` is $\epsilon$, used in $\sqrt{Var[x^{(k)}] + \epsilon}$ for numerical stability
|
||||
* `momentum` is the momentum in taking the exponential moving average
|
||||
* `estimate` is whether to use running mean and variance for batch norm
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
self.eps = eps
|
||||
self.momentum = momentum
|
||||
self.affine = affine
|
||||
self.channels = channels
|
||||
|
||||
# Channel wise transformation parameters
|
||||
if self.affine:
|
||||
self.scale = nn.Parameter(torch.ones(channels))
|
||||
self.shift = nn.Parameter(torch.zeros(channels))
|
||||
|
||||
# Tensors for $\hat{\mu}_C$ and $\hat{\sigma}^2_C$
|
||||
self.register_buffer('exp_mean', torch.zeros(channels))
|
||||
self.register_buffer('exp_var', torch.ones(channels))
|
||||
|
||||
def forward(self, x: torch.Tensor):
|
||||
"""
|
||||
`x` is a tensor of shape `[batch_size, channels, *]`.
|
||||
`*` denotes any number of (possibly 0) dimensions.
|
||||
For example, in an image (2D) convolution this will be
|
||||
`[batch_size, channels, height, width]`
|
||||
"""
|
||||
# Keep old shape
|
||||
x_shape = x.shape
|
||||
# Get the batch size
|
||||
batch_size = x_shape[0]
|
||||
|
||||
# Sanity check to make sure the number of features is correct
|
||||
assert self.channels == x.shape[1]
|
||||
|
||||
# Reshape into `[batch_size, channels, n]`
|
||||
x = x.view(batch_size, self.channels, -1)
|
||||
|
||||
# Update $\hat{\mu}_C$ and $\hat{\sigma}^2_C$ in training mode only
|
||||
if self.training:
|
||||
# No backpropagation through $\hat{\mu}_C$ and $\hat{\sigma}^2_C$
|
||||
with torch.no_grad():
|
||||
# Calculate the mean across first and last dimensions;
|
||||
# $$\frac{1}{B H W} \sum_{b,h,w} X_{b,c,h,w}$$
|
||||
mean = x.mean(dim=[0, 2])
|
||||
# Calculate the squared mean across first and last dimensions;
|
||||
# $$\frac{1}{B H W} \sum_{b,h,w} X^2_{b,c,h,w}$$
|
||||
mean_x2 = (x ** 2).mean(dim=[0, 2])
|
||||
# Variance for each feature
|
||||
# $$\frac{1}{B H W} \sum_{b,h,w} \big(X_{b,c,h,w} - \hat{\mu}_C \big)^2$$
|
||||
var = mean_x2 - mean ** 2
|
||||
|
||||
# Update exponential moving averages
|
||||
#
|
||||
# \begin{align}
|
||||
# \hat{\mu}_C &\longleftarrow (1 - r)\hat{\mu}_C + r \frac{1}{B H W} \sum_{b,h,w} X_{b,c,h,w} \\
|
||||
# \hat{\sigma}^2_C &\longleftarrow (1 - r)\hat{\sigma}^2_C + r \frac{1}{B H W} \sum_{b,h,w} \big(X_{b,c,h,w} - \hat{\mu}_C \big)^2
|
||||
# \end{align}
|
||||
self.exp_mean = (1 - self.momentum) * self.exp_mean + self.momentum * mean
|
||||
self.exp_var = (1 - self.momentum) * self.exp_var + self.momentum * var
|
||||
|
||||
# Normalize
|
||||
# $$\frac{X_{\cdot, C, \cdot, \cdot} - \hat{\mu}_C}{\hat{\sigma}_C}$$
|
||||
x_norm = (x - self.exp_mean.view(1, -1, 1)) / torch.sqrt(self.exp_var + self.eps).view(1, -1, 1)
|
||||
# Scale and shift
|
||||
# $$ \gamma_C
|
||||
# \frac{X_{\cdot, C, \cdot, \cdot} - \hat{\mu}_C}{\hat{\sigma}_C}
|
||||
# + \beta_C$$
|
||||
if self.affine:
|
||||
x_norm = self.scale.view(1, -1, 1) * x_norm + self.shift.view(1, -1, 1)
|
||||
|
||||
# Reshape to original and return
|
||||
return x_norm.view(x_shape)
|
||||
|
||||
|
||||
class ChannelNorm(nn.Module):
|
||||
"""
|
||||
## Channel Normalization
|
||||
|
||||
This is similar to [Group Normalization](../group_norm/index.html) but affine transform is done group wise.
|
||||
"""
|
||||
|
||||
def __init__(self, channels, groups,
|
||||
eps: float = 1e-5, affine: bool = True):
|
||||
"""
|
||||
* `groups` is the number of groups the features are divided into
|
||||
* `channels` is the number of features in the input
|
||||
* `eps` is $\epsilon$, used in $\sqrt{Var[x^{(k)}] + \epsilon}$ for numerical stability
|
||||
* `affine` is whether to scale and shift the normalized value
|
||||
"""
|
||||
super().__init__()
|
||||
self.channels = channels
|
||||
self.groups = groups
|
||||
self.eps = eps
|
||||
self.affine = affine
|
||||
# Parameters for affine transformation.
|
||||
#
|
||||
# *Note that these transforms are per group, unlike in group norm where
|
||||
# they are transformed channel-wise.*
|
||||
if self.affine:
|
||||
self.scale = nn.Parameter(torch.ones(groups))
|
||||
self.shift = nn.Parameter(torch.zeros(groups))
|
||||
|
||||
def forward(self, x: torch.Tensor):
|
||||
"""
|
||||
`x` is a tensor of shape `[batch_size, channels, *]`.
|
||||
`*` denotes any number of (possibly 0) dimensions.
|
||||
For example, in an image (2D) convolution this will be
|
||||
`[batch_size, channels, height, width]`
|
||||
"""
|
||||
|
||||
# Keep the original shape
|
||||
x_shape = x.shape
|
||||
# Get the batch size
|
||||
batch_size = x_shape[0]
|
||||
# Sanity check to make sure the number of features is the same
|
||||
assert self.channels == x.shape[1]
|
||||
|
||||
# Reshape into `[batch_size, groups, n]`
|
||||
x = x.view(batch_size, self.groups, -1)
|
||||
|
||||
# Calculate the mean across last dimension;
|
||||
# i.e. the means for each sample and channel group $\mathbb{E}[x_{(i_N, i_G)}]$
|
||||
mean = x.mean(dim=[-1], keepdim=True)
|
||||
# Calculate the squared mean across last dimension;
|
||||
# i.e. the means for each sample and channel group $\mathbb{E}[x^2_{(i_N, i_G)}]$
|
||||
mean_x2 = (x ** 2).mean(dim=[-1], keepdim=True)
|
||||
# Variance for each sample and feature group
|
||||
# $Var[x_{(i_N, i_G)}] = \mathbb{E}[x^2_{(i_N, i_G)}] - \mathbb{E}[x_{(i_N, i_G)}]^2$
|
||||
var = mean_x2 - mean ** 2
|
||||
|
||||
# Normalize
|
||||
# $$\hat{x}_{(i_N, i_G)} =
|
||||
# \frac{x_{(i_N, i_G)} - \mathbb{E}[x_{(i_N, i_G)}]}{\sqrt{Var[x_{(i_N, i_G)}] + \epsilon}}$$
|
||||
x_norm = (x - mean) / torch.sqrt(var + self.eps)
|
||||
|
||||
# Scale and shift group-wise
|
||||
# $$y_{i_G} =\gamma_{i_G} \hat{x}_{i_G} + \beta_{i_G}$$
|
||||
if self.affine:
|
||||
x_norm = self.scale.view(1, -1, 1) * x_norm + self.shift.view(1, -1, 1)
|
||||
|
||||
# Reshape to original and return
|
||||
return x_norm.view(x_shape)
|
||||
@@ -0,0 +1,225 @@
|
||||
"""
|
||||
---
|
||||
title: Batch Normalization
|
||||
summary: >
|
||||
A PyTorch implementation/tutorial of batch normalization.
|
||||
---
|
||||
|
||||
# Batch Normalization
|
||||
|
||||
This is a [PyTorch](https://pytorch.org) implementation of Batch Normalization from paper
|
||||
[Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift](https://arxiv.org/abs/1502.03167).
|
||||
|
||||
### Internal Covariate Shift
|
||||
|
||||
The paper defines *Internal Covariate Shift* as the change in the
|
||||
distribution of network activations due to the change in
|
||||
network parameters during training.
|
||||
For example, let's say there are two layers $l_1$ and $l_2$.
|
||||
During the beginning of the training $l_1$ outputs (inputs to $l_2$)
|
||||
could be in distribution $\mathcal{N}(0.5, 1)$.
|
||||
Then, after some training steps, it could move to $\mathcal{N}(0.6, 1.5)$.
|
||||
This is *internal covariate shift*.
|
||||
|
||||
Internal covariate shift will adversely affect training speed because the later layers
|
||||
($l_2$ in the above example) have to adapt to this shifted distribution.
|
||||
|
||||
By stabilizing the distribution, batch normalization minimizes the internal covariate shift.
|
||||
|
||||
## Normalization
|
||||
|
||||
It is known that whitening improves training speed and convergence.
|
||||
*Whitening* is linearly transforming inputs to have zero mean, unit variance,
|
||||
and be uncorrelated.
|
||||
|
||||
### Normalizing outside gradient computation doesn't work
|
||||
|
||||
Normalizing outside the gradient computation using pre-computed (detached)
|
||||
means and variances doesn't work. For instance. (ignoring variance), let
|
||||
$$\hat{x} = x - \mathbb{E}[x]$$
|
||||
where $x = u + b$ and $b$ is a trained bias
|
||||
and $\mathbb{E}[x]$ is an outside gradient computation (pre-computed constant).
|
||||
|
||||
Note that $\hat{x}$ has no effect on $b$.
|
||||
Therefore,
|
||||
$b$ will increase or decrease based
|
||||
$\frac{\partial{\mathcal{L}}}{\partial x}$,
|
||||
and keep on growing indefinitely in each training update.
|
||||
The paper notes that similar explosions happen with variances.
|
||||
|
||||
### Batch Normalization
|
||||
|
||||
Whitening is computationally expensive because you need to de-correlate and
|
||||
the gradients must flow through the full whitening calculation.
|
||||
|
||||
The paper introduces a simplified version which they call *Batch Normalization*.
|
||||
First simplification is that it normalizes each feature independently to have
|
||||
zero mean and unit variance:
|
||||
$$\hat{x}^{(k)} = \frac{x^{(k)} - \mathbb{E}[x^{(k)}]}{\sqrt{Var[x^{(k)}]}}$$
|
||||
where $x = (x^{(1)} ... x^{(d)})$ is the $d$-dimensional input.
|
||||
|
||||
The second simplification is to use estimates of mean $\mathbb{E}[x^{(k)}]$
|
||||
and variance $Var[x^{(k)}]$ from the mini-batch
|
||||
for normalization; instead of calculating the mean and variance across the whole dataset.
|
||||
|
||||
Normalizing each feature to zero mean and unit variance could affect what the layer
|
||||
can represent.
|
||||
As an example paper illustrates that, if the inputs to a sigmoid are normalized
|
||||
most of it will be within $[-1, 1]$ range where the sigmoid is linear.
|
||||
To overcome this each feature is scaled and shifted by two trained parameters
|
||||
$\gamma^{(k)}$ and $\beta^{(k)}$.
|
||||
$$y^{(k)} =\gamma^{(k)} \hat{x}^{(k)} + \beta^{(k)}$$
|
||||
where $y^{(k)}$ is the output of the batch normalization layer.
|
||||
|
||||
Note that when applying batch normalization after a linear transform
|
||||
like $Wu + b$ the bias parameter $b$ gets cancelled due to normalization.
|
||||
So you can and should omit bias parameter in linear transforms right before the
|
||||
batch normalization.
|
||||
|
||||
Batch normalization also makes the back propagation invariant to the scale of the weights
|
||||
and empirically it improves generalization, so it has regularization effects too.
|
||||
|
||||
## Inference
|
||||
|
||||
We need to know $\mathbb{E}[x^{(k)}]$ and $Var[x^{(k)}]$ in order to
|
||||
perform the normalization.
|
||||
So during inference, you either need to go through the whole (or part of) dataset
|
||||
and find the mean and variance, or you can use an estimate calculated during training.
|
||||
The usual practice is to calculate an exponential moving average of
|
||||
mean and variance during the training phase and use that for inference.
|
||||
|
||||
Here's [the training code](mnist.html) and a notebook for training
|
||||
a CNN classifier that uses batch normalization for MNIST dataset.
|
||||
|
||||
[](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/normalization/batch_norm/mnist.ipynb)
|
||||
"""
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
|
||||
|
||||
class BatchNorm(nn.Module):
|
||||
r"""
|
||||
## Batch Normalization Layer
|
||||
|
||||
Batch normalization layer $\text{BN}$ normalizes the input $X$ as follows:
|
||||
|
||||
When input $X \in \mathbb{R}^{B \times C \times H \times W}$ is a batch of image representations,
|
||||
where $B$ is the batch size, $C$ is the number of channels, $H$ is the height and $W$ is the width.
|
||||
$\gamma \in \mathbb{R}^{C}$ and $\beta \in \mathbb{R}^{C}$.
|
||||
$$\text{BN}(X) = \gamma
|
||||
\frac{X - \underset{B, H, W}{\mathbb{E}}[X]}{\sqrt{\underset{B, H, W}{Var}[X] + \epsilon}}
|
||||
+ \beta$$
|
||||
|
||||
When input $X \in \mathbb{R}^{B \times C}$ is a batch of embeddings,
|
||||
where $B$ is the batch size and $C$ is the number of features.
|
||||
$\gamma \in \mathbb{R}^{C}$ and $\beta \in \mathbb{R}^{C}$.
|
||||
$$\text{BN}(X) = \gamma
|
||||
\frac{X - \underset{B}{\mathbb{E}}[X]}{\sqrt{\underset{B}{Var}[X] + \epsilon}}
|
||||
+ \beta$$
|
||||
|
||||
When input $X \in \mathbb{R}^{B \times C \times L}$ is a batch of a sequence embeddings,
|
||||
where $B$ is the batch size, $C$ is the number of features, and $L$ is the length of the sequence.
|
||||
$\gamma \in \mathbb{R}^{C}$ and $\beta \in \mathbb{R}^{C}$.
|
||||
$$\text{BN}(X) = \gamma
|
||||
\frac{X - \underset{B, L}{\mathbb{E}}[X]}{\sqrt{\underset{B, L}{Var}[X] + \epsilon}}
|
||||
+ \beta$$
|
||||
"""
|
||||
|
||||
def __init__(self, channels: int, *,
|
||||
eps: float = 1e-5, momentum: float = 0.1,
|
||||
affine: bool = True, track_running_stats: bool = True):
|
||||
"""
|
||||
* `channels` is the number of features in the input
|
||||
* `eps` is $\epsilon$, used in $\sqrt{Var[x^{(k)}] + \epsilon}$ for numerical stability
|
||||
* `momentum` is the momentum in taking the exponential moving average
|
||||
* `affine` is whether to scale and shift the normalized value
|
||||
* `track_running_stats` is whether to calculate the moving averages or mean and variance
|
||||
|
||||
We've tried to use the same names for arguments as PyTorch `BatchNorm` implementation.
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
self.channels = channels
|
||||
|
||||
self.eps = eps
|
||||
self.momentum = momentum
|
||||
self.affine = affine
|
||||
self.track_running_stats = track_running_stats
|
||||
# Create parameters for $\gamma$ and $\beta$ for scale and shift
|
||||
if self.affine:
|
||||
self.scale = nn.Parameter(torch.ones(channels))
|
||||
self.shift = nn.Parameter(torch.zeros(channels))
|
||||
# Create buffers to store exponential moving averages of
|
||||
# mean $\mathbb{E}[x^{(k)}]$ and variance $Var[x^{(k)}]$
|
||||
if self.track_running_stats:
|
||||
self.register_buffer('exp_mean', torch.zeros(channels))
|
||||
self.register_buffer('exp_var', torch.ones(channels))
|
||||
|
||||
def forward(self, x: torch.Tensor):
|
||||
"""
|
||||
`x` is a tensor of shape `[batch_size, channels, *]`.
|
||||
`*` denotes any number of (possibly 0) dimensions.
|
||||
For example, in an image (2D) convolution this will be
|
||||
`[batch_size, channels, height, width]`
|
||||
"""
|
||||
# Keep the original shape
|
||||
x_shape = x.shape
|
||||
# Get the batch size
|
||||
batch_size = x_shape[0]
|
||||
# Sanity check to make sure the number of features is the same
|
||||
assert self.channels == x.shape[1]
|
||||
|
||||
# Reshape into `[batch_size, channels, n]`
|
||||
x = x.view(batch_size, self.channels, -1)
|
||||
|
||||
# We will calculate the mini-batch mean and variance
|
||||
# if we are in training mode or if we have not tracked exponential moving averages
|
||||
if self.training or not self.track_running_stats:
|
||||
# Calculate the mean across first and last dimension;
|
||||
# i.e. the means for each feature $\mathbb{E}[x^{(k)}]$
|
||||
mean = x.mean(dim=[0, 2])
|
||||
# Calculate the squared mean across first and last dimension;
|
||||
# i.e. the means for each feature $\mathbb{E}[(x^{(k)})^2]$
|
||||
mean_x2 = (x ** 2).mean(dim=[0, 2])
|
||||
# Variance for each feature $Var[x^{(k)}] = \mathbb{E}[(x^{(k)})^2] - \mathbb{E}[x^{(k)}]^2$
|
||||
var = mean_x2 - mean ** 2
|
||||
|
||||
# Update exponential moving averages
|
||||
if self.training and self.track_running_stats:
|
||||
self.exp_mean = (1 - self.momentum) * self.exp_mean + self.momentum * mean
|
||||
self.exp_var = (1 - self.momentum) * self.exp_var + self.momentum * var
|
||||
# Use exponential moving averages as estimates
|
||||
else:
|
||||
mean = self.exp_mean
|
||||
var = self.exp_var
|
||||
|
||||
# Normalize $$\hat{x}^{(k)} = \frac{x^{(k)} - \mathbb{E}[x^{(k)}]}{\sqrt{Var[x^{(k)}] + \epsilon}}$$
|
||||
x_norm = (x - mean.view(1, -1, 1)) / torch.sqrt(var + self.eps).view(1, -1, 1)
|
||||
# Scale and shift $$y^{(k)} =\gamma^{(k)} \hat{x}^{(k)} + \beta^{(k)}$$
|
||||
if self.affine:
|
||||
x_norm = self.scale.view(1, -1, 1) * x_norm + self.shift.view(1, -1, 1)
|
||||
|
||||
# Reshape to original and return
|
||||
return x_norm.view(x_shape)
|
||||
|
||||
|
||||
def _test():
|
||||
"""
|
||||
Simple test
|
||||
"""
|
||||
from labml.logger import inspect
|
||||
|
||||
x = torch.zeros([2, 3, 2, 4])
|
||||
inspect(x.shape)
|
||||
bn = BatchNorm(3)
|
||||
|
||||
x = bn(x)
|
||||
inspect(x.shape)
|
||||
inspect(bn.exp_var.shape)
|
||||
|
||||
|
||||
#
|
||||
if __name__ == '__main__':
|
||||
_test()
|
||||
@@ -0,0 +1,64 @@
|
||||
"""
|
||||
---
|
||||
title: CIFAR10 Experiment to try Group Normalization
|
||||
summary: >
|
||||
This trains is a simple convolutional neural network that uses group normalization
|
||||
to classify CIFAR10 images.
|
||||
---
|
||||
|
||||
# CIFAR10 Experiment for Group Normalization
|
||||
"""
|
||||
|
||||
import torch.nn as nn
|
||||
|
||||
from labml import experiment
|
||||
from labml.configs import option
|
||||
from labml_nn.experiments.cifar10 import CIFAR10Configs, CIFAR10VGGModel
|
||||
from labml_nn.normalization.batch_norm import BatchNorm
|
||||
|
||||
|
||||
class Model(CIFAR10VGGModel):
|
||||
"""
|
||||
### VGG 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:
|
||||
return nn.Sequential(
|
||||
nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
|
||||
BatchNorm(out_channels),
|
||||
nn.ReLU(inplace=True),
|
||||
)
|
||||
|
||||
def __init__(self):
|
||||
super().__init__([[64, 64], [128, 128], [256, 256, 256], [512, 512, 512], [512, 512, 512]])
|
||||
|
||||
|
||||
@option(CIFAR10Configs.model)
|
||||
def model(c: CIFAR10Configs):
|
||||
"""
|
||||
### Create model
|
||||
"""
|
||||
return Model().to(c.device)
|
||||
|
||||
|
||||
def main():
|
||||
# Create experiment
|
||||
experiment.create(name='cifar10', comment='batch norm')
|
||||
# Create configurations
|
||||
conf = CIFAR10Configs()
|
||||
# Load configurations
|
||||
experiment.configs(conf, {
|
||||
'optimizer.optimizer': 'Adam',
|
||||
'optimizer.learning_rate': 2.5e-4,
|
||||
'train_batch_size': 64,
|
||||
})
|
||||
# Start the experiment and run the training loop
|
||||
with experiment.start():
|
||||
conf.run()
|
||||
|
||||
|
||||
#
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,85 @@
|
||||
"""
|
||||
---
|
||||
title: MNIST Experiment to try Batch Normalization
|
||||
summary: >
|
||||
This trains is a simple convolutional neural network that uses batch normalization
|
||||
to classify MNIST digits.
|
||||
---
|
||||
|
||||
# MNIST Experiment for Batch Normalization
|
||||
"""
|
||||
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import torch.utils.data
|
||||
|
||||
from labml import experiment
|
||||
from labml.configs import option
|
||||
from labml_nn.experiments.mnist import MNISTConfigs
|
||||
from labml_nn.normalization.batch_norm import BatchNorm
|
||||
|
||||
|
||||
class Model(nn.Module):
|
||||
"""
|
||||
### Model definition
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
# Note that we omit the bias parameter
|
||||
self.conv1 = nn.Conv2d(1, 20, 5, 1, bias=False)
|
||||
# Batch normalization with 20 channels (output of convolution layer).
|
||||
# The input to this layer will have shape `[batch_size, 20, height(24), width(24)]`
|
||||
self.bn1 = BatchNorm(20)
|
||||
#
|
||||
self.conv2 = nn.Conv2d(20, 50, 5, 1, bias=False)
|
||||
# Batch normalization with 50 channels.
|
||||
# The input to this layer will have shape `[batch_size, 50, height(8), width(8)]`
|
||||
self.bn2 = BatchNorm(50)
|
||||
#
|
||||
self.fc1 = nn.Linear(4 * 4 * 50, 500, bias=False)
|
||||
# Batch normalization with 500 channels (output of fully connected layer).
|
||||
# The input to this layer will have shape `[batch_size, 500]`
|
||||
self.bn3 = BatchNorm(500)
|
||||
#
|
||||
self.fc2 = nn.Linear(500, 10)
|
||||
|
||||
def forward(self, x: torch.Tensor):
|
||||
x = F.relu(self.bn1(self.conv1(x)))
|
||||
x = F.max_pool2d(x, 2, 2)
|
||||
x = F.relu(self.bn2(self.conv2(x)))
|
||||
x = F.max_pool2d(x, 2, 2)
|
||||
x = x.view(-1, 4 * 4 * 50)
|
||||
x = F.relu(self.bn3(self.fc1(x)))
|
||||
return self.fc2(x)
|
||||
|
||||
|
||||
@option(MNISTConfigs.model)
|
||||
def model(c: MNISTConfigs):
|
||||
"""
|
||||
### Create model
|
||||
|
||||
We use [`MNISTConfigs`](../../experiments/mnist.html#MNISTConfigs) configurations
|
||||
and set a new function to calculate the model.
|
||||
"""
|
||||
return Model().to(c.device)
|
||||
|
||||
|
||||
def main():
|
||||
# Create experiment
|
||||
experiment.create(name='mnist_batch_norm')
|
||||
# Create configurations
|
||||
conf = MNISTConfigs()
|
||||
# Load configurations
|
||||
experiment.configs(conf, {
|
||||
'optimizer.optimizer': 'Adam',
|
||||
'optimizer.learning_rate': 0.001,
|
||||
})
|
||||
# Start the experiment and run the training loop
|
||||
with experiment.start():
|
||||
conf.run()
|
||||
|
||||
|
||||
#
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,87 @@
|
||||
# [Batch Normalization](https://nn.labml.ai/normalization/batch_norm/index.html)
|
||||
|
||||
This is a [PyTorch](https://pytorch.org) implementation of Batch Normalization from paper
|
||||
[Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift](https://arxiv.org/abs/1502.03167).
|
||||
|
||||
### Internal Covariate Shift
|
||||
|
||||
The paper defines *Internal Covariate Shift* as the change in the
|
||||
distribution of network activations due to the change in
|
||||
network parameters during training.
|
||||
For example, let's say there are two layers $l_1$ and $l_2$.
|
||||
During the beginning of the training $l_1$ outputs (inputs to $l_2$)
|
||||
could be in distribution $\mathcal{N}(0.5, 1)$.
|
||||
Then, after some training steps, it could move to $\mathcal{N}(0.6, 1.5)$.
|
||||
This is *internal covariate shift*.
|
||||
|
||||
Internal covariate shift will adversely affect training speed because the later layers
|
||||
($l_2$ in the above example) have to adapt to this shifted distribution.
|
||||
|
||||
By stabilizing the distribution, batch normalization minimizes the internal covariate shift.
|
||||
|
||||
## Normalization
|
||||
|
||||
It is known that whitening improves training speed and convergence.
|
||||
*Whitening* is linearly transforming inputs to have zero mean, unit variance,
|
||||
and be uncorrelated.
|
||||
|
||||
### Normalizing outside gradient computation doesn't work
|
||||
|
||||
Normalizing outside the gradient computation using pre-computed (detached)
|
||||
means and variances doesn't work. For instance. (ignoring variance), let
|
||||
$$\hat{x} = x - \mathbb{E}[x]$$
|
||||
where $x = u + b$ and $b$ is a trained bias
|
||||
and $\mathbb{E}[x]$ is an outside gradient computation (pre-computed constant).
|
||||
|
||||
Note that $\hat{x}$ has no effect on $b$.
|
||||
Therefore,
|
||||
$b$ will increase or decrease based
|
||||
$\frac{\partial{\mathcal{L}}}{\partial x}$,
|
||||
and keep on growing indefinitely in each training update.
|
||||
The paper notes that similar explosions happen with variances.
|
||||
|
||||
### Batch Normalization
|
||||
|
||||
Whitening is computationally expensive because you need to de-correlate and
|
||||
the gradients must flow through the full whitening calculation.
|
||||
|
||||
The paper introduces a simplified version which they call *Batch Normalization*.
|
||||
First simplification is that it normalizes each feature independently to have
|
||||
zero mean and unit variance:
|
||||
$$\hat{x}^{(k)} = \frac{x^{(k)} - \mathbb{E}[x^{(k)}]}{\sqrt{Var[x^{(k)}]}}$$
|
||||
where $x = (x^{(1)} ... x^{(d)})$ is the $d$-dimensional input.
|
||||
|
||||
The second simplification is to use estimates of mean $\mathbb{E}[x^{(k)}]$
|
||||
and variance $Var[x^{(k)}]$ from the mini-batch
|
||||
for normalization; instead of calculating the mean and variance across the whole dataset.
|
||||
|
||||
Normalizing each feature to zero mean and unit variance could affect what the layer
|
||||
can represent.
|
||||
As an example paper illustrates that, if the inputs to a sigmoid are normalized
|
||||
most of it will be within $[-1, 1]$ range where the sigmoid is linear.
|
||||
To overcome this each feature is scaled and shifted by two trained parameters
|
||||
$\gamma^{(k)}$ and $\beta^{(k)}$.
|
||||
$$y^{(k)} =\gamma^{(k)} \hat{x}^{(k)} + \beta^{(k)}$$
|
||||
where $y^{(k)}$ is the output of the batch normalization layer.
|
||||
|
||||
Note that when applying batch normalization after a linear transform
|
||||
like $Wu + b$ the bias parameter $b$ gets cancelled due to normalization.
|
||||
So you can and should omit bias parameter in linear transforms right before the
|
||||
batch normalization.
|
||||
|
||||
Batch normalization also makes the back propagation invariant to the scale of the weights
|
||||
and empirically it improves generalization, so it has regularization effects too.
|
||||
|
||||
## Inference
|
||||
|
||||
We need to know $\mathbb{E}[x^{(k)}]$ and $Var[x^{(k)}]$ in order to
|
||||
perform the normalization.
|
||||
So during inference, you either need to go through the whole (or part of) dataset
|
||||
and find the mean and variance, or you can use an estimate calculated during training.
|
||||
The usual practice is to calculate an exponential moving average of
|
||||
mean and variance during the training phase and use that for inference.
|
||||
|
||||
Here's [the training code](mnist.html) and a notebook for training
|
||||
a CNN classifier that uses batch normalization for MNIST dataset.
|
||||
|
||||
[](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/normalization/batch_norm/mnist.ipynb)
|
||||
@@ -0,0 +1,173 @@
|
||||
"""
|
||||
---
|
||||
title: DeepNorm
|
||||
summary: >
|
||||
A PyTorch implementation/tutorial of DeepNorm from paper DeepNet: Scaling Transformers to 1,000 Layers.
|
||||
---
|
||||
|
||||
# DeepNorm
|
||||
|
||||
[](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/normalization/deep_norm/experiment.ipynb)
|
||||
|
||||
This is a [PyTorch](https://pytorch.org) implementation of
|
||||
the DeepNorm from the paper
|
||||
[DeepNet: Scaling Transformers to 1,000 Layers](https://arxiv.org/abs/2203.00555).
|
||||
|
||||
The paper proposes a method to stabilize extremely deep transformers through a new normalizing function
|
||||
to replace LayerNorm and a weight initialization scheme.
|
||||
This combines the performance of Post-LayerNorm and the stability of Pre-LayerNorm.
|
||||
Transformers with DeepNorms are supposed to be stable even without a learning rate warm-up.
|
||||
|
||||
The paper first shows that the changes to layer outputs (for the same input)
|
||||
change gradually during stable training;
|
||||
when unstable it changes rapidly during the initial training steps.
|
||||
This happens with initializing weights to small values, and learning rate warm-ups where the
|
||||
training is stable.
|
||||
They use the idea of keeping the changes to layer outputs small to derive the new
|
||||
normalization and weight initialization mechanism.
|
||||
|
||||
## Weight Initializations
|
||||
|
||||
Usually, the weights are initialized with Xavier or Kaiming initializations.
|
||||
This paper scales (sets the gain) the weights by a constant $\beta$ depending on the size of the
|
||||
transformer.
|
||||
|
||||
DeepNorm suggests scaling the weights of the two linear transforms in the
|
||||
[Feed-Forward Network](../../transformers/feed_forward.html),
|
||||
the value projection transform, and the output projection transform of the
|
||||
attention layer.
|
||||
Weights of these transforms are scaled by (has a gain equal to) $\beta$.
|
||||
|
||||
The scaling is implemented in the
|
||||
|
||||
## Normalization Function
|
||||
|
||||
$$x_{l + 1} = \mathop{LN}\Big( \alpha x_l + \mathop{G}_l \big(x_l, \theta_l \big)\Big)$$
|
||||
|
||||
where $\alpha$ is a constant that depends on the depth of the transformer,
|
||||
$\mathop{LN}$ is [Layer Normalization](../layer_norm/index.html), and
|
||||
$\mathop{G}_l (x_l, \theta_l)$ is the function of the $l$-th transformer sub-layer (FFN or attention).
|
||||
|
||||
This function is used to replace Post-LayerNorm.
|
||||
|
||||
## $\alpha$ and $\beta$ constants
|
||||
|
||||
\begin{align}
|
||||
\begin{array} {c|cc|cc}
|
||||
\text{Type} & \text{Enc-} \alpha & \text{Enc-} \beta & \text{Dec-} \alpha & \text{Dec-} \beta \\
|
||||
\hline \\
|
||||
\text{Encoder only} & (2N)^{\frac{1}{4}} & (8N)^{-\frac{1}{4}} & - & - \\
|
||||
\text{Decoder only} & - & - & (2M)^{\frac{1}{4}} & (8M)^{-\frac{1}{4}} \\
|
||||
\text{Enc-Dec} & 0.81 (N^4M)^{\frac{1}{16}} & 0.87 (N^4 M)^{-\frac{1}{16}} &
|
||||
(3M)^{\frac{1}{4}} & (12M)^{-\frac{1}{4}} \\
|
||||
\end{array}
|
||||
\end{align}
|
||||
|
||||
Where $N$ is the number of layers in the encoder and $M$ is the number of layers in the decoder.
|
||||
|
||||
Refer to [the paper](https://arxiv.org/abs/2203.00555) for derivation.
|
||||
|
||||
[Here is an experiment implementation](experiment.html) that uses DeepNorm.
|
||||
"""
|
||||
|
||||
from typing import Union, List
|
||||
|
||||
import torch
|
||||
from torch import nn, Size
|
||||
|
||||
from labml_nn.normalization.layer_norm import LayerNorm
|
||||
from labml_nn.transformers import MultiHeadAttention
|
||||
from labml_nn.transformers.feed_forward import FeedForward
|
||||
from labml_nn.transformers.utils import subsequent_mask
|
||||
|
||||
|
||||
class DeepNorm(nn.Module):
|
||||
"""
|
||||
## DeepNorm Normalization
|
||||
|
||||
$$x_{l + 1} = \mathop{LN}\Big( \alpha x_l + \mathop{G}_l \big(x_l, \theta_l \big)\Big)$$
|
||||
"""
|
||||
|
||||
def __init__(self, alpha: float, normalized_shape: Union[int, List[int], Size], *,
|
||||
eps: float = 1e-5,
|
||||
elementwise_affine: bool = True):
|
||||
"""
|
||||
:param alpha: is $\alpha$
|
||||
:param normalized_shape: is the shape for LayerNorm $\mathop{LN}$
|
||||
:param eps: is $\epsilon$ for LayerNorm
|
||||
:param elementwise_affine: is a flag indicating whether to do an elementwise transformation in LayerNorm
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
self.alpha = alpha
|
||||
# Initialize $\mathop{LN}$
|
||||
self.layer_norm = LayerNorm(normalized_shape, eps=eps, elementwise_affine=elementwise_affine)
|
||||
|
||||
def forward(self, x: torch.Tensor, gx: torch.Tensor):
|
||||
"""
|
||||
:param x: is the output from the previous layer $x_l$
|
||||
:param gx: is the output of the current sub-layer $\mathop{G}_l (x_l, \theta_l)$
|
||||
"""
|
||||
# $$x_{l + 1} = \mathop{LN}\Big( \alpha x_l + \mathop{G}_l \big(x_l, \theta_l \big)\Big)$$
|
||||
return self.layer_norm(x + self.alpha * gx)
|
||||
|
||||
|
||||
class DeepNormTransformerLayer(nn.Module):
|
||||
"""
|
||||
## Transformer Decoder Layer with DeepNorm
|
||||
|
||||
This implements a transformer decoder layer with DeepNorm.
|
||||
Encoder layers will have a similar form.
|
||||
"""
|
||||
def __init__(self, *,
|
||||
d_model: int,
|
||||
self_attn: MultiHeadAttention,
|
||||
feed_forward: FeedForward,
|
||||
deep_norm_alpha: float,
|
||||
deep_norm_beta: float,
|
||||
):
|
||||
"""
|
||||
:param d_model: is the token embedding size
|
||||
:param self_attn: is the self attention module
|
||||
:param feed_forward: is the feed forward module
|
||||
:param deep_norm_alpha: is $\alpha$ coefficient in DeepNorm
|
||||
:param deep_norm_beta: is $\beta$ constant for scaling weights initialization
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
self.self_attn = self_attn
|
||||
self.feed_forward = feed_forward
|
||||
# DeepNorms after attention and feed forward network
|
||||
self.self_attn_norm = DeepNorm(deep_norm_alpha, [d_model])
|
||||
self.feed_forward_norm = DeepNorm(deep_norm_alpha, [d_model])
|
||||
|
||||
# Scale weights after initialization
|
||||
with torch.no_grad():
|
||||
# Feed forward network linear transformations
|
||||
feed_forward.layer1.weight *= deep_norm_beta
|
||||
feed_forward.layer2.weight *= deep_norm_beta
|
||||
|
||||
# Attention value projection
|
||||
self_attn.value.linear.weight *= deep_norm_beta
|
||||
# Attention output project
|
||||
self_attn.output.weight *= deep_norm_beta
|
||||
|
||||
# The mask will be initialized on the first call
|
||||
self.mask = None
|
||||
|
||||
def forward(self, x: torch.Tensor):
|
||||
"""
|
||||
:param x: are the embeddings of shape `[seq_len, batch_size, d_model]`
|
||||
"""
|
||||
# Create causal 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)
|
||||
|
||||
# Run through self attention, i.e. keys and values are from self
|
||||
x = self.self_attn_norm(x, self.self_attn(query=x, key=x, value=x, mask=self.mask))
|
||||
# Pass through the feed-forward network
|
||||
x = self.feed_forward_norm(x, self.feed_forward(x))
|
||||
|
||||
#
|
||||
return x
|
||||
@@ -0,0 +1,274 @@
|
||||
{
|
||||
"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/normalization/deep_norm/experiment.ipynb)\n",
|
||||
"\n",
|
||||
"## DeepNorm\n",
|
||||
"\n",
|
||||
"This is an experiment training Shakespeare dataset with a deep transformer using [DeepNorm](https://nn.labml.ai/normalization/deep_norm/index.html)."
|
||||
]
|
||||
},
|
||||
{
|
||||
"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": {
|
||||
"id": "0hJXx_g0wS2C",
|
||||
"pycharm": {
|
||||
"name": "#%%\n"
|
||||
}
|
||||
},
|
||||
"source": [
|
||||
"from labml import experiment\n",
|
||||
"from labml_nn.normalization.deep_norm.experiment import Configs"
|
||||
],
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "Lpggo0wM6qb-",
|
||||
"pycharm": {
|
||||
"name": "#%% md\n"
|
||||
}
|
||||
},
|
||||
"source": [
|
||||
"### Create an experiment"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "bFcr9k-l4cAg",
|
||||
"pycharm": {
|
||||
"name": "#%%\n"
|
||||
}
|
||||
},
|
||||
"source": [
|
||||
"experiment.create(name=\"deep_norm\", writers={'screen'})"
|
||||
],
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "-OnHLi626tJt",
|
||||
"pycharm": {
|
||||
"name": "#%% md\n"
|
||||
}
|
||||
},
|
||||
"source": [
|
||||
"### Configurations"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "Piz0c5f44hRo",
|
||||
"pycharm": {
|
||||
"name": "#%%\n"
|
||||
}
|
||||
},
|
||||
"source": [
|
||||
"conf = Configs()"
|
||||
],
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "wwMzCqpD6vkL",
|
||||
"pycharm": {
|
||||
"name": "#%% md\n"
|
||||
}
|
||||
},
|
||||
"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": "29634715-42f4-4405-fb11-fc9522608627",
|
||||
"pycharm": {
|
||||
"name": "#%%\n"
|
||||
}
|
||||
},
|
||||
"source": [
|
||||
"experiment.configs(conf, {\n",
|
||||
" # Use character level tokenizer\n",
|
||||
" 'tokenizer': 'character',\n",
|
||||
" # Prompt separator is blank\n",
|
||||
" 'prompt_separator': '',\n",
|
||||
" # Starting prompt for sampling\n",
|
||||
" 'prompt': 'It is ',\n",
|
||||
" # Use Tiny Shakespeare dataset\n",
|
||||
" 'text': 'tiny_shakespeare',\n",
|
||||
"\n",
|
||||
" # Use a context size of $256$\n",
|
||||
" 'seq_len': 256,\n",
|
||||
" # Train for 32 epochs\n",
|
||||
" 'epochs': 32,\n",
|
||||
" # Batch size $16$\n",
|
||||
" 'batch_size': 16,\n",
|
||||
" # Switch between training and validation for $10$ times per epoch\n",
|
||||
" 'inner_iterations': 10,\n",
|
||||
"\n",
|
||||
" # Adam optimizer with no warmup\n",
|
||||
" 'optimizer.optimizer': 'Adam',\n",
|
||||
" 'optimizer.learning_rate': 3e-4,\n",
|
||||
"})"
|
||||
],
|
||||
"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({'model': conf.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",
|
||||
" conf.run()"
|
||||
],
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"accelerator": "GPU",
|
||||
"colab": {
|
||||
"collapsed_sections": [],
|
||||
"name": "DeepNorm",
|
||||
"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,175 @@
|
||||
"""
|
||||
---
|
||||
title: DeepNorm Experiment
|
||||
summary: >
|
||||
Training a DeepNorm transformer on Tiny Shakespeare.
|
||||
---
|
||||
|
||||
# [DeepNorm](index.html) Experiment
|
||||
|
||||
[](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/normalization/deep_norm/experiment.ipynb)
|
||||
"""
|
||||
|
||||
import copy
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from labml import experiment
|
||||
from labml.configs import option
|
||||
from labml_nn.experiments.nlp_autoregression import NLPAutoRegressionConfigs
|
||||
from labml_nn.normalization.deep_norm import DeepNormTransformerLayer
|
||||
from labml_nn.transformers import MultiHeadAttention
|
||||
from labml_nn.transformers.feed_forward import FeedForward
|
||||
|
||||
|
||||
class AutoregressiveTransformer(nn.Module):
|
||||
"""
|
||||
## Auto-Regressive model
|
||||
|
||||
This is a autoregressive transformer model that uses DeepNorm.
|
||||
"""
|
||||
|
||||
def __init__(self, n_tokens: int, d_model: int, n_layers: int, layer: DeepNormTransformerLayer):
|
||||
"""
|
||||
: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 tranformer.
|
||||
"""
|
||||
super().__init__()
|
||||
# Transformer with `n_layers` layers
|
||||
self.transformer = nn.Sequential(*[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)
|
||||
|
||||
def forward(self, x: torch.Tensor):
|
||||
"""
|
||||
:param x: are the input tokens of shape `[seq_len, batch_size]`
|
||||
"""
|
||||
# Get the token embeddings
|
||||
x = self.emb(x)
|
||||
# Transformer encoder
|
||||
x = self.transformer(x)
|
||||
# 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 = 32
|
||||
|
||||
# $\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 = 64
|
||||
# Size of each attention head
|
||||
d_k: int = 16
|
||||
|
||||
|
||||
@option(Configs.deep_norm_alpha)
|
||||
def _deep_norm_alpha(c: Configs):
|
||||
"""
|
||||
#### Calculate $\alpha$
|
||||
|
||||
$\alpha = (2M)^{\frac{1}{4}}$
|
||||
"""
|
||||
return (2. * c.n_layers) ** (1. / 4.)
|
||||
|
||||
|
||||
@option(Configs.deep_norm_beta)
|
||||
def _deep_norm_beta(c: Configs):
|
||||
"""
|
||||
#### Calculate $\beta$
|
||||
|
||||
$\beta = (8M)^{-\frac{1}{4}}$
|
||||
"""
|
||||
return (8. * c.n_layers) ** -(1. / 4.)
|
||||
|
||||
|
||||
@option(Configs.model)
|
||||
def _model(c: Configs):
|
||||
"""
|
||||
#### Initialize the model
|
||||
"""
|
||||
m = AutoregressiveTransformer(c.n_tokens, c.d_model, c.n_layers,
|
||||
DeepNormTransformerLayer(d_model=c.d_model,
|
||||
deep_norm_alpha=c.deep_norm_alpha,
|
||||
deep_norm_beta=c.deep_norm_beta,
|
||||
feed_forward=FeedForward(d_model=c.d_model,
|
||||
d_ff=c.d_model * 4),
|
||||
self_attn=MultiHeadAttention(c.n_heads, c.d_model,
|
||||
dropout_prob=0.0)))
|
||||
|
||||
return m.to(c.device)
|
||||
|
||||
|
||||
def main():
|
||||
"""
|
||||
#### Create and run the experiment
|
||||
"""
|
||||
# Create experiment
|
||||
experiment.create(name="deep_norm", writers={'screen', 'web_api'})
|
||||
# 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,
|
||||
|
||||
# Number of layers
|
||||
'n_layers': 50,
|
||||
|
||||
|
||||
# Adam optimizer with no warmup
|
||||
'optimizer.optimizer': 'Adam',
|
||||
'optimizer.learning_rate': 1.25e-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,173 @@
|
||||
"""
|
||||
---
|
||||
title: Group Normalization
|
||||
summary: >
|
||||
A PyTorch implementation/tutorial of group normalization.
|
||||
---
|
||||
|
||||
# Group Normalization
|
||||
|
||||
This is a [PyTorch](https://pytorch.org) implementation of
|
||||
the [Group Normalization](https://arxiv.org/abs/1803.08494) paper.
|
||||
|
||||
[Batch Normalization](../batch_norm/index.html) works well for large enough batch sizes
|
||||
but not well for small batch sizes, because it normalizes over the batch.
|
||||
Training large models with large batch sizes is not possible due to the memory capacity of the
|
||||
devices.
|
||||
|
||||
This paper introduces Group Normalization, which normalizes a set of features together as a group.
|
||||
This is based on the observation that classical features such as
|
||||
[SIFT](https://en.wikipedia.org/wiki/Scale-invariant_feature_transform) and
|
||||
[HOG](https://en.wikipedia.org/wiki/Histogram_of_oriented_gradients) are group-wise features.
|
||||
The paper proposes dividing feature channels into groups and then separately normalizing
|
||||
all channels within each group.
|
||||
|
||||
## Formulation
|
||||
|
||||
All normalization layers can be defined by the following computation.
|
||||
|
||||
$$\hat{x}_i = \frac{1}{\sigma_i} (x_i - \mu_i)$$
|
||||
|
||||
where $x$ is the tensor representing the batch,
|
||||
and $i$ is the index of a single value.
|
||||
For instance, when it's 2D images
|
||||
$i = (i_N, i_C, i_H, i_W)$ is a 4-d vector for indexing
|
||||
image within batch, feature channel, vertical coordinate and horizontal coordinate.
|
||||
$\mu_i$ and $\sigma_i$ are mean and standard deviation.
|
||||
|
||||
\begin{align}
|
||||
\mu_i &= \frac{1}{m} \sum_{k \in \mathcal{S}_i} x_k \\
|
||||
\sigma_i &= \sqrt{\frac{1}{m} \sum_{k \in \mathcal{S}_i} (x_k - \mu_i)^2 + \epsilon}
|
||||
\end{align}
|
||||
|
||||
$\mathcal{S}_i$ is the set of indexes across which the mean and standard deviation
|
||||
are calculated for index $i$.
|
||||
$m$ is the size of the set $\mathcal{S}_i$ which is the same for all $i$.
|
||||
|
||||
The definition of $\mathcal{S}_i$ is different for
|
||||
[Batch normalization](../batch_norm/index.html),
|
||||
[Layer normalization](../layer_norm/index.html), and
|
||||
[Instance normalization](../instance_norm/index.html).
|
||||
|
||||
### [Batch Normalization](../batch_norm/index.html)
|
||||
|
||||
$$\mathcal{S}_i = \{k | k_C = i_C\}$$
|
||||
|
||||
The values that share the same feature channel are normalized together.
|
||||
|
||||
### [Layer Normalization](../layer_norm/index.html)
|
||||
|
||||
$$\mathcal{S}_i = \{k | k_N = i_N\}$$
|
||||
|
||||
The values from the same sample in the batch are normalized together.
|
||||
|
||||
### [Instance Normalization](../instance_norm/index.html)
|
||||
|
||||
$$\mathcal{S}_i = \{k | k_N = i_N, k_C = i_C\}$$
|
||||
|
||||
The values from the same sample and same feature channel are normalized together.
|
||||
|
||||
### Group Normalization
|
||||
|
||||
$$\mathcal{S}_i = \{k | k_N = i_N,
|
||||
\bigg \lfloor \frac{k_C}{C/G} \bigg \rfloor = \bigg \lfloor \frac{i_C}{C/G} \bigg \rfloor\}$$
|
||||
|
||||
where $G$ is the number of groups and $C$ is the number of channels.
|
||||
|
||||
Group normalization normalizes values of the same sample and the same group of channels together.
|
||||
|
||||
Here's a [CIFAR 10 classification model](experiment.html) that uses instance normalization.
|
||||
|
||||
[](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/normalization/group_norm/experiment.ipynb)
|
||||
"""
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
|
||||
|
||||
class GroupNorm(nn.Module):
|
||||
r"""
|
||||
## Group Normalization Layer
|
||||
"""
|
||||
|
||||
def __init__(self, groups: int, channels: int, *,
|
||||
eps: float = 1e-5, affine: bool = True):
|
||||
"""
|
||||
* `groups` is the number of groups the features are divided into
|
||||
* `channels` is the number of features in the input
|
||||
* `eps` is $\epsilon$, used in $\sqrt{Var[x^{(k)}] + \epsilon}$ for numerical stability
|
||||
* `affine` is whether to scale and shift the normalized value
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
assert channels % groups == 0, "Number of channels should be evenly divisible by the number of groups"
|
||||
self.groups = groups
|
||||
self.channels = channels
|
||||
|
||||
self.eps = eps
|
||||
self.affine = affine
|
||||
# Create parameters for $\gamma$ and $\beta$ for scale and shift
|
||||
if self.affine:
|
||||
self.scale = nn.Parameter(torch.ones(channels))
|
||||
self.shift = nn.Parameter(torch.zeros(channels))
|
||||
|
||||
def forward(self, x: torch.Tensor):
|
||||
"""
|
||||
`x` is a tensor of shape `[batch_size, channels, *]`.
|
||||
`*` denotes any number of (possibly 0) dimensions.
|
||||
For example, in an image (2D) convolution this will be
|
||||
`[batch_size, channels, height, width]`
|
||||
"""
|
||||
# Keep the original shape
|
||||
x_shape = x.shape
|
||||
# Get the batch size
|
||||
batch_size = x_shape[0]
|
||||
# Sanity check to make sure the number of features is the same
|
||||
assert self.channels == x.shape[1]
|
||||
|
||||
# Reshape into `[batch_size, groups, n]`
|
||||
x = x.view(batch_size, self.groups, -1)
|
||||
|
||||
# Calculate the mean across last dimension;
|
||||
# i.e. the means for each sample and channel group $\mathbb{E}[x_{(i_N, i_G)}]$
|
||||
mean = x.mean(dim=[-1], keepdim=True)
|
||||
# Calculate the squared mean across last dimension;
|
||||
# i.e. the means for each sample and channel group $\mathbb{E}[x^2_{(i_N, i_G)}]$
|
||||
mean_x2 = (x ** 2).mean(dim=[-1], keepdim=True)
|
||||
# Variance for each sample and feature group
|
||||
# $Var[x_{(i_N, i_G)}] = \mathbb{E}[x^2_{(i_N, i_G)}] - \mathbb{E}[x_{(i_N, i_G)}]^2$
|
||||
var = mean_x2 - mean ** 2
|
||||
|
||||
# Normalize
|
||||
# $$\hat{x}_{(i_N, i_G)} =
|
||||
# \frac{x_{(i_N, i_G)} - \mathbb{E}[x_{(i_N, i_G)}]}{\sqrt{Var[x_{(i_N, i_G)}] + \epsilon}}$$
|
||||
x_norm = (x - mean) / torch.sqrt(var + self.eps)
|
||||
|
||||
# Scale and shift channel-wise
|
||||
# $$y_{i_C} =\gamma_{i_C} \hat{x}_{i_C} + \beta_{i_C}$$
|
||||
if self.affine:
|
||||
x_norm = x_norm.view(batch_size, self.channels, -1)
|
||||
x_norm = self.scale.view(1, -1, 1) * x_norm + self.shift.view(1, -1, 1)
|
||||
|
||||
# Reshape to original and return
|
||||
return x_norm.view(x_shape)
|
||||
|
||||
|
||||
def _test():
|
||||
"""
|
||||
Simple test
|
||||
"""
|
||||
from labml.logger import inspect
|
||||
|
||||
x = torch.zeros([2, 6, 2, 4])
|
||||
inspect(x.shape)
|
||||
bn = GroupNorm(2, 6)
|
||||
|
||||
x = bn(x)
|
||||
inspect(x.shape)
|
||||
|
||||
|
||||
#
|
||||
if __name__ == '__main__':
|
||||
_test()
|
||||
@@ -0,0 +1,451 @@
|
||||
{
|
||||
"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/normalization/group_norm/experiment.ipynb) \n",
|
||||
"\n",
|
||||
"## Group Norm - CIFAR 10\n",
|
||||
"\n",
|
||||
"This is an experiment training a model with group norm to classify CIFAR-10 dataset."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "AahG_i2y5tY9"
|
||||
},
|
||||
"source": [
|
||||
"Install the `labml-nn` package."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "ZCzmCrAIVg0L"
|
||||
},
|
||||
"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",
|
||||
"import torch.nn as nn\n",
|
||||
"\n",
|
||||
"from labml import experiment\n",
|
||||
"from labml_nn.normalization.group_norm.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=\"cifar10\", comment=\"group norm\")"
|
||||
],
|
||||
"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": "50ad9e07-84f4-47cf-9d26-034448eb611b"
|
||||
},
|
||||
"source": [
|
||||
"experiment.configs(conf, {\n",
|
||||
" 'optimizer.optimizer': 'Adam',\n",
|
||||
" 'optimizer.learning_rate': 2.5e-4,\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": 882,
|
||||
"referenced_widgets": [
|
||||
"14841e99103e41f69dd9b709301d3204",
|
||||
"60f10c7bff5c4eea845a14f3f3075e8d",
|
||||
"cf2f7c0f10454901bc5b48872b364dbf",
|
||||
"8755fc8b0f6b40b3b08822e0a705d403",
|
||||
"c6c192b58fa242008fbc2983c7866c5f",
|
||||
"6e5c4becab6b40aaafce1a4575d3199c",
|
||||
"328dbbbc3cdb4163896913308059c23c",
|
||||
"3bff44b4205f40119715fae60d4a04a9"
|
||||
]
|
||||
},
|
||||
"id": "aIAWo7Fw5DR8",
|
||||
"outputId": "2e82cce8-eaad-4cab-88d3-efd5d095b5b7"
|
||||
},
|
||||
"source": [
|
||||
"with experiment.start():\n",
|
||||
" conf.run()"
|
||||
],
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "oBXXlP2b7XZO"
|
||||
},
|
||||
"source": [],
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"accelerator": "GPU",
|
||||
"colab": {
|
||||
"collapsed_sections": [],
|
||||
"name": "Group Norm",
|
||||
"provenance": []
|
||||
},
|
||||
"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"
|
||||
},
|
||||
"widgets": {
|
||||
"application/vnd.jupyter.widget-state+json": {
|
||||
"14841e99103e41f69dd9b709301d3204": {
|
||||
"model_module": "@jupyter-widgets/controls",
|
||||
"model_name": "HBoxModel",
|
||||
"state": {
|
||||
"_dom_classes": [],
|
||||
"_model_module": "@jupyter-widgets/controls",
|
||||
"_model_module_version": "1.5.0",
|
||||
"_model_name": "HBoxModel",
|
||||
"_view_count": null,
|
||||
"_view_module": "@jupyter-widgets/controls",
|
||||
"_view_module_version": "1.5.0",
|
||||
"_view_name": "HBoxView",
|
||||
"box_style": "",
|
||||
"children": [
|
||||
"IPY_MODEL_cf2f7c0f10454901bc5b48872b364dbf",
|
||||
"IPY_MODEL_8755fc8b0f6b40b3b08822e0a705d403"
|
||||
],
|
||||
"layout": "IPY_MODEL_60f10c7bff5c4eea845a14f3f3075e8d"
|
||||
}
|
||||
},
|
||||
"328dbbbc3cdb4163896913308059c23c": {
|
||||
"model_module": "@jupyter-widgets/controls",
|
||||
"model_name": "DescriptionStyleModel",
|
||||
"state": {
|
||||
"_model_module": "@jupyter-widgets/controls",
|
||||
"_model_module_version": "1.5.0",
|
||||
"_model_name": "DescriptionStyleModel",
|
||||
"_view_count": null,
|
||||
"_view_module": "@jupyter-widgets/base",
|
||||
"_view_module_version": "1.2.0",
|
||||
"_view_name": "StyleView",
|
||||
"description_width": ""
|
||||
}
|
||||
},
|
||||
"3bff44b4205f40119715fae60d4a04a9": {
|
||||
"model_module": "@jupyter-widgets/base",
|
||||
"model_name": "LayoutModel",
|
||||
"state": {
|
||||
"_model_module": "@jupyter-widgets/base",
|
||||
"_model_module_version": "1.2.0",
|
||||
"_model_name": "LayoutModel",
|
||||
"_view_count": null,
|
||||
"_view_module": "@jupyter-widgets/base",
|
||||
"_view_module_version": "1.2.0",
|
||||
"_view_name": "LayoutView",
|
||||
"align_content": null,
|
||||
"align_items": null,
|
||||
"align_self": null,
|
||||
"border": null,
|
||||
"bottom": null,
|
||||
"display": null,
|
||||
"flex": null,
|
||||
"flex_flow": null,
|
||||
"grid_area": null,
|
||||
"grid_auto_columns": null,
|
||||
"grid_auto_flow": null,
|
||||
"grid_auto_rows": null,
|
||||
"grid_column": null,
|
||||
"grid_gap": null,
|
||||
"grid_row": null,
|
||||
"grid_template_areas": null,
|
||||
"grid_template_columns": null,
|
||||
"grid_template_rows": null,
|
||||
"height": null,
|
||||
"justify_content": null,
|
||||
"justify_items": null,
|
||||
"left": null,
|
||||
"margin": null,
|
||||
"max_height": null,
|
||||
"max_width": null,
|
||||
"min_height": null,
|
||||
"min_width": null,
|
||||
"object_fit": null,
|
||||
"object_position": null,
|
||||
"order": null,
|
||||
"overflow": null,
|
||||
"overflow_x": null,
|
||||
"overflow_y": null,
|
||||
"padding": null,
|
||||
"right": null,
|
||||
"top": null,
|
||||
"visibility": null,
|
||||
"width": null
|
||||
}
|
||||
},
|
||||
"60f10c7bff5c4eea845a14f3f3075e8d": {
|
||||
"model_module": "@jupyter-widgets/base",
|
||||
"model_name": "LayoutModel",
|
||||
"state": {
|
||||
"_model_module": "@jupyter-widgets/base",
|
||||
"_model_module_version": "1.2.0",
|
||||
"_model_name": "LayoutModel",
|
||||
"_view_count": null,
|
||||
"_view_module": "@jupyter-widgets/base",
|
||||
"_view_module_version": "1.2.0",
|
||||
"_view_name": "LayoutView",
|
||||
"align_content": null,
|
||||
"align_items": null,
|
||||
"align_self": null,
|
||||
"border": null,
|
||||
"bottom": null,
|
||||
"display": null,
|
||||
"flex": null,
|
||||
"flex_flow": null,
|
||||
"grid_area": null,
|
||||
"grid_auto_columns": null,
|
||||
"grid_auto_flow": null,
|
||||
"grid_auto_rows": null,
|
||||
"grid_column": null,
|
||||
"grid_gap": null,
|
||||
"grid_row": null,
|
||||
"grid_template_areas": null,
|
||||
"grid_template_columns": null,
|
||||
"grid_template_rows": null,
|
||||
"height": null,
|
||||
"justify_content": null,
|
||||
"justify_items": null,
|
||||
"left": null,
|
||||
"margin": null,
|
||||
"max_height": null,
|
||||
"max_width": null,
|
||||
"min_height": null,
|
||||
"min_width": null,
|
||||
"object_fit": null,
|
||||
"object_position": null,
|
||||
"order": null,
|
||||
"overflow": null,
|
||||
"overflow_x": null,
|
||||
"overflow_y": null,
|
||||
"padding": null,
|
||||
"right": null,
|
||||
"top": null,
|
||||
"visibility": null,
|
||||
"width": null
|
||||
}
|
||||
},
|
||||
"6e5c4becab6b40aaafce1a4575d3199c": {
|
||||
"model_module": "@jupyter-widgets/base",
|
||||
"model_name": "LayoutModel",
|
||||
"state": {
|
||||
"_model_module": "@jupyter-widgets/base",
|
||||
"_model_module_version": "1.2.0",
|
||||
"_model_name": "LayoutModel",
|
||||
"_view_count": null,
|
||||
"_view_module": "@jupyter-widgets/base",
|
||||
"_view_module_version": "1.2.0",
|
||||
"_view_name": "LayoutView",
|
||||
"align_content": null,
|
||||
"align_items": null,
|
||||
"align_self": null,
|
||||
"border": null,
|
||||
"bottom": null,
|
||||
"display": null,
|
||||
"flex": null,
|
||||
"flex_flow": null,
|
||||
"grid_area": null,
|
||||
"grid_auto_columns": null,
|
||||
"grid_auto_flow": null,
|
||||
"grid_auto_rows": null,
|
||||
"grid_column": null,
|
||||
"grid_gap": null,
|
||||
"grid_row": null,
|
||||
"grid_template_areas": null,
|
||||
"grid_template_columns": null,
|
||||
"grid_template_rows": null,
|
||||
"height": null,
|
||||
"justify_content": null,
|
||||
"justify_items": null,
|
||||
"left": null,
|
||||
"margin": null,
|
||||
"max_height": null,
|
||||
"max_width": null,
|
||||
"min_height": null,
|
||||
"min_width": null,
|
||||
"object_fit": null,
|
||||
"object_position": null,
|
||||
"order": null,
|
||||
"overflow": null,
|
||||
"overflow_x": null,
|
||||
"overflow_y": null,
|
||||
"padding": null,
|
||||
"right": null,
|
||||
"top": null,
|
||||
"visibility": null,
|
||||
"width": null
|
||||
}
|
||||
},
|
||||
"8755fc8b0f6b40b3b08822e0a705d403": {
|
||||
"model_module": "@jupyter-widgets/controls",
|
||||
"model_name": "HTMLModel",
|
||||
"state": {
|
||||
"_dom_classes": [],
|
||||
"_model_module": "@jupyter-widgets/controls",
|
||||
"_model_module_version": "1.5.0",
|
||||
"_model_name": "HTMLModel",
|
||||
"_view_count": null,
|
||||
"_view_module": "@jupyter-widgets/controls",
|
||||
"_view_module_version": "1.5.0",
|
||||
"_view_name": "HTMLView",
|
||||
"description": "",
|
||||
"description_tooltip": null,
|
||||
"layout": "IPY_MODEL_3bff44b4205f40119715fae60d4a04a9",
|
||||
"placeholder": "",
|
||||
"style": "IPY_MODEL_328dbbbc3cdb4163896913308059c23c",
|
||||
"value": " 170499072/? [00:03<00:00, 54808451.08it/s]"
|
||||
}
|
||||
},
|
||||
"c6c192b58fa242008fbc2983c7866c5f": {
|
||||
"model_module": "@jupyter-widgets/controls",
|
||||
"model_name": "ProgressStyleModel",
|
||||
"state": {
|
||||
"_model_module": "@jupyter-widgets/controls",
|
||||
"_model_module_version": "1.5.0",
|
||||
"_model_name": "ProgressStyleModel",
|
||||
"_view_count": null,
|
||||
"_view_module": "@jupyter-widgets/base",
|
||||
"_view_module_version": "1.2.0",
|
||||
"_view_name": "StyleView",
|
||||
"bar_color": null,
|
||||
"description_width": "initial"
|
||||
}
|
||||
},
|
||||
"cf2f7c0f10454901bc5b48872b364dbf": {
|
||||
"model_module": "@jupyter-widgets/controls",
|
||||
"model_name": "FloatProgressModel",
|
||||
"state": {
|
||||
"_dom_classes": [],
|
||||
"_model_module": "@jupyter-widgets/controls",
|
||||
"_model_module_version": "1.5.0",
|
||||
"_model_name": "FloatProgressModel",
|
||||
"_view_count": null,
|
||||
"_view_module": "@jupyter-widgets/controls",
|
||||
"_view_module_version": "1.5.0",
|
||||
"_view_name": "ProgressView",
|
||||
"bar_style": "success",
|
||||
"description": "",
|
||||
"description_tooltip": null,
|
||||
"layout": "IPY_MODEL_6e5c4becab6b40aaafce1a4575d3199c",
|
||||
"max": 170498071,
|
||||
"min": 0,
|
||||
"orientation": "horizontal",
|
||||
"style": "IPY_MODEL_c6c192b58fa242008fbc2983c7866c5f",
|
||||
"value": 170498071
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 4
|
||||
}
|
||||
@@ -0,0 +1,68 @@
|
||||
"""
|
||||
---
|
||||
title: CIFAR10 Experiment to try Group Normalization
|
||||
summary: >
|
||||
This trains is a simple convolutional neural network that uses group normalization
|
||||
to classify CIFAR10 images.
|
||||
---
|
||||
|
||||
# CIFAR10 Experiment for Group Normalization
|
||||
"""
|
||||
|
||||
import torch.nn as nn
|
||||
|
||||
from labml import experiment
|
||||
from labml.configs import option
|
||||
from labml_nn.experiments.cifar10 import CIFAR10Configs, CIFAR10VGGModel
|
||||
|
||||
|
||||
class Model(CIFAR10VGGModel):
|
||||
"""
|
||||
### VGG 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:
|
||||
return nn.Sequential(
|
||||
nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
|
||||
fnorm.GroupNorm(self.groups, out_channels), # new
|
||||
nn.ReLU(inplace=True),
|
||||
)
|
||||
|
||||
def __init__(self, groups: int = 32):
|
||||
self.groups = groups # input param:groups to conv_block
|
||||
super().__init__([[64, 64], [128, 128], [256, 256, 256], [512, 512, 512], [512, 512, 512]])
|
||||
|
||||
|
||||
class Configs(CIFAR10Configs):
|
||||
# Number of groups
|
||||
groups: int = 16
|
||||
|
||||
|
||||
@option(Configs.model)
|
||||
def model(c: Configs):
|
||||
"""
|
||||
### Create model
|
||||
"""
|
||||
return Model(c.groups).to(c.device)
|
||||
|
||||
|
||||
def main():
|
||||
# Create experiment
|
||||
experiment.create(name='cifar10', comment='group norm')
|
||||
# Create configurations
|
||||
conf = Configs()
|
||||
# Load configurations
|
||||
experiment.configs(conf, {
|
||||
'optimizer.optimizer': 'Adam',
|
||||
'optimizer.learning_rate': 2.5e-4,
|
||||
})
|
||||
# Start the experiment and run the training loop
|
||||
with experiment.start():
|
||||
conf.run()
|
||||
|
||||
|
||||
#
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,20 @@
|
||||
# [Group Normalization](https://nn.labml.ai/normalization/group_norm/index.html)
|
||||
|
||||
This is a [PyTorch](https://pytorch.org) implementation of
|
||||
the [Group Normalization](https://arxiv.org/abs/1803.08494) paper.
|
||||
|
||||
[Batch Normalization](https://nn.labml.ai/normalization/batch_norm/index.html) works well for large enough batch sizes
|
||||
but not well for small batch sizes, because it normalizes over the batch.
|
||||
Training large models with large batch sizes is not possible due to the memory capacity of the
|
||||
devices.
|
||||
|
||||
This paper introduces Group Normalization, which normalizes a set of features together as a group.
|
||||
This is based on the observation that classical features such as
|
||||
[SIFT](https://en.wikipedia.org/wiki/Scale-invariant_feature_transform) and
|
||||
[HOG](https://en.wikipedia.org/wiki/Histogram_of_oriented_gradients) are group-wise features.
|
||||
The paper proposes dividing feature channels into groups and then separately normalizing
|
||||
all channels within each group.
|
||||
|
||||
Here's a [CIFAR 10 classification model](https://nn.labml.ai/normalization/group_norm/experiment.html) that uses group normalization.
|
||||
|
||||
[](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/normalization/group_norm/experiment.ipynb)
|
||||
@@ -0,0 +1,122 @@
|
||||
"""
|
||||
---
|
||||
title: Instance Normalization
|
||||
summary: >
|
||||
A PyTorch implementation/tutorial of instance normalization.
|
||||
---
|
||||
|
||||
# Instance Normalization
|
||||
|
||||
This is a [PyTorch](https://pytorch.org) implementation of
|
||||
[Instance Normalization: The Missing Ingredient for Fast Stylization](https://arxiv.org/abs/1607.08022).
|
||||
|
||||
Instance normalization was introduced to improve [style transfer](https://paperswithcode.com/task/style-transfer).
|
||||
It is based on the observation that stylization should not depend on the contrast of the content image.
|
||||
The "contrast normalization" is
|
||||
|
||||
$$y_{t,i,j,k} = \frac{x_{t,i,j,k}}{\sum_{l=1}^H \sum_{m=1}^W x_{t,i,l,m}}$$
|
||||
|
||||
where $x$ is a batch of images with dimensions image index $t$,
|
||||
feature channel $i$, and
|
||||
spatial position $j, k$.
|
||||
|
||||
Since it's hard for a convolutional network to learn "contrast normalization", this paper
|
||||
introduces instance normalization which does that.
|
||||
|
||||
Here's a [CIFAR 10 classification model](experiment.html) that uses instance normalization.
|
||||
"""
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
|
||||
|
||||
class InstanceNorm(nn.Module):
|
||||
r"""
|
||||
## Instance Normalization Layer
|
||||
|
||||
Instance normalization layer $\text{IN}$ normalizes the input $X$ as follows:
|
||||
|
||||
When input $X \in \mathbb{R}^{B \times C \times H \times W}$ is a batch of image representations,
|
||||
where $B$ is the batch size, $C$ is the number of channels, $H$ is the height and $W$ is the width.
|
||||
$\gamma \in \mathbb{R}^{C}$ and $\beta \in \mathbb{R}^{C}$. The affine transformation with $gamma$ and
|
||||
$beta$ are optional.
|
||||
|
||||
$$\text{IN}(X) = \gamma
|
||||
\frac{X - \underset{H, W}{\mathbb{E}}[X]}{\sqrt{\underset{H, W}{Var}[X] + \epsilon}}
|
||||
+ \beta$$
|
||||
"""
|
||||
|
||||
def __init__(self, channels: int, *,
|
||||
eps: float = 1e-5, affine: bool = True):
|
||||
"""
|
||||
* `channels` is the number of features in the input
|
||||
* `eps` is $\epsilon$, used in $\sqrt{Var[X] + \epsilon}$ for numerical stability
|
||||
* `affine` is whether to scale and shift the normalized value
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
self.channels = channels
|
||||
|
||||
self.eps = eps
|
||||
self.affine = affine
|
||||
# Create parameters for $\gamma$ and $\beta$ for scale and shift
|
||||
if self.affine:
|
||||
self.scale = nn.Parameter(torch.ones(channels))
|
||||
self.shift = nn.Parameter(torch.zeros(channels))
|
||||
|
||||
def forward(self, x: torch.Tensor):
|
||||
"""
|
||||
`x` is a tensor of shape `[batch_size, channels, *]`.
|
||||
`*` denotes any number of (possibly 0) dimensions.
|
||||
For example, in an image (2D) convolution this will be
|
||||
`[batch_size, channels, height, width]`
|
||||
"""
|
||||
# Keep the original shape
|
||||
x_shape = x.shape
|
||||
# Get the batch size
|
||||
batch_size = x_shape[0]
|
||||
# Sanity check to make sure the number of features is the same
|
||||
assert self.channels == x.shape[1]
|
||||
|
||||
# Reshape into `[batch_size, channels, n]`
|
||||
x = x.view(batch_size, self.channels, -1)
|
||||
|
||||
# Calculate the mean across last dimension
|
||||
# i.e. the means for each feature $\mathbb{E}[x_{t,i}]$
|
||||
mean = x.mean(dim=[-1], keepdim=True)
|
||||
# Calculate the squared mean across first and last dimension;
|
||||
# i.e. the means for each feature $\mathbb{E}[(x_{t,i}^2]$
|
||||
mean_x2 = (x ** 2).mean(dim=[-1], keepdim=True)
|
||||
# Variance for each feature $Var[x_{t,i}] = \mathbb{E}[x_{t,i}^2] - \mathbb{E}[x_{t,i}]^2$
|
||||
var = mean_x2 - mean ** 2
|
||||
|
||||
# Normalize $$\hat{x}_{t,i} = \frac{x_{t,i} - \mathbb{E}[x_{t,i}]}{\sqrt{Var[x_{t,i}] + \epsilon}}$$
|
||||
x_norm = (x - mean) / torch.sqrt(var + self.eps)
|
||||
x_norm = x_norm.view(batch_size, self.channels, -1)
|
||||
|
||||
# Scale and shift $$y_{t,i} =\gamma_i \hat{x}_{t,i} + \beta_i$$
|
||||
if self.affine:
|
||||
x_norm = self.scale.view(1, -1, 1) * x_norm + self.shift.view(1, -1, 1)
|
||||
|
||||
# Reshape to original and return
|
||||
return x_norm.view(x_shape)
|
||||
|
||||
|
||||
def _test():
|
||||
"""
|
||||
Simple test
|
||||
"""
|
||||
from labml.logger import inspect
|
||||
|
||||
x = torch.zeros([2, 6, 2, 4])
|
||||
inspect(x.shape)
|
||||
bn = InstanceNorm(6)
|
||||
|
||||
x = bn(x)
|
||||
inspect(x.shape)
|
||||
|
||||
|
||||
#
|
||||
if __name__ == '__main__':
|
||||
_test()
|
||||
@@ -0,0 +1,67 @@
|
||||
"""
|
||||
---
|
||||
title: CIFAR10 Experiment to try Instance Normalization
|
||||
summary: >
|
||||
This trains is a simple convolutional neural network that uses instance normalization
|
||||
to classify CIFAR10 images.
|
||||
---
|
||||
|
||||
# CIFAR10 Experiment for Instance Normalization
|
||||
|
||||
This demonstrates the use of an instance normalization layer in a convolutional
|
||||
neural network for classification. Not that instance normalization was designed for
|
||||
style transfer and this is only a demo.
|
||||
"""
|
||||
|
||||
import torch.nn as nn
|
||||
|
||||
from labml import experiment
|
||||
from labml.configs import option
|
||||
from labml_nn.experiments.cifar10 import CIFAR10Configs, CIFAR10VGGModel
|
||||
from labml_nn.normalization.instance_norm import InstanceNorm
|
||||
|
||||
|
||||
class Model(CIFAR10VGGModel):
|
||||
"""
|
||||
### VGG 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:
|
||||
return nn.Sequential(
|
||||
nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
|
||||
InstanceNorm(out_channels),
|
||||
nn.ReLU(inplace=True),
|
||||
)
|
||||
|
||||
def __init__(self):
|
||||
super().__init__([[64, 64], [128, 128], [256, 256, 256], [512, 512, 512], [512, 512, 512]])
|
||||
|
||||
|
||||
@option(CIFAR10Configs.model)
|
||||
def _model(c: CIFAR10Configs):
|
||||
"""
|
||||
### Create model
|
||||
"""
|
||||
return Model().to(c.device)
|
||||
|
||||
|
||||
def main():
|
||||
# Create experiment
|
||||
experiment.create(name='cifar10', comment='instance norm')
|
||||
# Create configurations
|
||||
conf = CIFAR10Configs()
|
||||
# Load configurations
|
||||
experiment.configs(conf, {
|
||||
'optimizer.optimizer': 'Adam',
|
||||
'optimizer.learning_rate': 2.5e-4,
|
||||
})
|
||||
# Start the experiment and run the training loop
|
||||
with experiment.start():
|
||||
conf.run()
|
||||
|
||||
|
||||
#
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,9 @@
|
||||
# [Instance Normalization](https://nn.labml.ai/normalization/instance_norm/index.html)
|
||||
|
||||
This is a [PyTorch](https://pytorch.org) implementation of
|
||||
[Instance Normalization: The Missing Ingredient for Fast Stylization](https://arxiv.org/abs/1607.08022).
|
||||
|
||||
Instance normalization was introduced to improve [style transfer](https://paperswithcode.com/task/style-transfer).
|
||||
It is based on the observation that stylization should not depend on the contrast of the content image.
|
||||
Since it's hard for a convolutional network to learn "contrast normalization", this paper
|
||||
introduces instance normalization which does that.
|
||||
@@ -0,0 +1,150 @@
|
||||
"""
|
||||
---
|
||||
title: Layer Normalization
|
||||
summary: >
|
||||
A PyTorch implementation/tutorial of layer normalization.
|
||||
---
|
||||
|
||||
# Layer Normalization
|
||||
|
||||
This is a [PyTorch](https://pytorch.org) implementation of
|
||||
[Layer Normalization](https://arxiv.org/abs/1607.06450).
|
||||
|
||||
### Limitations of [Batch Normalization](../batch_norm/index.html)
|
||||
|
||||
* You need to maintain running means.
|
||||
* Tricky for RNNs. Do you need different normalizations for each step?
|
||||
* Doesn't work with small batch sizes;
|
||||
large NLP models are usually trained with small batch sizes.
|
||||
* Need to compute means and variances across devices in distributed training.
|
||||
|
||||
## Layer Normalization
|
||||
|
||||
Layer normalization is a simpler normalization method that works
|
||||
on a wider range of settings.
|
||||
Layer normalization transforms the inputs to have zero mean and unit variance
|
||||
across the features.
|
||||
*Note that batch normalization fixes the zero mean and unit variance for each element.*
|
||||
Layer normalization does it for each batch across all elements.
|
||||
|
||||
Layer normalization is generally used for NLP tasks.
|
||||
|
||||
We have used layer normalization in most of the
|
||||
[transformer implementations](../../transformers/gpt/index.html).
|
||||
"""
|
||||
from typing import Union, List
|
||||
|
||||
import torch
|
||||
from torch import nn, Size
|
||||
|
||||
|
||||
|
||||
class LayerNorm(nn.Module):
|
||||
r"""
|
||||
## Layer Normalization
|
||||
|
||||
Layer normalization $\text{LN}$ normalizes the input $X$ as follows:
|
||||
|
||||
When input $X \in \mathbb{R}^{B \times C}$ is a batch of embeddings,
|
||||
where $B$ is the batch size and $C$ is the number of features.
|
||||
$\gamma \in \mathbb{R}^{C}$ and $\beta \in \mathbb{R}^{C}$.
|
||||
$$\text{LN}(X) = \gamma
|
||||
\frac{X - \underset{C}{\mathbb{E}}[X]}{\sqrt{\underset{C}{Var}[X] + \epsilon}}
|
||||
+ \beta$$
|
||||
|
||||
When input $X \in \mathbb{R}^{L \times B \times C}$ is a batch of a sequence of embeddings,
|
||||
where $B$ is the batch size, $C$ is the number of channels, $L$ is the length of the sequence.
|
||||
$\gamma \in \mathbb{R}^{C}$ and $\beta \in \mathbb{R}^{C}$.
|
||||
$$\text{LN}(X) = \gamma
|
||||
\frac{X - \underset{C}{\mathbb{E}}[X]}{\sqrt{\underset{C}{Var}[X] + \epsilon}}
|
||||
+ \beta$$
|
||||
|
||||
When input $X \in \mathbb{R}^{B \times C \times H \times W}$ is a batch of image representations,
|
||||
where $B$ is the batch size, $C$ is the number of channels, $H$ is the height and $W$ is the width.
|
||||
This is not a widely used scenario.
|
||||
$\gamma \in \mathbb{R}^{C \times H \times W}$ and $\beta \in \mathbb{R}^{C \times H \times W}$.
|
||||
$$\text{LN}(X) = \gamma
|
||||
\frac{X - \underset{C, H, W}{\mathbb{E}}[X]}{\sqrt{\underset{C, H, W}{Var}[X] + \epsilon}}
|
||||
+ \beta$$
|
||||
"""
|
||||
|
||||
def __init__(self, normalized_shape: Union[int, List[int], Size], *,
|
||||
eps: float = 1e-5,
|
||||
elementwise_affine: bool = True):
|
||||
"""
|
||||
* `normalized_shape` $S$ is the shape of the elements (except the batch).
|
||||
The input should then be
|
||||
$X \in \mathbb{R}^{* \times S[0] \times S[1] \times ... \times S[n]}$
|
||||
* `eps` is $\epsilon$, used in $\sqrt{Var[X] + \epsilon}$ for numerical stability
|
||||
* `elementwise_affine` is whether to scale and shift the normalized value
|
||||
|
||||
We've tried to use the same names for arguments as PyTorch `LayerNorm` implementation.
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
# Convert `normalized_shape` to `torch.Size`
|
||||
if isinstance(normalized_shape, int):
|
||||
normalized_shape = torch.Size([normalized_shape])
|
||||
elif isinstance(normalized_shape, list):
|
||||
normalized_shape = torch.Size(normalized_shape)
|
||||
assert isinstance(normalized_shape, torch.Size)
|
||||
|
||||
#
|
||||
self.normalized_shape = normalized_shape
|
||||
self.eps = eps
|
||||
self.elementwise_affine = elementwise_affine
|
||||
# Create parameters for $\gamma$ and $\beta$ for gain and bias
|
||||
if self.elementwise_affine:
|
||||
self.gain = nn.Parameter(torch.ones(normalized_shape))
|
||||
self.bias = nn.Parameter(torch.zeros(normalized_shape))
|
||||
|
||||
def forward(self, x: torch.Tensor):
|
||||
"""
|
||||
`x` is a tensor of shape `[*, S[0], S[1], ..., S[n]]`.
|
||||
`*` could be any number of dimensions.
|
||||
For example, in an NLP task this will be
|
||||
`[seq_len, batch_size, features]`
|
||||
"""
|
||||
# Sanity check to make sure the shapes match
|
||||
assert self.normalized_shape == x.shape[-len(self.normalized_shape):]
|
||||
|
||||
# The dimensions to calculate the mean and variance on
|
||||
dims = [-(i + 1) for i in range(len(self.normalized_shape))]
|
||||
|
||||
# Calculate the mean of all elements;
|
||||
# i.e. the means for each element $\mathbb{E}[X]$
|
||||
mean = x.mean(dim=dims, keepdim=True)
|
||||
# Calculate the squared mean of all elements;
|
||||
# i.e. the means for each element $\mathbb{E}[X^2]$
|
||||
mean_x2 = (x ** 2).mean(dim=dims, keepdim=True)
|
||||
# Variance of all element $Var[X] = \mathbb{E}[X^2] - \mathbb{E}[X]^2$
|
||||
var = mean_x2 - mean ** 2
|
||||
|
||||
# Normalize $$\hat{X} = \frac{X - \mathbb{E}[X]}{\sqrt{Var[X] + \epsilon}}$$
|
||||
x_norm = (x - mean) / torch.sqrt(var + self.eps)
|
||||
# Scale and shift $$\text{LN}(x) = \gamma \hat{X} + \beta$$
|
||||
if self.elementwise_affine:
|
||||
x_norm = self.gain * x_norm + self.bias
|
||||
|
||||
#
|
||||
return x_norm
|
||||
|
||||
|
||||
def _test():
|
||||
"""
|
||||
Simple test
|
||||
"""
|
||||
from labml.logger import inspect
|
||||
|
||||
x = torch.zeros([2, 3, 2, 4])
|
||||
inspect(x.shape)
|
||||
ln = LayerNorm(x.shape[2:])
|
||||
|
||||
x = ln(x)
|
||||
inspect(x.shape)
|
||||
inspect(ln.gain.shape)
|
||||
|
||||
|
||||
#
|
||||
if __name__ == '__main__':
|
||||
_test()
|
||||
@@ -0,0 +1,26 @@
|
||||
# [Layer Normalization](https://nn.labml.ai/normalization/layer_norm/index.html)
|
||||
|
||||
This is a [PyTorch](https://pytorch.org) implementation of
|
||||
[Layer Normalization](https://arxiv.org/abs/1607.06450).
|
||||
|
||||
### Limitations of [Batch Normalization](https://nn.labml.ai/normalization/batch_norm/index.html)
|
||||
|
||||
* You need to maintain running means.
|
||||
* Tricky for RNNs. Do you need different normalizations for each step?
|
||||
* Doesn't work with small batch sizes;
|
||||
large NLP models are usually trained with small batch sizes.
|
||||
* Need to compute means and variances across devices in distributed training.
|
||||
|
||||
## Layer Normalization
|
||||
|
||||
Layer normalization is a simpler normalization method that works
|
||||
on a wider range of settings.
|
||||
Layer normalization transforms the inputs to have zero mean and unit variance
|
||||
across the features.
|
||||
*Note that batch normalization fixes the zero mean and unit variance for each element.*
|
||||
Layer normalization does it for each batch across all elements.
|
||||
|
||||
Layer normalization is generally used for NLP tasks.
|
||||
|
||||
We have used layer normalization in most of the
|
||||
[transformer implementations](https://nn.labml.ai/transformers/gpt/index.html).
|
||||
@@ -0,0 +1,84 @@
|
||||
"""
|
||||
---
|
||||
title: Weight Standardization
|
||||
summary: >
|
||||
A PyTorch implementation/tutorial of Weight Standardization.
|
||||
---
|
||||
|
||||
# Weight Standardization
|
||||
|
||||
This is a [PyTorch](https://pytorch.org) implementation of Weight Standardization from the paper
|
||||
[Micro-Batch Training with Batch-Channel Normalization and Weight Standardization](https://arxiv.org/abs/1903.10520).
|
||||
We also have an [annotated implementation of Batch-Channel Normalization](../batch_channel_norm/index.html).
|
||||
|
||||
Batch normalization **gives a smooth loss landscape** and
|
||||
**avoids elimination singularities**.
|
||||
Elimination singularities are nodes of the network that become
|
||||
useless (e.g. a ReLU that gives 0 all the time).
|
||||
|
||||
However, batch normalization doesn't work well when the batch size is too small,
|
||||
which happens when training large networks because of device memory limitations.
|
||||
The paper introduces Weight Standardization with Batch-Channel Normalization as
|
||||
a better alternative.
|
||||
|
||||
Weight Standardization:
|
||||
1. Normalizes the gradients
|
||||
2. Smoothes the landscape (reduced Lipschitz constant)
|
||||
3. Avoids elimination singularities
|
||||
|
||||
The Lipschitz constant is the maximum slope a function has between two points.
|
||||
That is, $L$ is the Lipschitz constant where $L$ is the smallest value that satisfies,
|
||||
$\forall a,b \in A: \lVert f(a) - f(b) \rVert \le L \lVert a - b \rVert$
|
||||
where $f: A \rightarrow \mathbb{R}^m, A \in \mathbb{R}^n$.
|
||||
|
||||
Elimination singularities are avoided because it keeps the statistics of the outputs similar to the
|
||||
inputs. So as long as the inputs are normally distributed the outputs remain close to normal.
|
||||
This avoids outputs of nodes from always falling beyond the active range of the activation function
|
||||
(e.g. always negative input for a ReLU).
|
||||
|
||||
*[Refer to the paper for proofs](https://arxiv.org/abs/1903.10520)*.
|
||||
|
||||
Here is [the training code](experiment.html) for training
|
||||
a VGG network that uses weight standardization to classify CIFAR-10 data.
|
||||
This uses a [2D-Convolution Layer with Weight Standardization](conv2d.html).
|
||||
|
||||
[](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/normalization/weight_standardization/experiment.ipynb)
|
||||
"""
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
def weight_standardization(weight: torch.Tensor, eps: float):
|
||||
r"""
|
||||
## Weight Standardization
|
||||
|
||||
$$\hat{W}_{i,j} = \frac{W_{i,j} - \mu_{W_{i,\cdot}}} {\sigma_{W_{i,\cdot}}}$$
|
||||
|
||||
where,
|
||||
|
||||
\begin{align}
|
||||
W &\in \mathbb{R}^{O \times I} \\
|
||||
\mu_{W_{i,\cdot}} &= \frac{1}{I} \sum_{j=1}^I W_{i,j} \\
|
||||
\sigma_{W_{i,\cdot}} &= \sqrt{\frac{1}{I} \sum_{j=1}^I W^2_{i,j} - \mu^2_{W_{i,\cdot}} + \epsilon} \\
|
||||
\end{align}
|
||||
|
||||
for a 2D-convolution layer $O$ is the number of output channels ($O = C_{out}$)
|
||||
and $I$ is the number of input channels times the kernel size ($I = C_{in} \times k_H \times k_W$)
|
||||
"""
|
||||
|
||||
# Get $C_{out}$, $C_{in}$ and kernel shape
|
||||
c_out, c_in, *kernel_shape = weight.shape
|
||||
# Reshape $W$ to $O \times I$
|
||||
weight = weight.view(c_out, -1)
|
||||
# Calculate
|
||||
#
|
||||
# \begin{align}
|
||||
# \mu_{W_{i,\cdot}} &= \frac{1}{I} \sum_{j=1}^I W_{i,j} \\
|
||||
# \sigma^2_{W_{i,\cdot}} &= \frac{1}{I} \sum_{j=1}^I W^2_{i,j} - \mu^2_{W_{i,\cdot}}
|
||||
# \end{align}
|
||||
var, mean = torch.var_mean(weight, dim=1, keepdim=True)
|
||||
# Normalize
|
||||
# $$\hat{W}_{i,j} = \frac{W_{i,j} - \mu_{W_{i,\cdot}}} {\sigma_{W_{i,\cdot}}}$$
|
||||
weight = (weight - mean) / (torch.sqrt(var + eps))
|
||||
# Change back to original shape and return
|
||||
return weight.view(c_out, c_in, *kernel_shape)
|
||||
@@ -0,0 +1,60 @@
|
||||
"""
|
||||
---
|
||||
title: 2D Convolution Layer with Weight Standardization
|
||||
summary: >
|
||||
A PyTorch implementation/tutorial of a 2D Convolution Layer with Weight Standardization.
|
||||
---
|
||||
|
||||
# 2D Convolution Layer with Weight Standardization
|
||||
|
||||
This is an implementation of a 2 dimensional convolution layer with [Weight Standardization](./index.html)
|
||||
"""
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.nn import functional as F
|
||||
|
||||
from labml_nn.normalization.weight_standardization import weight_standardization
|
||||
|
||||
|
||||
class Conv2d(nn.Conv2d):
|
||||
"""
|
||||
## 2D Convolution Layer
|
||||
|
||||
This extends the standard 2D Convolution layer and standardize the weights before the convolution step.
|
||||
"""
|
||||
def __init__(self, in_channels, out_channels, kernel_size,
|
||||
stride=1,
|
||||
padding=0,
|
||||
dilation=1,
|
||||
groups: int = 1,
|
||||
bias: bool = True,
|
||||
padding_mode: str = 'zeros',
|
||||
eps: float = 1e-5):
|
||||
super(Conv2d, self).__init__(in_channels, out_channels, kernel_size,
|
||||
stride=stride,
|
||||
padding=padding,
|
||||
dilation=dilation,
|
||||
groups=groups,
|
||||
bias=bias,
|
||||
padding_mode=padding_mode)
|
||||
self.eps = eps
|
||||
|
||||
def forward(self, x: torch.Tensor):
|
||||
return F.conv2d(x, weight_standardization(self.weight, self.eps), self.bias, self.stride,
|
||||
self.padding, self.dilation, self.groups)
|
||||
|
||||
|
||||
def _test():
|
||||
"""
|
||||
A simple test to verify the tensor sizes
|
||||
"""
|
||||
conv2d = Conv2d(10, 20, 5)
|
||||
from labml.logger import inspect
|
||||
inspect(conv2d.weight)
|
||||
import torch
|
||||
inspect(conv2d(torch.zeros(10, 10, 100, 100)))
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
_test()
|
||||
@@ -0,0 +1,444 @@
|
||||
{
|
||||
"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/normalization/group_norm/experiment.ipynb) \n",
|
||||
"\n",
|
||||
"## Weight Standardization & Batch-Channel Normalization - CIFAR 10\n",
|
||||
"\n",
|
||||
"This is an experiment training a model with Weight Standardization & Batch-Channel Normalization to classify CIFAR-10 dataset."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "AahG_i2y5tY9"
|
||||
},
|
||||
"source": [
|
||||
"Install the `labml-nn` package."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "ZCzmCrAIVg0L",
|
||||
"outputId": "8332d030-2fab-4a24-876f-152b7f99a226"
|
||||
},
|
||||
"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.normalization.weight_standardization.experiment import CIFAR10Configs as 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=\"cifar10\", comment=\"WS + BCN\")"
|
||||
],
|
||||
"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": "33a5979f-70eb-4d19-e82a-a6b113316cca"
|
||||
},
|
||||
"source": [
|
||||
"experiment.configs(conf, {\n",
|
||||
" 'optimizer.optimizer': 'Adam',\n",
|
||||
" 'optimizer.learning_rate': 2.5e-4,\n",
|
||||
" 'train_batch_size': 64,\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": 933,
|
||||
"referenced_widgets": [
|
||||
"1dacd9c82caf40c8b2def0c5fe4f7643",
|
||||
"67e37901c44b4bde9062014542cae018",
|
||||
"9d57be286b7244ca9a5f74cff6c3cf4e",
|
||||
"9e511aaba1204ea9985fb6f55e024039",
|
||||
"61849c45ed0c4436995bb81d7efdac7e",
|
||||
"7ed3376588734d64a229fd8dd6bfe06e",
|
||||
"afd6865a483f4273addb343059632961",
|
||||
"7b765f7623474999a8c9815e02c1cc6b"
|
||||
]
|
||||
},
|
||||
"id": "aIAWo7Fw5DR8",
|
||||
"outputId": "cf3d16aa-a9a1-4fd4-cfb4-23e8ea5ea940"
|
||||
},
|
||||
"source": [
|
||||
"with experiment.start():\n",
|
||||
" conf.run()"
|
||||
],
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"accelerator": "GPU",
|
||||
"colab": {
|
||||
"collapsed_sections": [],
|
||||
"name": "Weight Standardization & Batch-Channel Normalization",
|
||||
"provenance": []
|
||||
},
|
||||
"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"
|
||||
},
|
||||
"widgets": {
|
||||
"application/vnd.jupyter.widget-state+json": {
|
||||
"1dacd9c82caf40c8b2def0c5fe4f7643": {
|
||||
"model_module": "@jupyter-widgets/controls",
|
||||
"model_name": "HBoxModel",
|
||||
"state": {
|
||||
"_dom_classes": [],
|
||||
"_model_module": "@jupyter-widgets/controls",
|
||||
"_model_module_version": "1.5.0",
|
||||
"_model_name": "HBoxModel",
|
||||
"_view_count": null,
|
||||
"_view_module": "@jupyter-widgets/controls",
|
||||
"_view_module_version": "1.5.0",
|
||||
"_view_name": "HBoxView",
|
||||
"box_style": "",
|
||||
"children": [
|
||||
"IPY_MODEL_9d57be286b7244ca9a5f74cff6c3cf4e",
|
||||
"IPY_MODEL_9e511aaba1204ea9985fb6f55e024039"
|
||||
],
|
||||
"layout": "IPY_MODEL_67e37901c44b4bde9062014542cae018"
|
||||
}
|
||||
},
|
||||
"61849c45ed0c4436995bb81d7efdac7e": {
|
||||
"model_module": "@jupyter-widgets/controls",
|
||||
"model_name": "ProgressStyleModel",
|
||||
"state": {
|
||||
"_model_module": "@jupyter-widgets/controls",
|
||||
"_model_module_version": "1.5.0",
|
||||
"_model_name": "ProgressStyleModel",
|
||||
"_view_count": null,
|
||||
"_view_module": "@jupyter-widgets/base",
|
||||
"_view_module_version": "1.2.0",
|
||||
"_view_name": "StyleView",
|
||||
"bar_color": null,
|
||||
"description_width": "initial"
|
||||
}
|
||||
},
|
||||
"67e37901c44b4bde9062014542cae018": {
|
||||
"model_module": "@jupyter-widgets/base",
|
||||
"model_name": "LayoutModel",
|
||||
"state": {
|
||||
"_model_module": "@jupyter-widgets/base",
|
||||
"_model_module_version": "1.2.0",
|
||||
"_model_name": "LayoutModel",
|
||||
"_view_count": null,
|
||||
"_view_module": "@jupyter-widgets/base",
|
||||
"_view_module_version": "1.2.0",
|
||||
"_view_name": "LayoutView",
|
||||
"align_content": null,
|
||||
"align_items": null,
|
||||
"align_self": null,
|
||||
"border": null,
|
||||
"bottom": null,
|
||||
"display": null,
|
||||
"flex": null,
|
||||
"flex_flow": null,
|
||||
"grid_area": null,
|
||||
"grid_auto_columns": null,
|
||||
"grid_auto_flow": null,
|
||||
"grid_auto_rows": null,
|
||||
"grid_column": null,
|
||||
"grid_gap": null,
|
||||
"grid_row": null,
|
||||
"grid_template_areas": null,
|
||||
"grid_template_columns": null,
|
||||
"grid_template_rows": null,
|
||||
"height": null,
|
||||
"justify_content": null,
|
||||
"justify_items": null,
|
||||
"left": null,
|
||||
"margin": null,
|
||||
"max_height": null,
|
||||
"max_width": null,
|
||||
"min_height": null,
|
||||
"min_width": null,
|
||||
"object_fit": null,
|
||||
"object_position": null,
|
||||
"order": null,
|
||||
"overflow": null,
|
||||
"overflow_x": null,
|
||||
"overflow_y": null,
|
||||
"padding": null,
|
||||
"right": null,
|
||||
"top": null,
|
||||
"visibility": null,
|
||||
"width": null
|
||||
}
|
||||
},
|
||||
"7b765f7623474999a8c9815e02c1cc6b": {
|
||||
"model_module": "@jupyter-widgets/base",
|
||||
"model_name": "LayoutModel",
|
||||
"state": {
|
||||
"_model_module": "@jupyter-widgets/base",
|
||||
"_model_module_version": "1.2.0",
|
||||
"_model_name": "LayoutModel",
|
||||
"_view_count": null,
|
||||
"_view_module": "@jupyter-widgets/base",
|
||||
"_view_module_version": "1.2.0",
|
||||
"_view_name": "LayoutView",
|
||||
"align_content": null,
|
||||
"align_items": null,
|
||||
"align_self": null,
|
||||
"border": null,
|
||||
"bottom": null,
|
||||
"display": null,
|
||||
"flex": null,
|
||||
"flex_flow": null,
|
||||
"grid_area": null,
|
||||
"grid_auto_columns": null,
|
||||
"grid_auto_flow": null,
|
||||
"grid_auto_rows": null,
|
||||
"grid_column": null,
|
||||
"grid_gap": null,
|
||||
"grid_row": null,
|
||||
"grid_template_areas": null,
|
||||
"grid_template_columns": null,
|
||||
"grid_template_rows": null,
|
||||
"height": null,
|
||||
"justify_content": null,
|
||||
"justify_items": null,
|
||||
"left": null,
|
||||
"margin": null,
|
||||
"max_height": null,
|
||||
"max_width": null,
|
||||
"min_height": null,
|
||||
"min_width": null,
|
||||
"object_fit": null,
|
||||
"object_position": null,
|
||||
"order": null,
|
||||
"overflow": null,
|
||||
"overflow_x": null,
|
||||
"overflow_y": null,
|
||||
"padding": null,
|
||||
"right": null,
|
||||
"top": null,
|
||||
"visibility": null,
|
||||
"width": null
|
||||
}
|
||||
},
|
||||
"7ed3376588734d64a229fd8dd6bfe06e": {
|
||||
"model_module": "@jupyter-widgets/base",
|
||||
"model_name": "LayoutModel",
|
||||
"state": {
|
||||
"_model_module": "@jupyter-widgets/base",
|
||||
"_model_module_version": "1.2.0",
|
||||
"_model_name": "LayoutModel",
|
||||
"_view_count": null,
|
||||
"_view_module": "@jupyter-widgets/base",
|
||||
"_view_module_version": "1.2.0",
|
||||
"_view_name": "LayoutView",
|
||||
"align_content": null,
|
||||
"align_items": null,
|
||||
"align_self": null,
|
||||
"border": null,
|
||||
"bottom": null,
|
||||
"display": null,
|
||||
"flex": null,
|
||||
"flex_flow": null,
|
||||
"grid_area": null,
|
||||
"grid_auto_columns": null,
|
||||
"grid_auto_flow": null,
|
||||
"grid_auto_rows": null,
|
||||
"grid_column": null,
|
||||
"grid_gap": null,
|
||||
"grid_row": null,
|
||||
"grid_template_areas": null,
|
||||
"grid_template_columns": null,
|
||||
"grid_template_rows": null,
|
||||
"height": null,
|
||||
"justify_content": null,
|
||||
"justify_items": null,
|
||||
"left": null,
|
||||
"margin": null,
|
||||
"max_height": null,
|
||||
"max_width": null,
|
||||
"min_height": null,
|
||||
"min_width": null,
|
||||
"object_fit": null,
|
||||
"object_position": null,
|
||||
"order": null,
|
||||
"overflow": null,
|
||||
"overflow_x": null,
|
||||
"overflow_y": null,
|
||||
"padding": null,
|
||||
"right": null,
|
||||
"top": null,
|
||||
"visibility": null,
|
||||
"width": null
|
||||
}
|
||||
},
|
||||
"9d57be286b7244ca9a5f74cff6c3cf4e": {
|
||||
"model_module": "@jupyter-widgets/controls",
|
||||
"model_name": "FloatProgressModel",
|
||||
"state": {
|
||||
"_dom_classes": [],
|
||||
"_model_module": "@jupyter-widgets/controls",
|
||||
"_model_module_version": "1.5.0",
|
||||
"_model_name": "FloatProgressModel",
|
||||
"_view_count": null,
|
||||
"_view_module": "@jupyter-widgets/controls",
|
||||
"_view_module_version": "1.5.0",
|
||||
"_view_name": "ProgressView",
|
||||
"bar_style": "success",
|
||||
"description": "",
|
||||
"description_tooltip": null,
|
||||
"layout": "IPY_MODEL_7ed3376588734d64a229fd8dd6bfe06e",
|
||||
"max": 170498071,
|
||||
"min": 0,
|
||||
"orientation": "horizontal",
|
||||
"style": "IPY_MODEL_61849c45ed0c4436995bb81d7efdac7e",
|
||||
"value": 170498071
|
||||
}
|
||||
},
|
||||
"9e511aaba1204ea9985fb6f55e024039": {
|
||||
"model_module": "@jupyter-widgets/controls",
|
||||
"model_name": "HTMLModel",
|
||||
"state": {
|
||||
"_dom_classes": [],
|
||||
"_model_module": "@jupyter-widgets/controls",
|
||||
"_model_module_version": "1.5.0",
|
||||
"_model_name": "HTMLModel",
|
||||
"_view_count": null,
|
||||
"_view_module": "@jupyter-widgets/controls",
|
||||
"_view_module_version": "1.5.0",
|
||||
"_view_name": "HTMLView",
|
||||
"description": "",
|
||||
"description_tooltip": null,
|
||||
"layout": "IPY_MODEL_7b765f7623474999a8c9815e02c1cc6b",
|
||||
"placeholder": "",
|
||||
"style": "IPY_MODEL_afd6865a483f4273addb343059632961",
|
||||
"value": " 170499072/? [00:04<00:00, 42242006.52it/s]"
|
||||
}
|
||||
},
|
||||
"afd6865a483f4273addb343059632961": {
|
||||
"model_module": "@jupyter-widgets/controls",
|
||||
"model_name": "DescriptionStyleModel",
|
||||
"state": {
|
||||
"_model_module": "@jupyter-widgets/controls",
|
||||
"_model_module_version": "1.5.0",
|
||||
"_model_name": "DescriptionStyleModel",
|
||||
"_view_count": null,
|
||||
"_view_module": "@jupyter-widgets/base",
|
||||
"_view_module_version": "1.2.0",
|
||||
"_view_name": "StyleView",
|
||||
"description_width": ""
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 4
|
||||
}
|
||||
@@ -0,0 +1,65 @@
|
||||
"""
|
||||
---
|
||||
title: CIFAR10 Experiment to try Weight Standardization and Batch-Channel Normalization
|
||||
summary: >
|
||||
This trains is a VGG net that uses weight standardization and batch-channel normalization
|
||||
to classify CIFAR10 images.
|
||||
---
|
||||
|
||||
# CIFAR10 Experiment to try Weight Standardization and Batch-Channel Normalization
|
||||
"""
|
||||
|
||||
import torch.nn as nn
|
||||
|
||||
from labml import experiment
|
||||
from labml.configs import option
|
||||
from labml_nn.experiments.cifar10 import CIFAR10Configs, CIFAR10VGGModel
|
||||
from labml_nn.normalization.batch_channel_norm import BatchChannelNorm
|
||||
from labml_nn.normalization.weight_standardization.conv2d import Conv2d
|
||||
|
||||
|
||||
class Model(CIFAR10VGGModel):
|
||||
"""
|
||||
### VGG 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:
|
||||
return nn.Sequential(
|
||||
Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
|
||||
BatchChannelNorm(out_channels, 32),
|
||||
nn.ReLU(inplace=True),
|
||||
)
|
||||
|
||||
def __init__(self):
|
||||
super().__init__([[64, 64], [128, 128], [256, 256, 256], [512, 512, 512], [512, 512, 512]])
|
||||
|
||||
|
||||
@option(CIFAR10Configs.model)
|
||||
def _model(c: CIFAR10Configs):
|
||||
"""
|
||||
### Create model
|
||||
"""
|
||||
return Model().to(c.device)
|
||||
|
||||
|
||||
def main():
|
||||
# Create experiment
|
||||
experiment.create(name='cifar10', comment='weight standardization')
|
||||
# Create configurations
|
||||
conf = CIFAR10Configs()
|
||||
# Load configurations
|
||||
experiment.configs(conf, {
|
||||
'optimizer.optimizer': 'Adam',
|
||||
'optimizer.learning_rate': 2.5e-4,
|
||||
'train_batch_size': 64,
|
||||
})
|
||||
# Start the experiment and run the training loop
|
||||
with experiment.start():
|
||||
conf.run()
|
||||
|
||||
|
||||
#
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,6 @@
|
||||
# [Weight Standardization](https://nn.labml.ai/normalization/weight_standardization/index.html)
|
||||
|
||||
This is a [PyTorch](https://pytorch.org) implementation of Weight Standardization from the paper
|
||||
[Micro-Batch Training with Batch-Channel Normalization and Weight Standardization](https://arxiv.org/abs/1903.10520).
|
||||
We also have an
|
||||
[annotated implementation of Batch-Channel Normalization](https://nn.labml.ai/normalization/batch_channel_norm/index.html).
|
||||
Reference in New Issue
Block a user