python神经网络模型怎么保存_Pytorch模型的保存和提取以及保存和提取神经网络 - pytorch中文网...

有时候我们训练了一个模型, 希望保存它下次直接使用,不需要下次再花时间去训练 ，本节我们来讲解一下pytorch序列化语义以及我们保存和提取回归的神经网络.中文文档地址为：序列化语义

基本的pytorch保存和加载模型

保存和提取主要使用torch.save和torch.load方法实现保存和提取

import torch

test_data = torch.FloatTensor(2,3)

# 保存数据

torch.save(test_data, "test_data.pkl")

print test_data

# 提取数据

print torch.load("test_data.pkl")

保存和提取神经网络

第一种：只保存和加载模型参数(推荐使用)

# 保存

torch.save(the_model.state_dict(), PATH)

# 提取

the_model = TheModelClass(*args, **kwargs)

the_model.load_state_dict(torch.load(PATH))

使用实例：

net = torch.nn.Sequential(

torch.nn.Linear(1, 10),

torch.nn.ReLU(),

torch.nn.Linear(10, 1)

)

torch.save(net.state_dict(), "net_params.pkl")

# 这种方式将会提取整个神经网络, 网络大的时候可能会比较慢.

print torch.load('net_pa