目录
我们之前讲过了如何部署一个别人已经训练好的AI模型、也学会了如何微调一个AI模型,也讲了预训练模型和微调模型的区别,那本文就聊聊如何从零训练一个语言模型吧!
收集或制造数据集
在机器学习中,数据集的收集是非常重要的一步,质量高或者相关性高的数据集对模型的训练有非常大的帮助。
如下两份数据集可供使用。
(1)15G gpt2数据集
链接:百度网盘 请输入提取码 提取码:khiq
(2)200G悟道数据集
链接:百度网盘 请输入提取码 提取码:sp3n
下载并解压
我使用的是15G gpt2数据集进行的预训练
安装python依赖
pip install gradio transformers datasets -i https://pypi.mirrors.ustc.edu.cn/simple/
编写训练脚本
from datasets import load_dataset, DatasetDict
from glob import glob
import random
random.seed(42)
print("开始加载包~")
all_file_list = glob(pathname="gpt2_data/*/**")
test_file_list = random.sample(all_file_list, 50)
train_file_list = [i for i in all_file_list if i not in test_file_list]
len(train_file_list), len(test_file_list)
print(len(train_file_list))
print(len(test_file_list))
print("加载包 ok")
print("开始解析和提取数据~")
raw_datasets =load_dataset("csv",data_files={'train':train_file_list,'valid':test_file_list}, cache_dir="cache_data")
print("解析和提取数据 ok")
raw_datasets
from transformers import AutoTokenizer, AutoConfig
context_length = 512
# 下载分词工具
print("下载分词工具~")
tokenizer = AutoTokenizer.from_pretrained("bert-base-chinese")
print("下载分词工具 ok")
print("对原始数据集进行分词和编码~")
# 将原始文本转换为模型可以理解的编码表示
outputs = tokenizer(
raw_datasets["train"][:2]["content"],
truncation=True,
max_length=context_length,
return_overflowing_tokens=True,
return_length=True,
)
print("对原始数据集进行分词和编码 ok")
print(f"Input IDs length: {len(outputs['input_ids'])}")
print(f"Input chunk lengths: {(outputs['length'])}")
print(f"Chunk mapping: {outputs['overflow_to_sample_mapping']}")
# 为 tokenizer 添加特殊标记
# <|endoftext|> 是模型在训练期间看到的文档分隔符
tokenizer.add_special_tokens(special_tokens_dict={'bos_token': '<|endoftext|>',
'eos_token': '<|endoftext|>',
'unk_token': '<|endoftext|>'})
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}
# 对原始数据集进行分词和编码,并生成经过处理的分词后的数据集
tokenized_datasets = raw_datasets.map(
tokenize, batched=True, remove_columns=raw_datasets["train"].column_names
)
tokenized_datasets
from transformers import AutoTokenizer, GPT2LMHeadModel, AutoConfig
# 创建和配置 GPT-2 模型
print("创建和配置 GPT-2 模型~")
config = AutoConfig.from_pretrained(
"gpt2",
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"GPT-2 size: {model_size/1000**2:.1f}M parameters")
from transformers import DataCollatorForLanguageModeling
# 为语言建模任务准备数据,并使用数据处理器对数据进行编码、填充和组织,生成适用于训练的批次数据
print("为语言建模任务准备数据,并使用数据处理器对数据进行编码、填充和组织,生成适用于训练的批次数据~")
tokenizer.pad_token = tokenizer.eos_token
data_collator = DataCollatorForLanguageModeling(tokenizer, mlm=False)
out = data_collator([tokenized_datasets["train"][i] for i in range(5)])
for key in out:
print(f"{key} shape: {out[key].shape}")
print("为语言建模任务准备数据,并使用数据处理器对数据进行编码、填充和组织,生成适用于训练的批次数据 ok")
from transformers import Trainer, TrainingArguments
print("创建并执行模型训练~")
args = TrainingArguments(
# 指定训练完成后保存模型的目录
output_dir="chinese_gpt2_big",
# 每个设备(GPU或CPU)上用于训练的批次大小
per_device_train_batch_size=20,
# 每个设备上用于评估的批次大小
per_device_eval_batch_size=16,
# 指定何时在训练过程中进行评估。这里设置为 "steps",表示将在指定的 eval_steps 处进行评估
evaluation_strategy="steps",
# 每次评估之间的训练步数
eval_steps=2_000,
# 每隔多少训练步数记录一次训练指标
logging_steps=2_000,
# 在执行反向传播/更新步骤之前累积梯度的步骤数
gradient_accumulation_steps=8,
# 遍历整个训练数据集的次数
num_train_epochs=2,
# 权重衰减正则化的系数
weight_decay=0.1,
# 学习率调度的预热步数
warmup_steps=1_000,
# 学习率调度器的类型。这里设置为 "cosine"
lr_scheduler_type="cosine",
# 优化器的初始学习率
learning_rate=5e-4,
# 每隔多少训练步数保存一次检查点
save_steps=2_000,
# 是否使用16位浮点精度(混合精度训练),以加快训练速度并减少内存使用
fp16=True,
# 是否将训练完成的模型推送到模型中心
push_to_hub=False,
)
trainer = Trainer(
# 要训练的模型
model=model,
# 用于处理输入数据的 tokenizer
tokenizer=tokenizer,
# 训练的参数配置,包括输出目录、批次大小、训练轮数等
args=args,
# 用于处理数据的数据处理器
data_collator=data_collator,
# 训练数据集,即经过处理的分词后的训练数据
train_dataset=tokenized_datasets["train"],
# 验证数据集,即经过处理的分词后的验证数据
eval_dataset=tokenized_datasets["valid"],
)
trainer.train()
print("创建并执行模型训练 ok")
执行训练脚本开始训练
python3 trainer.py
注意,训练底座大模型需要大量的计算资源和时间,所以一定要保证自己的硬件足够,我使用的是一张3090显卡,24G显存,训练15G数据集需要60个小时
60小时后,训练完成,如下:
编写推理脚本
vi chat.py
import gradio as gr
from transformers import GPT2Tokenizer, GPT2LMHeadModel, AutoTokenizer
model_name_or_path = "checkpoint-36000"
tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
model = GPT2LMHeadModel.from_pretrained(model_name_or_path, pad_token_id=tokenizer.eos_token_id)
def infer_model(txt: str) -> str:
# encode context the generation is conditioned on
input_ids = tokenizer.encode(txt, return_tensors='pt')
# set no_repeat_ngram_size to 2
beam_output = model.generate(
input_ids,
max_length=200,
num_beams=5,
no_repeat_ngram_size=2,
early_stopping=True
)
result = tokenizer.decode(beam_output[0], skip_special_tokens=True)
result = result.replace(" ", "")
return result
def predict(input, history=[]):
respose = infer_model(input)
res = [(input, respose)]
return res, history
def add_text(state, text):
res = infer_model(text)
state = state + [(text, res)]
return state, state
with gr.Blocks(css="#chatbot .overflow-y-auto{height:500px}") as demo:
chatbot = gr.Chatbot(elem_id="chatbot")
state = gr.State([])
with gr.Row():
txt = gr.Textbox(show_label=False, placeholder="输入文本").style(container=False)
txt.submit(add_text, [state, txt], [state, chatbot])
if __name__ == "__main__":
demo.launch(server_name='0.0.0.0',server_port=6006)
启动推理脚本
python chat.py
测试
注意
由于15G gpt2数据量太少,所以训练完成的AI模型对话效果并不理想,如想训练效果更好的模型,请使用200G悟道数据集。
以上便是本次训练的全部过程啦~~~