相信不少同学在学习了不少图像分类的模型后,都想将不同的经典网络作为backbone,使用在不同的任务中,比如:姿态估计、目标检测、语义分割等。
本文主要记载我将Light Weight OpenPose的backbone由原来的MobileNet修改成RepVGG,这里我们不针对修改的意义展开讨论。本次仅以姿态估计为例子,希望可以给大家为修改其他网络带来思考,后续如果修改YOLO等网络的时候,也会尝试发教程。
本次使用的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 这里我做简单介绍,这个姿态估计网络是由英特尔对原来的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 作为模型丢入到这个项目的主文件下(这不用纠结,如果你有良好的项目文件存放结构习惯,你也可以放在你喜欢的地方,因为我们只需要导入这个文件)
我们找到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的输入通道数,我这里纯粹就是为了耍酷。但是两个直接效果是什么样子的,你可以实验证明他们的区别。
·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 .
需要注意的是
–from-Repvgg这个需要在tain.py修改,原来名字是–from-MobileNet,这个就是为了触发我们使用RepVgg的预训练权重。
关于训练上的操作,建议大家可以详细看一下lw-openpose在github中给出的教程。根据不同需求去修改里面的参数。
官方给出是需要训练三次的。我也还在训练,训练结构后续发出,大家可以先尝试一下。
本文主要针对Light Weight OpenPose的backbone作出了修改,因为我发现网上很少教大家修改backbone的教程,当然不少的baseline的结构不一样,希望这个修改教程能引发你对其他网络的修改思路。
欢迎大家留言评论和讨论。欢迎一键三联!!!嘻嘻嘻。