轻量级的transformer代码

轻量级的transformer代码有多种实现方式，以下是其中一种示例代码：

import torch
import torch.nn as nn

class LightweightTransformer(nn.Module):
    def __init__(self, input_dim, hidden_dim, num_layers, num_heads):
        super(LightweightTransformer, self).__init__()
        self.embedding = nn.Linear(input_dim, hidden_dim)
        self.transformer = nn.TransformerEncoder(
            nn.TransformerEncoderLayer(hidden_dim, num_heads),
            num_layers
        )
        self.fc = nn.Linear(hidden_dim, input_dim)

    def forward(self, x):
        x = self.embedding(x)
        x = self.transformer(x)
        x = self.fc(x)
        return x

这是一个简单的轻量级transformer模型的代码示例。它包括一个线性层作为嵌入层，一个Transformer编码器作为主体，并在最后添加一个线性层进行预测。