保存加载用于推理的常规Checkpoint/或继续训练

正在搬砖嘤

已于 2022-08-14 20:16:12 修改

阅读量290

点赞数 1

分类专栏：深度学习文章标签： pytorch 人工智能 python

于 2021-10-22 15:54:08 首次发布

本文链接：https://blog.csdn.net/walkinging1/article/details/120907588

版权

深度学习专栏收录该内容

2 篇文章

订阅专栏

该博客介绍了如何在PyTorch中保存和加载模型的检查点，包括模型参数、优化器状态以及训练的epoch和loss。通过使用`torch.save()`和`torch.load()`，可以在训练过程中定期保存模型，并在之后的推理或继续训练时加载这些信息。检查点文件包含了模型的重要状态，而state_dict则用于存储模型参数。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

保存加载用于推理的常规Checkpoint/或继续训练

https://blog.csdn.net/qq_38765642/article/details/109784913

if (epoch+1) % checkpoint_interval == 0:

        checkpoint = {"model_state_dict": net.state_dict(),
                      "optimizer_state_dict": optimizer.state_dict(),
                      "epoch": epoch}
        path_checkpoint = "./checkpoint_{}_epoch.pkl".format(epoch)
        torch.save(checkpoint, path_checkpoint)
#或者
#保存
torch.save({
            'epoch': epoch,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'loss': loss,
            ...
            }, PATH)
#加载
model = TheModelClass(*args, **kwargs)
optimizer = TheOptimizerClass(*args, **kwargs)

checkpoint = torch.load(PATH)
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
epoch = checkpoint['epoch']
loss = checkpoint['loss']

model.eval()
# - 或者 -
model.train()

checkpoint检查点：不仅保存模型的参数，优化器参数，还有loss，epoch等（相当于一个保存模型的文件夹）
state_dict可以用于保存模型的参数及优化。