PyTorch保存和加载模型

最新推荐文章于 2025-04-08 10:39:55 发布

grllery

最新推荐文章于 2025-04-08 10:39:55 发布

阅读量1k

点赞数 1

分类专栏： PyTorch

PyTorch 专栏收录该内容

5 篇文章

订阅专栏

本文介绍了PyTorch中两种常见的模型保存与加载方法。一种仅保存模型参数，适用于模型迁移；另一种则保存整个模型，但可能受限于特定目录结构。文章详细展示了每种方法的实现代码。

转自：

http://www.cnblogs.com/darkknightzh/p/8108466.html

参考网址：

https://pytorch.org/tutorials/beginner/saving_loading_models.html

https://pytorch.org/docs/stable/notes/serialization.html#best-practices

https://github.com/clcarwin/sphereface_pytorch

有两种方式保存和载入模型

1. 只保存和载入模型参数

保存：

torch.save(the_model.state_dict(), PATH)

载入：

the_model = TheModelClass(*args, **kwargs)
the_model.load_state_dict(torch.load(PATH))

当model使用gpu训练时，可以将数据转换到cpu中，并保存（载入时，还是上面的方法。需要使用gpu时，加上.cuda()）：

def save_model(model, filename):
    state = model.state_dict()
    for key in state: state[key] = state[key].clone().cpu()
    torch.save(state, filename)

Remember that you must call model.eval() to set dropout and batch normalization layers to evaluation mode before running inference. Failing to do this will yield inconsistent inference results.

在加载已经保存好的模型, 务必调用model.eval()再进行前向计算.

A common PyTorch convention is to save models using either a .pt or .pth file extension.

2. 保存和载入整个模型

保存：

torch.save(the_model, PATH)

载入：

the_model = torch.load(PATH)

However in this case, the serialized data is bound to the specific classes and the exact directory structure used, so it can break in various ways when used in other projects, or after some serious refactors.

第二种方式，序列化后的数据使用特殊的结构，缺点就是当在其他工程中使用时，可能会碰到各种问题。因而，官方更建议使用第一种方式。