替换骨干网络之后使用预训练模型进行训练

前言

最近看了几篇使用transformer的文章，于是想用其中的一个transformer模块来替换另一个方法的骨干网络(backbone)，替换完之后跑起来感觉没有什么效果，想着可能是transformer模型要用预训练会好一些。但是。由于是自己把原来方法的backbone替换掉，因此没有现成的直接可以使用的预训练模型来使用，只能从两个方法中提取相应模块的权重然后整合起来当做预训练模型使用。

原理

不同的预训练模型之所以能够东拼一块、西拼一块成为一个可以用的预训练模型，是因为在预训练模型中有相应的键值对（key-value），只要把预训练模型中的对应到自己使用的网络中的键值对进行更新就好了。简单的举个例子：

# 使用的网络中有这样的减值对：
{
   'conv1': [1, 1, 1, 0, 0]}
# 加载的预训练模型中也有这样的键值对，但是值不同，这样你就可以通过字典更新的方式获得到预训练模型的训练权重了。
{
   'conv1': [1, 1, 1, 0, 0]} -> {
   'conv1': [0, 1, 0, 0, 1]}

但是值得注意的一点是，一定要根据自己构建的网络中的键值对来对预训练模型的键值对进行提取，否则将会更新失败。

更新backbone的预训练权重过程

1、查看网络的键值对和预训练模型的键值对

首先查看需要替换的backbone在原始的预训练模型中的键值对，因为这个是作为backbone使用，所以一般打印出来的信息会有‘backbone'几个字，挺好辨认的。代码和效果如下：

def extract_backbone():
	# 查看backbone部分的预训练模型键值对
    backbone_model_path = "backbone.pth"
    backbone_train_model = torch.load(backbone_model_path)
    print(backbone_train_model .keys())

打印结果如下：

odict_keys(['backbone.SA_modules.0.local_chunk.pe.0.conv.weight', 
'backbone.SA_modules.0.local_chunk.pe.0.bn.weight', ...., 
'bbox_head.vote_module.vote_conv.0.bn.bias']

可以看到在打印中的信息就能看到’backbone'几个字，也就能知道我们要提取的数据范围。
接下来查看网络的键值对：

    model = net()
    model_stat_dict = model.state_dict()
    print(model_stat_dict.keys())

打印结果：

odict_keys(['SA_modules.0.local_chunk.pe.0.conv.weight',
 'SA_modules.0.local_chunk.pe.0.bn.weight', ...]

可以看到里面的键的名字虽然不是完全一样'backbone.SA_modules.0.local_chunk.pe.0.conv.weight'对比'SA_modules.0.local_chunk.pe.0.conv.weight'，但是能够知道之间的对应关系，也就能知道该怎么样更新字典的键值对。

2、提取键值对并更新

方法：观察法，观察对应的键值对相差什么样的字符串，然后把多余的字符串去掉，如，把'backbone.SA_modules.0.local_chunk.pe.0.conv.weight'换为'SA_modules.0.local_chunk.pe.0.conv.weight'，

    # 提取backbone部分的权重
    backbone_stat_dict = {
   }
    for i in backbone_train_model.keys():
        if 'backbone' in i and 'FP_modules.1.mlps' not in i:
            backbone_stat_dict[i.replace('backbone.', '')] = backbone_train_model[i]

这里的方法是判断字符串是否在另一个字符串中来定位自己想要的键值对，另外一个就是，因为我把网络的最后一层的输出改变了，例如原始输出是256，现在改为288，那么最后一层的训练权重就不能用了，只能使用默认的值。这一步是需要自己debug或者慢慢观察的出来的，遇到什么bug就解决什么bug就好了。
接下来对自己新建的网络进行权值更新：

    # 更新网络权重
    model_stat_dict.update(backbone_stat_dict)
    model.load_state_dict(model_stat_dict)

更新非backbone部分的预训练权重过程

1、查看键值对

def extract_others():
    other_model_path = "xxx.pth"
    other_train_model = torch.load(other_model_path )
    print(other_train_model['model'].keys())

打印结果：

odict_keys(['module.backbone_net.sa1.mlp_module.layer0.conv.weight',...]

这里的代码跟提取backbone的差不多，多了一个[‘model’]是因为这个预训练模型的所有键值对放在一个叫model的字典中，保存的层次不一样而已，可以看到里面也有’backbone’几个字，但是这一次，提取的就不是‘backbone'部分的预训练权重了。这一次是要把‘backbone部分的训练权重丢掉，保留其他部分。
具体代码：

    other_state_dict = {
   }
    for i in other_train_model['model'].keys():
        if 'backbone_net' not in i and i.replace('module.', '') in model_stat_dict .keys():
            other_state_dict [i.replace('module.', '')] = other_train_model['model'][i]

这里提取的是除backbone部分和在新建的网络中的键值对。

2、更新键值对

    model_stat_dict.update(other_state_dict)
    model.load_state_dict(model_stat_dict)

结果对比

新建网络的默认权重：

('prediction_heads.5.bn1.weight', tensor([1., 1., 1.