基于pytorch简易实现gpt1网站首页 学无止境

基于pytorch简易实现gpt1

Rivennnnnnnnn 2026-07-09 00:01:05

简介基于pytorch简易实现gpt1

最近学习了GPT1模型，感觉模型相对来说比较简单，参数也不多。于是就根据模型结构用pytorch简单实现了一下，执行了一个机器翻译的任务。

GPT模型架构

在这里插入图片描述
上面就是GPT的架构了，基本上就是Transformer里面的解码器，可以看到多头注意力是用的掩码多头注意力模块。具体GPT和Transformer的学习可以参考别的视频和文章，Transformer可以看一下b站王木头的视频，再去看看原论文的解读。本文主要是给出基于pytorch的代码实现。

Embedding模块

首先是实现一下这个Embedding模块，根据架构图，GPT需要用到Position Embed和Text Embed也就是一个位置编码和词嵌入。位置编码用来提供序列的位置信息，词嵌入用来学习词和词之间的关系。

import torch
import torch.nn as nn
class GPT_Embedding(nn.Module):
    def __init__(self, maxlen,d_model,vocab_size):
        super(GPT_Embedding, self).__init__()
        self.positionEmbedding = nn.Embedding(maxlen, d_model)
        self.wordEmbedding = nn.Embedding(vocab_size, d_model)

    def forward(self, x):
        position=torch.arange(0, x.size(1), dtype=torch.long, device=x.device)
        position=self.positionEmbedding(position)
        x=self.wordEmbedding(x)
        return x+position

具体代码如上所示，GPT采用的是可学习的位置编码，代码实现的话就是一个Embedding层。maxlen是手动设置的最大序列长度，d_model是词向量的维度。wordEmbedding就是经典的词向量，学习词和词之间的关系，并用d_model维的向量表示，vocab_size是词汇表的长度，就是你的序列在进入GPT之前，就应该完成文字到数字之间的转换，比如你对模型输入"今天天气如何？"，"今天天气如何？"这句话要根据你的词汇表来进行文字到数字的转换，词汇表中记录的是文字对应的数字索引，比如"今"对应数字0，"天"对应1，这样。整句话转换成数字就可能是0，1，1，2，3，4，5这样一个序列。wordEmbedding学习完成后，记录的就是一张vocab_size*d_model的表，每一行代表对应索引的词向量。PositionEmbedding学习的就是位置信息的"词向量”。
然后看到forward函数，生成position，也就是0到输入序列长度的索引，比如输入序列长度是20，就生成0，1，2，3，4，5…20这样一个数组。然后放入positionEmbedding取出对应的向量。同理也需要取出输入对应的词向量表示，最后将词向量加上位置信息。
一般来说输入的长度就是最大的序列长度，长度不足的也会进行填充。

Decoder

这里实现的是解码器模块，代码如下所示

import torch
from torch import nn
from torch.cuda import device


class GPT_Attention(nn.Module):
    def __init__(self, d_model,n_head,d_ff,dropout):
        super(GPT_Attention, self).__init__()
        self.d_model = d_model
        self.n_head = n_head
        self.d_ff = d_ff
        self.dropout = dropout
        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.attn = nn.MultiheadAttention(d_model, n_head, dropout=dropout,batch_first=True)
        self.layer_norm1 = nn.LayerNorm(d_model)
        self.layer_norm2 = nn.LayerNorm(d_model)
        self.attn_weights=None
        self.feed_forward = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.Dropout(dropout),
            nn.Tanh(),
            nn.Linear(d_ff, d_model),
            nn.Dropout(dropout),
        )
    def forward(self, x):
        Q=self.wq(x)
        K=self.wk(x)
        V=self.wv(x)
        attn_mask=torch.triu(torch.ones(x.shape[1],x.shape[1]),diagonal=1).bool()
        attn_mask=attn_mask.to(x.device)
        attn_output,attn_weights=self.attn(Q,K,V,attn_mask=attn_mask,need_weights=True)
        self.attn_weights=attn_weights
        attn_output=attn_output+x#残差连接
        attn_output=self.layer_norm1(attn_output)
        attn_output=attn_output+self.feed_forward(attn_output)
        attn_output=self.layer_norm2(attn_output)
        return attn_output

这里我把注意力和前馈层放在一起实现了。d_model是词向量维度，n_head是注意力的头数，d_ff是前馈层中间输出的维度，dropout懂得都懂。wq,wk,wv是注意力机制里面Q，K，V的权重矩阵。这里的代码为什么不直接用一个d_model*d_model的矩阵来表示wq，wk，wv呢？显然nn.Linear本质就是学习一个系数矩阵，跟直接用一个矩阵和输入相乘区别不大，但又显然nn.Linear里面会有更多的策略，比如初始化，偏置项这些…
attn就是多头注意力模块，这里直接用torch实现的MultiheadAttention
feed_forward就是前馈层，由线性层，非线性函数，线性层组成。
然后就是forward函数，首先就是通过线性层，得到Q,K,V。生成掩码，然后将这些传入attn得到注意力得分。需要注意的是，torch实现的MultiheadAttention，掩码是为1的位置会屏蔽。所以这里生成的掩码是上三角矩阵，上三角区域全都1，其他为0。接下来就没什么好说的了，根据结构做好残差连接就行了。

生成器

首先提一嘴，按照架构图，最后得到的输出应该是seq_lend_model的矩阵，所以还需要做一次额外的映射，再乘一个d_modelvocab_size的矩阵，最后得到seq_len*vocab_size的矩阵。这样就相当于得到了所有词汇在相对位置的得分。比如有词汇0，1，2，3。在第一个位置，0的得分可能是20，1的得分为30，2的得分为99这样。当然最后还需要做一次softmax转成概率，但这样不需要做，因为后续的损失函数会自动进行。

