实战指南:基于PyTorch与Hugging Face Transformers微调中文GPT-2进行诗歌生成

在当今人工智能浪潮中,生成式预训练模型正深刻改变着自然语言处理(NLP)的格局。GPT-2作为其中的经典代表,以其强大的自回归生成能力,为中文文本创作任务提供了卓越的解决方案。本文将深入浅出,手把手带你完成一个完整的中文诗歌生成模型的微调实战,涵盖从数据处理、模型加载到训练优化的全链路细节,为你的AI创作之旅提供一份清晰的路线图。

一、环境搭建与核心工具

开启我们的AI创作之旅,首先需要搭建一个稳定高效的开发环境。本项目将依托两大核心框架:PyTorch作为深度学习的基础引擎,以及Hugging Face Transformers库,它为我们提供了预训练模型、分词器等一站式NLP工具,极大地简化了开发流程。通过以下命令,我们可以快速安装所有必要的依赖包。

pip install torch transformers

安装完成后,我们就拥有了构建和训练神经网络模型所需的一切工具。接下来,我们将目光投向数据的准备。

二、数据准备:构建诗歌数据集

高质量的模型始于高质量的数据。对于中文诗歌生成任务,我们需要一个格式规整的文本数据集。本实战使用 chinese_poems.txt 文件作为训练数据源,该文件应放置在 data 目录下。文件中的每一行代表一首独立的诗歌。

为了高效地加载和处理这些文本数据,我们需要自定义一个PyTorch的 Dataset 类。这允许我们以流式、批量的方式将数据喂给模型。创建 data.py 文件,并实现一个继承自 torch.utils.data.Dataset 的类,其核心职责包括:

  • 数据读取与清洗:在 __init__ 方法中,读取文件并利用 strip() 等方法去除每行首尾的空格和无效字符,确保输入模型的文本干净、一致。
  • 提供数据长度:实现 __len__ 方法,返回数据集的总条目数,这是 Dataset 抽象类的要求,也为后续的数据分批提供依据。
  • 按索引获取数据:实现 __getitem__ 方法,根据给定的索引返回单条诗歌文本。

我们可以通过实例化这个数据集类并打印其长度和第一条数据来验证数据加载是否正常。

from torch.utils.data import Dataset
class MyDataset(Dataset):
    def __init__(self):
        with open("data/chinese_poems.txt", encoding="utf-8") as f:
            lines = f.readlines()
        # 去除每数据的前后空格
        lines = [i.strip() for i in lines]
        self.lines = lines
    def __len__(self):
        return len(self.lines)
    def __getitem__(self, item):
        return self.lines[item]
if __name__ == '__main__':
    dataset = MyDataset()
    print(len(dataset), dataset[0])
[AFFILIATE_SLOT_1]

三、模型加载与数据流水线构建

数据就绪后,下一步是加载预训练模型并构建高效的数据处理流水线。这一切将在 train.py 文件中完成。

3.1 加载分词器与模型

首先,我们需要加载专门针对中文优化的GPT-2分词器(Tokenizer)和模型。这里以ModelScope上的 gpt2-chinese-cluecorpussmall 模型为例。

from torch.optim import AdamW
from transformers.optimization import get_scheduler
import torch
from data import MyDataset
from transformers import AutoTokenizer
from transformers import AutoModelForCausalLM, GPT2Model
dataset = MyDataset()
model_dir = r"D:\pyprojecgt\flaskProject\langchainstudy\modelscope\gpt2-chinese-cluecorpussmall";
# 加载编码器
tokenizer = AutoTokenizer.from_pretrained(model_dir)
# 加载模型
model = AutoModelForCausalLM.from_pretrained(model_dir)
print(model)

关键组件解析:

  • AutoTokenizer:负责将人类可读的中文文本(如“明月几时有”)转换为模型能够理解的数字ID序列(Token IDs),是连接自然语言与神经网络模型的桥梁。
  • AutoModelForCausalLM:这是用于因果语言建模(即文本生成)的GPT-2模型。它内部已经集成了计算语言模型损失(根据上文预测下一个词)的逻辑,因此我们无需手动定义损失函数,非常方便。
  • GPT2Model:基础模型类,如果需要进行更复杂的模型结构定制(如添加额外的网络层),可以从这个类继承。

打印模型结构(print(model))是一个好习惯,可以快速确认模型是否加载成功。

3.2 实现批处理函数(collate_fn)

在批量训练时,一个批次内的文本长度往往不一致。我们需要一个函数来统一处理:将短文本填充(Padding)到相同长度,将长文本截断(Truncation),并正确设置训练标签。这就是 collate_fn 函数的作用。

