pytorch下复现transformer完成分类

### 使用 PyTorch 实现 Transformer 模型完成文本分类任务 为了构建一个用于文本分类的Transformer模型，需要定义几个核心组件：嵌入层、位置编码、多头注意力机制、前馈神经网络以及最终的分类器。下面展示了一个简化版的实现方法。 #### 定义超参数 ```python import torch.nn as nn import math class Config: src_vocab_size = 10000 # 假设源词汇表大小为10000 tgt_vocab_size = 2 # 对于二元分类问题，目标词汇表大小通常设置为目标类别数 d_model = 512 # 模型尺寸 num_heads = 8 # 头的数量 num_layers = 6 # 编码器堆叠层数 d_ff = 2048 # FeedForward层中的隐藏单元数量 max_seq_length = 512 # 序列的最大长度 dropout_rate = 0.1 # Dropout比率 config = Config() ``` #### 构建Transformer类 ```python class PositionalEncoding(nn.Module): def __init__(self, d_model: int, dropout: float = 0.1, max_len: int = 5000): super().__init__() self.dropout = nn.Dropout(p=dropout) position = torch.arange(max_len).unsqueeze(1) div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model)) pe = torch.zeros(max_len, 1, d_model) pe[:, 0, 0::2] = torch.sin(position * div_term) pe[:, 0, 1::2] = torch.cos(position * div_term) self.register_buffer('pe', pe) def forward(self, x): """ Args: x: Tensor, shape [seq_len, batch_size, embedding_dim] """ x = x + self.pe[:x.size(0)] return self.dropout(x) class TransformerModel(nn.Module): def __init__(cls, config): super().__init__() cls.src_word_embedding = nn.Embedding(config.src_vocab_size, config.d_model) cls.pos_encoder = PositionalEncoding(d_model=config.d_model, dropout=config.dropout_rate, max_len=config.max_seq_length) cls.transformer = nn.Transformer( d_model=config.d_model, nhead=config.num_heads, num_encoder_layers=config.num_layers, dim_feedforward=config.d_ff, dropout=config.dropout_rate ) cls.fc_out = nn.Linear(config.d_model, config.tgt_vocab_size) cls.init_weights() def init_weights(cls): initrange = 0.1 cls.src_word_embedding.weight.data.uniform_(-initrange, initrange) cls.fc_out.bias.data.zero_() cls.fc_out.weight.data.uniform_(-initrange, initrange) def forward(cls, src, src_mask=None): embedded_src = cls.src_word_embedding(src) * math.sqrt(cls.config.d_model) encoded_src = cls.pos_encoder(embedded_src) output = cls.transformer.encoder(encoded_src, src_key_padding_mask=src_mask) pooled_output = output.mean(dim=1) # 平均池化作为句子表示 logits = cls.fc_out(pooled_output) return logits ``` 此代码片段展示了如何基于`nn.Transformer`来创建自定义的Transformer模型[^3]。注意这里只实现了编码端因为解码端对于分类任务来说不是必需的部分；另外通过平均池化的方式获取整个输入序列的一个固定长度向量表示，并将其传递给全连接层以获得最后的概率分布。