Spaces:
Running
on
T4
Running
on
T4
import torch | |
import torch.nn as nn | |
import torch.nn.functional as F | |
from .deform_conv import DeformConv2d | |
def add_conv(in_ch, out_ch, ksize, stride, leaky=True): | |
""" | |
Add a conv2d / batchnorm / leaky ReLU block. | |
Args: | |
in_ch (int): number of input channels of the convolution layer. | |
out_ch (int): number of output channels of the convolution layer. | |
ksize (int): kernel size of the convolution layer. | |
stride (int): stride of the convolution layer. | |
Returns: | |
stage (Sequential) : Sequential layers composing a convolution block. | |
""" | |
stage = nn.Sequential() | |
pad = (ksize - 1) // 2 | |
stage.add_module( | |
"conv", | |
nn.Conv2d(in_channels=in_ch, out_channels=out_ch, kernel_size=ksize, stride=stride, padding=pad, bias=False), | |
) | |
stage.add_module("batch_norm", nn.BatchNorm2d(out_ch)) | |
if leaky: | |
stage.add_module("leaky", nn.LeakyReLU(0.1)) | |
else: | |
stage.add_module("relu6", nn.ReLU6(inplace=True)) | |
return stage | |
class upsample(nn.Module): | |
__constants__ = ["size", "scale_factor", "mode", "align_corners", "name"] | |
def __init__(self, size=None, scale_factor=None, mode="nearest", align_corners=None): | |
super(upsample, self).__init__() | |
self.name = type(self).__name__ | |
self.size = size | |
self.scale_factor = scale_factor | |
self.mode = mode | |
self.align_corners = align_corners | |
def forward(self, input): | |
return F.interpolate(input, self.size, self.scale_factor, self.mode, self.align_corners) | |
def extra_repr(self): | |
if self.scale_factor is not None: | |
info = "scale_factor=" + str(self.scale_factor) | |
else: | |
info = "size=" + str(self.size) | |
info += ", mode=" + self.mode | |
return info | |
class SPPLayer(nn.Module): | |
def __init__(self): | |
super(SPPLayer, self).__init__() | |
def forward(self, x): | |
x_1 = x | |
x_2 = F.max_pool2d(x, 5, stride=1, padding=2) | |
x_3 = F.max_pool2d(x, 9, stride=1, padding=4) | |
x_4 = F.max_pool2d(x, 13, stride=1, padding=6) | |
out = torch.cat((x_1, x_2, x_3, x_4), dim=1) | |
return out | |
class DropBlock(nn.Module): | |
def __init__(self, block_size=7, keep_prob=0.9): | |
super(DropBlock, self).__init__() | |
self.block_size = block_size | |
self.keep_prob = keep_prob | |
self.gamma = None | |
self.kernel_size = (block_size, block_size) | |
self.stride = (1, 1) | |
self.padding = (block_size // 2, block_size // 2) | |
def reset(self, block_size, keep_prob): | |
self.block_size = block_size | |
self.keep_prob = keep_prob | |
self.gamma = None | |
self.kernel_size = (block_size, block_size) | |
self.stride = (1, 1) | |
self.padding = (block_size // 2, block_size // 2) | |
def calculate_gamma(self, x): | |
return ( | |
(1 - self.keep_prob) * x.shape[-1] ** 2 / (self.block_size**2 * (x.shape[-1] - self.block_size + 1) ** 2) | |
) | |
def forward(self, x): | |
if not self.training or self.keep_prob == 1: # set keep_prob=1 to turn off dropblock | |
return x | |
if self.gamma is None: | |
self.gamma = self.calculate_gamma(x) | |
if x.type() == "torch.cuda.HalfTensor": # TODO: not fully support for FP16 now | |
FP16 = True | |
x = x.float() | |
else: | |
FP16 = False | |
p = torch.ones_like(x) * (self.gamma) | |
mask = 1 - torch.nn.functional.max_pool2d(torch.bernoulli(p), self.kernel_size, self.stride, self.padding) | |
out = mask * x * (mask.numel() / mask.sum()) | |
if FP16: | |
out = out.half() | |
return out | |
class resblock(nn.Module): | |
""" | |
Sequential residual blocks each of which consists of \ | |
two convolution layers. | |
Args: | |
ch (int): number of input and output channels. | |
nblocks (int): number of residual blocks. | |
shortcut (bool): if True, residual tensor addition is enabled. | |
""" | |
def __init__(self, ch, nblocks=1, shortcut=True): | |
super().__init__() | |
self.shortcut = shortcut | |
self.module_list = nn.ModuleList() | |
for i in range(nblocks): | |
resblock_one = nn.ModuleList() | |
resblock_one.append(add_conv(ch, ch // 2, 1, 1)) | |
resblock_one.append(add_conv(ch // 2, ch, 3, 1)) | |
self.module_list.append(resblock_one) | |
def forward(self, x): | |
for module in self.module_list: | |
h = x | |
for res in module: | |
h = res(h) | |
x = x + h if self.shortcut else h | |
return x | |
class RFBblock(nn.Module): | |
def __init__(self, in_ch, residual=False): | |
super(RFBblock, self).__init__() | |
inter_c = in_ch // 4 | |
self.branch_0 = nn.Sequential( | |
nn.Conv2d(in_channels=in_ch, out_channels=inter_c, kernel_size=1, stride=1, padding=0), | |
) | |
self.branch_1 = nn.Sequential( | |
nn.Conv2d(in_channels=in_ch, out_channels=inter_c, kernel_size=1, stride=1, padding=0), | |
nn.Conv2d(in_channels=inter_c, out_channels=inter_c, kernel_size=3, stride=1, padding=1), | |
) | |
self.branch_2 = nn.Sequential( | |
nn.Conv2d(in_channels=in_ch, out_channels=inter_c, kernel_size=1, stride=1, padding=0), | |
nn.Conv2d(in_channels=inter_c, out_channels=inter_c, kernel_size=3, stride=1, padding=1), | |
nn.Conv2d(in_channels=inter_c, out_channels=inter_c, kernel_size=3, stride=1, dilation=2, padding=2), | |
) | |
self.branch_3 = nn.Sequential( | |
nn.Conv2d(in_channels=in_ch, out_channels=inter_c, kernel_size=1, stride=1, padding=0), | |
nn.Conv2d(in_channels=inter_c, out_channels=inter_c, kernel_size=5, stride=1, padding=2), | |
nn.Conv2d(in_channels=inter_c, out_channels=inter_c, kernel_size=3, stride=1, dilation=3, padding=3), | |
) | |
self.residual = residual | |
def forward(self, x): | |
x_0 = self.branch_0(x) | |
x_1 = self.branch_1(x) | |
x_2 = self.branch_2(x) | |
x_3 = self.branch_3(x) | |
out = torch.cat((x_0, x_1, x_2, x_3), 1) | |
if self.residual: | |
out += x | |
return out | |
class FeatureAdaption(nn.Module): | |
def __init__(self, in_ch, out_ch, n_anchors, rfb=False, sep=False): | |
super(FeatureAdaption, self).__init__() | |
if sep: | |
self.sep = True | |
else: | |
self.sep = False | |
self.conv_offset = nn.Conv2d( | |
in_channels=2 * n_anchors, | |
out_channels=2 * 9 * n_anchors, | |
groups=n_anchors, | |
kernel_size=1, | |
stride=1, | |
padding=0, | |
) | |
self.dconv = DeformConv2d( | |
in_channels=in_ch, out_channels=out_ch, kernel_size=3, stride=1, padding=1, deformable_groups=n_anchors | |
) | |
self.rfb = None | |
if rfb: | |
self.rfb = RFBblock(out_ch) | |
def forward(self, input, wh_pred): | |
# The RFB block is added behind FeatureAdaption | |
# For mobilenet, we currently don't support rfb and FeatureAdaption | |
if self.sep: | |
return input | |
if self.rfb is not None: | |
input = self.rfb(input) | |
wh_pred_new = wh_pred.detach() | |
offset = self.conv_offset(wh_pred_new) | |
out = self.dconv(input, offset) | |
return out | |
class ASFFmobile(nn.Module): | |
def __init__(self, level, rfb=False, vis=False): | |
super(ASFFmobile, self).__init__() | |
self.level = level | |
self.dim = [512, 256, 128] | |
self.inter_dim = self.dim[self.level] | |
if level == 0: | |
self.stride_level_1 = add_conv(256, self.inter_dim, 3, 2, leaky=False) | |
self.stride_level_2 = add_conv(128, self.inter_dim, 3, 2, leaky=False) | |
self.expand = add_conv(self.inter_dim, 1024, 3, 1, leaky=False) | |
elif level == 1: | |
self.compress_level_0 = add_conv(512, self.inter_dim, 1, 1, leaky=False) | |
self.stride_level_2 = add_conv(128, self.inter_dim, 3, 2, leaky=False) | |
self.expand = add_conv(self.inter_dim, 512, 3, 1, leaky=False) | |
elif level == 2: | |
self.compress_level_0 = add_conv(512, self.inter_dim, 1, 1, leaky=False) | |
self.compress_level_1 = add_conv(256, self.inter_dim, 1, 1, leaky=False) | |
self.expand = add_conv(self.inter_dim, 256, 3, 1, leaky=False) | |
compress_c = 8 if rfb else 16 # when adding rfb, we use half number of channels to save memory | |
self.weight_level_0 = add_conv(self.inter_dim, compress_c, 1, 1, leaky=False) | |
self.weight_level_1 = add_conv(self.inter_dim, compress_c, 1, 1, leaky=False) | |
self.weight_level_2 = add_conv(self.inter_dim, compress_c, 1, 1, leaky=False) | |
self.weight_levels = nn.Conv2d(compress_c * 3, 3, kernel_size=1, stride=1, padding=0) | |
self.vis = vis | |
def forward(self, x_level_0, x_level_1, x_level_2): | |
if self.level == 0: | |
level_0_resized = x_level_0 | |
level_1_resized = self.stride_level_1(x_level_1) | |
level_2_downsampled_inter = F.max_pool2d(x_level_2, 3, stride=2, padding=1) | |
level_2_resized = self.stride_level_2(level_2_downsampled_inter) | |
elif self.level == 1: | |
level_0_compressed = self.compress_level_0(x_level_0) | |
level_0_resized = F.interpolate(level_0_compressed, scale_factor=2, mode="nearest") | |
level_1_resized = x_level_1 | |
level_2_resized = self.stride_level_2(x_level_2) | |
elif self.level == 2: | |
level_0_compressed = self.compress_level_0(x_level_0) | |
level_0_resized = F.interpolate(level_0_compressed, scale_factor=4, mode="nearest") | |
level_1_compressed = self.compress_level_1(x_level_1) | |
level_1_resized = F.interpolate(level_1_compressed, scale_factor=2, mode="nearest") | |
level_2_resized = x_level_2 | |
level_0_weight_v = self.weight_level_0(level_0_resized) | |
level_1_weight_v = self.weight_level_1(level_1_resized) | |
level_2_weight_v = self.weight_level_2(level_2_resized) | |
levels_weight_v = torch.cat((level_0_weight_v, level_1_weight_v, level_2_weight_v), 1) | |
levels_weight = self.weight_levels(levels_weight_v) | |
levels_weight = F.softmax(levels_weight, dim=1) | |
fused_out_reduced = ( | |
level_0_resized * levels_weight[:, 0:1, :, :] | |
+ level_1_resized * levels_weight[:, 1:2, :, :] | |
+ level_2_resized * levels_weight[:, 2:, :, :] | |
) | |
out = self.expand(fused_out_reduced) | |
if self.vis: | |
return out, levels_weight, fused_out_reduced.sum(dim=1) | |
else: | |
return out | |
class ASFF(nn.Module): | |
def __init__(self, level, rfb=False, vis=False): | |
super(ASFF, self).__init__() | |
self.level = level | |
self.dim = [512, 256, 256] | |
self.inter_dim = self.dim[self.level] | |
if level == 0: | |
self.stride_level_1 = add_conv(256, self.inter_dim, 3, 2) | |
self.stride_level_2 = add_conv(256, self.inter_dim, 3, 2) | |
self.expand = add_conv(self.inter_dim, 1024, 3, 1) | |
elif level == 1: | |
self.compress_level_0 = add_conv(512, self.inter_dim, 1, 1) | |
self.stride_level_2 = add_conv(256, self.inter_dim, 3, 2) | |
self.expand = add_conv(self.inter_dim, 512, 3, 1) | |
elif level == 2: | |
self.compress_level_0 = add_conv(512, self.inter_dim, 1, 1) | |
self.expand = add_conv(self.inter_dim, 256, 3, 1) | |
compress_c = 8 if rfb else 16 # when adding rfb, we use half number of channels to save memory | |
self.weight_level_0 = add_conv(self.inter_dim, compress_c, 1, 1) | |
self.weight_level_1 = add_conv(self.inter_dim, compress_c, 1, 1) | |
self.weight_level_2 = add_conv(self.inter_dim, compress_c, 1, 1) | |
self.weight_levels = nn.Conv2d(compress_c * 3, 3, kernel_size=1, stride=1, padding=0) | |
self.vis = vis | |
def forward(self, x_level_0, x_level_1, x_level_2): | |
if self.level == 0: | |
level_0_resized = x_level_0 | |
level_1_resized = self.stride_level_1(x_level_1) | |
level_2_downsampled_inter = F.max_pool2d(x_level_2, 3, stride=2, padding=1) | |
level_2_resized = self.stride_level_2(level_2_downsampled_inter) | |
elif self.level == 1: | |
level_0_compressed = self.compress_level_0(x_level_0) | |
level_0_resized = F.interpolate(level_0_compressed, scale_factor=2, mode="nearest") | |
level_1_resized = x_level_1 | |
level_2_resized = self.stride_level_2(x_level_2) | |
elif self.level == 2: | |
level_0_compressed = self.compress_level_0(x_level_0) | |
level_0_resized = F.interpolate(level_0_compressed, scale_factor=4, mode="nearest") | |
level_1_resized = F.interpolate(x_level_1, scale_factor=2, mode="nearest") | |
level_2_resized = x_level_2 | |
level_0_weight_v = self.weight_level_0(level_0_resized) | |
level_1_weight_v = self.weight_level_1(level_1_resized) | |
level_2_weight_v = self.weight_level_2(level_2_resized) | |
levels_weight_v = torch.cat((level_0_weight_v, level_1_weight_v, level_2_weight_v), 1) | |
levels_weight = self.weight_levels(levels_weight_v) | |
levels_weight = F.softmax(levels_weight, dim=1) | |
fused_out_reduced = ( | |
level_0_resized * levels_weight[:, 0:1, :, :] | |
+ level_1_resized * levels_weight[:, 1:2, :, :] | |
+ level_2_resized * levels_weight[:, 2:, :, :] | |
) | |
out = self.expand(fused_out_reduced) | |
if self.vis: | |
return out, levels_weight, fused_out_reduced.sum(dim=1) | |
else: | |
return out | |
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 | |
:param v: | |
:param divisor: | |
:param min_value: | |
:return: | |
""" | |
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 ConvBNReLU(nn.Sequential): | |
def __init__(self, in_planes, out_planes, kernel_size=3, stride=1, groups=1): | |
padding = (kernel_size - 1) // 2 | |
super(ConvBNReLU, self).__init__( | |
nn.Conv2d(in_planes, out_planes, kernel_size, stride, padding, groups=groups, bias=False), | |
nn.BatchNorm2d(out_planes), | |
nn.ReLU6(inplace=True), | |
) | |
def add_sepconv(in_ch, out_ch, ksize, stride): | |
stage = nn.Sequential() | |
pad = (ksize - 1) // 2 | |
stage.add_module( | |
"sepconv", | |
nn.Conv2d( | |
in_channels=in_ch, | |
out_channels=in_ch, | |
kernel_size=ksize, | |
stride=stride, | |
padding=pad, | |
groups=in_ch, | |
bias=False, | |
), | |
) | |
stage.add_module("sepbn", nn.BatchNorm2d(in_ch)) | |
stage.add_module("seprelu6", nn.ReLU6(inplace=True)) | |
stage.add_module("ptconv", nn.Conv2d(in_ch, out_ch, 1, 1, 0, bias=False)) | |
stage.add_module("ptbn", nn.BatchNorm2d(out_ch)) | |
stage.add_module("ptrelu6", nn.ReLU6(inplace=True)) | |
return stage | |
class InvertedResidual(nn.Module): | |
def __init__(self, inp, oup, stride, expand_ratio): | |
super(InvertedResidual, 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 | |
layers = [] | |
if expand_ratio != 1: | |
# pw | |
layers.append(ConvBNReLU(inp, hidden_dim, kernel_size=1)) | |
layers.extend( | |
[ | |
# dw | |
ConvBNReLU(hidden_dim, hidden_dim, stride=stride, groups=hidden_dim), | |
# pw-linear | |
nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False), | |
nn.BatchNorm2d(oup), | |
] | |
) | |
self.conv = nn.Sequential(*layers) | |
def forward(self, x): | |
if self.use_res_connect: | |
return x + self.conv(x) | |
else: | |
return self.conv(x) | |
class ressepblock(nn.Module): | |
def __init__(self, ch, out_ch, in_ch=None, shortcut=True): | |
super().__init__() | |
self.shortcut = shortcut | |
self.module_list = nn.ModuleList() | |
in_ch = ch // 2 if in_ch == None else in_ch | |
resblock_one = nn.ModuleList() | |
resblock_one.append(add_conv(ch, in_ch, 1, 1, leaky=False)) | |
resblock_one.append(add_conv(in_ch, out_ch, 3, 1, leaky=False)) | |
self.module_list.append(resblock_one) | |
def forward(self, x): | |
for module in self.module_list: | |
h = x | |
for res in module: | |
h = res(h) | |
x = x + h if self.shortcut else h | |
return x | |