transformer与lstm 代码

以下是使用PyTorch实现Transformer和LSTM的代码示例：

Transformer:

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

class TransformerEncoder(nn.Module):
    def __init__(self, input_dim, hidden_dim, num_layers, num_heads, dropout):
        super(TransformerEncoder, self).__init__()
        self.input_dim = input_dim
        self.hidden_dim = hidden_dim
        self.num_layers = num_layers
        self.num_heads = num_heads
        self.dropout = dropout
        
        self.embedding = nn.Embedding(input_dim, hidden_dim)
        self.pos_embedding = nn.Embedding(1000, hidden_dim)
        self.layers = nn.ModuleList([TransformerEncoderLayer(hidden_dim, num_heads, dropout) for _ in range(num_layers)])
        self.norm = nn.LayerNorm(hidden_dim)
        
    def forward(self, x):
        seq_len = x.size(1)
        pos = torch.arange(0, seq_len).unsqueeze(0).repeat(x.size(0), 1).to(x.device)
        x = self.embedding(x) * math.sqrt(self.hidden_dim)
        pos = self.pos_embedding(pos)
        x = x + pos
        x = F.dropout(x, p=self.dropout, training=self.training)
        
        for layer in self.layers:
            x = layer(x)
            
        x = self.norm(x)
        
        return x

class TransformerEncoderLayer(nn.Module):
    def __init__(self, hidden_dim, num_heads, dropout):
        super(TransformerEncoderLayer, self).__init__()
        self.hidden_dim = hidden_dim
        self.num_heads = num_heads
        self.dropout = dropout
        
        self.self_attn = nn.MultiheadAttention(hidden_dim, num_heads)
        self.feed_forward = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim * 4),
            nn.ReLU(),
            nn.Linear(hidden_dim * 4, hidden_dim)
        )
        self.norm1 = nn.LayerNorm(hidden_dim)
        self.norm2 = nn.LayerNorm(hidden_dim)
        
    def forward(self, x):
        residual = x
        x, _ = self.self_attn(x, x, x)
        x = F.dropout(x, p=self.dropout, training=self.training)
        x = self.norm1(residual + x)
        
        residual = x
        x = self.feed_forward(x)
        x = F.dropout(x, p=self.dropout, training=self.training)
        x = self.norm2(residual + x)
        
        return x

LSTM:

import torch
import torch.nn as nn

class LSTMEncoder(nn.Module):
    def __init__(self, input_dim, hidden_dim, num_layers, dropout):
        super(LSTMEncoder, self).__init__()
        self.input_dim = input_dim
        self.hidden_dim = hidden_dim
        self.num_layers = num_layers
        self.dropout = dropout
        
        self.embedding = nn.Embedding(input_dim, hidden_dim)
        self.lstm = nn.LSTM(hidden_dim, hidden_dim, num_layers=num_layers, dropout=dropout, batch_first=True)
        
    def forward(self, x):
        x = self.embedding(x)
        output, _ = self.lstm(x)
        
        return output[:, -1, :]