Transformer实现的一个Demo
RT，直接上代码，可以跑通：
#encoding:utf-8
import torch
import torch.nn as nn
import numpy as np
import math
class Config(object):
def __init__(self):
self.vocab_size = 6
self.d_model = 512
self.n_heads = 4
assert self.d_model % self.n_heads == 0
self.dim_k = self.d_model // self.n_heads
self.dim_v = self.d_model // self.n_heads
self.padding_size = 30
self.UNK = 5
self.PAD = 4
self.N = 6
self.p = 0.1
config = Config()
class Embedding(nn.Module):
def __init__(self, vocab_size):
super(Embedding, self).__init__()
self.embedding = nn.Embedding(vocab_size, config.d_model, padding_idx=config.PAD)
def forward(self, x):
# print("type(x):", type(x))
# for i in range(len(x)):
# if len(x[i]) < config.padding_size:
# x[i].extend([config.UNK] * (config.padding_size - len(x[i])))
# else:
# tmp = x[i][:config.padding_size]
# print("tmp.shape:", tmp.shape)
# print("type(x[i]):", type(x[i]))
# print("x[i].shape:", x[i].shape)
# x[i,:] = x[i][:config.padding_size]
x = self.embedding(torch.tensor(x))
return x
class Positional_Encoding(nn.Module):
def __init__(self, d_model):
super(Positional_Encoding, self).__init__()
self.d_model = d_model
def forward(self, seq_len, embedding_dim):
positional_encoding = np.zeros((seq_len, embedding_dim))
for pos in range(positional_encoding.shape[0]):
for i in range(positional_encoding.shape[1]):
positional_encoding[pos][i] = math.sin(pos / (10000**(2*i/self.d_model))) if i%2==0 else math.cos(pos/(10000**(2*i/self.d_model)))
return torch.from_numpy(positional_encoding)
class Multihead_Attention(nn.Module):
def __init__(self, d_model, dim_k, dim_v, n_heads):
super(Multihead_Attention, self).__init__()
self.dim_v = dim_v
self.dim_k = dim_k
self.n_heads = n_heads
self.q = nn.Linear(d_model, dim_k)
self.k = nn.Linear(d_model, dim_k)
self.v = nn.Linear(d_model, dim_v)
self.o = nn.Linear(dim_v, d_model)
self.norm_fact = 1 / math.sqrt(d_model)
def generate_mask(self, dim):
matrix = np.ones((dim, dim))
mask = torch.BoolTensor(np.tril(matrix).astype(np.bool_))
return mask
def forward(self, x, y, requires_mask=False):
Q = self.q(x).reshape(-1, x.shape[0], x.shape[1], self.dim_k//self.n_heads)
K = self.k(x).reshape(-1, x.shape[0], x.shape[1], self.dim_k//self.n_heads)
V = self.v(y).reshape(-1, x.shape[0], x.shape[1], self.dim_v//self.n_heads)
attention_score = torch.matmul(Q, K.permute(0,1,3,2)) * self.norm_fact
if requires_mask:
mask = self.generate_mask(x.shape[1])
attention_score.masked_fill(mask, value=float("-inf"))
output = torch.matmul(attention_score, V).reshape(y.shape[0], y.shape[1], -1)
output = self.o(output)
return output
class Feed_Forward(nn.Module):
def __init__(self, input_dim, hidden_dim=2048):
super(Feed_Forward, self).__init__()
self.L1 = nn.Linear(input_dim, hidden_dim)
self.L2 = nn.Linear(hidden_dim, input_dim)
def forward(self, x):
output = nn.ReLU()(self.L1(x))
output = self.L2(output)
return output
class Add_Norm(nn.Module):
def __init__(self):
super(Add_Norm, self).__init__()
self.dropout = nn.Dropout(config.p)
def forward(self, x, sub_layer, **kwargs):
sub_output = sub_layer(x, **kwargs)
x = self.dropout(x+sub_output)
layer_norm = nn.LayerNorm(x.size()[1:])
out = layer_norm(x)
return out
class Encoder(nn.Module):
def __init__(self):
super(Encoder, self).__init__()
self.positional_encoding = Positional_Encoding(config.d_model)
self.muti_atten = Multihead_Attention(config.d_model, config.dim_k, config.dim_v, config.n_heads)
self.feed_forward = Feed_Forward(config.d_model)
self.add_norm = Add_Norm()
def forward(self, x):
x += self.positional_encoding(x.shape[1], config.d_model)
output = self.add_norm(x, self.muti_atten, y=x)
output = self.add_norm(output, self.feed_forward)
return output
class Decoder(nn.Module):
def __init__(self):
super(Decoder, self).__init__()
self.positional_encoding = Positional_Encoding(config.d_model)
self.muti_atten = Multihead_Attention(config.d_model, config.dim_k, config.dim_v, config.n_heads)
self.feed_forward = Feed_Forward(config.d_model)
self.add_norm = Add_Norm()
def forward(self, x, encoder_output):
x += self.positional_encoding(x.shape[1], config.d_model)
output = self.add_norm(x, self.muti_atten, y=x, requires_mask=True)
output = self.add_norm(output, self.muti_atten, y=encoder_output, requires_mask=True)
output = self.add_norm(output, self.feed_forward)
return output
class Transformer_layer(nn.Module):
def __init__(self):
super(Transformer_layer, self).__init__()
self.encoder = Encoder()
self.decoder = Decoder()
def forward(self, x):
x_input, x_output = x
encoder_output = self.encoder(x_input)
decoder_output = self.decoder(x_output, encoder_output)
return (encoder_output, decoder_output)
class Transformer(nn.Module):
def __init__(self, N, vocab_size, output_dim):
super(Transformer, self).__init__()
self.embedding_input = Embedding(vocab_size=vocab_size)
self.embedding_output = Embedding(vocab_size=vocab_size)
self.output_dim = output_dim
self.linear = nn.Linear(config.d_model, output_dim)
self.softmax = nn.Softmax(dim=-1)
self.model = nn.Sequential(*[Transformer_layer() for _ in range(N)])
def forward(self, x):
x_input, x_output = x, x
x_input = self.embedding_input(x_input)
x_output = self.embedding_output(x_output)
_, output = self.model((x_input, x_output))
output = self.linear(output)
output = self.softmax(output)
return output
def main():
transformer = Transformer(4, 1024, 512)
input = (np.random.rand(8, 512) * 512).astype(np.int64)
output = transformer(input)
print("output.shape:", output.shape)
if __name__ == '__main__':
main()
---------------------
输出：
output.shape: torch.Size([8, 512, 512])