能否生成transformer的torch代码

当然可以！以下是一个简单的示例代码，用于生成一个基本的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, num_heads, num_layers):
        super(Transformer, self).__init__()
        self.embedding = nn.Embedding(input_dim, hidden_dim)
        self.position_encoding = PositionalEncoding(hidden_dim)
        self.encoder = Encoder(hidden_dim, num_heads, num_layers)
        self.decoder = Decoder(hidden_dim, num_heads, num_layers)
        self.fc = nn.Linear(hidden_dim, input_dim)

    def forward(self, x):
        x = self.embedding(x)
        x = self.position_encoding(x)
        encoded = self.encoder(x)
        decoded = self.decoder(encoded)
        output = self.fc(decoded)
        return output


class PositionalEncoding(nn.Module):
    def __init__(self, hidden_dim, max_length=5000):
        super(PositionalEncoding, self).__init__()
        self.hidden_dim = hidden_dim
        pe = torch.zeros(max_length, hidden_dim)
        position = torch.arange(0, max_length, 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 * math.sqrt(self.hidden_dim)
        x = x + self.pe[:x.size(0), :]
        return x


class MultiHeadAttention(nn.Module):
    def __init__(self, hidden_dim, num_heads):
        super(MultiHeadAttention, self).__init__()
        self.hidden_dim = hidden_dim
        self.num_heads = num_heads
        self.head_dim = hidden_dim // num_heads

        self.query_linear = nn.Linear(hidden_dim, hidden_dim)
        self.key_linear = nn.Linear(hidden_dim, hidden_dim)
        self.value_linear = nn.Linear(hidden_dim, hidden_dim)
        self.fc = nn.Linear(hidden_dim, hidden_dim)

    def forward(self, query, key, value):
        batch_size = query.size(0)

        query = self.query_linear(query).view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)
        key = self.key_linear(key).view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)
        value = self.value_linear(value).view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)

        scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(self.head_dim)
        attention_weights = F.softmax(scores, dim=-1)

        attended_values = torch.matmul(attention_weights, value)
        attended_values = attended_values.transpose(1, 2).contiguous().view(batch_size, -1, self.hidden_dim)

        output = self.fc(attended_values)
        return output


class FeedForward(nn.Module):
    def __init__(self, hidden_dim, feedforward_dim):
        super(FeedForward, self).__init__()
        self.fc1 = nn.Linear(hidden_dim, feedforward_dim)
        self.fc2 = nn.Linear(feedforward_dim, hidden_dim)

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


class Encoder(nn.Module):
    def __init__(self, hidden_dim, num_heads, num_layers):
        super(Encoder, self).__init__()
        self.layers = nn.ModuleList([EncoderLayer(hidden_dim, num_heads) for _ in range(num_layers)])

    def forward(self, x):
        for layer in self.layers:
            x = layer(x)
        return x


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

    def forward(self, x):
        attended = self.multihead_attention(x, x, x)
        encoded = attended + x
        output = self.feedforward(encoded)
        output = output + encoded
        return output


class Decoder(nn.Module):
    def __init__(self, hidden_dim, num_heads, num_layers):
        super(Decoder, self).__init__()
        self.layers = nn.ModuleList([DecoderLayer(hidden_dim, num_heads) for _ in range(num_layers)])

    def forward(self, x):
        for layer in self.layers:
            x = layer(x)
        return x


class DecoderLayer(nn.Module):
    def __init__(self, hidden_dim, num_heads):
        super(DecoderLayer, self).__init__()
        self.masked_multihead_attention = MultiHeadAttention(hidden_dim, num_heads)
        self.multihead_attention = MultiHeadAttention(hidden_dim, num_heads)
        self.feedforward = FeedForward(hidden_dim, hidden_dim)

    def forward(self, x):
        masked_attended = self.masked_multihead_attention(x, x, x)
        attended = self.multihead_attention(masked_attended, x, x)
        encoded = attended + x
        output = self.feedforward(encoded)
        output = output + encoded
        return output

这是一个简单的Transformer模型实现，包括位置编码、多头注意力机制、前馈神经网络等组件。你可以根据需要进行修改和扩展。