Pytorch 模型保存与使用、Finetune参数冻结、train/eval

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
... ...

2. 自己定义好模型，直接加载模型；本部分讲述了两个相关方法，但是就我个人而言目前没有使用过相关的方法，对于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