前言
相信不少同学在学习了不少图像分类的模型后，都想将不同的经典网络作为backbone，使用在不同的任务中，比如：姿态估计、目标检测、语义分割等。
本文主要记载我将Light Weight OpenPose的backbone由原来的MobileNet修改成RepVGG，这里我们不针对修改的意义展开讨论。本次仅以姿态估计为例子，希望可以给大家为修改其他网络带来思考，后续如果修改YOLO等网络的时候，也会尝试发教程。
一、BackBone
本次使用的BackBone是前段时间小伙伴介绍的来自清华大学的RepVGG，这个网络的详细介绍大家可以去看原作者的知乎。我会把作者对RepVGG的介绍链接放在下方给到大家。
我们知道OpenPose的backbone原来使用的就是VGG，然后Light Weight OpenPose使用的是MobileNet，这里我就想尝试将RepVGG放到Light Weight OpenPose上。这样修改后可能就不再轻量了。但是我想看看这样精度上有什么变化，也许也可以运用在其他场景。
Rep-VGG：https://zhuanlan.zhihu.com/p/344324470
在作者RepVGG的github上下载我们需要的模型和权重文件。
github：https://github.com/DingXiaoH/RepVGG
本次使用的是作者的RepVGG-A0网络，可以通过repvgg.py查看，关于repvgg的代码，这里不做详细说明。大家自行详细阅读作者博文。
def create_RepVGG_A0(deploy=False):
return RepVGG(num_blocks=[2, 4, 14, 1], num_classes=1000,
width_multiplier=[0.75, 0.75, 0.75, 2.5], override_groups_map=None, deploy=deploy)
这里需要提醒的是，RepVGG的width_multiplier是输出Channel的影响因子，最后一层的输出为1280（512*2.5）
二、Light Weight OpenPose
关于Light Weight OpenPose 这里我做简单介绍，这个姿态估计网络是由英特尔对原来的OpenPose的一次升级，其中它们使用了较为轻量的MobileNet来代替原来的VGG作为backbone。
本文因为针对的是Light Weight OpenPose修改。它的源码的github也放在下方给到大家。大家可以自行下载。
Light-Weight-Openpose-github: https://github.com/Daniil-Osokin/lightweight-human-pose-estimation.pytorch#training
拥有了这个项目以后，我们将我们下载好的repvgg.py 作为模型丢入到这个项目的主文件下（这不用纠结，如果你有良好的项目文件存放结构习惯，你也可以放在你喜欢的地方，因为我们只需要导入这个文件）
1.修改backbone
我们找到Lw-openpose的模型文件，它存放在项目中的models文件夹下的with_mobilenet.py
1.导入我们的repvgg中的repvgg-A0
from repvgg import create_RepVGG_A0 as RepVGG
2.找到class PoseEstimationWithMobileNet这个类 ，我们将在里面修改我们的backbone
class PoseEstimationWithRepVgg(nn.Module):
def __init__(self, num_refinement_stages=1, num_channels=128, num_heatmaps=19, num_pafs=38):
super().__init__()
# self.model = nn.Sequential(
# conv( 3, 32, stride=2, bias=False),
# conv_dw( 32, 64),
# conv_dw( 64, 128, stride=2),
# conv_dw(128, 128),
# conv_dw(128, 256, stride=2),
# conv_dw(256, 256),
# conv_dw(256, 512), # conv4_2
# conv_dw(512, 512, dilation=2, padding=2),
# conv_dw(512, 512),
# conv_dw(512, 512),
# conv_dw(512, 512),
# conv_dw(512, 512) # conv5_5
# )
self.model = RepVGG() # 将backbone修改成我们导入的RepVGG
# 因为最后输出的一层是1280的channel 但是Cpm输入的是512的channel 我在这里用了一个1x1的卷积进行了一个下采样操作。当然你也可以尝试修改Cpm输入Channe为1280。
self.conv1x1 = nn.Sequential(
nn.Conv2d(1280,512,kernel_size=1,stride=1,bias=False)
)
self.cpm = Cpm(512, num_channels)
self.initial_stage = InitialStage(num_channels, num_heatmaps, num_pafs)
self.refinement_stages = nn.ModuleList()
for idx in range(num_refinement_stages):
self.refinement_stages.append(RefinementStage(num_channels + num_heatmaps + num_pafs, num_channels,
num_heatmaps, num_pafs))
代码中，我将model替换成了RepVGG，同时添加了一个1x1的卷积进行下采样，将原来的通道数从1280变成了512了。当然你可以直接修改Cmp的输入通道数，我这里纯粹就是为了耍酷。但是两个直接效果是什么样子的，你可以实验证明他们的区别。
2.修改train.py中的内容
·1. 修改优化器中需要优化的参数，主要针对RepVGG和新增的1x1卷积层
optimizer = optim.Adam([
# net.model 部分就是我们的backbone部分，我们将它的相关权重拿出来，给到优化器。这里使用的get_parameters_repvgg函数在后面告诉写在哪。
{'params': get_parameters_repvgg(net.model), 'lr': base_lr},
{'params': get_parameters_repvgg(net.conv1x1), 'lr': base_lr},
# {'params': get_parameters(net.model, 'weight')},
# {'params': get_parameters_conv_depthwise(net.model, 'weight'), 'weight_decay': 0},
# {'params': get_parameters_bn(net.model, 'weight'), 'weight_decay': 0},
# {'params': get_parameters_bn(net.model, 'bias'), 'lr': base_lr * 2, 'weight_decay': 0},
{'params': get_parameters_conv(net.cpm, 'weight'), 'lr': base_lr},
{'params': get_parameters_conv(net.cpm, 'bias'), 'lr': base_lr * 2, 'weight_decay': 0},
{'params': get_parameters_conv_depthwise(net.cpm, 'weight'), 'weight_decay': 0},
{'params': get_parameters_conv(net.initial_stage, 'weight'), 'lr': base_lr},
{'params': get_parameters_conv(net.initial_stage, 'bias'), 'lr': base_lr * 2, 'weight_decay': 0},
{'params': get_parameters_conv(net.refinement_stages, 'weight'), 'lr': base_lr * 4},
{'params': get_parameters_conv(net.refinement_stages, 'bias'), 'lr': base_lr * 8, 'weight_decay': 0},
{'params': get_parameters_bn(net.refinement_stages, 'weight'), 'weight_decay': 0},
{'params': get_parameters_bn(net.refinement_stages, 'bias'), 'lr': base_lr * 2, 'weight_decay': 0},
], lr=base_lr, weight_decay=5e-4)
net.model 部分就是我们的backbone部分，我们将它的相关权重拿出来，给到优化器。这里使用的get_parameters_repvgg函数见下面。
2.添加get_parameters_repvgg函数
get_parameters_repvgg（）函数主要是把每个参数读出来，在modules/get_parameters.py中添加这个函数。
def get_parameters_repvgg(model):
for m in model.modules():
for param_name, param in m.named_parameters():
yield param
3.删除RepVGG-A0权重中的FC部分（全连接层部分）
导入RepVGG的权重，需要删除FC部分的参数。需要注意的是，我们在load_from_RepVgg(）记得将net修改成net.model。仅仅针对backbone部分进行权重加载。load_from_RepVgg(）函数写入到modules/load_state.py中
if checkpoint_path:
checkpoint = torch.load(checkpoint_path)
# # 删除linear部分的weight和bias
del checkpoint['linear.weight']
del checkpoint['linear.bias']
if from_Repvgg:
load_from_RepVgg(net.model, checkpoint)
else:
load_state(net, checkpoint)
if not weights_only:
optimizer.load_state_dict(checkpoint['optimizer'])
scheduler.load_state_dict(checkpoint['scheduler'])
num_iter = checkpoint['iter']
current_epoch = checkpoint['current_epoch']
modules/load_state.py
加入这个代码
def load_from_RepVgg(net, checkpoint):
net.load_state_dict(checkpoint)
3.修改两个地方的stride
修改train中的stride为16，并且修改RepVgg的stage4的strid为1，这样他们最后算loss的时候高宽是一致的，变成23*23。效果有待验证…
4.训练命令
python train.py --train-images-folder .<COCO_HOME>/train2017/ --prepared-train-labels prepared_train_annotation.pkl --val-labels val_subset.json --val-images-folder .<COCO_HOME>/val2017/ --checkpoint-path RepVGG-A0-train.pth --from-Repvgg
需要注意的是<COCO_HOME>是你存放COCO数据集的目录地址
–from-Repvgg这个需要在tain.py修改，原来名字是–from-MobileNet，这个就是为了触发我们使用RepVgg的预训练权重。
关于训练上的操作，建议大家可以详细看一下lw-openpose在github中给出的教程。根据不同需求去修改里面的参数。
官方给出是需要训练三次的。我也还在训练，训练结构后续发出，大家可以先尝试一下。
总结
本文主要针对Light Weight OpenPose的backbone作出了修改，因为我发现网上很少教大家修改backbone的教程，当然不少的baseline的结构不一样，希望这个修改教程能引发你对其他网络的修改思路。
欢迎大家留言评论和讨论。欢迎一键三联！！！嘻嘻嘻。
3085
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
