transformer代码python

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

import torch
import torch.nn as nn
import torch.optim as optim

class TransformerModel(nn.Module):
    def __init__(self, vocab_size, embedding_dim, hidden_dim, num_layers):
        super(TransformerModel, self).__init__()
        self.embedding = nn.Embedding(vocab_size, embedding_dim)
        self.transformer = nn.Transformer(
            d_model=embedding_dim,
            nhead=2,
            num_encoder_layers=num_layers,
            num_decoder_layers=num_layers,
            dim_feedforward=hidden_dim,
        )
        self.fc = nn.Linear(embedding_dim, vocab_size)
    
    def forward(self, src, tgt):
        src_embed = self.embedding(src)
        tgt_embed = self.embedding(tgt)
        
        src_pad_mask = self._generate_padding_mask(src)
        tgt_pad_mask = self._generate_padding_mask(tgt)
        tgt_subseq_mask = self._generate_subsequent_mask(tgt)
        tgt_mask = tgt_pad_mask & tgt_subseq_mask
        
        src_encoded = self.transformer.encoder(src_embed)
        tgt_decoded = self.transformer.decoder(tgt_embed, src_encoded, tgt_mask=tgt_mask, src_key_padding_mask=src_pad_mask, tgt_key_padding_mask=tgt_pad_mask)
        
        output = self.fc(tgt_decoded)
        return output
    
    def _generate_padding_mask(self, seq):
        mask = (seq == 0)  # Pad token index is assumed to be 0
        return mask
    
    def _generate_subsequent_mask(self, seq):
        seq_len = seq.shape[1]
        mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1).bool()  # Upper triangular matrix
        return mask

# 示例用法
vocab_size = 10000
embedding_dim = 256
hidden_dim = 512
num_layers = 4

model = TransformerModel(vocab_size, embedding_dim, hidden_dim, num_layers)
src = torch.tensor([[1, 2, 3, 4, 5]])
tgt = torch.tensor([[6, 7, 8, 9, 10]])
output = model(src, tgt)
print(output.shape)  # 输出: torch.Size([1, 5, 10000])

# 损失函数和优化器的定义和训练过程略过

这是一个简单的 Transformer 模型，用于序列到序列的任务。它使用了 Embedding 层、Transformer 层和全连接层来实现。你可以根据自己的需求进行修改和扩展。