机器学习 - save和load训练好的模型

如果已经训练好了一个模型,你就可以save和load这模型。

For saving and loading models in PyTorch, there are three main methods you should be aware of.

PyTorch method What does it do?
torch.save Saves a serialized object to disk using Python's pickle utility. Models, tensors and various other Python objects like dictionaries can be saved using torch.save
torch.load Uses pickle's unpickling features to deserialize and load pickled Python object files (like models, tensors or dictionaries) into memory. You can also set which device to load the object to (CPU, GPU etc)
torch.nn.Module.load_state_dict Loads a model's parameter dictionary (model.state_dict()) using a saved state_dict() object

在 PyTorch 中,pickle 是一个用于序列化和反序列化Python对象的标准库模块。它可以将Python对象转换为字节流 (即序列化),并将字节流转换回Python对象 (即反序列化)。pickle模块在很多情况下都非常有用,特别是在保存和加载模型,保存训练中间状态等方面。

在深度学习中,经常需要保存训练好的模型或者训练过程中的中间结果,以便后续的使用或分析。PyTorch提高了方便的API来保存和加载模型,其中就包括了使用pickle模块进行对象的序列化和反序列化。


save model

python 复制代码
import torch
from pathlib import Path 

# 1. Create models directory
MODEL_PATH = Path("models")
MODEL_PATH.mkdir(parents = True, exist_ok = True)

# 2. Create model save path
MODEL_NAME = "trained_model.pth"
MODEL_SAVE_PATH = MODEL_PATH / MODEL_NAME

# 3. Save the model state dict 
print(f"Saving model to: {MODEL_SAVE_PATH}")
torch.save(obj = model_0.state_dict(),
			f = MODEL_SAVE_PATH)

就能看到 trained_model.pth 文件下载到所属的文件夹位置。


Load the saved PyTorch model

You can load it in using torch.nn.Module.load_state_dict(torch.load(f)) where f is the filepath of the saved model state_dict().

Why call torch.load() inside torch.nn.Module.load_state_dict()?

Because you only saved the model's state_dict() which is a dictionary of learned parameters and not the entire model, you first have to load the state_dict() with torch.load() and then pass that state_dict() to a new instance of the model (which is a subclass of nn.Module).

python 复制代码
# Instantiate a new instance of the model 
loaded_model_0 = LinearRegressionModel()

# Load the state_dict of the saved model
loaded_model_0.load_state_dict(torch.load(f=MODEL_SAVE_PATH))

# 结果如下
<All keys matched successfully>

测试 loaded model。

python 复制代码
# Put the loaded model into evaluation model 
loaded_model_0.eval() 

# 2. Use the inference mode context manager to make predictions
with torch.inference_mode():
  loaded_model_preds = loaded_model_0(X_test)

# Compare previous model predictions with loaded model predictions
print(y_preds == loaded_model_preds) 

# 结果如下
tensor([[True],
        [True],
        [True],
        [True],
        [True],
        [True],
        [True],
        [True],
        [True],
        [True]])

看到这了,点个赞呗~

相关推荐
Yo_Becky15 分钟前
【PyTorch】PyTorch预训练模型缓存位置迁移,也可拓展应用于其他文件的迁移
人工智能·pytorch·经验分享·笔记·python·程序人生·其他
DeepSeek-大模型系统教程17 分钟前
深入金融与多模态场景实战:金融文档分块技术与案例汇总
人工智能·ai·语言模型·程序员·大模型·大模型学习·大模型教程
xinxiangwangzhi_22 分钟前
pytorch底层原理学习--PyTorch 架构梳理
人工智能·pytorch·架构
yzx99101326 分钟前
关于网络协议
网络·人工智能·python·网络协议
AiTEN_Robot28 分钟前
AGV 无人叉车关键技术问题解析:精准定位算法 / 安全避障逻辑 / 系统对接协议全方案
人工智能·机器人·自动化·制造
云天徽上29 分钟前
【PaddleOCR】OCR常见关键信息抽取数据集,包含FUNSD、XFUND、WildReceipt等整理,持续更新中......
人工智能·计算机视觉·信息可视化·paddlepaddle·paddleocr·文本识别
zskj_zhyl29 分钟前
智绅科技:以科技为翼,构建养老安全守护网
人工智能·科技·安全
刘海东刘海东30 分钟前
结构型智能科技的关键可行性——信息型智能向结构型智能的转换(提纲)
人工智能
martian6651 小时前
支持向量机(SVM)深度解析:从数学根基到工程实践
算法·机器学习·支持向量机
Jay Kay1 小时前
TensorFlow源码深度阅读指南
人工智能·python·tensorflow