nn.moudle模型
所有神经网络模块的基类(linear,criterion,bn。。注意optimizer不是),我们自建的神经网络也要继承自他。
最著名的参数:
cuda() 所有模块的cuda都是调用的这个
forward()向前传播
state_dict()输出自己的状态字典
模型可以认为是键值对,是很多块(块的特点是继承了nn.moudle)组成的
模型保存
教程
源码详解Pytorch的state_dict和load_state_dict - 知乎
torch的save和load和pickle基本是一样的,都是数据存储
torch.save
torch.save({'epoch': epoch_log, 'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict(),
'scheduler': scheduler.state_dict(), 'best_iou': best_iou,
'is_best': is_best}
, filename)
filename为保存路径+保存文件名称(以pth结尾)
第一个参数传入需要保存的数据,这里是字典。
model.state_dict()是一个顺序字典,是模型的参数,分为需要更新的也就是parameter和不需要更新的buffer。优化器和scheduler也需要记录,因为他们也有参数要记录
torch.load
checkpoint = torch.load(args.model_path)
将刚刚的pth文件读入即可,checkpoint 就是刚刚存入的字典,也可以是任何数据结构,取决于存入什么。
state_dict
Pytorch:模型的保存与加载 torch.save()、torch.load()、torch.nn.Module.load_state_dict()_宁静致远*的博客-CSDN博客
非常好的知乎教程
Pytorch模型中的parameter与buffer - 知乎
state_dict可以由model.state_dict()得到,存储了所有要加载的数据,是一个OrderedDict。
OrderedDict:
相较于普通dict在于它里面的keys是有序的,取决于进dict的顺序,进dict的顺序就是你的模块在init里定义的顺序,这里在state_dict中的顺序取决于该模块在__init__中定义的顺序。
比如
class TheModelClass(nn.Module):
def __init__(self):
super(TheModelClass, self).__init__()
self.bn=nn.BatchNorm1d(10)
self.conv1 = nn.Conv2d(3, 6, 5)
self.pool = nn.MaxPool2d(2, 2)
self.conv2 = nn.Conv2d(6, 16, 5)
self.fc1 = nn.Linear(16 * 5 * 5, 120)
self.fc2 = nn.Linear(120, 84)
self.fc3 = nn.Linear(84, 10)
self.hahaTest=HahaClass()
class HahaClass(nn.Module):
def __init__(self):
super(HahaClass, self).__init__()
self.bn=nn.BatchNorm1d(10)
self.conv1 = nn.Conv2d(3, 6, 5)
>>>model.state_dict().keys()
odict_keys(['bn.weight', 'bn.bias', 'bn.running_mean', 'bn.running_var', 'bn.num_batches_tracked', 'conv1.weight', 'conv1.bias', 'conv2.weight', 'conv2.bias', 'fc1.weight', 'fc1.bias', 'fc2.weight', 'fc2.bias', 'fc3.weight', 'fc3.bias', 'hahaTest.bn.weight', 'hahaTest.bn.bias', 'hahaTest.bn.running_mean', 'hahaTest.bn.running_var', 'hahaTest.bn.num_batches_tracked', 'hahaTest.conv1.weight', 'hahaTest.conv1.bias'])
可以看到确实是bn层先被加入到字典中,hahaTest模块最后定义,所以在最后
还可以看出keys的命名规则:
就是最外层模块默认没有名字,然后内部的就是按照在init中定义的变量名来命名
state_dict存储了什么parameter和buffer
模型中需要保存下来的参数包括两种:
- 一种是反向传播需要被optimizer更新的,称之为 parameter
bias和weight为parameter
- 一种是反向传播不需要被optimizer更新,称之为 buffer
running_mean,running_var,num_batches_tracked等不需要方向传播更新也就是没有梯度的为buffer
parameter都是有grad的,需要更新,buffer都是没grad的,不需要更新,这就是最大的区别
如何创建这两种的参数?
见知乎教程
- 我们可以直接将模型的成员变量(http://self.xxx) 通过
nn.Parameter()
创建,会自动注册到parameters中,可以通过model.parameters() 返回,并且这样创建的参数会自动保存到OrderDict中去;
self.param = nn.Parameter(torch.randn(3, 3))
这样就行了很简单
- 通过
nn.Parameter()
创建普通Parameter对象,不作为模型的成员变量,然后将Parameter对象通过register_parameter()进行注册,可以通model.parameters()
返回,注册后的参数也会自动保存到OrderDict
中去;
model移动与参数的交互
model移动比如model.cuda()和model.todevice()的参数移动
model.cuda()就是将所有buffer和parameter的参数都放在cuda上,也就是最终操作的都是OrderDict,其他的变量都不动,
class MyModel(nn.Module):
def __init__(self):
super(MyModel, self).__init__()
self.my_tensor = torch.randn(1) # 参数直接作为模型类成员变量
self.register_buffer('my_buffer', torch.randn(1)) # 参数注册为 buffer
self.my_param = nn.Parameter(torch.randn(1))
def forward(self, x):
return x
model = MyModel()
print(model.state_dict())
model.cuda()
print(model.my_tensor)
print(model.my_buffer)
输出结果:
可以看到模型类的成员变量不在OrderDict中,不能进行保存,在的才能保存;
model.state_dict()
返回该模型的state_dict
model.load_state_dict
1. load_state_dict
这个函数会递归地对模型进行参数恢复。
首先我们需要明确state_dict这个变量表示你之前保存的模型参数序列,而_load_from_state_dict函数中的local_state表示你的代码中定义的模型的结构。
我的理解就是按照key去进行某一个元素的复制,加入有state_dict1要复制到state_dict2中,那就是按照key进行逐个的复制,key对上了就将value转移过去
2.if strict:
这个部分的作用是判断上面参数拷贝过程中是否有unexpected_keys
或者missing_keys
,如果有就报错,代码不能继续执行。当然,如果strict=False
,则会忽略这些细节。
model.load_state_dict(new_state_dict, strict=True)
加载状态字典,strict要求不能有对不上的keys,比如unexpected_keys或者missing_keys,也就是状态字典key和model必须完全一一对应才行,否则会直接报错,我们自己的模型肯定是要用到这个参数的。
对于分布式模型需要注意的点
比如我用的模型checkpoint在放入DistributedDataParallel之后,state_dict的keys的前面都会加上一个moudle,比如
odict_keys(['enc1.0.linear.weight', 'enc1.0.bn.weight', 'enc1.0.bn.bias', 'enc1.0.bn.running_mean'
变为了
odict_keys(['module.enc1.0.linear.weight', 'module.enc1.0.bn.weight', 'module.enc1.0.bn.bias', 'module.enc1.0.bn.running_mean',
此时我们在load时如果无法将key进行转换为同名就会报错,
解决方法:在新模型进行DistributedDataParallel之后,此时他的状态字典前面也会有moudle,再读入可以。
也可以将所有checkpoint的moudle都去掉
state_dict = checkpoint['state_dict'] #状态字典,也就是模型的权重,本质也是OrderedDict
new_state_dict = collections.OrderedDict() #顺序字典,按照放入字典顺序排序
for k, v in state_dict.items():#将老的字典的东西全都顺序放入新的字典,且key去掉了前面的moudle名
name = k[7:]
new_state_dict[name] = v
也就是对state_dict 的key进行统一的截取
model.named_parameters()
只返回parameters,buffer不返回,且返回的是每一次迭代时[k,v]的list,总体是一个迭代器,内部用yeild实现。
model.parameters()
返回的是上面named_parameters每一次迭代时[k,v]中的[v]
附实验代码
from __future__ import print_function, division
import torch
import torch.nn as nn
import torch.optim as optim
class HahaClass(nn.Module):
def __init__(self):
super(HahaClass, self).__init__()
self.bn=nn.BatchNorm1d(10)
self.conv1 = nn.Conv2d(3, 6, 5)
def forward(self, x):
x = self.pool(F.relu(self.conv1(x)))
x = self.pool(F.relu(self.conv2(x)))
x = x.view(-1, 16 * 5 * 5)
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
x = self.fc3(x)
x=self.bn(x)
return x
# 定义模型
class TheModelClass(nn.Module):
def __init__(self):
super(TheModelClass, self).__init__()
self.bn=nn.BatchNorm1d(10)
self.conv1 = nn.Conv2d(3, 6, 5)
self.pool = nn.MaxPool2d(2, 2)
self.conv2 = nn.Conv2d(6, 16, 5)
self.fc1 = nn.Linear(16 * 5 * 5, 120)
self.fc2 = nn.Linear(120, 84)
self.fc3 = nn.Linear(84, 10)
self.hahaTest=HahaClass()
def forward(self, x):
x = self.pool(F.relu(self.conv1(x)))
x = self.pool(F.relu(self.conv2(x)))
x = x.view(-1, 16 * 5 * 5)
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
x = self.fc3(x)
x=self.bn(x)
return x
# 初始化模型
model = TheModelClass()
# 初始化优化器
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
# 打印模型的 state_dict
print("Model's state_dict:")
for param_tensor in model.state_dict():
print(param_tensor, "\t", model.state_dict()[param_tensor].size())
# 打印优化器的 state_dict
print("Optimizer's state_dict:")
for var_name in optimizer.state_dict():
print(var_name, "\t", optimizer.state_dict()[var_name])