本节代码是一个完整的机器学习工作流程,用于训练一个基于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,TrainingArgumentsdef 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.2context_length = 128all_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()