Transformer编码器代码

Transformer编码器是一种用于自然语言处理任务的深度学习模型，它由多个相同结构的层组成，每个层都包含了自注意力机制和前馈神经网络。下面是一个简化的Transformer编码器的代码示例：

import torch
import torch.nn as nn

class TransformerEncoder(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers, num_heads, dropout):
        super(TransformerEncoder, self).__init__()
        self.embedding = nn.Embedding(input_size, hidden_size)
        self.position_encoding = PositionalEncoding(hidden_size)
        self.layers = nn.ModuleList([
            TransformerEncoderLayer(hidden_size, num_heads, dropout)
            for _ in range(num_layers)
        ])
        
    def forward(self, input):
        embedded = self.embedding(input)
        encoded = self.position_encoding(embedded)
        
        for layer in self.layers:
            encoded = layer(encoded)
        
        return encoded

class TransformerEncoderLayer(nn.Module):
    def __init__(self, hidden_size, num_heads, dropout):
        super(TransformerEncoderLayer, self).__init__()
        self.attention = MultiHeadAttention(hidden_size, num_heads)
        self.feed_forward = FeedForwardNetwork(hidden_size)
        self.dropout = nn.Dropout(dropout)
        self.layer_norm1 = nn.LayerNorm(hidden_size)
        self.layer_norm2 = nn.LayerNorm(hidden_size)
        
    def forward(self, input):
        attended = self.attention(input)
        attended = self.dropout(attended)
        attended = self.layer_norm1(input + attended)
        
        feed_forward_output = self.feed_forward(attended)
        feed_forward_output = self.dropout(feed_forward_output)
        output = self.layer_norm2(attended + feed_forward_output)
        
        return output

class MultiHeadAttention(nn.Module):
    def __init__(self, hidden_size, num_heads):
        super(MultiHeadAttention, self).__init__()
        self.hidden_size = hidden_size
        self.num_heads = num_heads
        self.head_size = hidden_size // num_heads
        
        self.query_projection = nn.Linear(hidden_size, hidden_size)
        self.key_projection = nn.Linear(hidden_size, hidden_size)
        self.value_projection = nn.Linear(hidden_size, hidden_size)
        self.output_projection = nn.Linear(hidden_size, hidden_size)
        
    def forward(self, input):
        batch_size, seq_len, hidden_size = input.size()
        
        query = self.query_projection(input).view(batch_size, seq_len, self.num_heads, self.head_size).transpose(1, 2)
        key = self.key_projection(input).view(batch_size, seq_len, self.num_heads, self.head_size).transpose(1, 2)
        value = self.value_projection(input).view(batch_size, seq_len, self.num_heads, self.head_size).transpose(1, 2)
        
        scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(self.head_size)
        attention_weights = nn.functional.softmax(scores, dim=-1)
        
        attended = torch.matmul(attention_weights, value).transpose(1, 2).contiguous().view(batch_size, seq_len, hidden_size)
        output = self.output_projection(attended)
        
        return output

class FeedForwardNetwork(nn.Module):
    def __init__(self, hidden_size):
        super(FeedForwardNetwork, self).__init__()
        self.hidden_size = hidden_size
        self.fc1 = nn.Linear(hidden_size, hidden_size * 4)
        self.fc2 = nn.Linear(hidden_size * 4, hidden_size)
        
    def forward(self, input):
        output = nn.functional.relu(self.fc1(input))
        output = self.fc2(output)
        
        return output

class PositionalEncoding(nn.Module):
    def __init__(self, hidden_size, max_seq_len=5000):
        super(PositionalEncoding, self).__init__()
        self.hidden_size = hidden_size
        
        position_encoding = torch.zeros(max_seq_len, hidden_size)
        position = torch.arange(0, max_seq_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, hidden_size, 2).float() * (-math.log(10000.0) / hidden_size))
        
        position_encoding[:, 0::2] = torch.sin(position * div_term)
        position_encoding[:, 1::2] = torch.cos(position * div_term)
        
        self.register_buffer('position_encoding', position_encoding)
        
    def forward(self, input):
        seq_len = input.size(1)
        position_encoding = self.position_encoding[:seq_len, :]
        
        return input + position_encoding

这段代码实现了一个简单的Transformer编码器，包括了多个TransformerEncoderLayer层，每个层由自注意力机制和前馈神经网络组成。同时还包括了位置编码(Positional Encoding)用于处理输入序列的位置信息。请注意，这只是一个简化的示例，实际应用中可能还需要进行更多的调整和优化。