Transformers是一个先进的NLP框架,适用于pytorch和tensorflow2.0,这边使用的GPT2框架是一个OpenAI的先进技术,在上下文连贯性和情感表达上有相当好的表现,实际使用时可以直接从transformer库导入:
from transformers.models.gpt2.modeling_gpt2 import GPT2Config, GPT2LMHeadModel
其实OpenAI还推出了GPT3——1750亿的参数和千万级别的算力需求就不是笔记本显卡带的动的了。GPT2有15亿参数,面对8M的语料就需要4h左右,还是比较可行的方案
定义参数,统计step,采用梯度累积和warmup策略,导入tensorboardX
def train(model, device, train_list, multi_gpu, args):
# model:被训练的模型
# device:使用的GPU
# train_list:经过划分后的训练集
# multi_gpu:多GPU设置参数
# args:代码参数管理器
train_dataset = MyDataset(train_list)
train_dataloader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers,
collate_fn=collate_fn)
model.train()
# 计算所有epoch进行参数优化的总步数total_steps
# 在一定条件下,branch_size效果越大越好,梯度累积就以通过积累多个branch的梯度,在有限内存的前提下达到大branch的效果
total_steps = int(train_dataset.__len__() * args.epochs / args.batch_size / args.gradient_accumulation)
logger.info('total training steps = {}'.format(total_steps))
# 设置优化器,并且在初始训练时,使用warmup策略
# 普遍观点认为在初始阶段以小于基础学习率的学习率进行warmup有助于避免局部过拟合,帮助模型熟悉数据,节约时间
optimizer = transformers.AdamW(model.parameters(), lr=args.lr