一个 Transformer 训练生成式模型的例子

最近在看chatGPT,想着chatGPT 是怎么训练出来的，不涉及神经网络算法，可以使用Transformer玩一下

 import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.data import Dataset, DataLoader # 构造词表 vocab = ["<PAD>", "<BOS>", "<EOS>", "明天", "天气", "很", "好"] word2idx = {w: i for i, w in enumerate(vocab)} idx2word = {i: w for w, i in word2idx.items()} # 示例训练数据：输入和标签都偏移一位 # 输入: <BOS> 明天 天气 很 # 输出: 明天 天气 很 好 inputs = torch.tensor([ [word2idx["<BOS>"], word2idx["明天"], word2idx["天气"], word2idx["很"]] ]) labels = torch.tensor([ [word2idx["明天"], word2idx["天气"], word2idx["很"], word2idx["好"]] ]) class TinyTransformerModel(nn.Module): def __init__(self, vocab_size, d_model=32, nhead=2, num_layers=2): super().__init__() self.embedding = nn.Embedding(vocab_size, d_model) self.pos_encoding = nn.Parameter(torch.randn(1, 100, d_model)) # 最多100个词 encoder_layer = nn.TransformerEncoderLayer(d_model, nhead) self.transformer = nn.TransformerEncoder(encoder_layer, num_layers) self.fc = nn.Linear(d_model, vocab_size) def forward(self, x): x = self.embedding(x) + self.pos_encoding[:, :x.size(1)] x = self.transformer(x) return self.fc(x) model = TinyTransformerModel(vocab_size=len(vocab)) optimizer = torch.optim.Adam(model.parameters(), lr=0.01) loss_fn = nn.CrossEntropyLoss() for epoch in range(50): model.train() out = model(inputs) # [batch_size, seq_len, vocab_size] loss = loss_fn(out.view(-1, len(vocab)), labels.view(-1)) optimizer.zero_grad() loss.backward() optimizer.step() if epoch % 20 == 0: print(f"Epoch {epoch}, loss: {loss.item():.4f}") def generate(model, start_tokens, max_len=5): model.eval() input_ids = torch.tensor([start_tokens]) for _ in range(max_len): with torch.no_grad(): logits = model(input_ids) next_token = logits[0, -1].argmax().item() input_ids = torch.cat([input_ids, torch.tensor([[next_token]])], dim=1) if next_token == word2idx["<EOS>"] or len(input_ids[0]) > max_len: break return [idx2word[i] for i in input_ids[0].tolist()] # 测试 generated = generate(model, [word2idx["<BOS>"]]) print("生成结果：", " ".join(generated))