def collate_fn(data):
    data = tokenizer.batch_encode_plus(data,
                                       padding=True,
                                       truncation=True,
                                       max_length=512,
                                       return_tensors='pt')
    data['labels'] = data['input_ids'].clone()
    return data

核心原理:

  • 使用分词器的 batch_encode_plus 方法对一批文本进行编码。padding=True 确保批次内长度一致,truncation=True 处理超长文本,return_tensors='pt' 返回PyTorch张量。
  • 标签设置:在自回归语言模型中,训练目标是“给定前文,预测下一个词”。因此,标签(labels)通常就是输入ID(input_ids)本身,只是模型在计算损失时会自动进行偏移对齐(例如,用 input_ids[:-1] 来预测 labels[1:])。

3.3 组装数据加载器(DataLoader)

最后,将我们自定义的数据集(Dataset)和批处理函数(collate_fn)结合起来,构建最终的数据加载器。

# 数据加载器
loader = torch.utils.data.DataLoader(
    dataset=dataset,
    batch_size=4,
    collate_fn=collate_fn,
    shuffle=True,
    drop_last=True,
)
print(len(loader))

参数说明:

  • batch_size=4:每个训练批次包含4首诗歌。
  • shuffle=True:在每个训练周期(Epoch)开始时打乱数据顺序,防止模型学到无关的序列信息。
  • drop_last=True:当数据集大小不能被批次大小整除时,丢弃最后一个不完整的批次,保证每次输入的张量形状一致。

打印 print(len(loader)) 可以确认数据加载器构建正确。

四、训练循环:核心逻辑深度解析

万事俱备,只欠训练。训练循环是机器学习项目的核心,它定义了模型如何从数据中学习。以下是完整的训练函数实现:

# 训练
def train():
    global model
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    model = model.to(device)
#加速训练
    optimizer = AdamW(model.parameters(), lr=2e-5)
    scheduler = get_scheduler(name='linear',
                              num_warmup_steps=0,
                              num_training_steps=len(loader),
                              optimizer=optimizer)
    model.train()
    for i, data in enumerate(loader):
        for k in data.keys():
            data[k] = data[k].to(device)
        out = model(**data)
        loss = out['loss']
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        optimizer.step()
        scheduler.step()
        optimizer.zero_grad()
        if i % 50 == 0:
            labels = data['labels'][:, 1:]
            out = out['logits'].argmax(dim=2)[:, :-1]
            select = labels != 0
            labels = labels[select]
            out = out[select]
            del select
            accuracy = (labels == out).sum().item() / labels.numel()
            lr = optimizer.state_dict()['param_groups'][0]['lr']
            print(i, loss.item(), lr, accuracy)
    # 保存模型参数,未保存模型结构
    torch.save(model.state_dict(), f='net.pt')
    print("权重保存成功!")
if __name__ == '__main__':
    for epoch in range(1000):
        train()

让我们深入拆解其中的关键技术点:

4.1 设备部署与优化器配置

首先,我们需要决定模型在GPU还是CPU上运行,以充分利用硬件资源。

device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = model.to(device)

接着,配置优化器和学习率调度器,这是控制模型学习步伐的关键。

  • 优化器:使用 AdamW,它在Adam算法的基础上加入了权重衰减(Weight Decay),能有效防止模型过拟合。初始学习率设为 2e-5,这是微调大型预训练模型的常用值。
  • 学习率调度:使用线性调度器(get_linear_schedule_with_warmup)。这里设置 num_warmup_steps=0,表示没有预热阶段,学习率从初始值开始,在总步数(num_training_steps=len(loader))内线性下降到0,有助于训练后期稳定收敛。

4.2 梯度计算与梯度裁剪

在每一批数据的前向传播计算损失后,需要进行反向传播来更新模型参数。

loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

loss.backward() 执行反向传播,计算损失相对于每个模型参数的梯度。然而,深度Transformer模型(如GPT-2)在训练时容易遇到梯度爆炸问题,即梯度值变得异常巨大,导致训练不稳定甚至失败。梯度裁剪(clip_grad_norm_)通过限制梯度向量的范数(这里限制L2范数不超过1.0)来有效缓解这一问题。

4.3 训练监控与模型保存

为了实时了解训练状态,我们需要定期打印关键指标。

if i % 50 == 0:
    labels = data['labels'][:, 1:]
    out = out['logits'].argmax(dim=2)[:, :-1]
    select = labels != 0
    labels = labels[select]
    out = out[select]
    del select
    accuracy = (labels == out).sum().item() / labels.numel()
    lr = optimizer.state_dict()['param_groups'][0]['lr']
    print(i, loss.item(), lr, accuracy)

