def __init__(self):
super(HR, self).__init__()
self.conv = ConvBlock(3, 64, 7, 2, add_relu=True)
self.pool = nn.MaxPool2d(3, 2, 1)
self.res2a = ResidualBlock('res2a', 64, 64, 256, downsample=True)
self.res2b = ResidualBlock('res2b', 256, 64, 256)
self.res2c = ResidualBlock('res2c', 256, 64, 256)
self.res3a = ResidualBlock('res3a', 256, 128, 512, downsample=True)
self.res3b1 = ResidualBlock('res3b1', 512, 128, 512)
self.res3b2 = ResidualBlock('res3b2', 512, 128, 512)
self.res3b3 = ResidualBlock('res3b3', 512, 128, 512)
self.res4a = ResidualBlock('res4a', 512, 256, 1024, downsample=True)
self.res4bX = nn.ModuleList()
for i in range(0, 22):
self.res4bX.append(ResidualBlock('res4bX', 1024, 256, 1024))
# Detection Head
self.score_res4 = ConvBlock(1024, 125, 1, bias=True, add_relu=False, add_bn=False)
self.score4 = ConvTransLayer(125, 125, 4, 2, 0)
self.score_res3 = ConvBlock(512, 125, 1, bias=True, add_relu=False, add_bn=False)
self.eltadd = nn.EltAdd()
def __init__(self):
super(SSH, self).__init__()
# backbone
self.vgg16 = nn.ModuleList(make_layers(vgg_cfgs['D']))
# SSH - M3
self.M3 = M_Module(512, 256, 128)
self.M3_bbox_pred = nn.Conv2d(512, 8, 1, 1, 0)
self.M3_cls_score = nn.Conv2d(512, 4, 1, 1, 0)
self.M3_cls_score_softmax = nn.Softmax(dim=1)
# SSH - M2
self.M2 = M_Module(512, 256, 128)
self.M2_bbox_pred = nn.Conv2d(512, 8, 1, 1, 0)
self.M2_cls_score = nn.Conv2d(512, 4, 1, 1, 0)
self.M2_cls_score_softmax = nn.Softmax(dim=1)
# SSH - M1
self.conv4_128 = nn.Conv2d(512, 128, 1, 1, 0)
self.conv4_128_relu = nn.ReLU(inplace=True)
self.conv5_128 = nn.Conv2d(512, 128, 1, 1, 0)
self.conv5_128_relu = nn.ReLU(inplace=True)
self.conv5_128_up = nn.ConvTranspose2d(128, 128, 4, 2, 1, groups=128, bias=False)
self.eltadd = nn.EltAdd()
self.conv4_fuse_final = nn.Conv2d(128, 128, 3, 1, 1)
self.conv4_fuse_final_relu = nn.ReLU(inplace=True)
self.M1 = M_Module(128, 128, 64)
self.M1_bbox_pred = nn.Conv2d(256, 8, 1, 1, 0)
self.M1_cls_score = nn.Conv2d(256, 4, 1, 1, 0)
self.M1_cls_score_softmax = nn.Softmax(dim=1)
def __init__(self, in_planes, out_planes, stride=1, scale=0.1, map_reduce=8, vision=1, groups=1):
super(BasicRFB, self).__init__()
self.scale = scale
self.out_channels = out_planes
inter_planes = in_planes // map_reduce
self.branch0 = nn.Sequential(
BasicConv(in_planes, inter_planes, kernel_size=1, stride=1, groups=groups, relu=False),
BasicConv(inter_planes, 2 * inter_planes, kernel_size=(3, 3), stride=stride, padding=(1, 1), groups=groups),
BasicConv(2 * inter_planes, 2 * inter_planes, kernel_size=3, stride=1, padding=vision + 1, dilation=vision + 1, relu=False, groups=groups)
)
self.branch1 = nn.Sequential(
BasicConv(in_planes, inter_planes, kernel_size=1, stride=1, groups=groups, relu=False),
BasicConv(inter_planes, 2 * inter_planes, kernel_size=(3, 3), stride=stride, padding=(1, 1), groups=groups),
BasicConv(2 * inter_planes, 2 * inter_planes, kernel_size=3, stride=1, padding=vision + 2, dilation=vision + 2, relu=False, groups=groups)
)
self.branch2 = nn.Sequential(
BasicConv(in_planes, inter_planes, kernel_size=1, stride=1, groups=groups, relu=False),
BasicConv(inter_planes, (inter_planes // 2) * 3, kernel_size=3, stride=1, padding=1, groups=groups),
BasicConv((inter_planes // 2) * 3, 2 * inter_planes, kernel_size=3, stride=stride, padding=1, groups=groups),
BasicConv(2 * inter_planes, 2 * inter_planes, kernel_size=3, stride=1, padding=vision + 4, dilation=vision + 4, relu=False, groups=groups)
)
self.ConvLinear = BasicConv(6 * inter_planes, out_planes, kernel_size=1, stride=1, relu=False)
self.shortcut = BasicConv(in_planes, out_planes, kernel_size=1, stride=stride, relu=False)
self.relu = nn.ReLU(inplace=False)
self.eltadd = nn.EltAdd()
def __init__(self, in_channels_list, out_channels):
super(FPN, self).__init__()
leaky = 0
if (out_channels <= 64):
leaky = 0.1
self.output1 = conv_bn1X1(in_channels_list[0],
out_channels,
stride=1,
leaky=leaky)
self.output2 = conv_bn1X1(in_channels_list[1],
out_channels,
stride=1,
leaky=leaky)
self.output3 = conv_bn1X1(in_channels_list[2],
out_channels,
stride=1,
leaky=leaky)
self.upsample = nn.Upsample(scale_factor=2,
mode='bilinear',
align_corners=False)
self.eltadd = nn.EltAdd()
self.merge1 = conv_bn(out_channels, out_channels, leaky=leaky)
self.merge2 = conv_bn(out_channels, out_channels, leaky=leaky)
def __init__(self, model_type='32'):
super(EXTD, self).__init__()
self.model_type = int(model_type)
self.mobilenet = MobileNetV2(model_type=self.model_type)
if self.model_type == 32:
self.base = nn.ModuleList(self.mobilenet.features)[:8]
elif self.model_type == 48 or self.model_type == 64:
self.base = nn.ModuleList(self.mobilenet.features)[:6]
else:
raise NotImplementedError
self.upsample = nn.Upsample(scale_factor=2,
mode='bilinear',
align_corners=False)
self.eltadd = nn.EltAdd()
self.upfeat = []
for it in range(5):
self.upfeat.append(
upsample(in_channels=self.model_type,
out_channels=self.model_type))
self.upfeat = nn.ModuleList(self.upfeat)
self.loc = []
self.conf = []
self.net_source = [4, 4, 4, 4]
self.feature_dim = []
if self.model_type == 32:
self.feature_dim += [self.base[4].conv[-3].out_channels]
for idx in self.net_source:
self.feature_dim += [self.base[idx].conv[-3].out_channels]
else:
self.feature_dim += [self.base[4].conv[-2].out_channels]
for idx in self.net_source:
self.feature_dim += [self.base[idx].conv[-2].out_channels]
self.loc += [
nn.Conv2d(self.feature_dim[0], 4, kernel_size=3, padding=1)
]
self.conf += [
nn.Conv2d(self.feature_dim[0], 4, kernel_size=3, padding=1)
]
for k, v in enumerate(self.net_source, 1):
self.loc += [
nn.Conv2d(self.feature_dim[k], 4, kernel_size=3, padding=1)
]
self.conf += [
nn.Conv2d(self.feature_dim[k], 2, kernel_size=3, padding=1)
]
self.loc += [nn.Conv2d(self.model_type, 4, kernel_size=3, padding=1)]
self.conf += [nn.Conv2d(self.model_type, 2, kernel_size=3, padding=1)]
self.loc = nn.ModuleList(self.loc)
self.conf = nn.ModuleList(self.conf)
self.softmax = nn.Softmax(dim=-1)
def __init__(self):
super(SFA, self).__init__()
# M0, M1, M2, M3
self.ssh_in_channels = [128, 128, 512, 512]
self.ssh_out_channels = [256, 256, 512, 512]
self.conv1 = self.build_conv_block(in_channels=3, out_channels=64, n_conv=2, with_pool=False)
self.conv2 = self.build_conv_block(in_channels=64, out_channels=128, n_conv=2, with_pool=True)
self.conv3 = self.build_conv_block(in_channels=128, out_channels=256, n_conv=3, with_pool=True)
self.conv4 = self.build_conv_block(in_channels=256, out_channels=512, n_conv=3, with_pool=True)
self.conv5 = self.build_conv_block(in_channels=512, out_channels=512, n_conv=3, with_pool=True)
# M3
self.pool6 = nn.MaxPool2d(2, 2)
self.m3 = SSH(self.ssh_in_channels[3], self.ssh_out_channels[3], index=3)
# M2
self.m2 = SSH(self.ssh_in_channels[2], self.ssh_out_channels[2], index=2)
# share by M1 and M2
self.conv4_128 = self.build_conv_block(in_channels=512, out_channels=128, kernel_size=1, padding=0, n_conv=1, with_pool=False)
# M1
self.conv5_128 = self.build_conv_block(in_channels=512, out_channels=128, kernel_size=1, padding=0, n_conv=1, with_pool=False)
self.conv5_128_up = nn.ConvTranspose2d(in_channels=128, out_channels=128, kernel_size=4, stride=2, padding=1, groups=128, bias=False)
self.conv4_fuse = nn.EltAdd()
self.conv4_fuse_final = self.build_conv_block(in_channels=128, out_channels=128, kernel_size=3, padding=1, n_conv=1, with_pool=False)
self.m1 = SSH(self.ssh_in_channels[1], self.ssh_out_channels[1], index=1)
# M0
self.conv3_128 = self.build_conv_block(in_channels=256, out_channels=128, kernel_size=1, padding=0, n_conv=1, with_pool=False)
self.conv4_128_up = nn.ConvTranspose2d(in_channels=128, out_channels=128, kernel_size=4, stride=2, padding=1, groups=128, bias=False)
self.conv3_fuse = nn.EltAdd()
self.conv3_fuse_final = self.build_conv_block(in_channels=128, out_channels=128, kernel_size=3, padding=1, n_conv=1, with_pool=False)
self.m0 = SSH(self.ssh_in_channels[0], self.ssh_out_channels[0], index=0)
# detection heads
self.bbox_head, self.cls_head = self.build_detect_head()
self.softmax = nn.Softmax(dim=-1)
def __init__(self, name, in_channel, mid_channel, out_channel, downsample=False):
super(ResidualBlock, self).__init__()
stride = 2 if name.startswith('res3a') or name.startswith('res4a') else 1
self.cb1 = ConvBlock(in_channel, mid_channel, 1, stride, 0, add_relu=True)
self.cb2 = ConvBlock(mid_channel, mid_channel, 3, padding=1, add_relu=True)
self.cb3 = ConvBlock(mid_channel, out_channel, 1, padding=0, add_relu=False)
self.branch = None
if downsample:
self.branch = ConvBlock(in_channel, out_channel, 1, stride, 0, add_relu=False)
self.eltadd = nn.EltAdd()
def __init__(self, inplanes, planes, stride=1, downsample=None):
super(Bottleneck, self).__init__()
self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
self.bn1 = nn.BatchNorm2d(planes)
self.conv2 = nn.Conv2d(planes,
planes,
kernel_size=3,
stride=stride,
padding=1,
bias=False)
self.bn2 = nn.BatchNorm2d(planes)
self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
self.bn3 = nn.BatchNorm2d(planes * 4)
self.relu = nn.ReLU(inplace=True)
self.downsample = downsample
self.stride = stride
self.eltadd = nn.EltAdd()
def __init__(self, inp, oup, stride, expand_ratio, model_type=32):
super(InvertedResidual_dwc, self).__init__()
self.stride = stride
assert stride in [1, 2]
hidden_dim = int(round(inp * expand_ratio))
self.use_res_connect = self.stride == 1 and inp == oup
self.conv = []
if expand_ratio == 1:
self.conv.append(
nn.Conv2d(inp,
hidden_dim,
kernel_size=(3, 3),
stride=stride,
padding=1,
groups=hidden_dim))
self.conv.append(nn.BatchNorm2d(hidden_dim))
self.conv.append(nn.PReLU())
self.conv.append(nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False))
self.conv.append(nn.BatchNorm2d(oup))
if model_type == 32:
self.conv.append(nn.PReLU())
else:
self.conv.append(nn.Conv2d(inp, hidden_dim, 1, 1, 0, bias=False))
self.conv.append(nn.BatchNorm2d(hidden_dim))
self.conv.append(nn.PReLU())
self.conv.append(
nn.Conv2d(hidden_dim,
hidden_dim,
kernel_size=(3, 3),
stride=stride,
padding=1,
groups=hidden_dim))
self.conv.append(nn.BatchNorm2d(hidden_dim))
self.conv.append(nn.PReLU())
self.conv.append(nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False))
self.conv.append(nn.BatchNorm2d(oup))
if model_type == 32:
self.conv.append(nn.PReLU())
self.conv = nn.Sequential(*self.conv)
if self.use_res_connect:
self.eltadd = nn.EltAdd()
def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1):
super(IdentityBlock, self).__init__()
out_channels_1, out_channels_2, out_channels_3 = out_channels//4, out_channels//4, out_channels
self.conv1 = nn.Conv2d(in_channels, out_channels_1, kernel_size=(1, 1))
self.bn1 = nn.BatchNorm2d(out_channels_1)
self.relu1 = nn.ReLU(inplace=True)
self.conv2 = nn.Conv2d(out_channels_1, out_channels_2, kernel_size=(kernel_size, kernel_size), padding=(padding, padding), dilation=(dilation, dilation))
self.bn2 = nn.BatchNorm2d(out_channels_2)
self.relu2 = nn.ReLU(inplace=True)
self.conv3 = nn.Conv2d(out_channels_2, out_channels_3, kernel_size=(1, 1))
self.bn3 = nn.BatchNorm2d(out_channels_3)
self.eltadd = nn.EltAdd()
self.relu_f = nn.ReLU(inplace=True)
def __init__(self, in_channels):
super(CPM, self).__init__()
# residual
self.branch1 = Conv_BN(in_channels, 1024, 1, 1, 0, act=None)
self.branch2a = Conv_BN(in_channels, 256, 1, 1, 0, act='relu')
self.branch2b = Conv_BN(256, 256, 3, 1, 1, act='relu')
self.branch2c = Conv_BN(256, 1024, 1, 1, 0, act=None)
self.eltadd = nn.EltAdd()
self.rescomb_relu = nn.ReLU(inplace=True)
# ssh
self.ssh_1_conv = nn.Conv2d(1024, 256, 3, 1, 1)
self.ssh_dimred_conv = nn.Conv2d(1024, 128, 3, 1, 1)
self.ssh_dimred_relu = nn.ReLU(inplace=True)
self.ssh_2_conv = nn.Conv2d(128, 128, 3, 1, 1)
self.ssh_3a_conv = nn.Conv2d(128, 128, 3, 1, 1)
self.ssh_3a_relu = nn.ReLU(inplace=True)
self.ssh_3b_conv = nn.Conv2d(128, 128, 3, 1, 1)
self.concat_relu = nn.ReLU(inplace=True)
def __init__(self, channels, det_out=False):
super(ResBlock, self).__init__()
self.channels = channels
self.det_out = det_out
self.relu = nn.ReLU()
self.block = nn.Sequential(
nn.Conv2d(in_channels=self.channels,
out_channels=self.channels,
kernel_size=3,
stride=1,
padding=1), nn.ReLU(),
nn.Conv2d(in_channels=self.channels,
out_channels=self.channels,
kernel_size=3,
stride=1,
padding=1))
self.eltadd = nn.EltAdd()
from _utils import test_on
import sys
sys.path.append('.')
from flops_counter import nn
from flops_counter.tensorsize import TensorSize
######
# test on EltAdd
######
eltadd = {
'layers': [
nn.EltAdd() # same shape
],
'ins': [
TensorSize([64, 112, 112])
],
'out_shape': [
TensorSize([64, 112, 112])
],
'out_flops': [
802816
]
}
test_on(eltadd)
######
# test on EltMul
######
eltmul = {
def __init__(self, in_planes=256, out_planes=256):
super(RFE, self).__init__()
self.out_channels = out_planes
self.inter_channels = int(in_planes / 4)
self.branch0 = nn.Sequential(
nn.Conv2d(in_planes,
self.inter_channels,
kernel_size=1,
stride=1,
padding=0), nn.ReLU(inplace=True),
nn.Conv2d(self.inter_channels,
self.inter_channels,
kernel_size=(1, 5),
stride=1,
padding=(0, 2)), nn.ReLU(inplace=True),
nn.Conv2d(self.inter_channels,
self.inter_channels,
kernel_size=1,
stride=1,
padding=0), nn.ReLU(inplace=True))
self.branch1 = nn.Sequential(
nn.Conv2d(in_planes,
self.inter_channels,
kernel_size=1,
stride=1,
padding=0), nn.ReLU(inplace=True),
nn.Conv2d(self.inter_channels,
self.inter_channels,
kernel_size=(5, 1),
stride=1,
padding=(2, 0)), nn.ReLU(inplace=True),
nn.Conv2d(self.inter_channels,
self.inter_channels,
kernel_size=1,
stride=1,
padding=0), nn.ReLU(inplace=True))
self.branch2 = nn.Sequential(
nn.Conv2d(in_planes,
self.inter_channels,
kernel_size=1,
stride=1,
padding=0), nn.ReLU(inplace=True),
nn.Conv2d(self.inter_channels,
self.inter_channels,
kernel_size=(1, 3),
stride=1,
padding=(0, 1)), nn.ReLU(inplace=True),
nn.Conv2d(self.inter_channels,
self.inter_channels,
kernel_size=1,
stride=1,
padding=0), nn.ReLU(inplace=True))
self.branch3 = nn.Sequential(
nn.Conv2d(in_planes,
self.inter_channels,
kernel_size=1,
stride=1,
padding=0), nn.ReLU(inplace=True),
nn.Conv2d(self.inter_channels,
self.inter_channels,
kernel_size=(3, 1),
stride=1,
padding=(1, 0)), nn.ReLU(inplace=True),
nn.Conv2d(self.inter_channels,
self.inter_channels,
kernel_size=1,
stride=1,
padding=0), nn.ReLU(inplace=True))
self.cated_conv = nn.Sequential(
nn.Conv2d(in_planes,
out_planes,
kernel_size=1,
stride=1,
padding=0), nn.ReLU(inplace=True))
self.eltadd = nn.EltAdd()
def __init__(self):
super(SRN, self).__init__()
block = Bottleneck
layers = [3, 4, 6, 3]
self.inplanes = 64
self.conv1 = nn.Conv2d(3,
64,
kernel_size=7,
stride=2,
padding=3,
bias=False)
self.bn1 = nn.BatchNorm2d(64)
self.relu = nn.ReLU(inplace=True)
self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
self.layer1 = self._make_layer(block, 64, layers[0])
self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
self.layer5 = nn.Conv2d(2048, 1024, kernel_size=3, stride=2, padding=1)
self.layer6 = nn.Conv2d(1024, 256, kernel_size=3, stride=2, padding=1)
self.c5_lateral = nn.Conv2d(2048,
256,
kernel_size=1,
stride=1,
padding=0)
self.c4_lateral = nn.Conv2d(1024,
256,
kernel_size=1,
stride=1,
padding=0)
self.c3_lateral = nn.Conv2d(512,
256,
kernel_size=1,
stride=1,
padding=0)
self.c2_lateral = nn.Conv2d(256,
256,
kernel_size=1,
stride=1,
padding=0)
self.eltadd = nn.EltAdd()
self.upsample = nn.Upsample(scale_factor=2,
mode='bilinear',
align_corners=False)
self.p7_conv = nn.Conv2d(256, 256, kernel_size=3, stride=2, padding=1)
self.p6_conv = nn.Conv2d(256, 256, kernel_size=3, stride=2, padding=1)
self.p5_conv = nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1)
self.p4_conv = nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1)
self.p3_conv = nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1)
self.p2_conv = nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1)
self.c7_conv = nn.Conv2d(256, 256, kernel_size=1, stride=1, padding=0)
self.c6_conv = nn.Conv2d(1024, 256, kernel_size=1, stride=1, padding=0)
self.c5_conv = nn.Conv2d(2048, 256, kernel_size=1, stride=1, padding=0)
self.c4_conv = nn.Conv2d(1024, 256, kernel_size=1, stride=1, padding=0)
self.c3_conv = nn.Conv2d(512, 256, kernel_size=1, stride=1, padding=0)
self.c2_conv = nn.Conv2d(256, 256, kernel_size=1, stride=1, padding=0)
# subnet_first stage
num_anchors = 2 * 1
self.cls_subnet = nn.Sequential(
nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1),
nn.ReLU(inplace=True),
RFE(256, 256),
nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1),
nn.ReLU(inplace=True),
)
self.box_subnet = nn.Sequential(
nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1),
nn.ReLU(inplace=True),
RFE(256, 256),
nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1),
nn.ReLU(inplace=True),
)
self.cls_subnet_pred = nn.Conv2d(256,
num_anchors * 1,
kernel_size=1,
stride=1,
padding=0)
self.box_subnet_pred = nn.Conv2d(256,
num_anchors * 4,
kernel_size=1,
stride=1,
padding=0)
self.sigmoid = nn.Sigmoid()