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_params.pkl')
# 打印输出我们会发现，上面只保存了模型参数

# 将保存的参数复制到 net,这种方式将会提取所有的参数, 然后再放到你的新建网络中.
net.load_state_dict(torch.load('net_params.pkl'))

第二种：保存和加载整个模型

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

# 保存
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)
)