### 构建Transformer模型的代码实现 以下是基于PyTorch框架从零开始实现一个简单的Transformer模型的完整代码示例。此实现包括了编码器（Encoder）和解码器（Decoder）的基本结构，以及自注意力机制（Self-Attention）、位置编码（Positional Encoding）等核心组件。 #### 1. 自注意力机制（Self-Attention） 自注意力机制是Transformer的核心组件之一，它允许模型在处理序列数据时关注不同位置的信息[^1]。 ```python import torch import torch.nn as nn import math class SelfAttention(nn.Module): def __init__(self, embed_size, heads): super(SelfAttention, self).__init__() self.embed_size = embed_size self.heads = heads self.head_dim = embed_size // heads assert ( self.head_dim * heads == embed_size ), "Embedding size needs to be divisible by heads" self.values = nn.Linear(self.head_dim, self.head_dim, bias=False) self.keys = nn.Linear(self.head_dim, self.head_dim, bias=False) self.queries = nn.Linear(self.head_dim, self.head_dim, bias=False) self.fc_out = nn.Linear(heads * self.head_dim, embed_size) def forward(self, values, keys, query, mask): N = query.shape[0] value_len, key_len, query_len = values.shape[1], keys.shape[1], query.shape[1] # Split embedding into self.heads different pieces values = values.reshape(N, value_len, self.heads, self.head_dim) keys = keys.reshape(N, key_len, self.heads, self.head_dim) queries = query.reshape(N, query_len, self.heads, self.head_dim) energy = torch.einsum("nqhd,nkhd->nhqk", [queries, keys]) if mask is not None: energy = energy.masked_fill(mask == 0, float("-1e20")) attention = torch.softmax(energy / (self.embed_size ** (1 / 2)), dim=3) out = torch.einsum("nhql,nlhd->nqhd", [attention, values]).reshape( N, query_len, self.heads * self.head_dim ) out = self.fc_out(out) return out ``` #### 2. Transformer块（Transformer Block） Transformer块将多头注意力（Multi-Head Attention）与前馈神经网络（Feed Forward Network）结合在一起[^1]。 ```python class TransformerBlock(nn.Module): def __init__(self, embed_size, heads, dropout, forward_expansion): super(TransformerBlock, self).__init__() self.attention = SelfAttention(embed_size, heads) self.norm1 = nn.LayerNorm(embed_size) self.norm2 = nn.LayerNorm(embed_size) self.feed_forward = nn.Sequential( nn.Linear(embed_size, forward_expansion * embed_size), nn.ReLU(), nn.Linear(forward_expansion * embed_size, embed_size), ) self.dropout = nn.Dropout(dropout) def forward(self, value, key, query, mask): attention = self.attention(value, key, query, mask) x = self.dropout(self.norm1(attention + query)) forward = self.feed_forward(x) out = self.dropout(self.norm2(forward + x)) return out ``` #### 3. 编码器（Encoder） 编码器由多个Transformer块堆叠而成，负责将输入序列转换为上下文表示[^1]。 ```python class Encoder(nn.Module): def __init__( self, src_vocab_size, embed_size, num_layers, heads, device, forward_expansion, dropout, max_length, ): super(Encoder, self).__init__() self.embed_size = embed_size self.device = device self.word_embedding = nn.Embedding(src_vocab_size, embed_size) self.position_embedding = nn.Embedding(max_length, embed_size) self.layers = nn.ModuleList( [ TransformerBlock( embed_size, heads, dropout=dropout, forward_expansion=forward_expansion ) for _ in range(num_layers) ] ) self.dropout = nn.Dropout(dropout) def forward(self, x, mask): N, seq_length = x.shape positions = torch.arange(0, seq_length).expand(N, seq_length).to(self.device) out = self.dropout( (self.word_embedding(x) + self.position_embedding(positions)) ) for layer in self.layers: out = layer(out, out, out, mask) return out ``` #### 4. 解码器（Decoder） 解码器也由多个Transformer块组成，但需要额外处理掩码操作以确保预测时不会泄露未来信息。 ```python class DecoderBlock(nn.Module): def __init__(self, embed_size, heads, forward_expansion, dropout, device): super(DecoderBlock, self).__init__() self.norm = nn.LayerNorm(embed_size) self.attention = SelfAttention(embed_size, heads=heads) self.transformer_block = TransformerBlock( embed_size, heads, dropout, forward_expansion ) self.dropout = nn.Dropout(dropout) def forward(self, x, value, key, src_mask, trg_mask): attention = self.attention(x, x, x, trg_mask) query = self.dropout(self.norm(attention + x)) out = self.transformer_block(value, key, query, src_mask) return out class Decoder(nn.Module): def __init__( self, trg_vocab_size, embed_size, num_layers, heads, forward_expansion, dropout, device, max_length, ): super(Decoder, self).__init__() self.device = device self.word_embedding = nn.Embedding(trg_vocab_size, embed_size) self.position_embedding = nn.Embedding(max_length, embed_size) self.layers = nn.ModuleList( [ DecoderBlock(embed_size, heads, forward_expansion, dropout, device) for _ in range(num_layers) ] ) self.fc_out = nn.Linear(embed_size, trg_vocab_size) self.dropout = nn.Dropout(dropout) def forward(self, x, enc_out, src_mask, trg_mask): N, seq_length = x.shape positions = torch.arange(0, seq_length).expand(N, seq_length).to(self.device) x = self.dropout((self.word_embedding(x) + self.position_embedding(positions))) for layer in self.layers: x = layer(x, enc_out, enc_out, src_mask, trg_mask) out = self.fc_out(x) return out ``` #### 5. 完整Transformer模型 将编码器和解码器组合成一个完整的Transformer模型[^1]。 ```python class Transformer(nn.Module): def __init__( self, src_vocab_size, trg_vocab_size, src_pad_idx, trg_pad_idx, embed_size=512, num_layers=6, forward_expansion=4, heads=8, dropout=0.1, device="cuda", max_length=100, ): super(Transformer, self).__init__() self.encoder = Encoder( src_vocab_size, embed_size, num_layers, heads, device, forward_expansion, dropout, max_length, ) self.decoder = Decoder( trg_vocab_size, embed_size, num_layers, heads, forward_expansion, dropout, device, max_length, ) self.src_pad_idx = src_pad_idx self.trg_pad_idx = trg_pad_idx self.device = device def make_src_mask(self, src): src_mask = (src != self.src_pad_idx).unsqueeze(1).unsqueeze(2) return src_mask.to(self.device) def make_trg_mask(self, trg): N, trg_len = trg.shape trg_mask = torch.tril(torch.ones((trg_len, trg_len))).expand( N, 1, trg_len, trg_len ) return trg_mask.to(self.device) def forward(self, src, trg): src_mask = self.make_src_mask(src) trg_mask = self.make_trg_mask(trg) enc_src = self.encoder(src, src_mask) out = self.decoder(trg, enc_src, src_mask, trg_mask) return out ``` ###