chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,295 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "AYV_dMVDxyc2",
|
||||
"pycharm": {
|
||||
"name": "#%% md\n"
|
||||
}
|
||||
},
|
||||
"source": [
|
||||
"[](https://github.com/labmlai/annotated_deep_learning_paper_implementations)\n",
|
||||
"[](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/transformers/basic/autoregressive_experiment.ipynb)\n",
|
||||
"\n",
|
||||
"## Transformer Experiment\n",
|
||||
"\n",
|
||||
"This trains a simple transformer with\n",
|
||||
"[multi headed attention](https://nn.labml.ai/transformers/mha.html)\n",
|
||||
"introduced in [Attention Is All You Need](https://arxiv.org/abs/1706.03762)\n",
|
||||
"on an NLP auto-regression task (with Tiny Shakespeare dataset)."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "AahG_i2y5tY9",
|
||||
"pycharm": {
|
||||
"name": "#%% md\n"
|
||||
}
|
||||
},
|
||||
"source": [
|
||||
"### Install the packages"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "ZCzmCrAIVg0L",
|
||||
"outputId": "cf107fb2-4d50-4c67-af34-367624553421",
|
||||
"pycharm": {
|
||||
"name": "#%%\n"
|
||||
}
|
||||
},
|
||||
"source": [
|
||||
"!pip install labml-nn --quiet"
|
||||
],
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "SE2VUQ6L5zxI",
|
||||
"pycharm": {
|
||||
"name": "#%% md\n"
|
||||
}
|
||||
},
|
||||
"source": [
|
||||
"### Imports"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "0hJXx_g0wS2C",
|
||||
"pycharm": {
|
||||
"name": "#%%\n"
|
||||
}
|
||||
},
|
||||
"source": [
|
||||
"from labml import experiment\n",
|
||||
"from labml_nn.transformers.basic.autoregressive_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=\"transformer\", 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': 512,\n",
|
||||
" # Train for 32 epochs\n",
|
||||
" 'epochs': 32,\n",
|
||||
" # Batch size $32$\n",
|
||||
" 'batch_size': 16,\n",
|
||||
" # Switch between training and validation for $10$ times\n",
|
||||
" # per epoch\n",
|
||||
" 'inner_iterations': 10,\n",
|
||||
"\n",
|
||||
" # Model size\n",
|
||||
" 'd_model': 256,\n",
|
||||
" 'transformer.n_heads': 16,\n",
|
||||
" 'transformer.ffn.d_ff': 1024,\n",
|
||||
"\n",
|
||||
" # Use [Noam optimizer](../../optimizers/noam.html)\n",
|
||||
" 'optimizer.optimizer': 'Noam',\n",
|
||||
" 'optimizer.learning_rate': 1.,\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
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"metadata": {
|
||||
"id": "oBXXlP2b7XZO",
|
||||
"pycharm": {
|
||||
"name": "#%%\n"
|
||||
}
|
||||
},
|
||||
"source": [],
|
||||
"outputs": [],
|
||||
"execution_count": null
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"accelerator": "GPU",
|
||||
"colab": {
|
||||
"collapsed_sections": [],
|
||||
"name": "Transformer",
|
||||
"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,155 @@
|
||||
"""
|
||||
---
|
||||
title: Transformer Auto-Regression Experiment
|
||||
summary: >
|
||||
This trains a simple transformer model on NLP auto-regression.
|
||||
---
|
||||
|
||||
# Transformer Auto-Regression Experiment
|
||||
|
||||
[](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/transformers/basic/autoregressive_experiment.ipynb)
|
||||
|
||||
This trains a simple transformer introduced in [Attention Is All You Need](https://arxiv.org/abs/1706.03762)
|
||||
on an NLP auto-regression task (with Tiny Shakespeare dataset).
|
||||
"""
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from labml import experiment
|
||||
from labml.configs import option
|
||||
from labml_nn.experiments.nlp_autoregression import NLPAutoRegressionConfigs
|
||||
from labml_nn.transformers import TransformerConfigs, Encoder
|
||||
from labml_nn.transformers.utils import subsequent_mask
|
||||
|
||||
|
||||
class AutoregressiveTransformer(nn.Module):
|
||||
"""
|
||||
## Auto-Regressive model
|
||||
"""
|
||||
def __init__(self, encoder: Encoder, src_embed: nn.Module, generator: nn.Module):
|
||||
"""
|
||||
* `encoder` is the transformer [Encoder](../models.html#Encoder)
|
||||
* `src_embed` is the token
|
||||
[embedding module (with positional encodings)](../models.html#EmbeddingsWithLearnedPositionalEncoding)
|
||||
* `generator` is the [final fully connected layer](../models.html#Generator) that gives the logits.
|
||||
"""
|
||||
super().__init__()
|
||||
self.src_embed = src_embed
|
||||
self.encoder = encoder
|
||||
self.generator = generator
|
||||
|
||||
# The mask will be initialized on the first call
|
||||
self.mask = None
|
||||
|
||||
def forward(self, x: torch.Tensor):
|
||||
# Create subsequent mask if mask is not initialized
|
||||
# or if the size of the mask is different
|
||||
if self.mask is None or self.mask.size(0) != len(x):
|
||||
# Subsequent mask, will mask out tokens from seeing future tokens
|
||||
self.mask = subsequent_mask(len(x)).to(x.device)
|
||||
# Get the token embeddings with positional encodings
|
||||
x = self.src_embed(x)
|
||||
# Transformer encoder
|
||||
x = self.encoder(x, self.mask)
|
||||
# Get logits
|
||||
x = self.generator(x)
|
||||
|
||||
# Return results
|
||||
# (second value is for state, since our trainer is used with RNNs also)
|
||||
return x, None
|
||||
|
||||
|
||||
class Configs(NLPAutoRegressionConfigs):
|
||||
"""
|
||||
## Configurations
|
||||
|
||||
This inherits from
|
||||
[`NLPAutoRegressionConfigs`](../../experiments/nlp_autoregression.html#NLPAutoRegressionConfigs)
|
||||
"""
|
||||
|
||||
# GPT model
|
||||
model: AutoregressiveTransformer
|
||||
# Transformer
|
||||
transformer: TransformerConfigs
|
||||
|
||||
|
||||
@option(Configs.transformer, 'Transformer')
|
||||
def _transformer_configs(c: Configs):
|
||||
"""
|
||||
### Transformer configurations
|
||||
"""
|
||||
|
||||
# We use our
|
||||
# [configurable transformer implementation](../configs.html#TransformerConfigs)
|
||||
conf = TransformerConfigs()
|
||||
# Set the vocabulary sizes for embeddings and generating logits
|
||||
conf.n_src_vocab = c.n_tokens
|
||||
conf.n_tgt_vocab = c.n_tokens
|
||||
#
|
||||
conf.d_model = c.d_model
|
||||
|
||||
#
|
||||
return conf
|
||||
|
||||
|
||||
@option(Configs.model)
|
||||
def _model(c: Configs):
|
||||
"""
|
||||
Create GPT model and initialize weights
|
||||
"""
|
||||
m = AutoregressiveTransformer(c.transformer.encoder,
|
||||
c.transformer.src_embed,
|
||||
c.transformer.generator).to(c.device)
|
||||
|
||||
return m
|
||||
|
||||
|
||||
def main():
|
||||
# Create experiment
|
||||
experiment.create(name="transformer")
|
||||
# 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': 512,
|
||||
# Train for 32 epochs
|
||||
'epochs': 32,
|
||||
# Batch size $32$
|
||||
'batch_size': 16,
|
||||
# Switch between training and validation for $10$ times
|
||||
# per epoch
|
||||
'inner_iterations': 10,
|
||||
|
||||
# Model size
|
||||
'd_model': 256,
|
||||
'transformer.n_heads': 16,
|
||||
'transformer.ffn.d_ff': 1024,
|
||||
|
||||
# Use [Noam optimizer](../../optimizers/noam.html)
|
||||
'optimizer.optimizer': 'Noam',
|
||||
'optimizer.learning_rate': 1.,
|
||||
})
|
||||
|
||||
# Set models 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,159 @@
|
||||
"""
|
||||
---
|
||||
title: Transformer Auto-Regression Experiment with [Sophia-G optimizer](../../optimizers/sophia.html)
|
||||
summary: >
|
||||
This trains a simple transformer model on NLP auto-regression with Sophia-G optimizer.
|
||||
---
|
||||
|
||||
# Transformer Auto-Regression Experiment with [Sophia-G optimizer](../../optimizers/sophia.html)
|
||||
|
||||
This trains a simple transformer introduced in [Attention Is All You Need](https://arxiv.org/abs/1706.03762)
|
||||
on an NLP auto-regression task (with Tiny Shakespeare dataset) with [Sophia-G optimizer](../../optimizers/sophia.html).
|
||||
"""
|
||||
import torch
|
||||
|
||||
from labml import experiment, tracker
|
||||
from labml_nn.helpers.trainer import BatchIndex
|
||||
from labml_nn.optimizers.sophia import Sophia
|
||||
from labml_nn.transformers.basic.autoregressive_experiment import Configs as TransformerAutoRegressionConfigs
|
||||
|
||||
|
||||
class Configs(TransformerAutoRegressionConfigs):
|
||||
"""
|
||||
## Configurations
|
||||
|
||||
This inherits from [`Configs`](autoregressive_experiment.html)
|
||||
"""
|
||||
|
||||
hess_interval: int = 10
|
||||
|
||||
optimizer: Sophia
|
||||
|
||||
def step(self, batch: any, batch_idx: BatchIndex):
|
||||
"""
|
||||
### Training or validation step with Gauss-Newton-Bartlett (GNB) Hessian diagonal estimator
|
||||
"""
|
||||
|
||||
# Set training/eval mode
|
||||
self.model.train(self.mode.is_train)
|
||||
|
||||
# Move data to the device
|
||||
data, target = batch[0].to(self.device), batch[1].to(self.device)
|
||||
|
||||
# Estimate the Hessian diagonal every $k$ steps
|
||||
if isinstance(self.optimizer, Sophia) and self.mode.is_train and batch_idx.idx % self.hess_interval == 0:
|
||||
# Get model outputs
|
||||
output, *_ = self.model(data)
|
||||
|
||||
# Create a categorical distribution from logits
|
||||
samp_dist = torch.distributions.Categorical(logits=output)
|
||||
# Sample $\hat{y}$
|
||||
y_sample = samp_dist.sample()
|
||||
|
||||
# Calculate and log loss
|
||||
loss = self.loss_func(output, y_sample)
|
||||
tracker.add("loss.hess.", loss)
|
||||
|
||||
# Calculate gradients
|
||||
loss.backward()
|
||||
# Clip gradients
|
||||
torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=self.grad_norm_clip)
|
||||
# Update EMA Hessian diagonal
|
||||
#
|
||||
# \begin{align}
|
||||
# \hat{h}_t &= B \cdot \nabla_\theta \hat{L} (\theta) \odot \nabla_\theta \hat{L} (\theta) \\
|
||||
# h_t &= \beta_2 h_{t-k} + (1 - \beta_2) \hat{h}_t
|
||||
# \end{align}
|
||||
self.optimizer.update_hessian(data.numel())
|
||||
# Clear the gradients
|
||||
self.optimizer.zero_grad()
|
||||
else:
|
||||
# Move data to the device
|
||||
data, target = batch[0].to(self.device), batch[1].to(self.device)
|
||||
|
||||
# Update global step (number of tokens processed) when in training mode
|
||||
if self.mode.is_train:
|
||||
tracker.add_global_step(data.shape[0] * data.shape[1])
|
||||
|
||||
# Get model outputs.
|
||||
# It's returning a tuple for states when using RNNs.
|
||||
# This is not implemented yet. 😜
|
||||
output, *_ = self.model(data)
|
||||
|
||||
# Calculate and log loss
|
||||
loss = self.loss_func(output, target)
|
||||
tracker.add("loss.", loss)
|
||||
|
||||
# Calculate and log accuracy
|
||||
self.accuracy(output, target)
|
||||
self.accuracy.track()
|
||||
|
||||
self.other_metrics(output, target)
|
||||
|
||||
# Train the model
|
||||
if self.mode.is_train:
|
||||
# Calculate gradients
|
||||
loss.backward()
|
||||
# Clip gradients
|
||||
torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=self.grad_norm_clip)
|
||||
# Take optimizer step
|
||||
self.optimizer.step()
|
||||
# Log the model parameters and gradients on last batch of every epoch
|
||||
if batch_idx.is_last and self.is_log_model_params_grads:
|
||||
tracker.add('model', self.model)
|
||||
# Clear the gradients
|
||||
self.optimizer.zero_grad()
|
||||
|
||||
# Save the tracked metrics
|
||||
tracker.save()
|
||||
|
||||
|
||||
def main():
|
||||
# Create experiment
|
||||
experiment.create(name="transformer")
|
||||
# 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': 512,
|
||||
# Train for 32 epochs
|
||||
'epochs': 32,
|
||||
# Batch size $32$
|
||||
'batch_size': 16,
|
||||
# Switch between training and validation for $10$ times
|
||||
# per epoch
|
||||
'inner_iterations': 10,
|
||||
|
||||
# Model size
|
||||
'd_model': 256,
|
||||
'transformer.n_heads': 16,
|
||||
'transformer.ffn.d_ff': 1024,
|
||||
|
||||
# Use [Sophia optimizer](../../optimizers/sophia.html)
|
||||
'optimizer.optimizer': 'Sophia',
|
||||
'optimizer.learning_rate': 3e-4,
|
||||
'optimizer.rho': 0.03,
|
||||
})
|
||||
|
||||
# Set models 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()
|
||||
Reference in New Issue
Block a user