mirror of
https://github.com/storytold/artcraft.git
synced 2026-11-03 08:57:54 +00:00
102 KiB
102 KiB
In [ ]:
"""
You can run either this notebook locally (if you have all the dependencies and a GPU) or on Google Colab.
Instructions for setting up Colab are as follows:
1. Open a new Python 3 notebook.
2. Import this notebook from GitHub (File -> Upload Notebook -> "GITHUB" tab -> copy/paste GitHub URL)
3. Connect to an instance with a GPU (Runtime -> Change runtime type -> select "GPU" for hardware accelerator)
4. Run this cell to set up dependencies.
"""
# If you're using Google Colab and not running locally, run this cell.
## Install dependencies
!pip install wget
!apt-get install sox libsndfile1 ffmpeg
!pip install unidecode
# ## Install NeMo
BRANCH = 'v1.0.2'
!python -m pip install git+https://github.com/NVIDIA/NeMo.git@$BRANCH#egg=nemo_toolkit[all]
## Install TorchAudio
!pip install torchaudio>=0.6.0 -f https://download.pytorch.org/whl/torch_stable.html
## Grab the config we'll use in this example
!mkdir configsIn [ ]:
import torch
import nemo
from nemo.core import NeuralModule
from nemo.core import typecheckIn [ ]:
class MyEmptyModule(NeuralModule):
def forward(self):
print("Neural Module ~ hello world!")In [ ]:
x = MyEmptyModule()
x()In [ ]:
# Case 1:
embedding = torch.nn.Embedding(num_embeddings=10, embedding_dim=30)
x = torch.randint(high=10, size=(1, 5))
print("x :", x)
print("embedding(x) :", embedding(x).shape)In [ ]:
# Case 2
lstm = torch.nn.LSTM(1, 30, batch_first=True)
x = torch.randn(1, 5, 1)
print("x :", x)
print("lstm(x) :", lstm(x)[0].shape) # Let's take all timestep outputs of the LSTMIn [ ]:
from nemo.core.neural_types import NeuralType
from nemo.core.neural_types import *In [ ]:
class EmbeddingModule(NeuralModule):
def __init__(self):
super().__init__()
self.embedding = torch.nn.Embedding(num_embeddings=10, embedding_dim=30)
@typecheck()
def forward(self, x):
return self.embedding(x)
@property
def input_types(self):
return {
'x': NeuralType(axes=('B', 'T'), elements_type=Index())
}
@property
def output_types(self):
return {
'y': NeuralType(axes=('B', 'T', 'C'), elements_type=EmbeddedTextType())
}In [ ]:
embedding_module = EmbeddingModule()In [ ]:
class LSTMModule(NeuralModule):
def __init__(self):
super().__init__()
self.lstm = torch.nn.LSTM(1, 30, batch_first=True)
@typecheck()
def forward(self, x):
return self.lstm(x)
@property
def input_types(self):
return {
'x': NeuralType(axes=('B', 'T', 'C'), elements_type=SpectrogramType())
}
@property
def output_types(self):
return {
'y': NeuralType(axes=('B', 'T', 'C'), elements_type=EncodedRepresentation())
}In [ ]:
lstm_module = LSTMModule()In [ ]:
# Case 1 [ERROR CELL]
x1 = torch.randint(high=10, size=(1, 5))
print("x :", x1)
print("embedding(x) :", embedding_module(x1).shape)In [ ]:
# Case 1
print("x :", x1)
print("embedding(x) :", embedding_module(x=x1).shape)In [ ]:
# Case 2 [ERROR CELL]
x2 = torch.randn(1, 5, 1) # Input = [B=1, T=5, C=1]
print("x :", x2)
print("lstm(x) :", lstm_module(x=x2)[0].shape) # Let's take all timestep outputs of the LSTMIn [ ]:
class CorrectLSTMModule(LSTMModule): # Let's inherit the wrong class to make it easy to override
@property
def output_types(self):
return {
'y': NeuralType(axes=('B', 'T', 'C'), elements_type=EncodedRepresentation()),
'h_c': [NeuralType(axes=('D', 'B', 'C'), elements_type=EncodedRepresentation())],
}In [ ]:
lstm_module = CorrectLSTMModule()In [ ]:
# Case 2
x2 = torch.randn(1, 5, 1)
y2, (h, c) = lstm_module(x=x2)
print("x :", x2)
print("lstm(x) :", y2.shape) # The output of the LSTM RNN
print("hidden state (h) :", h.shape) # The first hidden state of the LSTM RNN
print("hidden state (c) :", c.shape) # The second hidden state of the LSTM RNNIn [ ]:
emb_out = embedding_module(x=x1)
lstm_out = lstm_module(x=x2)[0]
assert hasattr(emb_out, 'neural_type')
assert hasattr(lstm_out, 'neural_type')In [ ]:
print("Embedding tensor :", emb_out.neural_type)
print("LSTM tensor :", lstm_out.neural_type)In [ ]:
emb_out.neural_type.compare(lstm_out.neural_type)In [ ]:
emb_out.neural_type == lstm_out.neural_typeIn [ ]:
embedding_module = EmbeddingModule()
x1 = torch.randint(high=10, size=(1, 5))
# Attach correct neural type
x1.neural_type = NeuralType(('B', 'T'), Index())
print("embedding(x) :", embedding_module(x=x1).shape)In [ ]:
# Attach wrong neural type [ERROR CELL]
x1.neural_type = NeuralType(('B', 'T'), LabelsType())
print("embedding(x) :", embedding_module(x=x1).shape)In [ ]:
import math
from typing import List, Set, Dict, Tuple, Optional
import torch
import torch.nn as nn
from torch.nn import functional as FIn [ ]:
class AttentionType(EncodedRepresentation):
"""Basic Attention Element Type"""
class SelfAttentionType(AttentionType):
"""Self Attention Element Type"""
class CausalSelfAttentionType(SelfAttentionType):
"""Causal Self Attention Element Type"""In [ ]:
class CausalSelfAttention(nn.Module):
"""
A vanilla multi-head masked self-attention layer with a projection at the end.
It is possible to use torch.nn.MultiheadAttention here but I am including an
explicit implementation here to show that there is nothing too scary here.
"""
def __init__(self, n_embd, block_size, n_head, attn_pdrop, resid_pdrop):
super().__init__()
assert n_embd % n_head == 0
self.n_head = n_head
# key, query, value projections for all heads
self.key = nn.Linear(n_embd, n_embd)
self.query = nn.Linear(n_embd, n_embd)
self.value = nn.Linear(n_embd, n_embd)
# regularization
self.attn_drop = nn.Dropout(attn_pdrop)
self.resid_drop = nn.Dropout(resid_pdrop)
# output projection
self.proj = nn.Linear(n_embd, n_embd)
# causal mask to ensure that attention is only applied to the left in the input sequence
self.register_buffer("mask", torch.tril(torch.ones(block_size, block_size))
.view(1, 1, block_size, block_size))
def forward(self, x, layer_past=None):
B, T, C = x.size()
# calculate query, key, values for all heads in batch and move head forward to be the batch dim
k = self.key(x).view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
q = self.query(x).view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
v = self.value(x).view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
# causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)
att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
att = att.masked_fill(self.mask[:,:,:T,:T] == 0, float('-inf'))
att = F.softmax(att, dim=-1)
att = self.attn_drop(att)
y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)
y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side
# output projection
y = self.resid_drop(self.proj(y))
return y
class Block(nn.Module):
""" an unassuming Transformer block """
def __init__(self, n_embd, block_size, n_head, attn_pdrop, resid_pdrop):
super().__init__()
self.ln1 = nn.LayerNorm(n_embd)
self.ln2 = nn.LayerNorm(n_embd)
self.attn = CausalSelfAttention(n_embd, block_size, n_head, attn_pdrop, resid_pdrop)
self.mlp = nn.Sequential(
nn.Linear(n_embd, 4 * n_embd),
nn.GELU(),
nn.Linear(4 * n_embd, n_embd),
nn.Dropout(resid_pdrop),
)
def forward(self, x):
x = x + self.attn(self.ln1(x))
x = x + self.mlp(self.ln2(x))
return xIn [ ]:
import pytorch_lightning as ptl
from nemo.core import ModelPT
from omegaconf import OmegaConfIn [ ]:
class PTLGPT(ptl.LightningModule):
def __init__(self,
# model definition args
vocab_size: int, # size of the vocabulary (number of possible tokens)
block_size: int, # length of the model's context window in time
n_layer: int, # depth of the model; number of Transformer blocks in sequence
n_embd: int, # the "width" of the model, number of channels in each Transformer
n_head: int, # number of heads in each multi-head attention inside each Transformer block
# model optimization args
learning_rate: float = 3e-4, # the base learning rate of the model
weight_decay: float = 0.1, # amount of regularizing L2 weight decay on MatMul ops
betas: Tuple[float, float] = (0.9, 0.95), # momentum terms (betas) for the Adam optimizer
embd_pdrop: float = 0.1, # \in [0,1]: amount of dropout on input embeddings
resid_pdrop: float = 0.1, # \in [0,1]: amount of dropout in each residual connection
attn_pdrop: float = 0.1, # \in [0,1]: amount of dropout on the attention matrix
):
super().__init__()
# save these for optimizer init later
self.learning_rate = learning_rate
self.weight_decay = weight_decay
self.betas = betas
# input embedding stem: drop(content + position)
self.tok_emb = nn.Embedding(vocab_size, n_embd)
self.pos_emb = nn.Parameter(torch.zeros(1, block_size, n_embd))
self.drop = nn.Dropout(embd_pdrop)
# deep transformer: just a sequence of transformer blocks
self.blocks = nn.Sequential(*[Block(n_embd, block_size, n_head, attn_pdrop, resid_pdrop) for _ in range(n_layer)])
# decoder: at the end one more layernorm and decode the answers
self.ln_f = nn.LayerNorm(n_embd)
self.head = nn.Linear(n_embd, vocab_size, bias=False) # no need for extra bias due to one in ln_f
self.block_size = block_size
self.apply(self._init_weights)
print("number of parameters: %e" % sum(p.numel() for p in self.parameters()))
def forward(self, idx):
b, t = idx.size()
assert t <= self.block_size, "Cannot forward, model block size is exhausted."
# forward the GPT model
token_embeddings = self.tok_emb(idx) # each index maps to a (learnable) vector
position_embeddings = self.pos_emb[:, :t, :] # each position maps to a (learnable) vector
x = self.drop(token_embeddings + position_embeddings)
x = self.blocks(x)
x = self.ln_f(x)
logits = self.head(x)
return logits
def get_block_size(self):
return self.block_size
def _init_weights(self, module):
"""
Vanilla model initialization:
- all MatMul weights \in N(0, 0.02) and biases to zero
- all LayerNorm post-normalization scaling set to identity, so weight=1, bias=0
"""
if isinstance(module, (nn.Linear, nn.Embedding)):
module.weight.data.normal_(mean=0.0, std=0.02)
if isinstance(module, nn.Linear) and module.bias is not None:
module.bias.data.zero_()
elif isinstance(module, nn.LayerNorm):
module.bias.data.zero_()
module.weight.data.fill_(1.0)In [ ]:
m = PTLGPT(vocab_size=100, block_size=32, n_layer=1, n_embd=32, n_head=4)In [ ]:
class GPTEmbedding(NeuralModule):
def __init__(self, vocab_size: int, n_embd: int, block_size: int, embd_pdrop: float = 0.0):
super().__init__()
# input embedding stem: drop(content + position)
self.tok_emb = nn.Embedding(vocab_size, n_embd)
self.pos_emb = nn.Parameter(torch.zeros(1, block_size, n_embd))
self.drop = nn.Dropout(embd_pdrop)
@typecheck()
def forward(self, idx):
b, t = idx.size()
# forward the GPT model
token_embeddings = self.tok_emb(idx) # each index maps to a (learnable) vector
position_embeddings = self.pos_emb[:, :t, :] # each position maps to a (learnable) vector
x = self.drop(token_embeddings + position_embeddings)
return x
@property
def input_types(self):
return {
'idx': NeuralType(('B', 'T'), Index())
}
@property
def output_types(self):
return {
'embeddings': NeuralType(('B', 'T', 'C'), EmbeddedTextType())
}In [ ]:
class GPTTransformerEncoder(NeuralModule):
def __init__(self, n_embd: int, block_size: int, n_head: int, n_layer: int, attn_pdrop: float = 0.0, resid_pdrop: float = 0.0):
super().__init__()
self.blocks = nn.Sequential(*[Block(n_embd, block_size, n_head, attn_pdrop, resid_pdrop)
for _ in range(n_layer)])
@typecheck()
def forward(self, embed):
return self.blocks(embed)
@property
def input_types(self):
return {
'embed': NeuralType(('B', 'T', 'C'), EmbeddedTextType())
}
@property
def output_types(self):
return {
'encoding': NeuralType(('B', 'T', 'C'), CausalSelfAttentionType())
}In [ ]:
class GPTDecoder(NeuralModule):
def __init__(self, n_embd: int, vocab_size: int):
super().__init__()
self.ln_f = nn.LayerNorm(n_embd)
self.head = nn.Linear(n_embd, vocab_size, bias=False) # no need for extra bias due to one in ln_f
@typecheck()
def forward(self, encoding):
x = self.ln_f(encoding)
logits = self.head(x)
return logits
@property
def input_types(self):
return {
'encoding': NeuralType(('B', 'T', 'C'), EncodedRepresentation())
}
@property
def output_types(self):
return {
'logits': NeuralType(('B', 'T', 'C'), LogitsType())
}
In [ ]:
class AbstractNeMoGPT(ModelPT):
def __init__(self, cfg: OmegaConf, trainer: ptl.Trainer = None):
super().__init__(cfg=cfg, trainer=trainer)
# input embedding stem: drop(content + position)
self.embedding = self.from_config_dict(self.cfg.embedding)
# deep transformer: just a sequence of transformer blocks
self.encoder = self.from_config_dict(self.cfg.encoder)
# decoder: at the end one more layernorm and decode the answers
self.decoder = self.from_config_dict(self.cfg.decoder)
self.block_size = self.cfg.embedding.block_size
self.apply(self._init_weights)
print("number of parameters: %e" % self.num_weights)
@typecheck()
def forward(self, idx):
b, t = idx.size()
assert t <= self.block_size, "Cannot forward, model block size is exhausted."
# forward the GPT model
# Remember: Only kwargs are allowed !
e = self.embedding(idx=idx)
x = self.encoder(embed=e)
logits = self.decoder(encoding=x)
return logits
def get_block_size(self):
return self.block_size
def _init_weights(self, module):
"""
Vanilla model initialization:
- all MatMul weights \in N(0, 0.02) and biases to zero
- all LayerNorm post-normalization scaling set to identity, so weight=1, bias=0
"""
if isinstance(module, (nn.Linear, nn.Embedding)):
module.weight.data.normal_(mean=0.0, std=0.02)
if isinstance(module, nn.Linear) and module.bias is not None:
module.bias.data.zero_()
elif isinstance(module, nn.LayerNorm):
module.bias.data.zero_()
module.weight.data.fill_(1.0)
@property
def input_types(self):
return {
'idx': NeuralType(('B', 'T'), Index())
}
@property
def output_types(self):
return {
'logits': NeuralType(('B', 'T', 'C'), LogitsType())
}In [ ]:
# model definition args (required)
# ================================
# vocab_size: int # size of the vocabulary (number of possible tokens)
# block_size: int # length of the model's context window in time
# n_layer: int # depth of the model; number of Transformer blocks in sequence
# n_embd: int # the "width" of the model, number of channels in each Transformer
# n_head: int # number of heads in each multi-head attention inside each Transformer block
# model definition args (optional)
# ================================
# embd_pdrop: float = 0.1, # \in [0,1]: amount of dropout on input embeddings
# resid_pdrop: float = 0.1, # \in [0,1]: amount of dropout in each residual connection
# attn_pdrop: float = 0.1, # \in [0,1]: amount of dropout on the attention matrixIn [ ]:
from omegaconf import MISSINGIn [ ]:
# Let's create a utility for building the class path
def get_class_path(cls):
return f'{cls.__module__}.{cls.__name__}'In [ ]:
common_config = OmegaConf.create({
'vocab_size': MISSING,
'block_size': MISSING,
'n_layer': MISSING,
'n_embd': MISSING,
'n_head': MISSING,
})In [ ]:
embedding_config = OmegaConf.create({
'_target_': get_class_path(GPTEmbedding),
'vocab_size': '${model.vocab_size}',
'n_embd': '${model.n_embd}',
'block_size': '${model.block_size}',
'embd_pdrop': 0.1
})
encoder_config = OmegaConf.create({
'_target_': get_class_path(GPTTransformerEncoder),
'n_embd': '${model.n_embd}',
'block_size': '${model.block_size}',
'n_head': '${model.n_head}',
'n_layer': '${model.n_layer}',
'attn_pdrop': 0.1,
'resid_pdrop': 0.1
})
decoder_config = OmegaConf.create({
'_target_': get_class_path(GPTDecoder),
# n_embd: int, vocab_size: int
'n_embd': '${model.n_embd}',
'vocab_size': '${model.vocab_size}'
})In [ ]:
model_config = OmegaConf.create({
'model': common_config
})
# Then let's attach the sub-module configs
model_config.model.embedding = embedding_config
model_config.model.encoder = encoder_config
model_config.model.decoder = decoder_configIn [ ]:
print(OmegaConf.to_yaml(model_config))In [ ]:
temp_config = OmegaConf.create(OmegaConf.to_container(model_config, resolve=True))
print(OmegaConf.to_yaml(temp_config))In [ ]:
import copyIn [ ]:
# Let's work on a copy of the model config and update it before we send it into the Model.
cfg = copy.deepcopy(model_config)In [ ]:
# Let's set the values of the config (for some plausible small model)
cfg.model.vocab_size = 100
cfg.model.block_size = 128
cfg.model.n_layer = 1
cfg.model.n_embd = 32
cfg.model.n_head = 4In [ ]:
print(OmegaConf.to_yaml(cfg))In [ ]:
# Try to create a model with this config [ERROR CELL]
m = AbstractNeMoGPT(cfg.model)In [ ]:
from nemo.core.classes.common import PretrainedModelInfoWarning:
Output truncated. This notebook contains too many cells to display efficiently.