指标计算逻辑:

  • 标签(shifted_labels)取输入的第2个token开始的部分(labels = data['labels'][:, 1:]),因为模型要预测的是“下一个词”。
  • 预测值(preds)取模型输出logits中概率最大的token ID(out = out['logits'].argmax(dim=2)[:, :-1]),并截断最后一个token以与标签对齐。
  • 准确率计算:关键一步是忽略掉用于填充的token(ID为0)。我们通过 select = labels != 0 创建一个掩码,只计算有效文本部分的预测正确率(accuracy),这能更真实地反映模型的语言建模能力。

最后,在每一轮训练结束后,保存模型的权重,以便后续使用或继续训练。

torch.save(model.state_dict(), f='net.pt')
print("权重保存成功!")
if __name__ == '__main__':
    for epoch in range(1000):
        train()

这里使用 torch.save(model.state_dict(), 'net.pt') 仅保存模型参数,文件较小。外层通过 for epoch in range(1000) 循环控制训练的总轮数。

[AFFILIATE_SLOT_2]

五、运行实践与效果评估

现在,让我们启动整个训练流程。

  1. 准备数据:确保 data/chinese_poems.txt 文件已就位。
  2. 配置模型路径:将代码中的 model_dir 修改为你本地或远程预训练模型的实际路径。
  3. 执行训练:直接运行 train.py 文件。

程序将自动执行:加载数据、初始化模型、开始多轮训练(每50步打印日志)、并在每轮结束后保存权重文件(如 net.pt)。

如何解读训练结果?

  • 损失值(loss):应随着训练步数增加而总体呈下降趋势并逐渐趋于稳定,这表明模型正在从数据中有效学习。如果损失值剧烈波动或上升,可能需要检查学习率是否过高。
  • 学习率(lr):在调度器作用下应平滑下降,符合“初期大胆探索,后期精细调整”的训练哲学。
  • 准确率(accuracy):反映模型预测下一个词的准确度。这个值会逐步提升,越高说明模型对诗歌的格律、用词规律掌握得越好。

六、总结与进阶优化方向

本文通过整合 data.py 和 train.py 的完整代码,系统性地演示了使用PyTorch和Hugging Face Transformers微调中文GPT-2模型的全过程。我们从自定义 Dataset 处理数据开始,到加载模型、构建训练流水线,最终实现了一个包含梯度裁剪、学习率调度和监控的稳健训练循环。

这个项目框架具有良好的扩展性。要将其应用于其他中文生成任务(如小说续写、对联生成、广告文案创作),你通常只需要:1)替换训练数据集;2)微调超参数。

进阶优化建议:

  • 引入早停机制:当前固定训练1000轮可能过拟合。可划分验证集,当验证集损失连续多轮不下降时自动停止训练。
  • 优化数据处理:将分词过程移至 Dataset 的 __getitem__ 方法中,减轻 collate_fn 函数的负担。
  • 完整模型保存:使用 model.save_pretrained('gpt2-poetry') 保存模型结构和权重,便于直接加载推理,无需重新定义模型。
  • 超参数调优:尝试调整批次大小(batch_size)、学习率、最大序列长度(max_length)等,以找到最适合你任务和数据集的配置。
  • 开发推理接口:基于训练好的模型,实现一个生成函数,输入“春江潮水连海平”等开头,让模型自动创作出完整的诗歌。

希望这份详实的实战指南,能帮助你顺利踏入AI文本生成的世界,并激发更多创作灵感。完整的项目代码附后,供你参考和实践。

附录:完整代码文件

data.py

from torch.utils.data import Dataset
class MyDataset(Dataset):
    def __init__(self):
        """
        数据集初始化函数
        负责读取数据文件并进行基础的数据清洗
        """
        # 以UTF-8编码打开中文诗歌数据文件
        with open("data/chinese_poems.txt", encoding="utf-8") as f:
            # 读取文件的所有行,每行作为列表中的一个元素
            lines = f.readlines()
        # 数据清洗:去除每条数据的前后空格和换行符
        # 使用列表推导式对每一行数据进行strip()处理
        lines = [i.strip() for i in lines]
        # 将处理后的数据保存为实例变量,供其他方法使用
        self.lines = lines
    def __len__(self):
        """
        返回数据集的长度
        这是PyTorch Dataset类必须实现的方法之一
        :return: 数据集中样本的总数
        """
        return len(self.lines)
    def __getitem__(self, item):
        """
        根据索引获取单个数据样本
        这是PyTorch Dataset类必须实现的方法之一
        :param item: 数据索引(整数)
        :return: 对应索引的数据样本
        """
        return self.lines[item]
if __name__ == '__main__':
    """
    主程序入口
    当直接运行此脚本时执行以下代码(用于测试数据集类)
    """
    # 创建数据集实例
    dataset = MyDataset()
    # 打印数据集大小和第一个样本,用于验证数据加载是否正确
    print(len(dataset), dataset[0])

