transformer编码器代码

以下是一个简单的 Transformer 编码器的示例代码：

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

class TransformerEncoderLayer(nn.Module):
    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1):
        super(TransformerEncoderLayer, self).__init__()
        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
        self.linear1 = nn.Linear(d_model, dim_feedforward)
        self.linear2 = nn.Linear(dim_feedforward, d_model)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, src, src_mask=None):
        src2 = self.self_attn(src, src, src, attn_mask=src_mask)[0]
        src = src + self.dropout(src2)
        src = self.norm1(src)
        src2 = self.linear2(F.relu(self.linear1(src)))
        src = src + self.dropout(src2)
        src = self.norm2(src)
        return src

class TransformerEncoder(nn.Module):
    def __init__(self, encoder_layer, num_layers):
        super(TransformerEncoder, self).__init__()
        self.layers = nn.ModuleList([encoder_layer for _ in range(num_layers)])

    def forward(self, src, mask=None):
        output = src
        for layer in self.layers:
            output = layer(output, mask)
        return output

# 使用示例
d_model = 512  # 模型维度
nhead = 8  # 多头注意力的头数
dim_feedforward = 2048  # 前馈网络隐藏层维度
dropout = 0.1  # 丢弃率
num_layers = 6  # 编码器层数

encoder_layer = TransformerEncoderLayer(d_model, nhead, dim_feedforward, dropout)
transformer_encoder = TransformerEncoder(encoder_layer, num_layers)

src = torch.randn(10, 32, d_model)  # 输入数据，形状为 (sequence_length, batch_size, d_model)
output = transformer_encoder(src)

这是一个简化的实现，仅包含 Transformer 编码器的前向传播。你可以根据自己的需要进行修改和扩展。