基于pytorch的保存和加载模型参数的方法

yipeiwu_com6年前Python基础

当我们花费大量的精力训练完网络,下次预测数据时不想再(有时也不必再)训练一次时,这时候torch.save(),torch.load()就要登场了。

保存和加载模型参数有两种方式:

方式一:

torch.save(net.state_dict(),path):

功能:保存训练完的网络的各层参数(即weights和bias)

其中:net.state_dict()获取各层参数,path是文件存放路径(通常保存文件格式为.pt或.pth)

net2.load_state_dict(torch.load(path)):

功能:加载保存到path中的各层参数到神经网络

注意:不可以直接为torch.load_state_dict(path),此函数不能直接接收字符串类型参数

方式二:

torch.save(net,path):

功能:保存训练完的整个网络模型(不止weights和bias)

net2=torch.load(path):

功能:加载保存到path中的整个神经网络

说明:官方推荐方式一,原因自然是保存的内容少,速度会更快。

以上这篇基于pytorch的保存和加载模型参数的方法就是小编分享给大家的全部内容了,希望能给大家一个参考,也希望大家多多支持【听图阁-专注于Python设计】。

相关文章

Python写的一个定时重跑获取数据库数据

Python写的一个定时重跑获取数据库数据

做大数据的童鞋经常会写定时任务跑数据,由于任务之间的依赖(一般都是下游依赖上游的数据产出),所以经常会导致数据获取失败,因为很多人发现数据失败后 都会去查看日志,然后手动去执行自己的任务...

对pandas写入读取h5文件的方法详解

1、引言 通过参考相关博客对hdf5格式简要介绍。 hdf5在存储的是支持压缩,使用的方式是blosc,这个是速度最快的也是pandas默认支持的。 使用压缩可以提磁盘利用率,节省空间。...

利用Pytorch实现简单的线性回归算法

利用Pytorch实现简单的线性回归算法

最近听了张江老师的深度学习课程,用Pytorch实现神经网络预测,之前做Titanic生存率预测的时候稍微了解过Tensorflow,听说Tensorflow能做的Pyorch都可以做,...

django写用户登录判定并跳转制定页面的实例

1. 首先看要设置登陆的界面 book/view.py @user_util.my_login #相当于 select_all=my_login(select_all) def se...

tensorflow ckpt模型和pb模型获取节点名称,及ckpt转pb模型实例

ckpt from tensorflow.python import pywrap_tensorflow checkpoint_path = 'model.ckpt-8000'...