写一个transformer模块代码

好的，以下是一个简单的Transformer模块的代码实现：

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


class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super().__init__()
        self.d_model = d_model
        self.num_heads = num_heads
        
        assert d_model % num_heads == 0
        
        self.depth = d_model // num_heads
        
        self.wq = nn.Linear(d_model, d_model)
        self.wk = nn.Linear(d_model, d_model)
        self.wv = nn.Linear(d_model, d_model)
        
        self.fc = nn.Linear(d_model, d_model)
        
    def forward(self, query, key, value, mask=None):
        batch_size = query.size(0)
        
        # Linearly project the queries, keys, and values
        Q = self.wq(query)
        K = self.wk(key)
        V = self.wv(value)
        
        # Split the queries, keys, and values into num_heads pieces
        Q = Q.view(batch_size * self.num_heads, -1, self.depth)
        K = K.view(batch_size * self.num_heads, -1, self.depth)
        V = V.view(batch_size * self.num_heads, -1, self.depth)
        
        # Calculate the dot products of the queries and keys
        scores = torch.bmm(Q, K.transpose(1, 2)) / torch.sqrt(torch.tensor(self.d_model).float())
        
        # Apply the mask (if provided)
        if mask is not None:
            mask = mask.unsqueeze(1)
            scores = scores.masked_fill(mask == 0, -1e9)
        
        # Apply softmax to the scores
        softmax_scores = F.softmax(scores, dim=-1)
        
        # Multiply the softmax scores by the values
        output = torch.bmm(softmax_scores, V)
        
        # Reshape the output and apply a linear layer
        output = output.view(batch_size, -1, self.num_heads * self.depth)
        output = self.fc(output)
        
        return output
        
        
class PositionwiseFeedForward(nn.Module):
    def __init__(self, d_model, d_ff):
        super().__init__()
        self.linear1 = nn.Linear(d_model, d_ff)
        self.linear2 = nn.Linear(d_ff, d_model)
        
    def forward(self, x):
        x = F.relu(self.linear1(x))
        x = self.linear2(x)
        return x
        
        
class EncoderLayer(nn.Module):
    def __init__(self, d_model, num_heads, d_ff):
        super().__init__()
        self.multihead_attention = MultiHeadAttention(d_model, num_heads)
        self.norm1 = nn.LayerNorm(d_model)
        
        self.positionwise_feedforward = PositionwiseFeedForward(d_model, d_ff)
        self.norm2 = nn.LayerNorm(d_model)
        
    def forward(self, x, mask=None):
        # Apply multi-head attention
        attention_output = self.multihead_attention(x, x, x, mask=mask)
        # Add residual connection and apply layer normalization
        output = self.norm1(x + attention_output)
        
        # Apply position-wise feedforward network
        ff_output = self.positionwise_feedforward(output)
        # Add residual connection and apply layer normalization
        output = self.norm2(output + ff_output)
        
        return output
        
        
class Transformer(nn.Module):
    def __init__(self, input_dim, output_dim, d_model, num_heads, d_ff, num_layers, dropout):
        super().__init__()
        self.input_dim = input_dim
        self.output_dim = output_dim
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_ff = d_ff
        self.num_layers = num_layers
        self.dropout = dropout
        
        self.embedding = nn.Embedding(input_dim, d_model)
        self.position_encoding = nn.Embedding(1000, d_model)
        
        self.layers = nn.ModuleList([EncoderLayer(d_model, num_heads, d_ff) for _ in range(num_layers)])
        
        self.fc_out = nn.Linear(d_model, output_dim)
        
        self.dropout = nn.Dropout(dropout)
        
    def forward(self, src, src_mask=None):
        # Create positions tensor
        positions = torch.arange(0, src.shape[1]).unsqueeze(0).repeat(src.shape[0], 1).to(src.device)
        pos_enc = self.position_encoding(positions)
        
        # Get the embedded input
        embedded_input = self.embedding(src)
        
        # Add the positional encoding to the embedded input
        x = embedded_input + pos_enc
        
        x = self.dropout(x)
        
        # Apply the transformer layers
        for layer in self.layers:
            x = layer(x, mask=src_mask)
            
        # Apply the final output layer
        x = self.fc_out(x)
        
        return x

这个代码实现了一个简单的Transformer模块。其中包括了Multi-Head Attention、Position-wise Feedforward等子模块。可以通过调整参数来修改模型的深度、宽度和dropout等超参数。