transformer代码实现例子

以下是一个使用PyTorch实现Transformer模型的简单示例代码：

import torch
import torch.nn as nn
import torch.nn.functional as F

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=5000):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=0.1)
        
        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() * (-math.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):
        x = x + self.pe[:x.size(0), :]
        return self.dropout(x)

class Transformer(nn.Module):
    def __init__(self, d_model=512, nhead=8, num_encoder_layers=6, num_decoder_layers=6,
                 dim_feedforward=2048, dropout=0.1):
        super(Transformer, self).__init__()
        encoder_layer = nn.TransformerEncoderLayer(d_model, nhead, dim_feedforward, dropout)
        encoder_norm = nn.LayerNorm(d_model)
        self.encoder = nn.TransformerEncoder(encoder_layer, num_encoder_layers, encoder_norm)
        
        decoder_layer = nn.TransformerDecoderLayer(d_model, nhead, dim_feedforward, dropout)
        decoder_norm = nn.LayerNorm(d_model)
        self.decoder = nn.TransformerDecoder(decoder_layer, num_decoder_layers, decoder_norm)
        
        self.fc = nn.Linear(d_model, vocab_size)
        
        self.source_embedding = nn.Embedding(vocab_size, d_model)
        self.target_embedding = nn.Embedding(vocab_size, d_model)
        self.positional_encoding = PositionalEncoding(d_model)
    
    def forward(self, source, target):
        source_embedded = self.source_embedding(source)
        target_embedded = self.target_embedding(target)
        
        source_embedded = self.positional_encoding(source_embedded)
        target_embedded = self.positional_encoding(target_embedded)
        
        encoder_output = self.encoder(source_embedded)
        decoder_output = self.decoder(target_embedded, encoder_output)
        
        output = self.fc(decoder_output)
        return F.log_softmax(output, dim=-1)

上述代码定义了一个简单的Transformer模型，其中使用了`nn.TransformerEncoder`和`nn.TransformerDecoder`来实现编码器和解码器。模型通过嵌入层将输入序列编码为嵌入向量，然后经过位置编码层进行位置编码。编码器和解码器分别通过多层的Transformer层进行处理，最后通过全连接层输出预测结果。

请注意，上述代码仅为示例，实际应用中可能需要根据具体任务进行修改和优化。