from torch import nn


class Generator(nn.Module):
    def __init__(self,d_model,vocab_size):
        super(Generator,self).__init__()
        self.d_model=d_model
        self.vocab_size=vocab_size
        self.linear=nn.Linear(d_model,vocab_size)
    def forward(self,x):
        return self.linear(x)

GPT组装

这里就直接贴代码了，我代码中是把层数写死了，大家可以自己改一下

import copy

from torch import nn

from GPT_Attention import GPT_Attention
from GPT_Embedding import GPT_Embedding
from Generator import Generator

def getGPT_Attention(d_model,n_head,d_ff,dropout):
    return GPT_Attention(d_model,n_head,d_ff,dropout)
class GPT(nn.Module):
    def __init__(self, d_model, vocab_size,n_head,d_ff,max_len,dropout):
        super().__init__()
        self.d_model = d_model
        self.vocab_size = vocab_size
        self.n_head = n_head
        self.d_ff = d_ff
        self.dropout = dropout
        self.max_len = max_len
        self.generator=Generator(d_model,vocab_size)
        self.embedding = GPT_Embedding(max_len,d_model,vocab_size)
        def gpt_attetion_block(n):
            model_list=list()
            for i in range(n):
                model_list.append(getGPT_Attention(d_model,n_head,d_ff,dropout))
            return model_list
        self.GPT_block=nn.Sequential(*gpt_attetion_block(12))
    def forward(self,x):
        x = self.embedding(x)
        x = self.GPT_block(x)
        x = self.generator(x)
        return x

训练测试

这里用的一个机器翻译的数据集，我是直接抄的另一篇博客的数据，源自这一篇博客包括数据处理也是直接复制了这位博主的代码，稍后我会放出源码，大家可以自己去运行。数据部分的代码我就不贴出来了，重点是看看训练的过程。

import torch
from torch import nn
from transformers import BertTokenizer

import dataHandle
from GPT import GPT
from Generator import Generator
from dataHandle import data_loader


def main():
    file_path=r"./textData/train.txt"
    tokenizer=BertTokenizer.from_pretrained('bert-base-multilingual-uncased')
    loader,dataset=data_loader(file_path,16,tokenizer,128)
    model=GPT(768,tokenizer.vocab_size,12,3072,128,0.4)
    optimizer1=torch.optim.Adam(model.parameters(),lr=0.0001)
    loss_fn=nn.CrossEntropyLoss()
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    loss_fn.to(device)
    model.to(device)
    for epoch in range(1,100):
        total_loss=0
        for data in loader:
            x,target=data
            x=x.to(device,dtype=torch.long)
            target=target.to(device,dtype=torch.long)
            output=model(x)
            loss=loss_fn(output.view(-1,tokenizer.vocab_size),target.view(-1))
            total_loss+=loss.item()
            optimizer1.zero_grad()
            loss.backward()
            optimizer1.step()
        print(f"第{epoch}轮训练已完成,总损失为{total_loss}")
        torch.save(model.state_dict(), 'last_model.pt')
if __name__=="__main__":
    main()

训练的时候是把输入一次性全部输入，然后得到所有词汇在每个位置上的得分，最后用目标序列和得分做交叉熵。这里的输入是英文，比如how are you。那么目标就是你咋地。

推理

import torch
from transformers import BertTokenizer

from GPT import GPT
from dataHandle import data_loader

def id2String(Tokenizer:BertTokenizer,Input_id):
    result=""
    for Id in Input_id:
        result=result+' '+Tokenizer.decode(Id)
    return result

def main():
    tokenizer=BertTokenizer.from_pretrained('bert-base-multilingual-uncased')
    model=GPT(768,tokenizer.vocab_size,12,3072,128,0.4)
    model.load_state_dict(torch.load("./last_model.pt"))
    model.eval()
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    model.to(device)
    file_path = r"./textData/test.txt"
    loader, dataset = data_loader(file_path, 16, tokenizer, 128)
    for data in loader:
        x,y = data
        x = x.to(device)
        y = y.to(device)
        output = model(x)
        Input_id=output.argmax(dim=-1)
        for x_ids,y_ids in zip(Input_id,x):
            print(id2String(tokenizer,y_ids))
            print(id2String(tokenizer,x_ids))
if __name__ == '__main__':
    main()

推理的代码也贴出来，需要把得到的索引转化成字符串。可以看看我自己运行的效果
在这里插入图片描述
第一句是原始的英文，第二句是给出的翻译，以此类推…可以看到确实是有在翻译，也别管效果好不好吧。
我是在笔记本上运行的，应该是跑了四个小时左右，100轮训练。显卡是4070laptop
需要提一嘴就是，GPT1推荐的是这样的推理过程，比如你想做一个翻译任务，就需要这样输入:“请帮我将这个句子翻译成中文:how are you"这样，第一次gpt会输出一堆概率，但只选取其中概率最大的一个，然后将这个输出拼接输入，比如说gpt输出了你，那么我们第二次输入就是"请帮我将这个句子翻译成中文:how are you[拼接符号]你”,第二次gpt输出好，那么第三次输入就是"请帮我将这个句子翻译成中文:how are you[拼接符号]你好"。
上述过程就是自回归生成了，我自己没有去实测过，但对于我用的数据集，肯定是不支持我这样去输出的。有条件的可以去试试。
代码以及训练样本都在这里