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

yipeiwu_com5年前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自定义进程池。分享给大家供大家参考,具体如下: 代码说明一切: #encoding=utf-8 #author: walker #date: 2014-05...

Python MySQLdb模块连接操作mysql数据库实例

mysql是一个优秀的开源数据库,它现在的应用非常的广泛,因此很有必要简单的介绍一下用python操作mysql数据库的方法。python操作数据库需要安装一个第三方的模块,在http:...

Android分包MultiDex策略详解

1.分包背景 这里首先介绍下MultiDex的产生背景。 当Android系统安装一个应用的时候,有一步是对Dex进行优化,这个过程有一个专门的工具来处理,叫DexOpt。DexOpt的...

python使用xmlrpc实例讲解

RPC是Remote Procedure Call的缩写,翻译成中文就是远程方法调用,是一种在本地的机器上调用远端机器上的一个过程(方法)的技术,这个过程也被大家称为“分布式计算”,是为...

Pytorch 保存模型生成图片方式

三通道数组转成彩色图片 img=np.array(img1) img=img.reshape(3,img1.shape[2],img1.shape[3])...