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