介绍
虽然使用的大语言模型和数据集都相对较小,但训练时间仍然较长,如果数据量大,需要使用价格高昂的GPU,所以训练大语言模型代价非常高昂。
因此,保存大语言模型的参数非常重要,这样就不必每次使用它时都重新运行训练。
本文实现对模型进行保存,通常能为以后的继续训练带来方便(第(6)步)
保存模型中的参数
推荐使用 torch.save() 函数保存模型的 state_dict,即将每个层映射到其参数的字典
torch.save(model.state_dict(), "model.pth")
model.pth 是保存 state_dict 的文件名。.pth 扩展名是PyTorch文件的规范,尽管从技术上讲可以使用任何文件扩展名。
加载模型中的参数
在通过 state_dict 保存模型权重之后,可以将模型权重加载到一个新的 GPTModel 模型实例中:
model = GPTModel(GPT_CONFIG_124M)
model.load_state_dict(torch.load("model.pth", map_location=device))
model.eval()
dropout 通过在训练过程中 随机 “丢弃”一层的神经元,有助于防止模型对训练数据过拟合。
但是,在推断过程中,要使用网络学习到的所有信息。因此,使用 model.eval() 将模型 切换到推断模式,这样就会 禁用模型的 dropout 层。
如果继续训练模型,可以使用前面定义的 train_model_simple() 函数,建议同时保存优化器状态。
保存优化器中的参数
像AdamW这样的自适应优化器可以为每个模型权重存储额外的参数。AdamW可以使用历史数据动态地调整每个模型参数的学习率。如果没有它,那么优化器就会重置,模型可能学习效果不佳,甚至无法正确收敛,这意味着模型将失去生成连贯文本的能力。可以使用 torch.save() 保存模型和优化器的 state_dict 内容:
torch.save({
"model_state_dict": model.state_dict(),
"optimizer_state_dict": optimizer.state_dict(),
},
"model_and_optimizer.pth"
)
加载优化器中的参数
checkpoint = torch.load("model_and_optimizer.pth", map_location=device)
model = GPTModel(GPT_CONFIG_124M)
model.load_state_dict(checkpoint["model_state_dict"])
optimizer = torch.optim.AdamW(model.parameters(), lr=5e-4, weight_decay=0.1)
optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
model.train();
