pytorch学习1.0系列------04 保存模型

1. 基本的pytorch 的 Tensor保存和加载

保存和提取主要使用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")

2. 保存和提取神经网络

2.1 只保存和加载模型参数（推荐使用）

# 保存
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_params.pkl')
# 打印输出我们会发现，上面只保存了模型参数
 
# 将保存的参数复制到 net,这种方式将会提取所有的参数, 然后再放到你的新建网络中.
net.load_state_dict(torch.load('net_params.pkl'))

2.2 保存和加载整个模型

如下会保存整个网络，如果数据量比较大，会消耗大量时间。占用的内存也比较高，所以不推荐使用

# 保存
torch.save(the_model, PATH)
# 提取
the_model = torch.load(PATH)

实例：

#保存提取一个神经网络
net = torch.nn.Sequential(
        torch.nn.Linear(1, 10),
        torch.nn.ReLU(),
        torch.nn.Linear(10, 1)
    )
 
torch.save(net, "net_params.pkl")
print torch.load('net_params.pkl')

输出：

Sequential (
  (0): Linear (1 -> 10)
  (1): ReLU ()
  (2): Linear (10 -> 1)
)