train.py

# 导入必要的库和模块
from torch.optim import AdamW  # AdamW优化器
from transformers.optimization import get_scheduler  # 学习率调度器
import torch  # PyTorch深度学习框架
from data import MyDataset  # 自定义数据集类
from transformers import AutoTokenizer  # 自动分词器
from transformers import AutoModelForCausalLM, GPT2Model  # 预训练模型
# 创建数据集实例
dataset = MyDataset()
# 指定预训练模型的本地路径
model_dir = r"D:\pyprojecgt\flaskProject\langchainstudy\modelscope\gpt2-chinese-cluecorpussmall"
# 加载分词器(用于文本编码和解码)
tokenizer = AutoTokenizer.from_pretrained(model_dir)
# 加载预训练的语言模型(用于因果语言建模,如文本生成)
model = AutoModelForCausalLM.from_pretrained(model_dir)
print(model)  # 打印模型结构,便于调试和了解模型组成
def collate_fn(data):
    """
    数据批处理函数
    将原始文本数据转换为模型可接受的张量格式
    :param 一个批次的文本数据列表
    :return: 编码后的张量字典
    """
    # 使用分词器批量编码文本
    data = tokenizer.batch_encode_plus(data,
                                       padding=True,  # 填充到相同长度
                                       truncation=True,  # 截断超过最大长度的文本
                                       max_length=512,  # 最大序列长度
                                       return_tensors='pt')  # 返回PyTorch张量
    # 为语言模型训练创建标签(GPT是自回归模型,输入即标签)
    data['labels'] = data['input_ids'].clone()
    return data
# 创建数据加载器,用于批量加载和预处理数据
loader = torch.utils.data.DataLoader(
    dataset=dataset,  # 使用的数据集
    batch_size=4,  # 每个批次的样本数量
    collate_fn=collate_fn,  # 批处理函数
    shuffle=True,  # 每个epoch打乱数据顺序
    drop_last=True,  # 丢弃最后一个不完整的批次
)
print(len(loader))  # 打印总批次数
# 定义训练函数
def train():
    """
    模型训练函数
    包含完整的前向传播、反向传播和参数更新流程
    """
    global model  # 声明使用全局变量model
    # 自动检测并设置计算设备(GPU优先)
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    model = model.to(device)  # 将模型移动到指定设备
    # 创建优化器(AdamW,适合Transformer模型)
    optimizer = AdamW(model.parameters(), lr=2e-5)
    # 创建学习率调度器(线性衰减)
    scheduler = get_scheduler(name='linear',
                              num_warmup_steps=0,  # 预热步数为0
                              num_training_steps=len(loader),  # 总训练步数
                              optimizer=optimizer)
    # 设置模型为训练模式(启用dropout等训练专用层)
    model.train()
    # 遍历数据加载器中的每个批次
    for i, data in enumerate(loader):
        # 将批次数据中的所有张量移动到指定设备
        for k in data.keys():
            data[k] = data[k].to(device)
        # 前向传播:将数据传入模型计算损失和输出
        out = model(**data)
        loss = out['loss']  # 提取损失值
        # 反向传播:计算梯度
        loss.backward()
        # 梯度裁剪:防止梯度爆炸,提高训练稳定性
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        # 参数更新:使用优化器更新模型参数
        optimizer.step()
        # 更新学习率
        scheduler.step()
        # 清空梯度,为下一批次做准备
        optimizer.zero_grad()
        # 每50个批次打印一次训练信息
        if i % 50 == 0:
            # 准备标签和预测结果(忽略序列的第一个token)
            labels = data['labels'][:, 1:]  # 标签,去掉第一个token
            # 获取预测结果(取logits中概率最大的索引),去掉最后一个预测
            out = out['logits'].argmax(dim=2)[:, :-1]
            # 创建掩码,忽略填充token(ID为0的token)
            select = labels != 0
            # 应用掩码,只计算非填充token的准确率
            labels = labels[select]
            out = out[select]
            del select  # 释放掩码张量内存
            # 计算准确率:预测正确的token比例
            accuracy = (labels == out).sum().item() / labels.numel()
            # 获取当前学习率
            lr = optimizer.state_dict()['param_groups'][0]['lr']
            # 打印训练信息:批次索引、损失值、学习率、准确率
            print(i, loss.item(), lr, accuracy)
    # 保存模型权重(不包含模型结构,需要原模型结构加载)
    torch.save(model.state_dict(), f='net.pt')
    print("权重保存成功!")
# 主程序入口
if __name__ == '__main__':
    # 进行1000个epoch的训练
    for epoch in range(1000):
        train()  # 调用训练函数
posted @ 2026-03-02 18:27  yangykaifa  阅读(82)  评论(0)    收藏  举报