Spaces:
Running
Running
# code taken from https://towardsdatascience.com/how-to-code-the-transformer-in-pytorch-24db27c8f9ec | |
# and https://pytorch.org/tutorials/beginner/transformer_tutorial.html | |
import torch | |
import math | |
import copy | |
class Embedder(torch.nn.Module): | |
def __init__(self, vocab_size, d_model): | |
super().__init__() | |
self.embed = torch.nn.Embedding(vocab_size, d_model) | |
def forward(self, x): | |
return self.embed(x) | |
class PositionalEncoder(torch.nn.Module): | |
def __init__(self, d_model, dropout=0.1, max_seq_len=80): | |
super().__init__() | |
self.dropout = torch.nn.Dropout(p=dropout) | |
position = torch.arange(max_seq_len).unsqueeze(1) | |
div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model)) | |
pe = torch.zeros(max_seq_len, 1, d_model) | |
pe[:, 0, 0::2] = torch.sin(position * div_term) | |
pe[:, 0, 1::2] = torch.cos(position * div_term) | |
self.register_buffer('pe', | |
pe) # notifies PyTorch that this value should be saved like a model parameter but should not have gradients | |
def forward(self, x): | |
x = x + self.pe[:x.size(0)] | |
return self.dropout(x) | |
class MultiHeadAttention(torch.nn.Module): | |
def __init__(self, heads, d_model, dropout=0.1): | |
super().__init__() | |
self.d_model = d_model | |
self.d_k = d_model // heads | |
self.h = heads | |
self.q_linear = torch.nn.Linear(d_model, d_model) | |
self.v_linear = torch.nn.Linear(d_model, d_model) | |
self.k_linear = torch.nn.Linear(d_model, d_model) | |
self.dropout = torch.nn.Dropout(dropout) | |
self.out = torch.nn.Linear(d_model, d_model) | |
def forward(self, q, k, v, mask=None): | |
bs = q.size(0) | |
# perform linear operation and split into h heads | |
k = self.k_linear(k).view(bs, -1, self.h, self.d_k) | |
q = self.q_linear(q).view(bs, -1, self.h, self.d_k) | |
v = self.v_linear(v).view(bs, -1, self.h, self.d_k) | |
# transpose to get dimensions bs * h * sl * d_model | |
k = k.transpose(1, 2) | |
q = q.transpose(1, 2) | |
v = v.transpose(1, 2) | |
# calculate attention using function we will define next | |
scores = attention(q, k, v, self.d_k, mask, self.dropout) | |
# concatenate heads and put through final linear layer | |
concat = scores.transpose(1, 2).contiguous().view(bs, -1, self.d_model) | |
output = self.out(concat) | |
return output | |
def attention(q, k, v, d_k, mask=None, dropout=None): | |
scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(d_k) | |
if mask is not None: | |
mask = mask.unsqueeze(1) | |
scores = scores.masked_fill(mask == 0, -1e9) | |
scores = torch.nn.functional.softmax(scores, dim=-1) | |
if dropout is not None: | |
scores = dropout(scores) | |
output = torch.matmul(scores, v) | |
return output | |
class FeedForward(torch.nn.Module): | |
def __init__(self, d_model, d_ff=2048, dropout=0.1): | |
super().__init__() | |
# We set d_ff as a default to 2048 | |
self.linear_1 = torch.nn.Linear(d_model, d_ff) | |
self.dropout = torch.nn.Dropout(dropout) | |
self.linear_2 = torch.nn.Linear(d_ff, d_model) | |
def forward(self, x): | |
x = self.dropout(torch.nn.functional.relu(self.linear_1(x))) | |
x = self.linear_2(x) | |
return x | |
class Norm(torch.nn.Module): | |
def __init__(self, d_model, eps=1e-6): | |
super().__init__() | |
self.size = d_model | |
# create two learnable parameters to calibrate normalization | |
self.alpha = torch.nn.Parameter(torch.ones(self.size)) | |
self.bias = torch.nn.Parameter(torch.zeros(self.size)) | |
self.eps = eps | |
def forward(self, x): | |
norm = self.alpha * (x - x.mean(dim=-1, keepdim=True)) / (x.std(dim=-1, keepdim=True) + self.eps) + self.bias | |
return norm | |
# build an encoder layer with one multi-head attention layer and one # feed-forward layer | |
class EncoderLayer(torch.nn.Module): | |
def __init__(self, d_model, heads, dropout=0.1): | |
super().__init__() | |
self.norm_1 = Norm(d_model) | |
self.norm_2 = Norm(d_model) | |
self.attn = MultiHeadAttention(heads, d_model) | |
self.ff = FeedForward(d_model) | |
self.dropout_1 = torch.nn.Dropout(dropout) | |
self.dropout_2 = torch.nn.Dropout(dropout) | |
def forward(self, x, mask): | |
x2 = self.norm_1(x) | |
x = x + self.dropout_1(self.attn(x2, x2, x2, mask)) | |
x2 = self.norm_2(x) | |
x = x + self.dropout_2(self.ff(x2)) | |
return x | |
# build a decoder layer with two multi-head attention layers and | |
# one feed-forward layer | |
class DecoderLayer(torch.nn.Module): | |
def __init__(self, d_model, heads, dropout=0.1): | |
super().__init__() | |
self.norm_1 = Norm(d_model) | |
self.norm_2 = Norm(d_model) | |
self.norm_3 = Norm(d_model) | |
self.dropout_1 = torch.nn.Dropout(dropout) | |
self.dropout_2 = torch.nn.Dropout(dropout) | |
self.dropout_3 = torch.nn.Dropout(dropout) | |
self.attn_1 = MultiHeadAttention(heads, d_model) | |
self.attn_2 = MultiHeadAttention(heads, d_model) | |
self.ff = FeedForward(d_model) | |
def forward(self, x, e_outputs, src_mask, trg_mask): | |
x2 = self.norm_1(x) | |
x = x + self.dropout_1(self.attn_1(x2, x2, x2, trg_mask)) | |
x2 = self.norm_2(x) | |
x = x + self.dropout_2(self.attn_2(x2, e_outputs, e_outputs, | |
src_mask)) | |
x2 = self.norm_3(x) | |
x = x + self.dropout_3(self.ff(x2)) | |
return x | |
# We can then build a convenient cloning function that can generate multiple layers: | |
def get_clones(module, N): | |
return torch.nn.ModuleList([copy.deepcopy(module) for i in range(N)]) | |
class Encoder(torch.nn.Module): | |
def __init__(self, vocab_size, d_model, N, heads): | |
super().__init__() | |
self.N = N | |
self.embed = Embedder(vocab_size, d_model) | |
self.pe = PositionalEncoder(d_model) | |
self.layers = get_clones(EncoderLayer(d_model, heads), N) | |
self.norm = Norm(d_model) | |
def forward(self, src, mask): | |
x = self.embed(src) | |
x = self.pe(x) | |
for i in range(self.N): | |
x = self.layers[i](x, mask) | |
return self.norm(x) | |
class Decoder(torch.nn.Module): | |
def __init__(self, vocab_size, d_model, N, heads): | |
super().__init__() | |
self.N = N | |
self.embed = Embedder(vocab_size, d_model) | |
self.pe = PositionalEncoder(d_model) | |
self.layers = get_clones(DecoderLayer(d_model, heads), N) | |
self.norm = Norm(d_model) | |
def forward(self, trg, e_outputs, src_mask, trg_mask): | |
x = self.embed(trg) | |
x = self.pe(x) | |
for i in range(self.N): | |
x = self.layers[i](x, e_outputs, src_mask, trg_mask) | |
return self.norm(x) | |
class Transformer(torch.nn.Module): | |
def __init__(self, src_vocab, trg_vocab, d_model, N, heads): | |
super().__init__() | |
self.encoder = Encoder(src_vocab, d_model, N, heads) | |
self.decoder = Decoder(trg_vocab, d_model, N, heads) | |
self.out = torch.nn.Linear(d_model, trg_vocab) | |
def forward(self, src, trg, src_mask, trg_mask): | |
e_outputs = self.encoder(src, src_mask) | |
d_output = self.decoder(trg, e_outputs, src_mask, trg_mask) | |
output = self.out(d_output) | |
return output | |
# we don't perform softmax on the output as this will be handled | |
# automatically by our loss function |