pytorch 加载模型_【pytorch】模型的搭建保存加载

使用pytorch进行网络模型的搭建、保存与加载，是非常快速、方便的、妙不可言的。

搭建ConvNet

所有的网络都要继承torch.nn.Module，然后在构造函数中使用torch.nn中的提供的接口定义layer的属性，最后，在forward函数中将各个layer连接起来。

下面，以LeNet为例：

class 

这样一来，我们就搭建好了网络模型，是不是很简洁明了呢？此外，还可以使用torch.nn.Sequential，更方便进行模块化的定义，如下：

class 

Module有很多属性，可以查看权重、参数等等；如下：

net = lenet.LeNet()
print(net)

for param in net.parameters():
     print(type(param.data), param.size())
     print(list(param.data)) 

print(net.state_dict().keys())
#参数的keys

for key in net.state_dict():#模型参数
    print key, 'corresponds to', list(net.state_dict()[key])

那么，如何进行参数初始化呢？使用 torch.nn.init ，如下：

def 

保存ConvNet

使用torch.save()对网络结构和模型参数的保存，有两种保存方式：

• 保存整个神经网络的的结构信息和模型参数信息，save的对象是网络net；
• 保存神经网络的训练模型参数，save的对象是net.state_dict()。

torch.save(net1, 'net.pkl')  # 保存整个神经网络的结构和模型参数    
torch.save(net1.state_dict(), 'net_params.pkl') # 只保存神经网络的模型参数    

加载ConvNet

对应上面两种保存方式，重载方式也有两种。

• 对应第一种完整网络结构信息，重载的时候通过torch.load(‘.pth’)直接初始化新的神经网络对象即可。
• 对应第二种只保存模型参数信息，需要首先导入对应的网络，通过net.load_state_dict(torch.load('.pth'))完成模型参数的重载。

在网络比较大的时候，第一种方法会花费较多的时间，所占的存储空间也比较大。

# 保存和加载整个模型  
torch.save(model_object, 'model.pth')  
model = torch.load('model.pth')  

# 仅保存和加载模型参数  
torch.save(model_object.state_dict(), 'params.pth')  
model_object.load_state_dict(torch.load('params.pth'))  

---------------------------------------------------------------------------------------------------------------------------

在各方小伙伴的努力和支持下，pytorch中文文档第一版终于上线啦！！！（鼓掌）文档还有很多小瑕疵，但是大体可以放心使用了～我们遵循快速迭代的原则，所以赶紧上线第一版来接受广大开源社区的意见和建议。欢迎加入我们！

pytorch中文文档：http://pytorch-cn.readthedocs.io/zh/latest/

github项目地址：awfssv/pytorch-cn

中文翻译组QQ群：628478868

还有pytorch项目交流群：613523596

欢迎关注！