pytorch是一个非常好用的模型,值得花时间去理解去搞一搞。
1 模型保存与使用
torch.save
:将对象序列化到硬盘上,该对象可以是 Models, tensors和 dictionaries 等。实际上是使用了python的 pickle方法。torch.load
:将硬盘上序列化的对象加载设备中。实际是使用了pickle的解包方法。torch.nn.Module.load_state_dict
:通过反序列化state_dict加载模型参数字典。
1.1 模型参数
在pytorch中torch.nn.Module模型的参数存放在模型的parameters中(model.parameters()),而state_dict是参数tensor的字典。仅当该层有可学习的参数才会有属性state_dict。torch.optim也有state_dict属性,其中包含的是优化器的参数,即网络所使用的超参数。
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
class TheModelClass(nn.Module):
def __init__(self):
super(TheModelClass, self).__init__()
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)
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)
return x
# 初始化模型
model = TheModelClass()
# 初始化优化器
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
# 打印模型参数
print("Model's state_dict:")
for param_tensor in model.state_dict():
print(param_tensor, "\t", model.state_dict()[param_tensor].size())
# 打印优化器参数
print("Optimizer's state_dict:")
for var_name in optimizer.state_dict():
print(var_name, "\t", optimizer.state_dict()[var_name])
1.2 保存和加载模型
保存的文件结尾一般使用.pt或.pth。最后的model.eval()
表示将drop和batch nromalization层设置为测试模式。可通过mode.train()
转化为训练模式。
model.eval(),pytorch会自动把BN和DropOut固定住,不会取平均,而是用训练好的值。不然的话,一旦test的batch_size过小,很容易就会被BN层导致生成图片颜色失真极大;在模型测试阶段使用model.train() 让model变成训练模式,此时 dropout和batch normalization的操作在训练过程中起到防止网络过拟合的问题
总结: model.train() 和 model.eval() 一般在模型训练和评价的时候会加上这两句,主要是针对由于model 在训练时和评价时 Batch Normalization 和 Dropout 方法模式不同;因此,在使用PyTorch进行训练和测试时一定注意要把实例化的model指定train/eval;model.eval() 和 model.train()这两句话的重要性就在这儿了
保存模型参数:
#保存模型参数
torch.save(model.state_dict(), PATH)
#加载模型参数
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.eval()
保存模型:
#保存模型
torch.save(model, PATH)
#加载模型
model = torch.load(PATH)
model.eval()
加载不同的模型:
在很多时候,我们加载的是已经训练好的模型,而训练好的模型可能与我们定义的模型不完全一样,而我们只想使用一样的那些层的参数。这个时候其实finetune里面常常需要用的技巧。
- 直接在训练好的模型开始搭建自己的模型,就是先加载训练好的模型,然后再它基础上定义自己的模型;
model_ft = models.resnet18(pretrained=use_pretrained)
self.conv1 = model_ft.conv1
self.bn = model_ft.bn
... ...
- 自己定义好模型,直接加载模型;本部分讲述了两个相关方法,但是就我个人而言目前没有使用过相关的方法,对于strict参数也是第一次使用。
#第一种方法:
mymodelB = TheModelBClass(*args, **kwargs)
# strict=False,设置为false,只保留键值相同的参数
mymodelB.load_state_dict(model_zoo.load_url(model_urls['resnet18']), strict=False)
#第二种方法:
# 加载模型
model_pretrained = models.resnet18(pretrained=use_pretrained)
# mymodel's state_dict,
# 如: conv1.weight
# conv1.bias
mymodelB_dict = mymodelB.state_dict()
# 将model_pretrained的建与自定义模型的建进行比较,剔除不同的
pretrained_dict = {k: v for k, v in model_pretrained.items() if k in mymodelB_dict}
# 更新现有的model_dict
mymodelB_dict.update(pretrained_dict)
# 加载我们真正需要的state_dict
mymodelB.load_state_dict(mymodelB_dict)
加载不同设备的模型
将由GPU保存的模型加载到CPU上。
将torch.load()函数中的map_location参数设置为torch.device(‘cpu’);博客中说可以使用如下方法,感觉是比较冗长的语句。我这边直接:torch.load("modle_weights",map_location='cpu')
device = torch.device('cpu')
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH, map_location=device))
将由GPU保存的模型加载到GPU上。确保对输入的tensors调用input = input.to(device)方法。。map_location是将模型加载到GPU上,model.to(torch.device(‘cuda’))是将模型参数加载为CUDA的tensor。最后保证使用.to(torch.device(‘cuda’))方法将需要使用的参数放入CUDA。
device = torch.device("cuda")
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.to(device)
将CPU传递的模型传输入指定的GPU模块,然后进行使用。
device = torch.device("cuda")
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH, map_location="cuda:0")) # Choose selected GPU device number you want
model.to(device)
2 参数冻结
也是在Finetune工程中总结的相关技巧,FInetune技巧集中在参数的冻结与学习率的改变上,目前微调的原网络模型较小,因此没有用到渐变式的学习率,但是确实用到了参数冻结技巧,在此进行总结。
2.1 requires_grad
如果你的模型不是nn.sequential,即你使用的是nn.module方法。此时你可以直接在模型定义部分,想要冻结的参数后加上对梯度求解参数的定义。
for p in self.parameters():
p.requires_grad = False
在文件中的使用可能是如下形态:
class RESNET_MF(nn.Module):
def __init__(self, model, pretrained):
super(RESNET_MF, self).__init__()
self.resnet = model(pretrained)
for p in self.parameters():
p.requires_grad = False #预训练模型加载进来后全部设置为不更新参数,然后再后面加层
self.f = SpectralNorm(nn.Conv2d(2048, 512, 1))
self.g = SpectralNorm(nn.Conv2d(2048, 512, 1))
self.h = SpectralNorm(nn.Conv2d(2048, 2048, 1))
同时在优化器中添加:filter(lambda p: p.requires_grad, model.parameters())
optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=0.001, \
betas=(0.9, 0.999), eps=1e-08, weight_decay=1e-5)
2.2 For loop
以上方法针对的是非sequential定义的方式,如果是使用nn.sequential句式定义的网络结构,这个时候是无法插入上一方法关于requires_grad 的定义的。此时可以使用字典循环查找,赋值的方式来进行参数冻结。
在load完model之后,可以查看模型参数,并对想要冻结的网络层进行设定。
最要同样需要在优化器定义部分加上filter实现对参数的屏蔽。
for i,p in enumerate(net.parameters()):
print(i,p)
if i < 5:
p.requires_grad = False