chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,260 @@
|
||||
"""
|
||||
---
|
||||
title: Arithmetic Dataset
|
||||
summary: >
|
||||
This creates arithmetic problems.
|
||||
---
|
||||
|
||||
*This is based on code by [Georges Harik (@gharik)](https://twitter.com/gharik).*
|
||||
"""
|
||||
|
||||
import random
|
||||
import string
|
||||
from typing import List
|
||||
|
||||
import torch
|
||||
from labml.logger import Text
|
||||
from torch.utils.data import DataLoader, Dataset
|
||||
|
||||
from labml import monit, logger, tracker
|
||||
from labml.configs import option
|
||||
from labml_nn.experiments.nlp_autoregression import NLPAutoRegressionConfigs, transpose_batch
|
||||
|
||||
|
||||
class ArithmeticDataset(Dataset):
|
||||
"""
|
||||
## Arithmetic Dataset
|
||||
|
||||
This creates arithmetic addition problems and solutions with workings.
|
||||
We've only implemented addition so far.
|
||||
|
||||
It's based on a character level tokenization.
|
||||
"""
|
||||
|
||||
def __init__(self, seq_len: int, max_digits: int, n_sequences: int):
|
||||
"""
|
||||
:param seq_len: is the sequence length of generated math problems.
|
||||
We fill as many problems as possible upto this length
|
||||
:max_digits: is the maximum number of digits in the operand integers
|
||||
:n_sequences: is the number of sequences per epoch
|
||||
"""
|
||||
self.n_sequences = n_sequences
|
||||
self.max_digits = max_digits
|
||||
self.seq_len = seq_len
|
||||
# Token id to string
|
||||
self.itos = list(string.digits + 'xe =\n?+;')
|
||||
# Character to token id
|
||||
self.stoi = {c: i for i, c in enumerate(self.itos)}
|
||||
|
||||
@staticmethod
|
||||
def make_int(n_digits: int):
|
||||
"""
|
||||
Generates an integer with `n_digit` number of digits
|
||||
"""
|
||||
res = 0
|
||||
for i in range(n_digits):
|
||||
d = random.randrange(1, 11) if i == 0 else random.randrange(0, 11)
|
||||
res = res * 10 + d
|
||||
|
||||
return res
|
||||
|
||||
@staticmethod
|
||||
def get_add_explanation(x: int, y: int):
|
||||
"""
|
||||
Generates the workings for `x + y`.
|
||||
For example for `11+29` it generates
|
||||
`1e0+9e0+0e0=10e0 1e0+2e0+1e0=4e0`.
|
||||
"""
|
||||
|
||||
carry = 0
|
||||
e = 0
|
||||
explanation = []
|
||||
while x > 0 or y > 0 or carry > 0:
|
||||
rx, ry = x % 10, y % 10
|
||||
total = rx + ry + carry
|
||||
explanation.append(f"{rx}e{e}+{ry}e{e}+{carry}e{e}=={total}e{e}")
|
||||
x, y, carry = x // 10, y // 10, total // 10
|
||||
e += 1
|
||||
|
||||
return ' '.join(explanation)
|
||||
|
||||
# Make a problem with a pre_explanation or not
|
||||
def make_add_problem(self):
|
||||
"""
|
||||
Creates an arithmetic addition problem with workings and answer.
|
||||
"""
|
||||
x = self.make_int(n_digits=random.randrange(1, self.max_digits + 1))
|
||||
y = self.make_int(n_digits=random.randrange(1, self.max_digits + 1))
|
||||
|
||||
explanation = self.get_add_explanation(x, y)
|
||||
return f"x={x}+{y}; {explanation} x=={x + y}\n"
|
||||
|
||||
def get_qa(self):
|
||||
"""
|
||||
Get arithmetic problem and answer. This is used for evaluation.
|
||||
"""
|
||||
x = self.make_int(n_digits=random.randrange(1, self.max_digits + 1))
|
||||
y = self.make_int(n_digits=random.randrange(1, self.max_digits + 1))
|
||||
|
||||
return f'x={x}+{y};', f'{x + y}'
|
||||
|
||||
def get_packed_math_input(self):
|
||||
"""
|
||||
Generate multiple problems and pack them into a sequence.
|
||||
"""
|
||||
s_enc = []
|
||||
while len(s_enc) <= self.seq_len:
|
||||
s_part = self.make_add_problem()
|
||||
s_part_enc = self.encode('?' + s_part)
|
||||
s_enc = s_enc + s_part_enc
|
||||
return s_enc
|
||||
|
||||
def encode(self, s: str):
|
||||
"""
|
||||
Encode a given string
|
||||
"""
|
||||
return [self.stoi[c] for c in s]
|
||||
|
||||
def decode(self, arr: List[int]):
|
||||
"""
|
||||
Decode a list of token ids
|
||||
"""
|
||||
return ''.join([self.itos[c] for c in arr])
|
||||
|
||||
def __getitem__(self, idx: int):
|
||||
"""
|
||||
Get a input and target pair for auto-regressive modelling
|
||||
"""
|
||||
s = torch.tensor(self.get_packed_math_input())
|
||||
return s[:self.seq_len], s[1:self.seq_len + 1]
|
||||
|
||||
def __len__(self):
|
||||
"""
|
||||
Number of sequences per epoch
|
||||
"""
|
||||
return self.n_sequences
|
||||
|
||||
|
||||
class ArithmeticAutoregression(NLPAutoRegressionConfigs):
|
||||
"""
|
||||
## Arithmetic Task Experiment Configurations
|
||||
"""
|
||||
# Maximum number of digits per operand integer
|
||||
max_digits: int = 4
|
||||
# Number of training sequences per epoch
|
||||
train_sequences_per_epoch: int = 2 ** 12
|
||||
# Training data loader
|
||||
train_loader: DataLoader = 'arithmetic_train_loader'
|
||||
# Number of problems in evaluation
|
||||
n_tests: int = 64
|
||||
# No need of a validation dataset
|
||||
validator = None
|
||||
# Number of times to run evaluations per epoch
|
||||
inner_iterations = 4
|
||||
# Number of tokens in the vocabulary
|
||||
n_tokens = len(ArithmeticDataset(1, 1, 1).itos)
|
||||
|
||||
@torch.no_grad()
|
||||
def sample(self):
|
||||
"""
|
||||
### Evaluation
|
||||
|
||||
We use the sampling function to evaluate the model on a set of problems
|
||||
"""
|
||||
|
||||
# Skip in the first epoch
|
||||
if self.training_loop.idx < 1:
|
||||
return
|
||||
|
||||
# Create a dataset to generate problems
|
||||
dataset = ArithmeticDataset(self.seq_len, self.max_digits, 1)
|
||||
# Get a set of problems and answers
|
||||
qa = [dataset.get_qa() for _ in range(self.n_tests)]
|
||||
# Collect the problems only
|
||||
questions = [p[0] for p in qa]
|
||||
|
||||
# Create a tensor with only the initial token
|
||||
data = torch.tensor([[dataset.stoi[p[0]] for p in questions]])
|
||||
# Move to device
|
||||
data = data.to(self.device)
|
||||
|
||||
# Number of sequences that have completed
|
||||
finished = torch.zeros((len(questions),)).bool().to(self.device)
|
||||
# Token id of the new line character - this marks end of the answer
|
||||
new_line = dataset.stoi['\n']
|
||||
|
||||
# Sampled results
|
||||
results = [p[0] for p in questions]
|
||||
|
||||
# Sample upto sequence length
|
||||
for i in monit.iterate('Sample', self.seq_len - 1):
|
||||
# If all the sequences have completed we skip this
|
||||
if finished.sum() == len(finished):
|
||||
continue
|
||||
|
||||
# Get the model output
|
||||
output, *_ = self.model(data)
|
||||
# Get the model prediction (greedy)
|
||||
output = output[-1].argmax(dim=-1)
|
||||
|
||||
# Find which sequences have finished
|
||||
finished = finished | (output == new_line)
|
||||
# Skip if all have finished
|
||||
if finished.sum() == len(finished):
|
||||
continue
|
||||
|
||||
# Override with the question
|
||||
for j, p in enumerate(questions):
|
||||
if len(p) > i + 1:
|
||||
output[j] = dataset.stoi[p[i + 1]]
|
||||
|
||||
# Add the next token to the input
|
||||
data = torch.cat([data, output[None, :]], dim=0)
|
||||
|
||||
# Get the sampled results
|
||||
for j, c in enumerate(output):
|
||||
results[j] += dataset.itos[c]
|
||||
|
||||
# Discard everything after the answer in the results
|
||||
results = [r.split('\n')[0] for r in results]
|
||||
|
||||
# Log a sample
|
||||
res_sample = results[0].split(';')
|
||||
logger.log([(res_sample[0], Text.key), (';', Text.subtle), (';'.join(res_sample[1:]), Text.none)])
|
||||
|
||||
# Get the answers
|
||||
results = [r.split('x==')[-1] for r in results]
|
||||
|
||||
# Count the number of correct answers
|
||||
correct = 0
|
||||
for r, _qa in zip(results, qa):
|
||||
if r == _qa[1]:
|
||||
correct += 1
|
||||
|
||||
# Log the score
|
||||
tracker.save('score', correct / len(results))
|
||||
|
||||
|
||||
@option(ArithmeticAutoregression.train_loader)
|
||||
def arithmetic_train_loader(c: ArithmeticAutoregression):
|
||||
"""
|
||||
Training data loader
|
||||
"""
|
||||
return DataLoader(ArithmeticDataset(c.seq_len, c.max_digits, c.train_sequences_per_epoch),
|
||||
batch_size=c.batch_size,
|
||||
collate_fn=transpose_batch,
|
||||
num_workers=4)
|
||||
|
||||
|
||||
def _test():
|
||||
"""
|
||||
Code to test generated problems
|
||||
"""
|
||||
dataset = ArithmeticDataset(256, 8, 10)
|
||||
|
||||
print(dataset.decode(dataset.get_packed_math_input()))
|
||||
|
||||
|
||||
#
|
||||
if __name__ == '__main__':
|
||||
_test()
|
||||
@@ -0,0 +1,111 @@
|
||||
"""
|
||||
---
|
||||
title: CIFAR10 Experiment
|
||||
summary: >
|
||||
This is a reusable trainer for CIFAR10 dataset
|
||||
---
|
||||
|
||||
# CIFAR10 Experiment
|
||||
"""
|
||||
from typing import List
|
||||
|
||||
import torch.nn as nn
|
||||
|
||||
from labml import lab
|
||||
from labml.configs import option
|
||||
from labml_nn.helpers.datasets import CIFAR10Configs as CIFAR10DatasetConfigs
|
||||
from labml_nn.experiments.mnist import MNISTConfigs
|
||||
|
||||
|
||||
class CIFAR10Configs(CIFAR10DatasetConfigs, MNISTConfigs):
|
||||
"""
|
||||
## Configurations
|
||||
|
||||
This extends from [CIFAR 10 dataset configurations](../helpers/datasets.html)
|
||||
and [`MNISTConfigs`](mnist.html).
|
||||
"""
|
||||
# Use CIFAR10 dataset by default
|
||||
dataset_name: str = 'CIFAR10'
|
||||
|
||||
|
||||
@option(CIFAR10Configs.train_dataset)
|
||||
def cifar10_train_augmented():
|
||||
"""
|
||||
### Augmented CIFAR 10 train dataset
|
||||
"""
|
||||
from torchvision.datasets import CIFAR10
|
||||
from torchvision.transforms import transforms
|
||||
return CIFAR10(str(lab.get_data_path()),
|
||||
train=True,
|
||||
download=True,
|
||||
transform=transforms.Compose([
|
||||
# Pad and crop
|
||||
transforms.RandomCrop(32, padding=4),
|
||||
# Random horizontal flip
|
||||
transforms.RandomHorizontalFlip(),
|
||||
#
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
|
||||
]))
|
||||
|
||||
|
||||
@option(CIFAR10Configs.valid_dataset)
|
||||
def cifar10_valid_no_augment():
|
||||
"""
|
||||
### Non-augmented CIFAR 10 validation dataset
|
||||
"""
|
||||
from torchvision.datasets import CIFAR10
|
||||
from torchvision.transforms import transforms
|
||||
return CIFAR10(str(lab.get_data_path()),
|
||||
train=False,
|
||||
download=True,
|
||||
transform=transforms.Compose([
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
|
||||
]))
|
||||
|
||||
|
||||
class CIFAR10VGGModel(nn.Module):
|
||||
"""
|
||||
### VGG model for CIFAR-10 classification
|
||||
"""
|
||||
|
||||
def conv_block(self, in_channels, out_channels) -> nn.Module:
|
||||
"""
|
||||
Convolution and activation combined
|
||||
"""
|
||||
return nn.Sequential(
|
||||
nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
|
||||
nn.ReLU(inplace=True),
|
||||
)
|
||||
|
||||
def __init__(self, blocks: List[List[int]]):
|
||||
super().__init__()
|
||||
|
||||
# 5 $2 \times 2$ pooling layers will produce a output of size $1 \ times 1$.
|
||||
# CIFAR 10 image size is $32 \times 32$
|
||||
assert len(blocks) == 5
|
||||
layers = []
|
||||
# RGB channels
|
||||
in_channels = 3
|
||||
# Number of channels in each layer in each block
|
||||
for block in blocks:
|
||||
# Convolution, Normalization and Activation layers
|
||||
for channels in block:
|
||||
layers += self.conv_block(in_channels, channels)
|
||||
in_channels = channels
|
||||
# Max pooling at end of each block
|
||||
layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
|
||||
|
||||
# Create a sequential model with the layers
|
||||
self.layers = nn.Sequential(*layers)
|
||||
# Final logits layer
|
||||
self.fc = nn.Linear(in_channels, 10)
|
||||
|
||||
def forward(self, x):
|
||||
# The VGG layers
|
||||
x = self.layers(x)
|
||||
# Reshape for classification layer
|
||||
x = x.view(x.shape[0], -1)
|
||||
# Final linear layer
|
||||
return self.fc(x)
|
||||
@@ -0,0 +1,111 @@
|
||||
"""
|
||||
---
|
||||
title: MNIST Experiment
|
||||
summary: >
|
||||
This is a reusable trainer for MNIST dataset
|
||||
---
|
||||
|
||||
# MNIST Experiment
|
||||
"""
|
||||
|
||||
import torch.nn as nn
|
||||
import torch.utils.data
|
||||
|
||||
from labml import tracker
|
||||
from labml.configs import option
|
||||
from labml_nn.helpers.datasets import MNISTConfigs as MNISTDatasetConfigs
|
||||
from labml_nn.helpers.device import DeviceConfigs
|
||||
from labml_nn.helpers.metrics import Accuracy
|
||||
from labml_nn.helpers.trainer import TrainValidConfigs, BatchIndex
|
||||
from labml_nn.optimizers.configs import OptimizerConfigs
|
||||
|
||||
|
||||
class MNISTConfigs(MNISTDatasetConfigs, TrainValidConfigs):
|
||||
"""
|
||||
<a id="MNISTConfigs"></a>
|
||||
|
||||
## Trainer configurations
|
||||
"""
|
||||
|
||||
# Optimizer
|
||||
optimizer: torch.optim.Adam
|
||||
# Training device
|
||||
device: torch.device = DeviceConfigs()
|
||||
|
||||
# Classification model
|
||||
model: nn.Module
|
||||
# Number of epochs to train for
|
||||
epochs: int = 10
|
||||
|
||||
# Number of times to switch between training and validation within an epoch
|
||||
inner_iterations = 10
|
||||
|
||||
# Accuracy function
|
||||
accuracy = Accuracy()
|
||||
# Loss function
|
||||
loss_func = nn.CrossEntropyLoss()
|
||||
|
||||
def init(self):
|
||||
"""
|
||||
### Initialization
|
||||
"""
|
||||
# Set tracker configurations
|
||||
tracker.set_scalar("loss.*", True)
|
||||
tracker.set_scalar("accuracy.*", True)
|
||||
# Add accuracy as a state module.
|
||||
# The name is probably confusing, since it's meant to store
|
||||
# states between training and validation for RNNs.
|
||||
# This will keep the accuracy metric stats separate for training and validation.
|
||||
self.state_modules = [self.accuracy]
|
||||
|
||||
def step(self, batch: any, batch_idx: BatchIndex):
|
||||
"""
|
||||
### Training or validation step
|
||||
"""
|
||||
|
||||
# Training/Evaluation mode
|
||||
self.model.train(self.mode.is_train)
|
||||
|
||||
# Move data to the device
|
||||
data, target = batch[0].to(self.device), batch[1].to(self.device)
|
||||
|
||||
# Update global step (number of samples processed) when in training mode
|
||||
if self.mode.is_train:
|
||||
tracker.add_global_step(len(data))
|
||||
|
||||
# Get model outputs.
|
||||
output = self.model(data)
|
||||
|
||||
# Calculate and log loss
|
||||
loss = self.loss_func(output, target)
|
||||
tracker.add("loss.", loss)
|
||||
|
||||
# Calculate and log accuracy
|
||||
self.accuracy(output, target)
|
||||
self.accuracy.track()
|
||||
|
||||
# Train the model
|
||||
if self.mode.is_train:
|
||||
# Calculate gradients
|
||||
loss.backward()
|
||||
# Take optimizer step
|
||||
self.optimizer.step()
|
||||
# Log the model parameters and gradients on last batch of every epoch
|
||||
if batch_idx.is_last:
|
||||
tracker.add('model', self.model)
|
||||
# Clear the gradients
|
||||
self.optimizer.zero_grad()
|
||||
|
||||
# Save the tracked metrics
|
||||
tracker.save()
|
||||
|
||||
|
||||
@option(MNISTConfigs.optimizer)
|
||||
def _optimizer(c: MNISTConfigs):
|
||||
"""
|
||||
### Default optimizer configurations
|
||||
"""
|
||||
opt_conf = OptimizerConfigs()
|
||||
opt_conf.parameters = c.model.parameters()
|
||||
opt_conf.optimizer = 'Adam'
|
||||
return opt_conf
|
||||
@@ -0,0 +1,331 @@
|
||||
"""
|
||||
---
|
||||
title: NLP auto-regression trainer
|
||||
summary: >
|
||||
This is a reusable trainer for auto-regressive tasks
|
||||
---
|
||||
|
||||
# Auto-regressive NLP model trainer
|
||||
"""
|
||||
|
||||
from typing import Callable
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from labml import lab, monit, logger, tracker
|
||||
from labml.configs import option
|
||||
from labml.logger import Text
|
||||
from labml_nn.helpers.datasets import TextDataset, SequentialDataLoader, SequentialUnBatchedDataset, TextFileDataset
|
||||
from labml_nn.helpers.device import DeviceConfigs
|
||||
from labml_nn.helpers.metrics import Accuracy
|
||||
from labml_nn.helpers.trainer import TrainValidConfigs, BatchIndex
|
||||
from labml_nn.optimizers.configs import OptimizerConfigs
|
||||
from torch.utils.data import DataLoader, RandomSampler
|
||||
|
||||
|
||||
class CrossEntropyLoss(nn.Module):
|
||||
"""
|
||||
### Cross entropy loss
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.loss = nn.CrossEntropyLoss()
|
||||
|
||||
def forward(self, outputs, targets):
|
||||
return self.loss(outputs.view(-1, outputs.shape[-1]), targets.view(-1))
|
||||
|
||||
|
||||
class NLPAutoRegressionConfigs(TrainValidConfigs):
|
||||
"""
|
||||
<a id="NLPAutoRegressionConfigs"></a>
|
||||
|
||||
## Trainer configurations
|
||||
|
||||
This has the basic configurations for NLP auto-regressive task training.
|
||||
All the properties are configurable.
|
||||
"""
|
||||
|
||||
# Optimizer
|
||||
optimizer: torch.optim.Adam
|
||||
# Training device
|
||||
device: torch.device = DeviceConfigs()
|
||||
|
||||
# Autoregressive model
|
||||
model: nn.Module
|
||||
# Text dataset
|
||||
text: TextDataset
|
||||
# Batch size
|
||||
batch_size: int = 16
|
||||
# Length of the sequence, or context size
|
||||
seq_len: int = 512
|
||||
# Number of token in vocabulary
|
||||
n_tokens: int
|
||||
# Tokenizer
|
||||
tokenizer: Callable = 'character'
|
||||
|
||||
# Text prompt to start sampling (for illustration)
|
||||
prompt: str
|
||||
# The token separator when sampling (blank for character level tokenization)
|
||||
prompt_separator: str
|
||||
|
||||
# Whether to periodically save models
|
||||
is_save_models = True
|
||||
|
||||
# Loss function
|
||||
loss_func = CrossEntropyLoss()
|
||||
# Accuracy function
|
||||
accuracy = Accuracy()
|
||||
# Model embedding size
|
||||
d_model: int = 512
|
||||
# Gradient clipping
|
||||
grad_norm_clip: float = 1.0
|
||||
|
||||
# Training data loader
|
||||
train_loader: DataLoader = 'shuffled_train_loader'
|
||||
# Validation data loader
|
||||
valid_loader: DataLoader = 'shuffled_valid_loader'
|
||||
|
||||
# Data loaders shuffle with replacement
|
||||
dataloader_shuffle_with_replacement: bool = False
|
||||
|
||||
# Whether to log model parameters and gradients (once per epoch).
|
||||
# These are summarized stats per layer, but it could still lead
|
||||
# to many indicators for very deep networks.
|
||||
is_log_model_params_grads: bool = False
|
||||
|
||||
# Whether to log model activations (once per epoch).
|
||||
# These are summarized stats per layer, but it could still lead
|
||||
# to many indicators for very deep networks.
|
||||
is_log_model_activations: bool = False
|
||||
|
||||
def init(self):
|
||||
"""
|
||||
### Initialization
|
||||
"""
|
||||
# Set tracker configurations
|
||||
tracker.set_scalar("accuracy.*", True)
|
||||
tracker.set_scalar("loss.*", True)
|
||||
tracker.set_text("sampled", False)
|
||||
# Add accuracy as a state module.
|
||||
# The name is probably confusing, since it's meant to store
|
||||
# states between training and validation for RNNs.
|
||||
# This will keep the accuracy metric stats separate for training and validation.
|
||||
self.state_modules = [self.accuracy]
|
||||
|
||||
def other_metrics(self, output: torch.Tensor, target: torch.Tensor):
|
||||
"""Override to calculate and log other metrics"""
|
||||
pass
|
||||
|
||||
def step(self, batch: any, batch_idx: BatchIndex):
|
||||
"""
|
||||
### Training or validation step
|
||||
"""
|
||||
|
||||
# Set training/eval mode
|
||||
self.model.train(self.mode.is_train)
|
||||
|
||||
# Move data to the device
|
||||
data, target = batch[0].to(self.device), batch[1].to(self.device)
|
||||
|
||||
# Update global step (number of tokens processed) when in training mode
|
||||
if self.mode.is_train:
|
||||
tracker.add_global_step(data.shape[0] * data.shape[1])
|
||||
|
||||
# Get model outputs.
|
||||
# It's returning a tuple for states when using RNNs.
|
||||
# This is not implemented yet. 😜
|
||||
output, *_ = self.model(data)
|
||||
|
||||
# Calculate and log loss
|
||||
loss = self.loss_func(output, target)
|
||||
tracker.add("loss.", loss)
|
||||
|
||||
# Calculate and log accuracy
|
||||
self.accuracy(output, target)
|
||||
self.accuracy.track()
|
||||
|
||||
self.other_metrics(output, target)
|
||||
|
||||
# Train the model
|
||||
if self.mode.is_train:
|
||||
# Calculate gradients
|
||||
loss.backward()
|
||||
# Clip gradients
|
||||
torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=self.grad_norm_clip)
|
||||
# Take optimizer step
|
||||
self.optimizer.step()
|
||||
# Log the model parameters and gradients on last batch of every epoch
|
||||
if batch_idx.is_last and self.is_log_model_params_grads:
|
||||
tracker.add('model', self.model)
|
||||
# Clear the gradients
|
||||
self.optimizer.zero_grad()
|
||||
|
||||
# Save the tracked metrics
|
||||
tracker.save()
|
||||
|
||||
def sample(self):
|
||||
"""
|
||||
### Sampling function to generate samples periodically while training
|
||||
"""
|
||||
|
||||
# Starting prompt
|
||||
prompt = self.prompt
|
||||
# Collect output for printing
|
||||
log = [(prompt, Text.subtle)]
|
||||
# Sample 25 tokens
|
||||
for i in monit.iterate('Sample', 25):
|
||||
# Tokenize the prompt
|
||||
data = self.text.text_to_i(prompt).unsqueeze(-1)
|
||||
data = data.to(self.device)
|
||||
# Get the model output
|
||||
output, *_ = self.model(data)
|
||||
# Get the model prediction (greedy)
|
||||
output = output.argmax(dim=-1).squeeze()
|
||||
# Add the prediction to prompt
|
||||
prompt += self.prompt_separator + self.text.itos[output[-1]]
|
||||
# Add the prediction for logging
|
||||
log += [(self.prompt_separator + self.text.itos[output[-1]], Text.value)]
|
||||
|
||||
tracker.add({'sampled': prompt})
|
||||
# Print the sampled output
|
||||
logger.log(log)
|
||||
|
||||
|
||||
@option(NLPAutoRegressionConfigs.optimizer)
|
||||
def _optimizer(c: NLPAutoRegressionConfigs):
|
||||
"""
|
||||
### Default [optimizer configurations](../optimizers/configs.html)
|
||||
"""
|
||||
|
||||
optimizer = OptimizerConfigs()
|
||||
optimizer.parameters = c.model.parameters()
|
||||
optimizer.optimizer = 'Adam'
|
||||
optimizer.d_model = c.d_model
|
||||
|
||||
return optimizer
|
||||
|
||||
|
||||
@option(NLPAutoRegressionConfigs.n_tokens)
|
||||
def _n_tokens(c: NLPAutoRegressionConfigs):
|
||||
"""
|
||||
Get number of tokens
|
||||
"""
|
||||
return c.text.n_tokens
|
||||
|
||||
|
||||
@option(NLPAutoRegressionConfigs.tokenizer)
|
||||
def basic_english():
|
||||
"""
|
||||
### Basic english tokenizer
|
||||
|
||||
We use character level tokenizer in this experiment.
|
||||
You can switch by setting,
|
||||
|
||||
```
|
||||
'tokenizer': 'basic_english',
|
||||
```
|
||||
|
||||
in the configurations dictionary when starting the experiment.
|
||||
"""
|
||||
|
||||
from torchtext.data import get_tokenizer
|
||||
return get_tokenizer('basic_english')
|
||||
|
||||
|
||||
def character_tokenizer(x: str):
|
||||
"""
|
||||
### Character level tokenizer
|
||||
"""
|
||||
return list(x)
|
||||
|
||||
|
||||
@option(NLPAutoRegressionConfigs.tokenizer)
|
||||
def character():
|
||||
"""
|
||||
### Character level tokenizer configuration
|
||||
"""
|
||||
return character_tokenizer
|
||||
|
||||
|
||||
@option(NLPAutoRegressionConfigs.text)
|
||||
def tiny_shakespeare(c: NLPAutoRegressionConfigs):
|
||||
"""
|
||||
### Tiny Shakespeare dataset
|
||||
|
||||
It will download from the url if not present
|
||||
"""
|
||||
return TextFileDataset(
|
||||
lab.get_data_path() / 'tiny_shakespeare.txt',
|
||||
c.tokenizer,
|
||||
url='https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt')
|
||||
|
||||
|
||||
@option(NLPAutoRegressionConfigs.train_loader)
|
||||
def sequential_train_loader(c: NLPAutoRegressionConfigs):
|
||||
"""
|
||||
### Sequential training data loader
|
||||
"""
|
||||
return SequentialDataLoader(text=c.text.train,
|
||||
dataset=c.text,
|
||||
batch_size=c.batch_size,
|
||||
seq_len=c.seq_len)
|
||||
|
||||
|
||||
@option(NLPAutoRegressionConfigs.valid_loader)
|
||||
def sequential_valid_loader(c: NLPAutoRegressionConfigs):
|
||||
"""
|
||||
### Sequential validation data loader
|
||||
"""
|
||||
return SequentialDataLoader(text=c.text.valid,
|
||||
dataset=c.text,
|
||||
batch_size=c.batch_size,
|
||||
seq_len=c.seq_len)
|
||||
|
||||
|
||||
def transpose_batch(batch):
|
||||
"""
|
||||
### Transpose batch
|
||||
|
||||
`DataLoader` collects the batches on the first dimension.
|
||||
We need to transpose it to be sequence first.
|
||||
"""
|
||||
|
||||
transposed_data = list(zip(*batch))
|
||||
# Stack the batch along the second dimension `dim=1`
|
||||
src = torch.stack(transposed_data[0], dim=1)
|
||||
tgt = torch.stack(transposed_data[1], dim=1)
|
||||
|
||||
return src, tgt
|
||||
|
||||
|
||||
@option(NLPAutoRegressionConfigs.train_loader)
|
||||
def shuffled_train_loader(c: NLPAutoRegressionConfigs):
|
||||
"""
|
||||
### Shuffled training data loader
|
||||
"""
|
||||
dataset = SequentialUnBatchedDataset(text=c.text.train,
|
||||
dataset=c.text,
|
||||
seq_len=c.seq_len)
|
||||
sampler = RandomSampler(dataset, replacement=c.dataloader_shuffle_with_replacement)
|
||||
|
||||
return DataLoader(dataset,
|
||||
batch_size=c.batch_size,
|
||||
collate_fn=transpose_batch,
|
||||
sampler=sampler)
|
||||
|
||||
|
||||
@option(NLPAutoRegressionConfigs.valid_loader)
|
||||
def shuffled_valid_loader(c: NLPAutoRegressionConfigs):
|
||||
"""
|
||||
### Shuffled validation data loader
|
||||
"""
|
||||
dataset = SequentialUnBatchedDataset(text=c.text.valid,
|
||||
dataset=c.text,
|
||||
seq_len=c.seq_len)
|
||||
sampler = RandomSampler(dataset, replacement=c.dataloader_shuffle_with_replacement)
|
||||
|
||||
return DataLoader(dataset,
|
||||
batch_size=c.batch_size,
|
||||
collate_fn=transpose_batch,
|
||||
sampler=sampler)
|
||||
@@ -0,0 +1,289 @@
|
||||
"""
|
||||
---
|
||||
title: NLP classification trainer
|
||||
summary: >
|
||||
This is a reusable trainer for classification tasks
|
||||
---
|
||||
|
||||
# NLP model trainer for classification
|
||||
"""
|
||||
|
||||
from collections import Counter
|
||||
from typing import Callable
|
||||
|
||||
import torchtext
|
||||
import torchtext.vocab
|
||||
from torchtext.vocab import Vocab
|
||||
|
||||
import torch
|
||||
from labml import lab, tracker, monit
|
||||
from labml.configs import option
|
||||
from labml_nn.helpers.device import DeviceConfigs
|
||||
from labml_nn.helpers.metrics import Accuracy
|
||||
from labml_nn.helpers.trainer import TrainValidConfigs, BatchIndex
|
||||
from labml_nn.optimizers.configs import OptimizerConfigs
|
||||
from torch import nn
|
||||
from torch.utils.data import DataLoader
|
||||
|
||||
|
||||
class NLPClassificationConfigs(TrainValidConfigs):
|
||||
"""
|
||||
<a id="NLPClassificationConfigs"></a>
|
||||
|
||||
## Trainer configurations
|
||||
|
||||
This has the basic configurations for NLP classification task training.
|
||||
All the properties are configurable.
|
||||
"""
|
||||
|
||||
# Optimizer
|
||||
optimizer: torch.optim.Adam
|
||||
# Training device
|
||||
device: torch.device = DeviceConfigs()
|
||||
|
||||
# Autoregressive model
|
||||
model: nn.Module
|
||||
# Batch size
|
||||
batch_size: int = 16
|
||||
# Length of the sequence, or context size
|
||||
seq_len: int = 512
|
||||
# Vocabulary
|
||||
vocab: Vocab = 'ag_news'
|
||||
# Number of token in vocabulary
|
||||
n_tokens: int
|
||||
# Number of classes
|
||||
n_classes: int = 'ag_news'
|
||||
# Tokenizer
|
||||
tokenizer: Callable = 'character'
|
||||
|
||||
# Whether to periodically save models
|
||||
is_save_models = True
|
||||
|
||||
# Loss function
|
||||
loss_func = nn.CrossEntropyLoss()
|
||||
# Accuracy function
|
||||
accuracy = Accuracy()
|
||||
# Model embedding size
|
||||
d_model: int = 512
|
||||
# Gradient clipping
|
||||
grad_norm_clip: float = 1.0
|
||||
|
||||
# Training data loader
|
||||
train_loader: DataLoader = 'ag_news'
|
||||
# Validation data loader
|
||||
valid_loader: DataLoader = 'ag_news'
|
||||
|
||||
# Whether to log model parameters and gradients (once per epoch).
|
||||
# These are summarized stats per layer, but it could still lead
|
||||
# to many indicators for very deep networks.
|
||||
is_log_model_params_grads: bool = False
|
||||
|
||||
# Whether to log model activations (once per epoch).
|
||||
# These are summarized stats per layer, but it could still lead
|
||||
# to many indicators for very deep networks.
|
||||
is_log_model_activations: bool = False
|
||||
|
||||
def init(self):
|
||||
"""
|
||||
### Initialization
|
||||
"""
|
||||
# Set tracker configurations
|
||||
tracker.set_scalar("accuracy.*", True)
|
||||
tracker.set_scalar("loss.*", True)
|
||||
# Add accuracy as a state module.
|
||||
# The name is probably confusing, since it's meant to store
|
||||
# states between training and validation for RNNs.
|
||||
# This will keep the accuracy metric stats separate for training and validation.
|
||||
self.state_modules = [self.accuracy]
|
||||
|
||||
def step(self, batch: any, batch_idx: BatchIndex):
|
||||
"""
|
||||
### Training or validation step
|
||||
"""
|
||||
|
||||
# Move data to the device
|
||||
data, target = batch[0].to(self.device), batch[1].to(self.device)
|
||||
|
||||
# Update global step (number of tokens processed) when in training mode
|
||||
if self.mode.is_train:
|
||||
tracker.add_global_step(data.shape[1])
|
||||
|
||||
# Get model outputs.
|
||||
# It's returning a tuple for states when using RNNs.
|
||||
# This is not implemented yet. 😜
|
||||
output, *_ = self.model(data)
|
||||
|
||||
# Calculate and log loss
|
||||
loss = self.loss_func(output, target)
|
||||
tracker.add("loss.", loss)
|
||||
|
||||
# Calculate and log accuracy
|
||||
self.accuracy(output, target)
|
||||
self.accuracy.track()
|
||||
|
||||
# Train the model
|
||||
if self.mode.is_train:
|
||||
# Calculate gradients
|
||||
loss.backward()
|
||||
# Clip gradients
|
||||
torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=self.grad_norm_clip)
|
||||
# Take optimizer step
|
||||
self.optimizer.step()
|
||||
# Log the model parameters and gradients on last batch of every epoch
|
||||
if batch_idx.is_last and self.is_log_model_params_grads:
|
||||
tracker.add('model', self.model)
|
||||
# Clear the gradients
|
||||
self.optimizer.zero_grad()
|
||||
|
||||
# Save the tracked metrics
|
||||
tracker.save()
|
||||
|
||||
|
||||
@option(NLPClassificationConfigs.optimizer)
|
||||
def _optimizer(c: NLPClassificationConfigs):
|
||||
"""
|
||||
### Default [optimizer configurations](../optimizers/configs.html)
|
||||
"""
|
||||
|
||||
optimizer = OptimizerConfigs()
|
||||
optimizer.parameters = c.model.parameters()
|
||||
optimizer.optimizer = 'Adam'
|
||||
optimizer.d_model = c.d_model
|
||||
|
||||
return optimizer
|
||||
|
||||
|
||||
@option(NLPClassificationConfigs.tokenizer)
|
||||
def basic_english():
|
||||
"""
|
||||
### Basic english tokenizer
|
||||
|
||||
We use character level tokenizer in this experiment.
|
||||
You can switch by setting,
|
||||
|
||||
```
|
||||
'tokenizer': 'basic_english',
|
||||
```
|
||||
|
||||
in the configurations dictionary when starting the experiment.
|
||||
|
||||
"""
|
||||
from torchtext.data import get_tokenizer
|
||||
return get_tokenizer('basic_english')
|
||||
|
||||
|
||||
def character_tokenizer(x: str):
|
||||
"""
|
||||
### Character level tokenizer
|
||||
"""
|
||||
return list(x)
|
||||
|
||||
|
||||
@option(NLPClassificationConfigs.tokenizer)
|
||||
def character():
|
||||
"""
|
||||
Character level tokenizer configuration
|
||||
"""
|
||||
return character_tokenizer
|
||||
|
||||
|
||||
@option(NLPClassificationConfigs.n_tokens)
|
||||
def _n_tokens(c: NLPClassificationConfigs):
|
||||
"""
|
||||
Get number of tokens
|
||||
"""
|
||||
return len(c.vocab) + 2
|
||||
|
||||
|
||||
class CollateFunc:
|
||||
"""
|
||||
## Function to load data into batches
|
||||
"""
|
||||
|
||||
def __init__(self, tokenizer, vocab: Vocab, seq_len: int, padding_token: int, classifier_token: int):
|
||||
"""
|
||||
* `tokenizer` is the tokenizer function
|
||||
* `vocab` is the vocabulary
|
||||
* `seq_len` is the length of the sequence
|
||||
* `padding_token` is the token used for padding when the `seq_len` is larger than the text length
|
||||
* `classifier_token` is the `[CLS]` token which we set at end of the input
|
||||
"""
|
||||
self.classifier_token = classifier_token
|
||||
self.padding_token = padding_token
|
||||
self.seq_len = seq_len
|
||||
self.vocab = vocab
|
||||
self.tokenizer = tokenizer
|
||||
|
||||
def __call__(self, batch):
|
||||
"""
|
||||
* `batch` is the batch of data collected by the `DataLoader`
|
||||
"""
|
||||
|
||||
# Input data tensor, initialized with `padding_token`
|
||||
data = torch.full((self.seq_len, len(batch)), self.padding_token, dtype=torch.long)
|
||||
# Empty labels tensor
|
||||
labels = torch.zeros(len(batch), dtype=torch.long)
|
||||
|
||||
# Loop through the samples
|
||||
for (i, (_label, _text)) in enumerate(batch):
|
||||
# Set the label
|
||||
labels[i] = int(_label) - 1
|
||||
# Tokenize the input text
|
||||
_text = [self.vocab[token] for token in self.tokenizer(_text)]
|
||||
# Truncate upto `seq_len`
|
||||
_text = _text[:self.seq_len]
|
||||
# Transpose and add to data
|
||||
data[:len(_text), i] = data.new_tensor(_text)
|
||||
|
||||
# Set the final token in the sequence to `[CLS]`
|
||||
data[-1, :] = self.classifier_token
|
||||
|
||||
#
|
||||
return data, labels
|
||||
|
||||
|
||||
@option([NLPClassificationConfigs.n_classes,
|
||||
NLPClassificationConfigs.vocab,
|
||||
NLPClassificationConfigs.train_loader,
|
||||
NLPClassificationConfigs.valid_loader])
|
||||
def ag_news(c: NLPClassificationConfigs):
|
||||
"""
|
||||
### AG News dataset
|
||||
|
||||
This loads the AG News dataset and the set the values for
|
||||
`n_classes`, `vocab`, `train_loader`, and `valid_loader`.
|
||||
"""
|
||||
|
||||
# Get training and validation datasets
|
||||
train, valid = torchtext.datasets.AG_NEWS(root=str(lab.get_data_path() / 'ag_news'), split=('train', 'test'))
|
||||
|
||||
# Load data to memory
|
||||
with monit.section('Load data'):
|
||||
from labml_nn.utils import MapStyleDataset
|
||||
|
||||
# Create [map-style datasets](../utils.html#map_style_dataset)
|
||||
train, valid = MapStyleDataset(train), MapStyleDataset(valid)
|
||||
|
||||
# Get tokenizer
|
||||
tokenizer = c.tokenizer
|
||||
|
||||
# Create a counter
|
||||
counter = Counter()
|
||||
# Collect tokens from training dataset
|
||||
for (label, line) in train:
|
||||
counter.update(tokenizer(line))
|
||||
# Collect tokens from validation dataset
|
||||
for (label, line) in valid:
|
||||
counter.update(tokenizer(line))
|
||||
# Create vocabulary
|
||||
vocab = torchtext.vocab.vocab(counter, min_freq=1)
|
||||
|
||||
# Create training data loader
|
||||
train_loader = DataLoader(train, batch_size=c.batch_size, shuffle=True,
|
||||
collate_fn=CollateFunc(tokenizer, vocab, c.seq_len, len(vocab), len(vocab) + 1))
|
||||
# Create validation data loader
|
||||
valid_loader = DataLoader(valid, batch_size=c.batch_size, shuffle=True,
|
||||
collate_fn=CollateFunc(tokenizer, vocab, c.seq_len, len(vocab), len(vocab) + 1))
|
||||
|
||||
# Return `n_classes`, `vocab`, `train_loader`, and `valid_loader`
|
||||
return 4, vocab, train_loader, valid_loader
|
||||
Reference in New Issue
Block a user