def __init__(self, nclass, backbone='', aux=False, jpu=False, pretrained_base=False, **kwargs): super(ESPNetV2, self).__init__() self.pretrained = eespnet(pretrained=pretrained_base, **kwargs) self.proj_L4_C = _ConvBNPReLU(256, 128, 1, **kwargs) self.pspMod = nn.Sequential( EESP(256, 128, stride=1, k=4, r_lim=7, **kwargs), _PSPModule(128, 128, **kwargs)) self.project_l3 = nn.Sequential(nn.Dropout2d(0.1), nn.Conv2d(128, nclass, 1, bias=False)) self.act_l3 = _BNPReLU(nclass, **kwargs) self.project_l2 = _ConvBNPReLU(64 + nclass, nclass, 1, **kwargs) self.project_l1 = nn.Sequential( nn.Dropout2d(0.1), nn.Conv2d(32 + nclass, nclass, 1, bias=False)) self.aux = aux self.__setattr__('exclusive', [ 'proj_L4_C', 'pspMod', 'project_l3', 'act_l3', 'project_l2', 'project_l1' ])
def __init__(self, nclass, backbone='', aux=False, jpu=False, pretrained_base=True, M=3, N=21, **kwargs): super(CGNet, self).__init__() # stage 1 self.stage1_0 = _ConvBNPReLU(3, 32, 3, 2, 1, **kwargs) self.stage1_1 = _ConvBNPReLU(32, 32, 3, 1, 1, **kwargs) self.stage1_2 = _ConvBNPReLU(32, 32, 3, 1, 1, **kwargs) self.sample1 = _InputInjection(1) self.sample2 = _InputInjection(2) self.bn_prelu1 = _BNPReLU(32 + 3, **kwargs) # stage 2 self.stage2_0 = ContextGuidedBlock(32 + 3, 64, dilation=2, reduction=8, down=True, residual=False, **kwargs) self.stage2 = nn.ModuleList() for i in range(0, M - 1): self.stage2.append( ContextGuidedBlock(64, 64, dilation=2, reduction=8, **kwargs)) self.bn_prelu2 = _BNPReLU(128 + 3, **kwargs) # stage 3 self.stage3_0 = ContextGuidedBlock(128 + 3, 128, dilation=4, reduction=16, down=True, residual=False, **kwargs) self.stage3 = nn.ModuleList() for i in range(0, N - 1): self.stage3.append( ContextGuidedBlock(128, 128, dilation=4, reduction=16, **kwargs)) self.bn_prelu3 = _BNPReLU(256, **kwargs) self.head = nn.Sequential(nn.Dropout2d(0.1, False), nn.Conv2d(256, nclass, 1)) self.__setattr__('exclusive', [ 'stage1_0', 'stage1_1', 'stage1_2', 'sample1', 'sample2', 'bn_prelu1', 'stage2_0', 'stage2', 'bn_prelu2', 'stage3_0', 'stage3', 'bn_prelu3', 'head' ])
def __init__( self, in_channels, out_channels, k=4, r_lim=9, reinf=True, inp_reinf=3, norm_layer=None, ): super(DownSampler, self).__init__() channels_diff = out_channels - in_channels self.eesp = EESP( in_channels, channels_diff, stride=2, k=k, r_lim=r_lim, down_method="avg", norm_layer=norm_layer, ) self.avg = nn.AvgPool2d(kernel_size=3, padding=1, stride=2) if reinf: self.inp_reinf = nn.Sequential( _ConvBNPReLU(inp_reinf, inp_reinf, 3, 1, 1), _ConvBN(inp_reinf, out_channels, 1, 1), ) self.act = nn.PReLU(out_channels)
def __init__(self, in_channels, out_channels, stride=1, k=4, r_lim=7, down_method='esp', norm_layer=nn.BatchNorm2d): super(EESP, self).__init__() self.stride = stride n = int(out_channels / k) n1 = out_channels - (k - 1) * n assert down_method in ['avg', 'esp' ], 'One of these is suppported (avg or esp)' assert n == n1, "n(={}) and n1(={}) should be equal for Depth-wise Convolution ".format( n, n1) self.proj_1x1 = _ConvBNPReLU(in_channels, n, 1, stride=1, groups=k, norm_layer=norm_layer) map_receptive_ksize = { 3: 1, 5: 2, 7: 3, 9: 4, 11: 5, 13: 6, 15: 7, 17: 8 } self.k_sizes = list() for i in range(k): ksize = int(3 + 2 * i) ksize = ksize if ksize <= r_lim else 3 self.k_sizes.append(ksize) self.k_sizes.sort() self.spp_dw = nn.ModuleList() for i in range(k): dilation = map_receptive_ksize[self.k_sizes[i]] self.spp_dw.append( nn.Conv2d(n, n, 3, stride, dilation, dilation=dilation, groups=n, bias=False)) self.conv_1x1_exp = _ConvBN(out_channels, out_channels, 1, 1, groups=k, norm_layer=norm_layer) self.br_after_cat = _BNPReLU(out_channels, norm_layer) self.module_act = nn.PReLU(out_channels) self.downAvg = True if down_method == 'avg' else False
def __init__(self, in_channels, out_channels, dilation=2, reduction=16, down=False, residual=True, norm_layer=nn.BatchNorm2d, **kwargs): super(ContextGuidedBlock, self).__init__() inter_channels = out_channels // 2 if not down else out_channels if down: self.conv = _ConvBNPReLU(in_channels, inter_channels, 3, 2, 1, norm_layer=norm_layer, **kwargs) self.reduce = nn.Conv2d(inter_channels * 2, out_channels, 1, bias=False) else: self.conv = _ConvBNPReLU(in_channels, inter_channels, 1, 1, 0, norm_layer=norm_layer, **kwargs) self.f_loc = _ChannelWiseConv(inter_channels, inter_channels, **kwargs) self.f_sur = _ChannelWiseConv(inter_channels, inter_channels, dilation, **kwargs) self.bn = norm_layer(inter_channels * 2) self.prelu = nn.PReLU(inter_channels * 2) self.f_glo = _FGlo(out_channels, reduction, **kwargs) self.down = down self.residual = residual
def __init__(self, in_channels, out_channels=1024, sizes=(1, 2, 4, 8), **kwargs): super(_PSPModule, self).__init__() self.stages = nn.ModuleList([ nn.Conv2d(in_channels, in_channels, 3, 1, 1, groups=in_channels, bias=False) for _ in sizes ]) self.project = _ConvBNPReLU(in_channels * (len(sizes) + 1), out_channels, 1, 1, **kwargs)
def __init__(self, num_classes=1000, scale=1, reinf=True, norm_layer=nn.BatchNorm2d): super(EESPNet, self).__init__() inp_reinf = 3 if reinf else None reps = [0, 3, 7, 3] r_lim = [13, 11, 9, 7, 5] K = [4] * len(r_lim) # set out_channels base, levels, base_s = 32, 5, 0 out_channels = [base] * levels for i in range(levels): if i == 0: base_s = int(base * scale) base_s = math.ceil(base_s / K[0]) * K[0] out_channels[i] = base if base_s > base else base_s else: out_channels[i] = base_s * pow(2, i) if scale <= 1.5: out_channels.append(1024) elif scale in [1.5, 2]: out_channels.append(1280) else: raise ValueError("Unknown scale value.") self.level1 = _ConvBNPReLU(3, out_channels[0], 3, 2, 1, norm_layer=norm_layer) self.level2_0 = DownSampler( out_channels[0], out_channels[1], k=K[0], r_lim=r_lim[0], reinf=reinf, inp_reinf=inp_reinf, norm_layer=norm_layer, ) self.level3_0 = DownSampler( out_channels[1], out_channels[2], k=K[1], r_lim=r_lim[1], reinf=reinf, inp_reinf=inp_reinf, norm_layer=norm_layer, ) self.level3 = nn.ModuleList() for i in range(reps[1]): self.level3.append( EESP( out_channels[2], out_channels[2], k=K[2], r_lim=r_lim[2], norm_layer=norm_layer, )) self.level4_0 = DownSampler( out_channels[2], out_channels[3], k=K[2], r_lim=r_lim[2], reinf=reinf, inp_reinf=inp_reinf, norm_layer=norm_layer, ) self.level4 = nn.ModuleList() for i in range(reps[2]): self.level4.append( EESP( out_channels[3], out_channels[3], k=K[3], r_lim=r_lim[3], norm_layer=norm_layer, )) self.level5_0 = DownSampler( out_channels[3], out_channels[4], k=K[3], r_lim=r_lim[3], reinf=reinf, inp_reinf=inp_reinf, norm_layer=norm_layer, ) self.level5 = nn.ModuleList() for i in range(reps[2]): self.level5.append( EESP( out_channels[4], out_channels[4], k=K[4], r_lim=r_lim[4], norm_layer=norm_layer, )) self.level5.append( _ConvBNPReLU( out_channels[4], out_channels[4], 3, 1, 1, groups=out_channels[4], norm_layer=norm_layer, )) self.level5.append( _ConvBNPReLU( out_channels[4], out_channels[5], 1, 1, 0, groups=K[4], norm_layer=norm_layer, )) self.fc = nn.Linear(out_channels[5], num_classes) for m in self.modules(): if isinstance(m, nn.Conv2d): nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu") if m.bias is not None: nn.init.constant_(m.bias, 0) elif isinstance(m, nn.BatchNorm2d): nn.init.constant_(m.weight, 1) nn.init.constant_(m.bias, 0) elif isinstance(m, nn.Linear): nn.init.normal_(m.weight, std=0.001) if m.bias is not None: nn.init.constant_(m.bias, 0)