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里面常常需要用的技巧。

  1. 直接在训练好的模型开始搭建自己的模型,就是先加载训练好的模型,然后再它基础上定义自己的模型;
model_ft = models.resnet18(pretrained=use_pretrained)
self.conv1 = model_ft.conv1
self.bn = model_ft.bn
... ...
  1. 自己定义好模型,直接加载模型;本部分讲述了两个相关方法,但是就我个人而言目前没有使用过相关的方法,对于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