A minimal Transformer Implementation¶
https://towardsdatascience.com/build-your-own-transformer-from-scratch-using-pytorch-84c850470dcb
In [1]:
import torch
import torch.nn as nn
import torch.optim as optim
import torch.utils.data as data
import math
import copy
In [207]:
class MultiHeadAttention(nn.Module):
def __init__(self, d_model, num_heads):
super(MultiHeadAttention, self).__init__()
assert d_model % num_heads == 0, "d_model must be divisible by num_heads"
self.d_model = d_model
self.num_heads = num_heads
self.d_k = d_model // num_heads
self.W_q = nn.Linear(d_model, d_model)
self.W_k = nn.Linear(d_model, d_model)
self.W_v = nn.Linear(d_model, d_model)
self.W_o = nn.Linear(d_model, d_model)
def scaled_dot_product_attention(self, Q, K, V, mask=None):
attn_scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
if mask is not None:
attn_scores = attn_scores.masked_fill(mask == 0, -1e9)
attn_probs = torch.softmax(attn_scores, dim=-1)
output = torch.matmul(attn_probs, V)
return output
def split_heads(self, x):
batch_size, seq_length, d_model = x.size()
return x.view(batch_size, seq_length, self.num_heads, self.d_k).transpose(1, 2)
def combine_heads(self, x):
batch_size, _, seq_length, d_k = x.size()
return x.transpose(1, 2).contiguous().view(batch_size, seq_length, self.d_model)
def forward(self, Q, K, V, mask=None):
Q = self.split_heads(self.W_q(Q))
K = self.split_heads(self.W_k(K))
V = self.split_heads(self.W_v(V))
attn_output = self.scaled_dot_product_attention(Q, K, V, mask)
output = self.W_o(self.combine_heads(attn_output))
return output
In [208]:
class PositionWiseFeedForward(nn.Module):
def __init__(self, d_model, d_ff):
super(PositionWiseFeedForward, self).__init__()
self.fc1 = nn.Linear(d_model, d_ff)
self.fc2 = nn.Linear(d_ff, d_model)
self.relu = nn.ReLU()
def forward(self, x):
return self.fc2(self.relu(self.fc1(x)))
In [209]:
class PositionalEncoding(nn.Module):
def __init__(self, d_model, max_seq_length):
super(PositionalEncoding, self).__init__()
pe = torch.zeros(max_seq_length, d_model)
position = torch.arange(0, max_seq_length, dtype=torch.float).unsqueeze(1)
div_term = torch.exp(torch.arange(0, d_model, 2).float()
* -(math.log(10000.0) / d_model))
pe[:, 0::2] = torch.sin(position * div_term)
pe[:, 1::2] = torch.cos(position * div_term)
self.register_buffer('pe', pe.unsqueeze(0))
def forward(self, x):
return x + self.pe[:, :x.size(1)]
In [210]:
class EncoderLayer(nn.Module):
def __init__(self, d_model, num_heads, d_ff, dropout):
super(EncoderLayer, self).__init__()
self.self_attn = MultiHeadAttention(d_model, num_heads)
self.feed_forward = PositionWiseFeedForward(d_model, d_ff)
self.norm1 = nn.LayerNorm(d_model)
self.norm2 = nn.LayerNorm(d_model)
self.dropout = nn.Dropout(dropout)
def forward(self, x, mask):
attn_output = self.self_attn(x, x, x, mask)
x = self.norm1(x + self.dropout(attn_output))
ff_output = self.feed_forward(x)
x = self.norm2(x + self.dropout(ff_output))
return x
In [211]:
class DecoderLayer(nn.Module):
def __init__(self, d_model, num_heads, d_ff, dropout):
super(DecoderLayer, self).__init__()
self.self_attn = MultiHeadAttention(d_model, num_heads)
self.cross_attn = MultiHeadAttention(d_model, num_heads)
self.feed_forward = PositionWiseFeedForward(d_model, d_ff)
self.norm1 = nn.LayerNorm(d_model)
self.norm2 = nn.LayerNorm(d_model)
self.norm3 = nn.LayerNorm(d_model)
self.dropout = nn.Dropout(dropout)
def forward(self, x, enc_output, src_mask, tgt_mask):
attn_output = self.self_attn(x, x, x, tgt_mask)
x = self.norm1(x + self.dropout(attn_output))
attn_output = self.cross_attn(x, enc_output, enc_output, src_mask)
x = self.norm2(x + self.dropout(attn_output))
ff_output = self.feed_forward(x)
x = self.norm3(x + self.dropout(ff_output))
return x
In [212]:
class Transformer(nn.Module):
def __init__(self, src_vocab_size, tgt_vocab_size, d_model, num_heads,
num_layers, d_ff, max_seq_length, dropout):
super(Transformer, self).__init__()
self.encoder_embedding = nn.Embedding(src_vocab_size, d_model)
self.decoder_embedding = nn.Embedding(tgt_vocab_size, d_model)
self.positional_encoding = PositionalEncoding(d_model, max_seq_length)
self.encoder_layers = nn.ModuleList([EncoderLayer(d_model, num_heads, d_ff, dropout)
for _ in range(num_layers)])
self.decoder_layers = nn.ModuleList([DecoderLayer(d_model, num_heads, d_ff, dropout)
for _ in range(num_layers)])
self.fc = nn.Linear(d_model, tgt_vocab_size)
self.dropout = nn.Dropout(dropout)
def generate_mask(self, src, tgt):
src_mask = (src != 0).unsqueeze(1).unsqueeze(2)
tgt_mask = (tgt != 0).unsqueeze(1).unsqueeze(3)
seq_length = tgt.size(1)
nopeak_mask = (1 - torch.triu(torch.ones(1, seq_length, seq_length),
diagonal=1)).bool()
tgt_mask = tgt_mask & nopeak_mask
return src_mask, tgt_mask
def forward(self, src, tgt):
src_mask, tgt_mask = self.generate_mask(src, tgt)
src_embedded = self.dropout(self.positional_encoding(self.encoder_embedding(src)))
tgt_embedded = self.dropout(self.positional_encoding(self.decoder_embedding(tgt)))
enc_output = src_embedded
for enc_layer in self.encoder_layers:
enc_output = enc_layer(enc_output, src_mask)
dec_output = tgt_embedded
for dec_layer in self.decoder_layers:
dec_output = dec_layer(dec_output, enc_output, src_mask, tgt_mask)
output = self.fc(dec_output)
return output
In [213]:
#src_vocab_size = 30
#nc = 2
#batch_size = 50
#d_model = 32
#num_heads = 2
#num_layers = 3
#d_ff = 128
#max_seq_length = 20
#dropout = 0.1
#transformer = Transformer(src_vocab_size, tgt_vocab_size, d_model, num_heads,
# num_layers, d_ff, max_seq_length, dropout)
# Generate random sample data
#src_data = torch.randint(1, src_vocab_size, (batch_size, max_seq_length))
#tgt_data = torch.randint(1, nc, (batch_size, nc))
#print(src_data.shape)
#print(tgt_data.shape)
In [214]:
def mytf(src, tgt, transformer):
criterion = nn.CrossEntropyLoss(ignore_index=0)
optimizer = optim.Adam(transformer.parameters(), lr=0.01, betas=(0.9, 0.98), eps=1e-9)
transformer.train()
n = int(len(src) * 0.8)
for ep in range(5):
optimizer.zero_grad()
output = transformer(src[:n,:], tgt[:n, :-1])
loss = criterion(output.contiguous().view(-1, tgt_vocab_size),
tgt[:n, 1:].contiguous().view(-1))
loss.backward()
optimizer.step()
#print('ep: %2d' % ep, 'loss: %.3f' % loss.item())
with torch.no_grad():
outp = transformer(src[n:], tgt[n:])
print('test: n:', n, 'outp:', outp.shape, 'tgt:', tgt[n:].shape)
print('outp[0]:', outp[0])
print('tgt[n]:', tgt[n])
accte = sum([np.argmax(outp[i][-1]) == np.argmax(tgt[n+i,:])
for i in range(len(outp))])
print('accte:', accte)
print('epoch: %2d loss: %5.1f acc_test: %.2f' %
(ep, loss, accte/(len(src)-n)))
#mytf(src_data, tgt_data)
Apply to IMDB Dataset¶
In [215]:
import numpy as np
import random
data = [ line.split() for line in open('imdb_labelled.txt').readlines() ]
random.seed(13)
random.shuffle(data)
rev = [ [ w.strip('.,').lower() for w in line[:-1] ] for line in data ]
sen = np.asarray([ int(line[-1]) for line in data ])
print(rev[:5], sen[:5])
print('mean length:', np.mean([len(s) for s in rev]))
print('reviews:', len(sen), 'positive:', sum(sen))
[['all', 'the', 'actors', 'give', 'a', 'wonderful', 'performance', 'especially', 'jennifer', 'rubin', 'as', 'jamie', 'harris', 'who', 'changes', 'from', 'the', 'nervous', 'starlet', 'in', 'the', 'beginning', 'through', 'the', 'strange', 'events', 'she', 'is', 'part', 'of', 'to', 'the', 'cool', 'star'], ['babie', 'bop', 'is', 'very', 'cute'], ['hackneyed', 'writing', 'certainly', 'but', 'made', 'even', 'worse', 'by', 'the', 'bad', 'directing'], ['it', 'is', 'indescribably', 'the', 'most', 'annoying', 'and', 'idiotic', 'show', 'i', 'have', 'ever', 'seen'], ['it', 'is', 'zillion', 'times', 'away', 'from', 'reality']] [1 1 0 0 0] mean length: 14.355 reviews: 1000 positive: 500
In [216]:
src_vocab_size = 500
cnt = {}
for r in rev:
for w in r:
if w in cnt: cnt[w] += 1
else: cnt[w] = 1
voc = list(sorted(cnt, key=cnt.get, reverse=True)[:src_vocab_size])
print(len(voc), voc[:50])
500 ['the', 'a', 'and', 'of', 'is', 'this', 'i', 'to', 'it', 'in', 'was', 'movie', 'that', 'film', 'for', 'as', 'but', 'with', 'are', 'you', 'on', 'one', 'not', 'bad', 'very', 'just', "it's", 'all', 'so', 'good', 'an', 'at', 'there', 'be', 'by', 'have', 'about', 'like', 'from', 'if', 'his', 'really', 'acting', 'time', 'even', 'out', 'great', 'were', 'or', 'has']
In [217]:
src = []
tgt = []
n = max_seq_length
for i in range(len(rev)):
src.append( ([ voc.index(w) for w in rev[i] if w in voc ] + n * [0] )[:n] )
tgt.append([[0,1],[1,0]][sen[i]])
print(src[:3])
print(tgt[:3])
src = torch.tensor(np.array(src))
tgt = torch.tensor(np.array(tgt))
print(src.shape, tgt.shape)
#mytf(src, tgt)
[[27, 0, 109, 161, 1, 106, 149, 277, 15, 52, 38, 0, 9, 0, 389, 177, 0, 162, 4, 163], [4, 24, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [136, 223, 16, 66, 44, 164, 34, 0, 23, 199, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]] [[1, 0], [1, 0], [0, 1]] torch.Size([1000, 20]) torch.Size([1000, 2])
In [218]:
nc = 2
d_model = 32
num_heads = 2
num_layers = 3
d_ff = 128
dropout = 0.1
transformer = Transformer(src_vocab_size, nc, d_model, num_heads,
num_layers, d_ff, max_seq_length, dropout)
# Generate random sample data
#src_data = torch.randint(1, src_vocab_size, (batch_size, max_seq_length))
#tgt_data = torch.randint(1, nc, (batch_size, nc))
#print(src_data.shape)
#print(tgt_data.shape)
mytf(src, tgt, transformer)
test: n: 800 outp: torch.Size([200, 2, 2]) tgt: torch.Size([200, 2]) outp[0]: tensor([[-1.9856, 1.6884], [-1.9457, 0.9406]]) tgt[n]: tensor([0, 1]) accte: tensor(105) epoch: 0 loss: 0.1 acc_test: 0.52 test: n: 800 outp: torch.Size([200, 2, 2]) tgt: torch.Size([200, 2]) outp[0]: tensor([[-2.3293, 2.7362], [-2.4762, 2.3961]]) tgt[n]: tensor([0, 1]) accte: tensor(105) epoch: 1 loss: 0.0 acc_test: 0.52 test: n: 800 outp: torch.Size([200, 2, 2]) tgt: torch.Size([200, 2]) outp[0]: tensor([[-2.3842, 3.0599], [-2.5533, 3.1412]]) tgt[n]: tensor([0, 1]) accte: tensor(105) epoch: 2 loss: 0.0 acc_test: 0.52 test: n: 800 outp: torch.Size([200, 2, 2]) tgt: torch.Size([200, 2]) outp[0]: tensor([[-2.7192, 3.3567], [-2.6117, 3.4418]]) tgt[n]: tensor([0, 1]) accte: tensor(105) epoch: 3 loss: 0.0 acc_test: 0.52 test: n: 800 outp: torch.Size([200, 2, 2]) tgt: torch.Size([200, 2]) outp[0]: tensor([[-2.8421, 3.6896], [-2.7736, 3.6028]]) tgt[n]: tensor([0, 1]) accte: tensor(105) epoch: 4 loss: 0.0 acc_test: 0.52
In [3]:
S, T, B, E = 20, 1, 100, 50
tf = torch.nn.Transformer(d_model=E, nhead=2, num_encoder_layers=3)
src = torch.rand((S, B, E))
tgt = torch.rand((T, B, E))
out = tf(src, tgt)
out.shape
Out[3]:
torch.Size([1, 100, 50])
In [ ]:
In [ ]: