Transformer模型代码实现

以下是一个简单的Transformer模型的代码实现：

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

class Transformer(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim, num_layers, num_heads, dropout):
        super().__init__()
        self.embedding = nn.Embedding(input_dim, hidden_dim)
        self.pos_encoding = PositionalEncoding(hidden_dim, dropout)
        self.encoder_layers = nn.ModuleList([EncoderLayer(hidden_dim, num_heads, dropout) for _ in range(num_layers)])
        self.decoder_layers = nn.ModuleList([DecoderLayer(hidden_dim, num_heads, dropout) for _ in range(num_layers)])
        self.fc_out = nn.Linear(hidden_dim, output_dim)
        self.dropout = nn.Dropout(dropout)
        self.scale = torch.sqrt(torch.FloatTensor([hidden_dim])).to(device)

    def forward(self, src, trg, src_mask, trg_mask):
        src_emb = self.dropout(self.embedding(src) * self.scale)
        trg_emb = self.dropout(self.embedding(trg) * self.scale)
        src_emb = self.pos_encoding(src_emb)
        trg_emb = self.pos_encoding(trg_emb)
        for layer in self.encoder_layers:
            src_emb = layer(src_emb, src_mask)
        for layer in self.decoder_layers:
            trg_emb = layer(trg_emb, src_emb, trg_mask, src_mask)
        output = self.fc_out(trg_emb)
        return output

这个模型包含了一个嵌入层、位置编码层、多层编码器和多层解码器。它可以用于序列到序列的任务，如机器翻译。