以下是添加了详细注释的代码和参数介绍:
Transformer 实现及自回归推理
本文展示了如何手动实现一个简化版的Transformer模型,并用自回归方式实现一个seq2seq任务,例如机器翻译。
导入必要的库
python
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import numpy as np
定义位置编码
Transformer 使用位置编码来捕捉序列中的位置信息。
python
class PositionalEncoding(nn.Module):
def __init__(self, d_model, max_len=5000):
super(PositionalEncoding, self).__init__()
pe = torch.zeros(max_len, d_model)
position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-np.log(10000.0) / d_model))
pe[:, 0::2] = torch.sin(position * div_term)
pe[:, 1::2] = torch.cos(position * div_term)
pe = pe.unsqueeze(0).transpose(0, 1)
self.register_buffer('pe', pe)
def forward(self, x):
return x + self.pe[:x.size(0), :]
参数介绍:
d_model
: 词嵌入和位置编码的维度。max_len
: 序列的最大长度。
定义自注意力机制
python
class MultiHeadAttention(nn.Module):
def __init__(self, d_model, num_heads):
super(MultiHeadAttention, self).__init__()
self.num_heads = num_heads
self.d_model = d_model
assert d_model % num_heads == 0
self.depth = d_model // num_heads
self.wq = nn.Linear(d_model, d_model)
self.wk = nn.Linear(d_model, d_model)
self.wv = nn.Linear(d_model, d_model)
self.dense = nn.Linear(d_model, d_model)
def split_heads(self, x, batch_size):
x = x.view(batch_size, -1, self.num_heads, self.depth)
return x.permute(0, 2, 1, 3)
def forward(self, v, k, q, mask):
batch_size = q.size(0)
q = self.split_heads(self.wq(q), batch_size)
k = self.split_heads(self.wk(k), batch_size)
v = self.split_heads(self.wv(v), batch_size)
matmul_qk = torch.matmul(q, k.transpose(-1, -2))
dk = torch.tensor(k.size(-1)).float()
scaled_attention_logits = matmul_qk / torch.sqrt(dk)
if mask is not None:
scaled_attention_logits += (mask * -1e9)
attention_weights = F.softmax(scaled_attention_logits, dim=-1)
output = torch.matmul(attention_weights, v)
output = output.permute(0, 2, 1, 3).contiguous()
output = output.view(batch_size, -1, self.d_model)
return self.dense(output)
参数介绍:
d_model
: 词嵌入和注意力机制的维度。num_heads
: 注意力头的数量。
定义前馈神经网络
python
class FeedForward(nn.Module):
def __init__(self, d_model, dff):
super(FeedForward, self).__init__()
self.linear1 = nn.Linear(d_model, dff)
self.linear2 = nn.Linear(dff, d_model)
def forward(self, x):
return self.linear2(F.relu(self.linear1(x)))
参数介绍:
d_model
: 词嵌入的维度。dff
: 前馈神经网络的隐藏层维度。
定义编码器层
python
class EncoderLayer(nn.Module):
def __init__(self, d_model, num_heads, dff, dropout=0.1):
super(EncoderLayer, self).__init__()
self.mha = MultiHeadAttention(d_model, num_heads)
self.ffn = FeedForward(d_model, dff)
self.layernorm1 = nn.LayerNorm(d_model)
self.layernorm2 = nn.LayerNorm(d_model)
self.dropout1 = nn.Dropout(dropout)
self.dropout2 = nn.Dropout(dropout)
def forward(self, x, mask):
attn_output = self.mha(x, x, x, mask)
out1 = self.layernorm1(x + self.dropout1(attn_output))
ffn_output = self.ffn(out1)
out2 = self.layernorm2(out1 + self.dropout2(ffn_output))
return out2
参数介绍:
d_model
: 词嵌入和注意力机制的维度。num_heads
: 注意力头的数量。dff
: 前馈神经网络的隐藏层维度。dropout
: Dropout 概率。
定义编码器
python
class Encoder(nn.Module):
def __init__(self, num_layers, d_model, num_heads, dff, input_vocab_size, max_len, dropout=0.1):
super(Encoder, self).__init__()
self.d_model = d_model
self.num_layers = num_layers
self.embedding = nn.Embedding(input_vocab_size, d_model)
self.pos_encoding = PositionalEncoding(d_model, max_len)
self.enc_layers = nn.ModuleList([EncoderLayer(d_model, num_heads, dff, dropout) for _ in range(num_layers)])
self.dropout = nn.Dropout(dropout)
def forward(self, x, mask):
seq_len = x.size(1)
x = self.embedding(x)
x *= torch.sqrt(torch.tensor(self.d_model, dtype=torch.float32))
x = self.pos_encoding(x.permute(1, 0, 2))
x = x.permute(1, 0, 2)
x = self.dropout(x)
for i in range(self.num_layers):
x = self.enc_layers[i](x, mask)
return x
参数介绍:
num_layers
: 编码器层的数量。d_model
: 词嵌入和注意力机制的维度。num_heads
: 注意力头的数量。dff
: 前馈神经网络的隐藏层维度。input_vocab_size
: 输入词汇表大小。max_len
: 序列的最大长度。dropout
: Dropout 概率。
定义解码器层
python
class DecoderLayer(nn.Module):
def __init__(self, d_model, num_heads, dff, dropout=0.1):
super(DecoderLayer, self).__init__()
self.mha1 = MultiHeadAttention(d_model, num_heads)
self.mha2 = MultiHeadAttention(d_model, num_heads)
self.ffn = FeedForward(d_model, dff)
self.layernorm1 = nn.LayerNorm(d_model)
self.layernorm2 = nn.LayerNorm(d_model)
self.layernorm3 = nn.LayerNorm(d_model)
self.dropout1 = nn.Dropout(dropout)
self.dropout2 = nn.Dropout(dropout)
self.dropout3 = nn.Dropout(dropout)
def forward(self, x, enc_output, look_ahead_mask, padding_mask):
attn1 = self.mha1(x, x, x, look_ahead_mask)
attn1 = self.dropout1(attn1)
out1 = self.layernorm1(attn1 + x)
attn2 = self.mha2(enc_output, enc_output, out1, padding_mask)
attn2 = self.dropout2(attn2)
out2 = self.layernorm2(attn2 + out1)
ffn_output = self.ffn(out2)
ffn_output = self.dropout3(ffn_output)
out3 = self.layernorm3(ffn_output + out2)
return out3
参数介绍:
d_model
: 词嵌入和注意力机制的维度。num_heads
: 注意力头的数量。dff
: 前馈神经网络的隐藏层维度。dropout
: Dropout 概率。
定义解码器
python
class Decoder(nn.Module):
def __init__(self, num_layers, d_model, num_heads, dff, target_vocab_size, max_len, dropout=0.1):
super(Decoder, self).__init__()
self.d_model = d_model
self.num_layers = num_layers
self.embedding = nn.Embedding(target_vocab_size, d_model)
self.pos_encoding = PositionalEncoding(d_model, max_len)
self.dec_layers = nn.ModuleList([DecoderLayer(d_model, num_heads, dff, dropout) for _ in range(num_layers)])
self.dropout = nn.Dropout(dropout)
def forward(self, x, enc_output, look_ahead_mask, padding_mask):
seq_len = x.size(1)
attention_weights = {}
x = self.embedding(x)
x *= torch.sqrt(torch.tensor(self.d_model, dtype=torch.float32))
x = self.pos_encoding(x.permute(1, 0, 2))
x = x.permute(1, 0, 2)
x = self.dropout(x)
for i in range(self.num_layers):
x = self.dec_layers[i](x, enc_output, look_ahead_mask, padding_mask)
return x
参数介绍:
num_layers
: 解码器层的数量。d_model
: 词嵌入和注意力机制的维度。num_heads
: 注意力头的数量。dff
: 前馈神经网络的隐藏层维度。target_vocab_size
: 目标词汇表大小。max_len
: 序列的最大长度。dropout
: Dropout 概率。
定义Transformer模型
python
class Transformer(nn.Module):
def __init__(self, num_layers, d_model, num_heads, dff, input_vocab_size, target_vocab_size, pe_input, pe_target, dropout=0.1):
super(Transformer, self).__init__()
self.encoder = Encoder(num_layers, d_model, num_heads, dff, input_vocab_size, pe_input, dropout)
self.decoder = Decoder(num_layers, d_model, num_heads, dff, target_vocab_size, pe_target, dropout)
self.final_layer = nn.Linear(d_model, target_vocab_size)
def forward(self, inp, tar, enc_padding_mask, look_ahead_mask, dec_padding_mask):
enc_output = self.encoder(inp, enc_padding_mask)
dec_output = self.decoder(tar, enc_output, look_ahead_mask, dec_padding_mask)
final_output = self.final_layer(dec_output)
return final_output
参数介绍:
num_layers
: 编码器和解码器层的数量。d_model
: 词嵌入和注意力机制的维度。num_heads
: 注意力头的数量。dff
: 前馈神经网络的隐藏层维度。input_vocab_size
: 输入词汇表大小。target_vocab_size
: 目标词汇表大小。pe_input
: 输入序列的最大长度。pe_target
: 目标序列的最大长度。dropout
: Dropout 概率。
创建掩码
python
def create_padding_mask(seq):
seq = torch.eq(seq, 0)
return seq[:, None, None, :]
def create_look_ahead_mask(size):
mask = torch.triu(torch.ones((size, size)), 1)
return mask
自回归推理
实现一个简化的自回归推理过程:
python
def generate_text(model, input_sequence, start_token, max_length, target_vocab_size):
generated = [start_token]
model.eval()
enc_padding_mask = create_padding_mask(input_sequence)
with torch.no_grad():
enc_output = model.encoder(input_sequence, enc_padding_mask)
for _ in range(max_length):
dec_input = torch.tensor(generated).unsqueeze(0)
look_ahead_mask = create_look_ahead_mask(dec_input.size(1))
dec_padding_mask = create_padding_mask(dec_input)
with torch.no_grad():
output = model.decoder(dec_input, enc_output, look_ahead_mask, dec_padding_mask)
output = model.final_layer(output)
next_token = torch.argmax(output[:, -1, :], dim=-1).item()
generated.append(next_token)
if next_token == eos_token:
break
return generated
参数介绍:
model
: 训练好的Transformer模型。input_sequence
: 输入的序列张量。start_token
: 生成序列的开始标记。max_length
: 生成序列的最大长度。target_vocab_size
: 目标词汇表大小。
使用示例
创建一个简单的模型并进行文本生成:
python
input_vocab_size = 1000 # 输入词汇表大小
target_vocab_size = 1000 # 目标词汇表大小
max_len = 50 # 序列最大长度
num_layers = 2 # 编码器和解码器层的数量
d_model = 512 # 词嵌入和注意力机制的维度
num_heads = 8 # 注意力头的数量
dff = 2048 # 前馈神经网络的隐藏层维度
# 创建Transformer模型
transformer = Transformer(num_layers, d_model, num_heads, dff, input_vocab_size, target_vocab_size, max_len, max_len)
# 输入序列,假设输入序列为[1, 2, 3, 4, 0, 0, 0]
input_sequence = torch.tensor([[1, 2, 3, 4, 0, 0, 0]])
# 假设开始标记为1,结束标记为2
start_token = 1
eos_token = 2
# 生成序列
generated_sequence = generate_text(transformer, input_sequence, start_token, max_length=20, target_vocab_size=target_vocab_size)
print("Generated sequence:", generated_sequence)
以上代码展示了一个简化的Transformer模型的实现,包括位置编码、自注意力机制、前馈神经网络、编码器层、解码器层、编码器和解码器整体的实现,以及一个基本的自回归推理过程。你可以根据需要进行调整和扩展。
关于mask的解释
以下关于掩码函数 create_padding_mask
和 create_look_ahead_mask
的详细介绍以及示例。
create_padding_mask
该函数用于生成填充掩码,以忽略序列中的填充值(通常是0)。在Transformer模型中,填充掩码用于屏蔽掉填充值在计算注意力时的影响。
代码实现
python
def create_padding_mask(seq):
seq = torch.eq(seq, 0) # 查找填充值(假设填充值为0),返回一个布尔张量
return seq[:, None, None, :] # 扩展维度以适配注意力机制中的广播
示例
假设我们有一个输入序列,其中0是填充值:
python
seq = torch.tensor([[7, 6, 0, 0, 0], [1, 2, 3, 0, 0]])
padding_mask = create_padding_mask(seq)
print(padding_mask)
输出
plaintext
tensor([[[[False, False, True, True, True]]],
[[[False, False, False, True, True]]]])
在输出中,True
表示填充值的位置,这些位置将在计算注意力时被忽略。
create_look_ahead_mask
该函数用于生成前瞻掩码,以确保解码器中的每个位置只能看到该位置之前的序列,不能看到未来的信息。在自回归生成中,前瞻掩码用于防止解码器在生成下一个标记时看到未来的标记。
代码实现
python
def create_look_ahead_mask(size):
mask = torch.triu(torch.ones((size, size)), 1) # 生成上三角矩阵,主对角线以上的元素为1
return mask # 返回前瞻掩码
示例
假设我们有一个序列长度为5:
python
size = 5
look_ahead_mask = create_look_ahead_mask(size)
print(look_ahead_mask)
输出
plaintext
tensor([[0., 1., 1., 1., 1.],
[0., 0., 1., 1., 1.],
[0., 0., 0., 1., 1.],
[0., 0., 0., 0., 1.],
[0., 0., 0., 0., 0.]])
在输出中,1
表示被掩盖的位置,这些位置在计算注意力时将被屏蔽。
综合示例
结合以上两种掩码,假设我们有以下输入序列:
python
seq = torch.tensor([[7, 6, 0, 0, 0], [1, 2, 3, 0, 0]])
size = seq.size(1)
padding_mask = create_padding_mask(seq)
look_ahead_mask = create_look_ahead_mask(size)
print("Padding Mask:\n", padding_mask)
print("Look Ahead Mask:\n", look_ahead_mask)
输出
plaintext
Padding Mask:
tensor([[[[False, False, True, True, True]]],
[[[False, False, False, True, True]]]])
Look Ahead Mask:
tensor([[0., 1., 1., 1., 1.],
[0., 0., 1., 1., 1.],
[0., 0., 0., 1., 1.],
[0., 0., 0., 0., 1.],
[0., 0., 0., 0., 0.]])
在实际使用中,编码器使用 padding_mask
来屏蔽填充值的影响,解码器则同时使用 look_ahead_mask
和 padding_mask
来屏蔽未来标记和填充值的影响。