chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,15 @@
|
||||
"""
|
||||
---
|
||||
title: Neural Networks with Adaptive Computation
|
||||
summary: >
|
||||
A set of PyTorch implementations/tutorials related to adaptive computation
|
||||
---
|
||||
|
||||
# Neural Networks with Adaptive Computation
|
||||
|
||||
These are neural network architectures that change the computation complexity based on the
|
||||
complexity of the input sample.
|
||||
|
||||
* 🚧 TODO: Adaptive Computation Time for Recurrent Neural Networks
|
||||
* [PonderNet: Learning to Ponder](ponder_net/index.html)
|
||||
"""
|
||||
@@ -0,0 +1,62 @@
|
||||
"""
|
||||
---
|
||||
title: "Parity Task"
|
||||
summary: >
|
||||
This creates data for Parity Task from the paper Adaptive Computation Time
|
||||
for Recurrent Neural Networks
|
||||
---
|
||||
|
||||
# Parity Task
|
||||
|
||||
This creates data for Parity Task from the paper
|
||||
[Adaptive Computation Time for Recurrent Neural Networks](https://arxiv.org/abs/1603.08983).
|
||||
|
||||
The input of the parity task is a vector with $0$'s $1$'s and $-1$'s.
|
||||
The output is the parity of $1$'s - one if there is an odd number of $1$'s and zero otherwise.
|
||||
The input is generated by making a random number of elements in the vector either $1$ or $-1$'s.
|
||||
"""
|
||||
|
||||
from typing import Tuple
|
||||
|
||||
import torch
|
||||
from torch.utils.data import Dataset
|
||||
|
||||
|
||||
class ParityDataset(Dataset):
|
||||
"""
|
||||
### Parity dataset
|
||||
"""
|
||||
|
||||
def __init__(self, n_samples: int, n_elems: int = 64):
|
||||
"""
|
||||
* `n_samples` is the number of samples
|
||||
* `n_elems` is the number of elements in the input vector
|
||||
"""
|
||||
self.n_samples = n_samples
|
||||
self.n_elems = n_elems
|
||||
|
||||
def __len__(self):
|
||||
"""
|
||||
Size of the dataset
|
||||
"""
|
||||
return self.n_samples
|
||||
|
||||
def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Generate a sample
|
||||
"""
|
||||
|
||||
# Empty vector
|
||||
x = torch.zeros((self.n_elems,))
|
||||
# Number of non-zero elements - a random number between $1$ and total number of elements
|
||||
n_non_zero = torch.randint(1, self.n_elems + 1, (1,)).item()
|
||||
# Fill non-zero elements with $1$'s and $-1$'s
|
||||
x[:n_non_zero] = torch.randint(0, 2, (n_non_zero,)) * 2 - 1
|
||||
# Randomly permute the elements
|
||||
x = x[torch.randperm(self.n_elems)]
|
||||
|
||||
# The parity
|
||||
y = (x == 1.).sum() % 2
|
||||
|
||||
#
|
||||
return x, y
|
||||
@@ -0,0 +1,267 @@
|
||||
"""
|
||||
---
|
||||
title: "PonderNet: Learning to Ponder"
|
||||
summary: >
|
||||
A PyTorch implementation/tutorial of PonderNet: Learning to Ponder.
|
||||
---
|
||||
|
||||
# PonderNet: Learning to Ponder
|
||||
|
||||
This is a [PyTorch](https://pytorch.org) implementation of the paper
|
||||
[PonderNet: Learning to Ponder](https://arxiv.org/abs/2107.05407).
|
||||
|
||||
PonderNet adapts the computation based on the input.
|
||||
It changes the number of steps to take on a recurrent network based on the input.
|
||||
PonderNet learns this with end-to-end gradient descent.
|
||||
|
||||
PonderNet has a step function of the form
|
||||
|
||||
$$\hat{y}_n, h_{n+1}, \lambda_n = s(x, h_n)$$
|
||||
|
||||
where $x$ is the input, $h_n$ is the state, $\hat{y}_n$ is the prediction at step $n$,
|
||||
and $\lambda_n$ is the probability of halting (stopping) at current step.
|
||||
|
||||
$s$ can be any neural network (e.g. LSTM, MLP, GRU, Attention layer).
|
||||
|
||||
The unconditioned probability of halting at step $n$ is then,
|
||||
|
||||
$$p_n = \lambda_n \prod_{j=1}^{n-1} (1 - \lambda_j)$$
|
||||
|
||||
That is the probability of not being halted at any of the previous steps and halting at step $n$.
|
||||
|
||||
During inference, we halt by sampling based on the halting probability $\lambda_n$
|
||||
and get the prediction at the halting layer $\hat{y}_n$ as the final output.
|
||||
|
||||
During training, we get the predictions from all the layers and calculate the losses for each of them.
|
||||
And then take the weighted average of the losses based on the probabilities of getting halted at each layer
|
||||
$p_n$.
|
||||
|
||||
The step function is applied to a maximum number of steps donated by $N$.
|
||||
|
||||
The overall loss of PonderNet is
|
||||
|
||||
\begin{align}
|
||||
L &= L_{Rec} + \beta L_{Reg} \\
|
||||
L_{Rec} &= \sum_{n=1}^N p_n \mathcal{L}(y, \hat{y}_n) \\
|
||||
L_{Reg} &= \mathop{KL} \Big(p_n \Vert p_G(\lambda_p) \Big)
|
||||
\end{align}
|
||||
|
||||
$\mathcal{L}$ is the normal loss function between target $y$ and prediction $\hat{y}_n$.
|
||||
|
||||
$\mathop{KL}$ is the [Kullback–Leibler divergence](https://en.wikipedia.org/wiki/Kullback%E2%80%93Leibler_divergence).
|
||||
|
||||
$p_G$ is the [Geometric distribution](https://en.wikipedia.org/wiki/Geometric_distribution) parameterized by
|
||||
$\lambda_p$. *$\lambda_p$ has nothing to do with $\lambda_n$; we are just sticking to same notation as the paper*.
|
||||
$$Pr_{p_G(\lambda_p)}(X = k) = (1 - \lambda_p)^k \lambda_p$$.
|
||||
|
||||
The regularization loss biases the network towards taking $\frac{1}{\lambda_p}$ steps and incentivizes
|
||||
non-zero probabilities for all steps; i.e. promotes exploration.
|
||||
|
||||
Here is the [training code `experiment.py`](experiment.html) to train a PonderNet on [Parity Task](../parity.html).
|
||||
"""
|
||||
|
||||
from typing import Tuple
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
|
||||
|
||||
class ParityPonderGRU(nn.Module):
|
||||
"""
|
||||
## PonderNet with GRU for Parity Task
|
||||
|
||||
This is a simple model that uses a [GRU Cell](https://pytorch.org/docs/stable/generated/torch.nn.GRUCell.html)
|
||||
as the step function.
|
||||
|
||||
This model is for the [Parity Task](../parity.html) where the input is a vector of `n_elems`.
|
||||
Each element of the vector is either `0`, `1` or `-1` and the output is the parity
|
||||
- a binary value that is true if the number of `1`s is odd and false otherwise.
|
||||
|
||||
The prediction of the model is the log probability of the parity being $1$.
|
||||
"""
|
||||
|
||||
def __init__(self, n_elems: int, n_hidden: int, max_steps: int):
|
||||
"""
|
||||
* `n_elems` is the number of elements in the input vector
|
||||
* `n_hidden` is the state vector size of the GRU
|
||||
* `max_steps` is the maximum number of steps $N$
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
self.max_steps = max_steps
|
||||
self.n_hidden = n_hidden
|
||||
|
||||
# GRU
|
||||
# $$h_{n+1} = s_h(x, h_n)$$
|
||||
self.gru = nn.GRUCell(n_elems, n_hidden)
|
||||
# $$\hat{y}_n = s_y(h_n)$$
|
||||
# We could use a layer that takes the concatenation of $h$ and $x$ as input
|
||||
# but we went with this for simplicity.
|
||||
self.output_layer = nn.Linear(n_hidden, 1)
|
||||
# $$\lambda_n = s_\lambda(h_n)$$
|
||||
self.lambda_layer = nn.Linear(n_hidden, 1)
|
||||
self.lambda_prob = nn.Sigmoid()
|
||||
# An option to set during inference so that computation is actually halted at inference time
|
||||
self.is_halt = False
|
||||
|
||||
def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
* `x` is the input of shape `[batch_size, n_elems]`
|
||||
|
||||
This outputs a tuple of four tensors:
|
||||
|
||||
1. $p_1 \dots p_N$ in a tensor of shape `[N, batch_size]`
|
||||
2. $\hat{y}_1 \dots \hat{y}_N$ in a tensor of shape `[N, batch_size]` - the log probabilities of the parity being $1$
|
||||
3. $p_m$ of shape `[batch_size]`
|
||||
4. $\hat{y}_m$ of shape `[batch_size]` where the computation was halted at step $m$
|
||||
"""
|
||||
|
||||
#
|
||||
batch_size = x.shape[0]
|
||||
|
||||
# We get initial state $h_1 = s_h(x)$
|
||||
h = x.new_zeros((x.shape[0], self.n_hidden))
|
||||
h = self.gru(x, h)
|
||||
|
||||
# Lists to store $p_1 \dots p_N$ and $\hat{y}_1 \dots \hat{y}_N$
|
||||
p = []
|
||||
y = []
|
||||
# $\prod_{j=1}^{n-1} (1 - \lambda_j)$
|
||||
un_halted_prob = h.new_ones((batch_size,))
|
||||
|
||||
# A vector to maintain which samples has halted computation
|
||||
halted = h.new_zeros((batch_size,))
|
||||
# $p_m$ and $\hat{y}_m$ where the computation was halted at step $m$
|
||||
p_m = h.new_zeros((batch_size,))
|
||||
y_m = h.new_zeros((batch_size,))
|
||||
|
||||
# Iterate for $N$ steps
|
||||
for n in range(1, self.max_steps + 1):
|
||||
# The halting probability $\lambda_N = 1$ for the last step
|
||||
if n == self.max_steps:
|
||||
lambda_n = h.new_ones(h.shape[0])
|
||||
# $\lambda_n = s_\lambda(h_n)$
|
||||
else:
|
||||
lambda_n = self.lambda_prob(self.lambda_layer(h))[:, 0]
|
||||
# $\hat{y}_n = s_y(h_n)$
|
||||
y_n = self.output_layer(h)[:, 0]
|
||||
|
||||
# $$p_n = \lambda_n \prod_{j=1}^{n-1} (1 - \lambda_j)$$
|
||||
p_n = un_halted_prob * lambda_n
|
||||
# Update $\prod_{j=1}^{n-1} (1 - \lambda_j)$
|
||||
un_halted_prob = un_halted_prob * (1 - lambda_n)
|
||||
|
||||
# Halt based on halting probability $\lambda_n$
|
||||
halt = torch.bernoulli(lambda_n) * (1 - halted)
|
||||
|
||||
# Collect $p_n$ and $\hat{y}_n$
|
||||
p.append(p_n)
|
||||
y.append(y_n)
|
||||
|
||||
# Update $p_m$ and $\hat{y}_m$ based on what was halted at current step $n$
|
||||
p_m = p_m * (1 - halt) + p_n * halt
|
||||
y_m = y_m * (1 - halt) + y_n * halt
|
||||
|
||||
# Update halted samples
|
||||
halted = halted + halt
|
||||
# Get next state $h_{n+1} = s_h(x, h_n)$
|
||||
h = self.gru(x, h)
|
||||
|
||||
# Stop the computation if all samples have halted
|
||||
if self.is_halt and halted.sum() == batch_size:
|
||||
break
|
||||
|
||||
#
|
||||
return torch.stack(p), torch.stack(y), p_m, y_m
|
||||
|
||||
|
||||
class ReconstructionLoss(nn.Module):
|
||||
"""
|
||||
## Reconstruction loss
|
||||
|
||||
$$L_{Rec} = \sum_{n=1}^N p_n \mathcal{L}(y, \hat{y}_n)$$
|
||||
|
||||
$\mathcal{L}$ is the normal loss function between target $y$ and prediction $\hat{y}_n$.
|
||||
"""
|
||||
|
||||
def __init__(self, loss_func: nn.Module):
|
||||
"""
|
||||
* `loss_func` is the loss function $\mathcal{L}$
|
||||
"""
|
||||
super().__init__()
|
||||
self.loss_func = loss_func
|
||||
|
||||
def forward(self, p: torch.Tensor, y_hat: torch.Tensor, y: torch.Tensor):
|
||||
"""
|
||||
* `p` is $p_1 \dots p_N$ in a tensor of shape `[N, batch_size]`
|
||||
* `y_hat` is $\hat{y}_1 \dots \hat{y}_N$ in a tensor of shape `[N, batch_size, ...]`
|
||||
* `y` is the target of shape `[batch_size, ...]`
|
||||
"""
|
||||
|
||||
# The total $\sum_{n=1}^N p_n \mathcal{L}(y, \hat{y}_n)$
|
||||
total_loss = p.new_tensor(0.)
|
||||
# Iterate upto $N$
|
||||
for n in range(p.shape[0]):
|
||||
# $p_n \mathcal{L}(y, \hat{y}_n)$ for each sample and the mean of them
|
||||
loss = (p[n] * self.loss_func(y_hat[n], y)).mean()
|
||||
# Add to total loss
|
||||
total_loss = total_loss + loss
|
||||
|
||||
#
|
||||
return total_loss
|
||||
|
||||
|
||||
class RegularizationLoss(nn.Module):
|
||||
"""
|
||||
## Regularization loss
|
||||
|
||||
$$L_{Reg} = \mathop{KL} \Big(p_n \Vert p_G(\lambda_p) \Big)$$
|
||||
|
||||
$\mathop{KL}$ is the [Kullback–Leibler divergence](https://en.wikipedia.org/wiki/Kullback%E2%80%93Leibler_divergence).
|
||||
|
||||
$p_G$ is the [Geometric distribution](https://en.wikipedia.org/wiki/Geometric_distribution) parameterized by
|
||||
$\lambda_p$. *$\lambda_p$ has nothing to do with $\lambda_n$; we are just sticking to same notation as the paper*.
|
||||
$$Pr_{p_G(\lambda_p)}(X = k) = (1 - \lambda_p)^k \lambda_p$$.
|
||||
|
||||
The regularization loss biases the network towards taking $\frac{1}{\lambda_p}$ steps and incentivies non-zero probabilities
|
||||
for all steps; i.e. promotes exploration.
|
||||
"""
|
||||
|
||||
def __init__(self, lambda_p: float, max_steps: int = 1_000):
|
||||
"""
|
||||
* `lambda_p` is $\lambda_p$ - the success probability of geometric distribution
|
||||
* `max_steps` is the highest $N$; we use this to pre-compute $p_G(\lambda_p)$
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
# Empty vector to calculate $p_G(\lambda_p)$
|
||||
p_g = torch.zeros((max_steps,))
|
||||
# $(1 - \lambda_p)^k$
|
||||
not_halted = 1.
|
||||
# Iterate upto `max_steps`
|
||||
for k in range(max_steps):
|
||||
# $$Pr_{p_G(\lambda_p)}(X = k) = (1 - \lambda_p)^k \lambda_p$$
|
||||
p_g[k] = not_halted * lambda_p
|
||||
# Update $(1 - \lambda_p)^k$
|
||||
not_halted = not_halted * (1 - lambda_p)
|
||||
|
||||
# Save $Pr_{p_G(\lambda_p)}$
|
||||
self.p_g = nn.Parameter(p_g, requires_grad=False)
|
||||
|
||||
# KL-divergence loss
|
||||
self.kl_div = nn.KLDivLoss(reduction='batchmean')
|
||||
|
||||
def forward(self, p: torch.Tensor):
|
||||
"""
|
||||
* `p` is $p_1 \dots p_N$ in a tensor of shape `[N, batch_size]`
|
||||
"""
|
||||
# Transpose `p` to `[batch_size, N]`
|
||||
p = p.transpose(0, 1)
|
||||
# Get $Pr_{p_G(\lambda_p)}$ upto $N$ and expand it across the batch dimension
|
||||
p_g = self.p_g[None, :p.shape[1]].expand_as(p)
|
||||
|
||||
# Calculate the KL-divergence.
|
||||
# *The [PyTorch KL-divergence](https://pytorch.org/docs/stable/generated/torch.nn.KLDivLoss.html)
|
||||
# implementation accepts log probabilities.*
|
||||
return self.kl_div(p.log(), p_g)
|
||||
@@ -0,0 +1,161 @@
|
||||
"""
|
||||
---
|
||||
title: "PonderNet Parity Task Experiment"
|
||||
summary: >
|
||||
This trains is a PonderNet on Parity Task
|
||||
---
|
||||
|
||||
# [PonderNet](index.html) [Parity Task](../parity.html) Experiment
|
||||
|
||||
This trains a [PonderNet](index.html) on [Parity Task](../parity.html).
|
||||
"""
|
||||
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.utils.data import DataLoader
|
||||
|
||||
from labml import tracker, experiment
|
||||
from labml_nn.helpers.metrics import AccuracyDirect
|
||||
from labml_nn.helpers.trainer import SimpleTrainValidConfigs, BatchIndex
|
||||
from labml_nn.adaptive_computation.parity import ParityDataset
|
||||
from labml_nn.adaptive_computation.ponder_net import ParityPonderGRU, ReconstructionLoss, RegularizationLoss
|
||||
|
||||
|
||||
class Configs(SimpleTrainValidConfigs):
|
||||
"""
|
||||
Configurations with a
|
||||
[simple training loop](../../helpers/trainer.html)
|
||||
"""
|
||||
|
||||
# Number of epochs
|
||||
epochs: int = 100
|
||||
# Number of batches per epoch
|
||||
n_batches: int = 500
|
||||
# Batch size
|
||||
batch_size: int = 128
|
||||
|
||||
# Model
|
||||
model: ParityPonderGRU
|
||||
|
||||
# $L_{Rec}$
|
||||
loss_rec: ReconstructionLoss
|
||||
# $L_{Reg}$
|
||||
loss_reg: RegularizationLoss
|
||||
|
||||
# The number of elements in the input vector.
|
||||
# *We keep it low for demonstration; otherwise, training takes a lot of time.
|
||||
# Although the parity task seems simple, figuring out the pattern by looking at samples
|
||||
# is quite hard.*
|
||||
n_elems: int = 8
|
||||
# Number of units in the hidden layer (state)
|
||||
n_hidden: int = 64
|
||||
# Maximum number of steps $N$
|
||||
max_steps: int = 20
|
||||
|
||||
# $\lambda_p$ for the geometric distribution $p_G(\lambda_p)$
|
||||
lambda_p: float = 0.2
|
||||
# Regularization loss $L_{Reg}$ coefficient $\beta$
|
||||
beta: float = 0.01
|
||||
|
||||
# Gradient clipping by norm
|
||||
grad_norm_clip: float = 1.0
|
||||
|
||||
# Training and validation loaders
|
||||
train_loader: DataLoader
|
||||
valid_loader: DataLoader
|
||||
|
||||
# Accuracy calculator
|
||||
accuracy = AccuracyDirect()
|
||||
|
||||
def init(self):
|
||||
# Print indicators to screen
|
||||
tracker.set_scalar('loss.*', True)
|
||||
tracker.set_scalar('loss_reg.*', True)
|
||||
tracker.set_scalar('accuracy.*', True)
|
||||
tracker.set_scalar('steps.*', True)
|
||||
|
||||
# We need to set the metrics to calculate them for the epoch for training and validation
|
||||
self.state_modules = [self.accuracy]
|
||||
|
||||
# Initialize the model
|
||||
self.model = ParityPonderGRU(self.n_elems, self.n_hidden, self.max_steps).to(self.device)
|
||||
# $L_{Rec}$
|
||||
self.loss_rec = ReconstructionLoss(nn.BCEWithLogitsLoss(reduction='none')).to(self.device)
|
||||
# $L_{Reg}$
|
||||
self.loss_reg = RegularizationLoss(self.lambda_p, self.max_steps).to(self.device)
|
||||
|
||||
# Training and validation loaders
|
||||
self.train_loader = DataLoader(ParityDataset(self.batch_size * self.n_batches, self.n_elems),
|
||||
batch_size=self.batch_size)
|
||||
self.valid_loader = DataLoader(ParityDataset(self.batch_size * 32, self.n_elems),
|
||||
batch_size=self.batch_size)
|
||||
|
||||
def step(self, batch: Any, batch_idx: BatchIndex):
|
||||
"""
|
||||
This method gets called by the trainer for each batch
|
||||
"""
|
||||
# Set the model mode
|
||||
self.model.train(self.mode.is_train)
|
||||
|
||||
# Get the input and labels and move them to the model's device
|
||||
data, target = batch[0].to(self.device), batch[1].to(self.device)
|
||||
|
||||
# Increment step in training mode
|
||||
if self.mode.is_train:
|
||||
tracker.add_global_step(len(data))
|
||||
|
||||
# Run the model
|
||||
p, y_hat, p_sampled, y_hat_sampled = self.model(data)
|
||||
|
||||
# Calculate the reconstruction loss
|
||||
loss_rec = self.loss_rec(p, y_hat, target.to(torch.float))
|
||||
tracker.add("loss.", loss_rec)
|
||||
|
||||
# Calculate the regularization loss
|
||||
loss_reg = self.loss_reg(p)
|
||||
tracker.add("loss_reg.", loss_reg)
|
||||
|
||||
# $L = L_{Rec} + \beta L_{Reg}$
|
||||
loss = loss_rec + self.beta * loss_reg
|
||||
|
||||
# Calculate the expected number of steps taken
|
||||
steps = torch.arange(1, p.shape[0] + 1, device=p.device)
|
||||
expected_steps = (p * steps[:, None]).sum(dim=0)
|
||||
tracker.add("steps.", expected_steps)
|
||||
|
||||
# Call accuracy metric
|
||||
self.accuracy(y_hat_sampled > 0, target)
|
||||
|
||||
if self.mode.is_train:
|
||||
# Compute gradients
|
||||
loss.backward()
|
||||
# Clip gradients
|
||||
torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=self.grad_norm_clip)
|
||||
# Optimizer
|
||||
self.optimizer.step()
|
||||
# Clear gradients
|
||||
self.optimizer.zero_grad()
|
||||
#
|
||||
tracker.save()
|
||||
|
||||
|
||||
def main():
|
||||
"""
|
||||
Run the experiment
|
||||
"""
|
||||
experiment.create(name='ponder_net')
|
||||
|
||||
conf = Configs()
|
||||
experiment.configs(conf, {
|
||||
'optimizer.optimizer': 'Adam',
|
||||
'optimizer.learning_rate': 0.0003,
|
||||
})
|
||||
|
||||
with experiment.start():
|
||||
conf.run()
|
||||
|
||||
#
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,8 @@
|
||||
# [PonderNet: Learning to Ponder](https://nn.labml.ai/adaptive_computation/ponder_net/index.html)
|
||||
|
||||
This is a [PyTorch](https://pytorch.org) implementation of the paper
|
||||
[PonderNet: Learning to Ponder](https://arxiv.org/abs/2107.05407).
|
||||
|
||||
PonderNet adapts the computation based on the input.
|
||||
It changes the number of steps to take on a recurrent network based on the input.
|
||||
PonderNet learns this with end-to-end gradient descent.
|
||||
@@ -0,0 +1,7 @@
|
||||
# [Neural Networks with Adaptive Computation](https://nn.labml.ai/adaptive_computation/index.html)
|
||||
|
||||
These are neural network architectures that change the computation complexity based on the
|
||||
complexity of the input sample.
|
||||
|
||||
* 🚧 TODO: Adaptive Computation Time for Recurrent Neural Networks
|
||||
* [PonderNet: Learning to Ponder](https://nn.labml.ai/adaptive_computation/ponder_net/index.html)
|
||||
Reference in New Issue
Block a user