Pytorch中保存模型的两种方法

目录

一、保存整个模型

这种方法会保存模型的整个定义(包括网络结构和参数)。这样保存的模型可以直接被加载并使用,但前提是保存和加载模型的代码环境需要完全一致。

保存模型时使用代码:

python 复制代码
torch.save(model, 'model.pth')

加载模型时使用代码:

python 复制代码
model = torch.load('model.pth')

二、只保存模型参数

这种方法只会保存模型的状态字典(state_dict),即模型的参数。加载这样的模型需要先定义模型的结构,然后加载参数。

保存模型使用代码:

python 复制代码
torch.save(model.state_dict(), 'params.pth')

加载模型使用代码:

python 复制代码
model = Model()  # 需要先定义模型结构
model.load_state_dict(torch.load('params.pth'))
相关推荐
cwn_1 分钟前
回归(多项式回归)
人工智能·机器学习·数据挖掘·回归
聚客AI32 分钟前
🔥 大模型开发进阶:基于LangChain的异步流式响应与性能优化
人工智能·langchain·agent
cooldream20091 小时前
Python 包管理新时代:深入了解 `uv` 的使用与实践
python·uv·包管理器
之歆1 小时前
Python-魔术方法-创建、初始化与销毁-hash-bool-可视化-运算符重载-容器和大小-可调用对象-上下文管理-反射-描述器-二分-学习笔记
笔记·python·学习
CareyWYR1 小时前
每周AI论文速递(250707-250711)
人工智能
AI街潜水的八角1 小时前
深度学习图像分类数据集—五种电器识别分类
人工智能·深度学习·分类
众链网络1 小时前
AI进化论08:机器学习的崛起——数据和算法的“二人转”,AI“闷声发大财”
人工智能·算法·机器学习
生命是有光的1 小时前
【机器学习】机器学习基础
人工智能·机器学习
平和男人杨争争1 小时前
机器学习13——支持向量机下
人工智能·机器学习·支持向量机
胖达不服输1 小时前
「日拱一码」025 机器学习——评价指标
人工智能·python·机器学习·评价指标