【PyTorch】保存和载入模型的两种方法

import torch
import argparseparser = argparse.ArgumentParser("-")
parser.add_argument("--pre_trained_model_type", type=str, choices=["model", "params"], default="model")
args = parser.parse_args()# 载入模型和载入参数
if torch.cuda.is_available():if opt.pre_trained_model_type == "model":model = torch.load(opt.pre_trained_model_path)if opt.pre_trained_model_type == "params":model = m()model.load_state_dict(torch.load(opt.pre_trained_model_path))
else:if opt.pre_trained_model_type == "model":model = torch.load(opt.pre_trained_model_path, map_location=lambda storage, loc: storage)if opt.pre_trained_model_type == "params":model = m()model.load_state_dict(torch.load(opt.pre_trained_model_path, map_location=lambda storage, loc: storage))# 保存模型和保存参数
if torch.cuda.is_available():if opt.pre_trained_model_type == "model":model = m()torch.save(model,cuda(), opt.pre_trained_model_path)if opt.pre_trained_model_type == "params":model = m()torch.save(model.cuda().state_dict(), opt.pre_trained_model_path)
else:if opt.pre_trained_model_type == "model":model = m()torch.save(model,cpu(), opt.pre_trained_model_path)if opt.pre_trained_model_type == "params":model = m()torch.save(model.cpu().state_dict(), opt.pre_trained_model_path)

【PyTorch】保存和载入模型的两种方法相关推荐

加载dict_PyTorch 7.保存和加载pytorch模型的两种方法
众所周知,python的对象都可以通过torch.save和torch.load函数进行保存和加载(不知道?那你现在知道了(*^_^*)),比如: x1 = {"d":" ...
pytorch保存模型的两种方法
文章目录前言一.保存整个模型二.只保存参数模型不同后缀名的区别总结前言模型的本质是一堆用某种结构存储起来的参数用数据对模型进行训练后得到了比较理想的模型,就需要将其存储起来,然后在需要 ...
pytorch加载自己的图片数据集的两种方法
目录 ImageFolder 加载数据集使用pytorch提供的Dataset类创建自己的数据集. Dataset加载数据集接下来我们就可以构建我们的网络架构: 训练我们的网络: 保存网络模型(这 ...
PDF体积太大怎么缩小？这两种方法轻松解决
在我们日常处理的文件中,PDF文件的体积已经算是比较小的文件了,但是随着工作时间增加,我们用到的PDF文件也越来越多,而且有些PDF文件的内容非常丰富,文件体积变得更大,这就不利于我们将文件传输给别人 ...
将pdf转换成jpg格式的两种方法
文件格式转换很难?很多人有这样的想法是因为转换时用了不正确的方法导致转换失败而又找不到好用的方法.其实只要用对了转换方法转换文件是一件非常轻松的事情,下面就以pdf转换成jpg为例,说说文件转换的方法 ...
TensorFlow 模型保存/载入的两种方法
TensorFlow 模型保存/载入我们在上线使用一个算法模型的时候,首先必须将已经训练好的模型保存下来.tensorflow保存模型的方式与sklearn不太一样,sklearn很直接,一个skl ...
matlab 函数到空间,Matlab中保存函数内部中间变量到工作空间的两种方法
比如函数 function y=myfunt1(x) a=5e-6; d=20e-6; boc=0.6328e-7; f=linspace(10e-3,10e-3,100); z=2; y=exp(- ...
DL之DNN：自定义2层神经网络TwoLayerNet模型(计算梯度两种方法)利用MNIST数据集进行训练、预测
DL之DNN:自定义2层神经网络TwoLayerNet模型(计算梯度两种方法)利用MNIST数据集进行训练.预测导读利用python的numpy计算库,进行自定义搭建2层神经网络TwoLayerN ...
NLP：两种方法(自定义函数和封装函数)实现提取两人对话内容(***分隔txt文档)，并各自保存为txt文档
NLP:两种方法(自定义函数和封装函数)实现提取两人对话内容(***分隔txt文档),并各自保存为txt文档目录问题探究实现代码问题探究实现代码 f=open("niu.txt&q ...

【PyTorch】保存和载入模型的两种方法

【PyTorch】保存和载入模型的两种方法相关推荐

最新文章

热门文章