使用pytorch写一段transformer的代码

我可以回答这个问题。以下是一个简单的Transformer模型的PyTorch代码示例：

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

class Transformer(nn.Module):
    def __init__(self, input_dim, hidden_dim, num_layers, num_heads):
        super(Transformer, self).__init__()
        self.embedding = nn.Embedding(input_dim, hidden_dim)
        self.pos_encoding = PositionalEncoding(hidden_dim)
        self.encoder_layers = nn.ModuleList([EncoderLayer(hidden_dim, num_heads) for _ in range(num_layers)])
        self.decoder_layers = nn.ModuleList([DecoderLayer(hidden_dim, num_heads) for _ in range(num_layers)])
        self.fc = nn.Linear(hidden_dim, input_dim)

    def forward(self, src, trg):
        src = self.embedding(src)
        trg = self.embedding(trg)
        src = self.pos_encoding(src)
        trg = self.pos_encoding(trg)
        for layer in self.encoder_layers:
            src = layer(src)
        for layer in self.decoder_layers:
            trg = layer(trg, src)
        output = self.fc(trg)
        return output

class PositionalEncoding(nn.Module):
    def __init__(self, hidden_dim, max_len=5000):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=0.1)
        pe = torch.zeros(max_len, hidden_dim)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, hidden_dim, 2).float() * (-math.log(10000.0) / hidden_dim))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0).transpose(0, 1)
        self.register_buffer('pe', pe)

    def forward(self, x):
        x = x + self.pe[:x.size(0), :]
        return self.dropout(x)

class EncoderLayer(nn.Module):
    def __init__(self, hidden_dim, num_heads):
        super(EncoderLayer, self).__init__()
        self.self_attn = MultiHeadAttention(hidden_dim, num_heads)
        self.feed_forward = FeedForward(hidden_dim)

    def forward(self, x):
        x = x + self.self_attn(x)
        x = x + self.feed_forward(x)
        return x

class DecoderLayer(nn.Module):
    def __init__(self, hidden_dim, num_heads):
        super(DecoderLayer, self).__init__()
        self.self_attn = MultiHeadAttention(hidden_dim, num_heads)
        self.src_attn = MultiHeadAttention(hidden_dim, num_heads)
        self.feed_forward = FeedForward(hidden_dim)

    def forward(self, x, src):
        x = x + self.self_attn(x)
        x = x + self.src_attn(x, src)
        x = x + self.feed_forward(x)
        return x

class MultiHeadAttention(nn.Module):
    def __init__(self, hidden_dim, num_heads):
        super(MultiHeadAttention, self).__init__()
        self.num_heads = num_heads
        self.head_dim = hidden_dim // num_heads
        self.query = nn.Linear(hidden_dim, hidden_dim)
        self.key = nn.Linear(hidden_dim, hidden_dim)
        self.value = nn.Linear(hidden_dim, hidden_dim)
        self.fc = nn.Linear(hidden_dim, hidden_dim)

    def forward(self, x):
        batch_size = x.size(0)
        query = self.query(x).view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)
        key = self.key(x).view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)
        value = self.value(x).view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)
        attn_weights = F.softmax(torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(self.head_dim), dim=-1)
        attn_output = torch.matmul(attn_weights, value)
        attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, -1, self.num_heads * self.head_dim)
        output = self.fc(attn_output)
        return output

class FeedForward(nn.Module):
    def __init__(self, hidden_dim, dropout=0.1):
        super(FeedForward, self).__init__()
        self.fc1 = nn.Linear(hidden_dim, hidden_dim * 4)
        self.fc2 = nn.Linear(hidden_dim * 4, hidden_dim)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = self.dropout(x)
        x = self.fc2(x)
        return x

这段代码实现了一个Transformer模型，包括编码器和解码器。其中，编码器由多个EncoderLayer组成，解码器由多个DecoderLayer组成。每个EncoderLayer和DecoderLayer包含了多头自注意力和前馈网络。在模型的前向传播过程中，输入先经过嵌入层和位置编码层，然后分别经过编码器和解码器，最后通过全连接层输出结果。