### PyTorch 内置 Transformer 模型的使用 PyTorch 提供了内置的 `nn.Transformer` 模块，可以直接用于构建 Transformer 模型[^2]。此模块实现了标准的 Transformer 架构，包括编码器和解码器部分。此外，PyTorch 还提供了更高级别的接口，例如 `torch.nn.TransformerEncoder` 和 `torch.nn.TransformerDecoder`，允许用户更灵活地设计模型。 #### 内置 Transformer 模型的结构 `nn.Transformer` 的主要组件包括多头注意力机制（Multi-Head Attention）、前馈神经网络（Feed-Forward Neural Network）以及位置编码（Positional Encoding）。这些组件共同构成了完整的 Transformer 模型架构[^2]。 以下是一个简单的代码示例，展示如何使用 PyTorch 内置的 `nn.Transformer` 模块： ```python import torch import torch.nn as nn class TransformerModel(nn.Module): def __init__(self, vocab_size, d_model, nhead, num_encoder_layers, num_decoder_layers, dim_feedforward, dropout=0.1): super(TransformerModel, self).__init__() self.embedding = nn.Embedding(vocab_size, d_model) self.transformer = nn.Transformer(d_model, nhead, num_encoder_layers, num_decoder_layers, dim_feedforward, dropout) self.fc_out = nn.Linear(d_model, vocab_size) def forward(self, src, tgt): src_embed = self.embedding(src) * torch.sqrt(torch.tensor(src.shape[-1], dtype=torch.float32)) tgt_embed = self.embedding(tgt) * torch.sqrt(torch.tensor(tgt.shape[-1], dtype=torch.float32)) output = self.transformer(src_embed, tgt_embed) output = self.fc_out(output) return output ``` #### 参数说明 - `vocab_size`: 词汇表大小。 - `d_model`: 输入和输出张量的特征维度。 - `nhead`: 多头注意力机制中的头数。 - `num_encoder_layers`: 编码器堆叠的层数。 - `num_decoder_layers`: 解码器堆叠的层数。 - `dim_feedforward`: 前馈神经网络的隐藏层维度。 - `dropout`: Dropout 的概率[^2]。 #### 训练与推理 在训练过程中，目标序列需要进行移位操作，以确保模型不会看到当前时间步之后的信息。推理时，则可以通过自回归方式逐步生成目标序列。 ```python # 示例训练代码 src = torch.randint(0, vocab_size, (seq_len, batch_size)) tgt = torch.randint(0, vocab_size, (seq_len, batch_size)) model = TransformerModel(vocab_size, d_model, nhead, num_encoder_layers, num_decoder_layers, dim_feedforward) criterion = nn.CrossEntropyLoss() optimizer = torch.optim.Adam(model.parameters(), lr=0.0001) output = model(src, tgt[:-1]) loss = criterion(output.view(-1, vocab_size), tgt[1:].reshape(-1)) loss.backward() optimizer.step() ``` #### 高级功能 除了基础的 Transformer 模型外，PyTorch 还支持多种变体，例如： - **Transformer Encoder-Only**: 仅使用编码器部分，适用于任务如文本分类或序列标注[^2]。 - **Pretrained Models**: 使用 Hugging Face 的 `transformers` 库可以加载预训练的 Transformer 模型，如 BERT、GPT 等[^5]。 ```python from transformers import AutoTokenizer, AutoModel tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased") model = AutoModel.from_pretrained("bert-base-uncased") ```