CenterTrack论文:Tracking Objects as Points
CenterTrack源码:CenterTrack
CenterTrack网络结构
CenterTrack的输入是将当前帧和前一帧的图片,以及前一帧的heatmap一起输入网络。默认使用的主干网络是DLA34,融合模块也是用的DLAUP结构,输出的是主要信息就是中心点热力图、中心点偏移量、宽高、预测框偏移。
图片参考于知乎文章
CenterTrack替换主干网络
CenterTrack代码的src/lib/model/neckworks/backbone目录下存放的为主干网络的代码文件,里面除了dla34外,还有resnet、mobilenetv2,本文的目的就加一个ghostnet.py
这个代码是从别人哪找的,将其修改为了适合CenterTrack的形式,主要就是添加了不同特征层的多阶段输出。
"""
Creates a GhostNet Model as defined in:
GhostNet: More Features from Cheap Operations By Kai Han, Yunhe Wang, Qi Tian, Jianyuan Guo, Chunjing Xu, Chang Xu.
https://arxiv.org/abs/1911.11907
Modified from https://github.com/d-li14/mobilenetv3.pytorch
"""
import torch
import torch.nn as nn
import math
cfgs = [
# k, t, c, SE, s
[[3, 16, 16, 0, 1]],
[[3, 48, 32, 0, 2]],
[[3, 72, 32, 0, 1]],
[[5, 72, 64, 1, 2]],
[[5, 120, 64, 1, 1]],
[[3, 240, 96, 0, 2]],
[[3, 200, 96, 0, 1],
[3, 184, 96, 0, 1],
[3, 184, 96, 0, 1],
[3, 480, 120, 1, 1],
[3, 672, 120, 1, 1]],
[[5, 672, 320, 1, 2]],
[[5, 960, 320, 0, 1],
[5, 960, 320, 1, 1],
[5, 960, 320, 0, 1],
[5, 960, 320, 1, 1]]
]
def _make_divisible(v, divisor, min_value=None):
"""
This function is taken from the original tf repo.
It ensures that all layers have a channel number that is divisible by 8
It can be seen here:
https://github.com/tensorflow/models/blob/master/research/slim/nets/mobilenet/mobilenet.py
"""
if min_value is None:
min_value = divisor
new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)
# Make sure that round down does not go down by more than 10%.
if new_v < 0.9 * v:
new_v += divisor
return new_v
class SELayer(nn.Module):
def __init__(self, channel, reduction=4):
super(SELayer, self).__init__()
self.avg_pool = nn.AdaptiveAvgPool2d(1)
self.fc = nn.Sequential(
nn.Linear(channel, channel // reduction),
nn.ReLU(inplace=True),
nn.Linear(channel // reduction, channel), )
def forward(self, x):
b, c, _, _ = x.size()
y = self.avg_pool(x).view(b, c)
y = self.fc(y).view(b, c, 1, 1)
y = torch.clamp(y, 0, 1)
return x * y
def depthwise_conv(inp, oup, kernel_size=3, stride=1, relu=False):
return nn.Sequential(
nn.Conv2d(inp, oup, kernel_size, stride, kernel_size // 2, groups=inp, bias=False),
nn.BatchNorm2d(oup),
nn.ReLU(inplace=True) if relu else nn.Sequential(),
)
class GhostModule(nn.Module):
def __init__(self, inp, oup, kernel_size=1, ratio=2, dw_size=3, stride=1, relu=True):
super(GhostModule, self).__init__()
self.oup = oup
init_channels = math.ceil(oup / ratio)
new_channels = init_channels * (ratio - 1)
self.primary_conv = nn.Sequential(
nn.Conv2d(inp, init_channels, kernel_size, stride, kernel_size // 2, bias=False),
nn.BatchNorm2d(init_channels),
nn.ReLU(inplace=True) if relu else nn.Sequential(),
)
self.cheap_operation = nn.Sequential(
nn.Conv2d(init_channels, new_channels, dw_size, 1, dw_size // 2, groups=init_channels, bias=False),
nn.BatchNorm2d(new_channels),
nn.ReLU(inplace=True) if relu else nn.Sequential(),
)
def forward(self, x):
x1 = self.primary_conv(x)
x2 = self.cheap_operation(x1)
out = torch.cat([x1, x2], dim=1)
return out[:, :self.oup, :, :]
class GhostBottleneck(nn.Module):
def __init__(self, inp, hidden_dim, oup, kernel_size, stride, use_se):
super(GhostBottleneck, self).__init__()
assert stride in [1, 2]
self.conv = nn.Sequential(
# pw
GhostModule(inp, hidden_dim, kernel_size=1, relu=True),
# dw
depthwise_conv(hidden_dim, hidden_dim, kernel_size, stride, relu=False) if stride == 2 else nn.Sequential(),
# Squeeze-and-Excite
SELayer(hidden_dim) if use_se else nn.Sequential(),
# pw-linear
GhostModule(hidden_dim, oup, kernel_size=1, relu=False),
)
if stride == 1 and inp == oup:
self.shortcut = nn.Sequential()
else:
self.shortcut = nn.Sequential(
depthwise_conv(inp, inp, kernel_size, stride, relu=False),
nn.Conv2d(inp, oup, 1, 1, 0, bias=False),
nn.BatchNorm2d(oup),
)
def forward(self, x):
return self.conv(x) + self.shortcut(x)
class GhostNet(nn.Module):
def __init__(self, opt, channel_stage=None, width_mult=1.):
super(GhostNet, self).__init__()
# setting of inverted residual blocks
self.cfgs = cfgs
# building first layer
# output_channel = _make_divisible(16 * width_mult, 4)
self.channels = [32, 16, 32, 64, 120, 320]
output_channel = 32
self.con_stem = nn.Sequential(
nn.Conv2d(3, output_channel, 3, 2, 1, bias=False),
nn.BatchNorm2d(output_channel),
nn.ReLU(inplace=True)
)
input_channel = output_channel
if opt.pre_img:
print('adding pre_img layer...')
self.pre_img_layer = nn.Sequential(
nn.Conv2d(3, input_channel, kernel_size=3, padding=1, stride=2, bias=False),
nn.BatchNorm2d(input_channel))
if opt.pre_hm:
print('adding pre_hm layer...')
self.pre_hm_layer = nn.Sequential(
nn.Conv2d(1, input_channel, kernel_size=3, padding=1, stride=2, bias=False),
nn.BatchNorm2d(input_channel))
# building inverted residual blocks
stages = []
block = GhostBottleneck
for cfg in self.cfgs:
layers = []
for k, exp_size, c, use_se, s in cfg:
output_channel = _make_divisible(c * width_mult, 4)
hidden_channel = _make_divisible(exp_size * width_mult, 4)
layers.append(block(input_channel, hidden_channel, output_channel, k, s, use_se))
input_channel = output_channel
stages.append(nn.Sequential(*layers))
self.blocks1 = nn.Sequential(
stages[0]
)
self.blocks2 = nn.Sequential(
stages[1],
stages[2]
)
self.blocks3 = nn.Sequential(
stages[3],
stages[4],
)
self.blocks4 = nn.Sequential(
stages[5],
stages[6]
)
self.blocks5 = nn.Sequential(
stages[7],
stages[8]
)
self._initialize_weights()
def forward(self, x, pre_img=None, pre_hm=None):
x = self.con_stem(x)
if pre_img is not None:
x = x + self.pre_img_layer(pre_img)
if pre_hm is not None:
x = x + self.pre_hm_layer(pre_hm)
y = [x]
x = self.blocks1(x)
y.append(x)
x = self.blocks2(x)
y.append(x)
x = self.blocks3(x)
y.append(x)
x = self.blocks4(x)
y.append(x)
x = self.blocks5(x)
y.append(x)
return y
def _initialize_weights(self):
for m in self.modules():
if isinstance(m, nn.Conv2d):
nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
elif isinstance(m, nn.BatchNorm2d):
m.weight.data.fill_(1)
m.bias.data.zero_()
# def ghost_net(**kwargs):
# """
# Constructs a GhostNet model
# """
# cfgs = [
# # k, t, c, SE, s
# [[3, 16, 16, 0, 1]],
#
# [[3, 48, 24, 0, 2]],
# [[3, 72, 24, 0, 1]],
#
# [[5, 72, 32, 1, 2]],
# [[5, 120, 32, 1, 1]],
#
# [[3, 240, 64, 0, 2]],
# [[3, 200, 64, 0, 1],
# [3, 184, 64, 0, 1],
# [3, 184, 64, 0, 1],
# [3, 480, 112, 1, 1],
# [3, 672, 112, 1, 1]],
#
# [[5, 672, 320, 1, 2]],
# [[5, 960, 320, 0, 1],
# [5, 960, 320, 1, 1],
# [5, 960, 320, 0, 1],
# [5, 960, 320, 1, 1]]
# ]
# return GhostNet(cfgs, **kwargs)
if __name__ == '__main__':
model = GhostNet()
model.eval()
print(model)
input = torch.randn(8, 3, 544, 960)
y = model(input)
print(y)
在CenterTrack代码的src/lib/model/neckworks目录下有一个generate_network.py文件,其功能就是根据前面有的backbone创建模型。在这个文件里面需要导入并添加ghostnet,
#导入ghostnet
from lib.model.networks.backbones.ghostnet import GhostNet
backbone_factory = {
'dla34': dla34,
'resnet': Resnet,
'mobilenetv2': MobileNetV2,
'ghostnet': GhostNet # 添加ghostnet
}
最后只需要在src/lib/opt.py中修改‘arch’参数改为‘generic’,‘backbone’参数为‘ghostnet’即可
实验结果
速度方面相比,相比dla34能快将近一倍作用,mota自然嘎嘎下降,哈哈哈