transformer完整的时序程序pytorch

以下是一个完整的Transformer模型的时序程序示例，使用PyTorch实现：

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

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, n_heads):
        super(MultiHeadAttention, self).__init__()
        self.n_heads = n_heads
        self.d_model = d_model
        
        self.q_linear = nn.Linear(d_model, d_model)
        self.v_linear = nn.Linear(d_model, d_model)
        self.k_linear = nn.Linear(d_model, d_model)
        self.out_linear = nn.Linear(d_model, d_model)
        
    def split_heads(self, x, batch_size):
        """
        Split the last dimension into (n_heads, d_model) and transpose the result.
        """
        x = x.view(batch_size, -1, self.n_heads, self.d_model // self.n_heads)
        return x.transpose(1,2)
    
    def forward(self, q, k, v):
        batch_size = q.size(0)
        
        q = self.q_linear(q)
        k = self.k_linear(k)
        v = self.v_linear(v)
        
        q = self.split_heads(q, batch_size)
        k = self.split_heads(k, batch_size)
        v = self.split_heads(v, batch_size)
        
        scores = torch.matmul(q, k.transpose(-2, -1)) / torch.sqrt(torch.tensor(self.d_model / self.n_heads))
        scores = F.softmax(scores, dim=-1)
        x = torch.matmul(scores, v)
        
        x = x.transpose(1,2).contiguous().view(batch_size, -1, self.d_model)
        x = self.out_linear(x)
        
        return x

class PositionwiseFeedforward(nn.Module):
    def __init__(self, d_model, d_ff):
        super(PositionwiseFeedforward, self).__init__()
        
        self.linear1 = nn.Linear(d_model, d_ff)
        self.linear2 = nn.Linear(d_ff, d_model)
        
    def forward(self, x):
        x = self.linear1(x)
        x = F.relu(x)
        x = self.linear2(x)
        
        return x

class EncoderLayer(nn.Module):
    def __init__(self, d_model, n_heads, d_ff):
        super(EncoderLayer, self).__init__()
        
        self.self_attn = MultiHeadAttention(d_model, n_heads)
        self.ffn = PositionwiseFeedforward(d_model, d_ff)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        
    def forward(self, x):
        x_norm = self.norm1(x)
        attn_output = self.self_attn(x_norm, x_norm, x_norm)
        x = x + attn_output
        x_norm = self.norm2(x)
        ffn_output = self.ffn(x_norm)
        x = x + ffn_output
        
        return x

class TransformerEncoder(nn.Module):
    def __init__(self, vocab_size, max_len, d_model, n_layers, n_heads, d_ff):
        super(TransformerEncoder, self).__init__()
        
        self.d_model = d_model
        self.n_layers = n_layers
        
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.position_encoding = nn.Embedding(max_len, d_model)
        
        self.layers = nn.ModuleList([EncoderLayer(d_model, n_heads, d_ff) for _ in range(n_layers)])
        
    def forward(self, x):
        positions = torch.arange(0, x.size(1), device=x.device).unsqueeze(0).repeat(x.size(0), 1)
        x = self.embedding(x) * math.sqrt(self.d_model) + self.position_encoding(positions)
        
        for i in range(self.n_layers):
            x = self.layers[i](x)
            
        return x

在这个示例中，我们定义了一个TransformerEncoder模型，它由多个EncoderLayer组成。每个EncoderLayer包括一个MultiHeadAttention层和一个PositionwiseFeedforward层。我们还定义了一个MultiHeadAttention模块，它用于计算注意力分数，并将结果与值向量相乘，产生注意力输出。我们还定义了一个PositionwiseFeedforward模块，它由两个线性层组成，用于在每个位置上对特征进行前向传递。

在模型的前向传递中，我们首先对输入进行嵌入，并添加位置编码。然后，我们按顺序对每个EncoderLayer进行前向传递。最后，我们返回编码器的输出。

这个示例中使用了PyTorch的nn.Module类来定义模型，并提供了一些方便的函数和类来构建模型。关键的组件包括nn.Linear，nn.Embedding，nn.LayerNorm和nn.ModuleList。为了提高效率，我们还使用了PyTorch的自动求导功能来计算梯度。