用pytorch搭建Transformer，完成一个简单任务，要求代码可运行，不要用torchtext库包，写出python代码

以下是使用PyTorch搭建Transformer的一个简单例子。 这个例子将通过Transformer模型将一系列输入整数加倍。

import math
import torch
import torch.nn as nn
import torch.optim as optim

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=5000):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=0.1)
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        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 TransformerModel(nn.Module):
    def __init__(self, ntoken, ninp, nhead, nhid, nlayers):
        super(TransformerModel, self).__init__()
        self.embedding = nn.Embedding(ntoken, ninp)
        self.pos_encoder = PositionalEncoding(ninp)
        encoder_layers = nn.TransformerEncoderLayer(ninp, nhead, nhid)
        self.transformer_encoder = nn.TransformerEncoder(encoder_layers, nlayers)
        self.decoder = nn.Linear(ninp, ntoken)
        self.init_weights()

    def init_weights(self):
        init_range = 0.1
        self.embedding.weight.data.uniform_(-init_range, init_range)
        self.decoder.bias.data.zero_()
        self.decoder.weight.data.uniform_(-init_range, init_range)

    def forward(self, x):
        x = self.embedding(x)
        x = self.pos_encoder(x)
        x = self.transformer_encoder(x)
        x = self.decoder(x)
        return x

def data_gen(batch_size=20, seq_len=10, limit=500):
    for _ in range(limit):
        data = torch.randint(1, 10, (batch_size, seq_len))
        targets = data * 2
        yield data, targets

if __name__ == "__main__":
    ntokens = 20
    emsize = 200
    nhead = 2
    nhid = 200
    nlayers = 2
    model = TransformerModel(ntokens, emsize, nhead, nhid, nlayers)

    criterion = nn.CrossEntropyLoss()
    lr = 0.001
    optimizer = optim.Adam(model.parameters(), lr=lr)

    num_epochs = 5
    for epoch in range(num_epochs):
        for i, (data, targets) in enumerate(data_gen()):
            optimizer.zero_grad()
            output = model(data)
            loss = criterion(output.view(-1, ntokens), targets.view(-1))
            loss.backward()
            optimizer.step()

            if i % 50 == 0:
                print(f"Epoch: {epoch}, Loss: {loss.item():.6f}")

    # Testing on some data
    test_data = torch.tensor([[3, 6, 9], [2, 4, 6]])
    print("Test input:", test_data)
    test_output = torch.argmax(model(test_data), dim=2)
    print("Test output:", test_output)

这个例子实现了一个简单的Transformer模型，该模型将输入序列中的每个整数乘以2。数据生成器data_gen函数生成用于训练的随机序列。在这个例子中，设置了一个较小的词汇表大小为20，表示单词无需太多。训练迭代次数为5，每50个批次输出损失。

在训练完成后，给出一个简单的测试样例。测试数据包括两个序列[3, 6, 9]和[2, 4, 6]。模型的输出是输入整数加倍的序列。请注意，这个模型并不是一个完美的解决方案，因为它只是识别这些特定的输入与目标之间的模式，并且可能对其他输入不太准确。这个例子只是为了展示如何搭建一个Transformer模型，并在实际任务上完成简单的训练。