当前位置：首页 > news >正文

GPT-2 语言模型 - 模型训练

news 来源：原创 2025/5/14 19:24:47

本节代码是一个完整的机器学习工作流程，用于训练一个基于GPT-2的语言模型。下面是对这段代码的详细解释：

文件目录如下

1. 初始化和数据准备

设置随机种子
```
random.seed(1002)
```
确保结果的可重复性。
定义参数
```
test_rate = 0.2
context_length = 128
```
- test_rate：测试集占总数据集的比例。
- context_length：模型处理的文本长度。
获取数据文件
```
all_files = glob(pathname=os.path.join("data","*"))
```
使用 glob 获取 data 目录下的所有文件。

划分数据集

test_file_list = random.sample(all_files, int(len(all_files) * test_rate))
train_file_list = [i for i in all_files if i not in test_file_list]

将数据集随机划分为训练集和测试集。

加载数据集

raw_datasets = load_dataset("csv", data_files={"train": train_file_list, "vaild": test_file_list}, cache_dir="cache_data")

使用 datasets 库加载 CSV 格式的数据集，并缓存到 cache_data 目录。

2. 数据预处理

初始化分词器

tokenizer = BertTokenizerFast.from_pretrained("D:/bert-base-chinese")
tokenizer.add_special_tokens({"bos_token":"[begin]","eos_token":"[end]"})

从本地路径加载预训练的 BERT 分词器，并添加自定义的开始和结束标记。

数据预处理
```
tokenize_datasets = raw_datasets.map(tokenize, batched=True, remove_columns=raw_datasets["train"].column_names)
```
使用 map 方法对数据集进行预处理，将文本转换为模型可接受的格式。
- tokenize 函数对文本进行分词和截断。
- batched=True 表示批量处理数据。
- remove_columns 删除原始数据集中的列。

3. 模型配置和初始化

模型配置

config = GPT2Config.from_pretrained("config",
                                    vocab_size=len(tokenizer),
                                    n_ctx=context_length,
                                    bos_token_id=tokenizer.bos_token_id,
                                    eos_token_id=tokenizer.eos_token_id,
                                    )

加载预训练的 GPT-2 配置，并根据分词器的词汇表大小和上下文长度进行调整。

初始化模型
```
model = GPT2LMHeadModel(config)
model_size = sum([t.numel() for t in model.parameters()])
print(f"model_size: {model_size/1000/1000} M")
```
根据配置初始化 GPT-2 语言模型，并计算模型参数的总数，打印模型大小（以兆字节为单位）。

4. 训练设置

数据整理器
```
data_collator = DataCollatorForLanguageModeling(tokenizer, mlm=False)
```
使用 DataCollatorForLanguageModeling 整理训练数据，设置 mlm=False 表示不使用掩码语言模型。

训练参数

args = TrainingArguments(
    learning_rate=1e-5,
    num_train_epochs=100,
    per_device_train_batch_size=10,
    per_device_eval_batch_size=10,
    eval_steps=2000,
    logging_steps=2000,
    gradient_accumulation_steps=5,
    weight_decay=0.1,
    warmup_steps=1000,
    lr_scheduler_type="cosine",
    save_steps=100,
    output_dir="model_output",
    fp16=True,
)

配置训练参数，包括学习率、训练轮数、批大小、评估间隔等。

初始化训练器

trianer = Trainer(
    model=model,
    args=args,
    tokenizer=tokenizer,
    data_collator=data_collator,
    train_dataset=tokenize_datasets["train"],
    eval_dataset=tokenize_datasets["vaild"]
)

启动训练
```
trianer.train()
```
使用 Trainer 类启动模型训练。

需复现完整代码

from glob import glob
import os
from torch.utils.data import Dataset
from datasets import load_dataset
import random
from transformers import BertTokenizerFast
from transformers import GPT2Config
from transformers import GPT2LMHeadModel
from transformers import DataCollatorForLanguageModeling
from transformers import Trainer,TrainingArguments

def tokenize(element):
    outputs = tokenizer(element["content"],truncation=True,max_length=context_length,return_overflowing_tokens=True,return_length=True)

    input_batch = []

    for length,input_ids in zip(outputs["length"],outputs["input_ids"]):

        if length == context_length:
            input_batch.append(input_ids)

    return {"input_ids":input_batch}

if __name__ == "__main__":
    random.seed(1002)
    test_rate = 0.2
    context_length = 128

    all_files = glob(pathname=os.path.join("data","*"))

    test_file_list = random.sample(all_files,int(len(all_files)*test_rate))
    train_file_list = [i for i in all_files if i not in test_file_list]

    raw_datasets = load_dataset("csv",data_files={"train":train_file_list,"vaild":test_file_list},cache_dir="cache_data")


    tokenizer = BertTokenizerFast.from_pretrained("D:/bert-base-chinese")
    tokenizer.add_special_tokens({"bos_token":"[begin]","eos_token":"[end]"})

    tokenize_datasets = raw_datasets.map(tokenize,batched=True,remove_columns=raw_datasets["train"].column_names)

    config = GPT2Config.from_pretrained("config",
                                        vocab_size=len(tokenizer),
                                        n_ctx=context_length,
                                        bos_token_id = tokenizer.bos_token_id,
                                        eos_token_id = tokenizer.eos_token_id,
                                        )

    model = GPT2LMHeadModel(config)
    model_size = sum([ t.numel() for t in model.parameters()])
    print(f"model_size: {model_size/1000/1000} M")

    data_collator = DataCollatorForLanguageModeling(tokenizer,mlm=False)

    args = TrainingArguments(
        learning_rate=1e-5,
        num_train_epochs=100,
        per_device_train_batch_size=10,
        per_device_eval_batch_size=10,
        eval_steps=2000,
        logging_steps=2000,
        gradient_accumulation_steps=5,
        weight_decay=0.1,
        warmup_steps=1000,
        lr_scheduler_type="cosine",
        save_steps=100,
        output_dir="model_output",
        fp16=True,
    )

    trianer = Trainer(
        model=model,
        args=args,
        tokenizer=tokenizer,
        data_collator=data_collator,
        train_dataset=tokenize_datasets["train"],
        eval_dataset=tokenize_datasets["vaild"]
    )

    trianer.train()