def __init__(self, in_channels, key_channels, value_channels, out_channels=None, scale=1, bn_type=None): super(_SelfAttentionBlock, self).__init__() self.scale = scale self.in_channels = in_channels self.out_channels = out_channels self.key_channels = key_channels self.value_channels = value_channels if out_channels == None: self.out_channels = in_channels self.pool = nn.MaxPool2d(kernel_size=(scale, scale)) self.f_key = nn.Sequential( nn.Conv2d(in_channels=self.in_channels, out_channels=self.key_channels, kernel_size=1, stride=1, padding=0), ModuleHelper.BNReLU(self.key_channels, bn_type=bn_type), nn.Conv2d(in_channels=self.key_channels, out_channels=self.key_channels, kernel_size=1, stride=1, padding=0), ModuleHelper.BNReLU(self.key_channels, bn_type=bn_type), ) self.f_query = nn.Sequential( nn.Conv2d(in_channels=self.in_channels, out_channels=self.key_channels, kernel_size=1, stride=1, padding=0), ModuleHelper.BNReLU(self.key_channels, bn_type=bn_type), nn.Conv2d(in_channels=self.key_channels, out_channels=self.key_channels, kernel_size=1, stride=1, padding=0), ModuleHelper.BNReLU(self.key_channels, bn_type=bn_type), ) self.f_value = nn.Conv2d(in_channels=self.in_channels, out_channels=self.value_channels, kernel_size=1, stride=1, padding=0) self.W = nn.Conv2d(in_channels=self.value_channels, out_channels=self.out_channels, kernel_size=1, stride=1, padding=0) nn.init.constant_(self.W.weight, 0) nn.init.constant_(self.W.bias, 0)
def __init__(self, configer): self.inplanes = 128 super(BaseOCNet, self).__init__() self.configer = configer self.num_classes = self.configer.get('data', 'num_classes') self.backbone = BackboneSelector(configer).get_backbone() # extra added layers if "wide_resnet38" in self.configer.get('network', 'backbone'): in_channels = [2048, 4096] else: in_channels = [1024, 2048] self.oc_module_pre = nn.Sequential( nn.Conv2d(in_channels[1], 512, kernel_size=3, stride=1, padding=1), ModuleHelper.BNReLU(512, bn_type=self.configer.get('network', 'bn_type')), ) from lib.models.modules.base_oc_block import BaseOC_Module self.oc_module = BaseOC_Module(in_channels=512, out_channels=512, key_channels=256, value_channels=256, dropout=0.05, sizes=([1]), bn_type=self.configer.get('network', 'bn_type')) self.cls = nn.Conv2d(512, self.num_classes, kernel_size=1, stride=1, padding=0, bias=True) self.dsn = nn.Sequential( nn.Conv2d(in_channels[0], 512, kernel_size=3, stride=1, padding=1), ModuleHelper.BNReLU(512, bn_type=self.configer.get('network', 'bn_type')), nn.Conv2d(512, self.num_classes, kernel_size=1, stride=1, padding=0, bias=True) )
def __init__(self, in_channels, out_channels, dropout, sizes=([1]), bn_type=None): super(Pyramid_OC_Module, self).__init__() self.group = len(sizes) self.stages = [] self.stages = nn.ModuleList([ self._make_stage(in_channels, out_channels, in_channels // 2, in_channels, size, bn_type) for size in sizes ]) self.conv_bn_dropout = nn.Sequential( nn.Conv2d(2 * in_channels * self.group, out_channels, kernel_size=1, padding=0), ModuleHelper.BNReLU(out_channels, bn_type=bn_type), nn.Dropout2d(dropout)) self.up_dr = nn.Sequential( nn.Conv2d(in_channels, in_channels * self.group, kernel_size=1, padding=0), ModuleHelper.BNReLU(in_channels * self.group, bn_type=bn_type))
def __init__(self, num_class=150, fc_dim=4096, bn_type=None): super(PPMBilinearDeepsup, self).__init__() self.bn_type = bn_type pool_scales = (1, 2, 3, 6) self.ppm = [] # assert bn_type == 'syncbn' or not self.training # Torch BN can't handle feature map size with 1x1. for scale in pool_scales: self.ppm.append(nn.Sequential( nn.AdaptiveAvgPool2d(scale), nn.Conv2d(fc_dim, 512, kernel_size=1, bias=False), ModuleHelper.BatchNorm2d(bn_type=bn_type)(512), nn.ReLU(inplace=True) )) self.ppm = nn.ModuleList(self.ppm) self.cbr_deepsup = _ConvBatchNormReluBlock(fc_dim // 2, fc_dim // 4, 3, 1, bn_type=bn_type) self.conv_last = nn.Sequential( nn.Conv2d(fc_dim+len(pool_scales)*512, 512, kernel_size=3, padding=1, bias=False), ModuleHelper.BatchNorm2d(bn_type=bn_type)(512), nn.ReLU(inplace=True), nn.Dropout2d(0.1), nn.Conv2d(512, num_class, kernel_size=1) ) self.conv_last_deepsup = nn.Conv2d(fc_dim // 4, num_class, 1, 1, 0) self.dropout_deepsup = nn.Dropout2d(0.1)
def __init__(self, inplanes, planes, stride=1, downsample=None, bn_type=None, bn_momentum=0.1): super(Bottleneck, self).__init__() self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False) self.bn1 = ModuleHelper.BatchNorm2d(bn_type=bn_type)( planes, momentum=bn_momentum) self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, padding=1, bias=False) self.bn2 = ModuleHelper.BatchNorm2d(bn_type=bn_type)( planes, momentum=bn_momentum) self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False) self.bn3 = ModuleHelper.BatchNorm2d(bn_type=bn_type)( planes * 4, momentum=bn_momentum) self.relu = nn.ReLU(inplace=False) self.relu_in = nn.ReLU(inplace=True) self.downsample = downsample self.stride = stride
def __init__(self, configer): self.inplanes = 128 super(SpatialOCRNet, self).__init__() self.configer = configer self.num_classes = self.configer.get('data', 'num_classes') self.backbone = BackboneSelector(configer).get_backbone() # extra added layers if "wide_resnet38" in self.configer.get('network', 'backbone'): in_channels = [2048, 4096] else: in_channels = [1024, 2048] self.conv_3x3 = nn.Sequential( nn.Conv2d(in_channels[1], 512, kernel_size=3, stride=1, padding=1), ModuleHelper.BNReLU(512, bn_type=self.configer.get('network', 'bn_type')), ) from lib.models.modules.spatial_ocr_block import SpatialGather_Module, SpatialOCR_Module self.spatial_context_head = SpatialGather_Module(self.num_classes) self.spatial_ocr_head = SpatialOCR_Module(in_channels=512, key_channels=256, out_channels=512, scale=1, dropout=0.05, bn_type=self.configer.get('network', 'bn_type')) self.head = nn.Conv2d(512, self.num_classes, kernel_size=1, stride=1, padding=0, bias=True) self.dsn_head = nn.Sequential( nn.Conv2d(in_channels[0], 512, kernel_size=3, stride=1, padding=1), ModuleHelper.BNReLU(512, bn_type=self.configer.get('network', 'bn_type')), nn.Dropout2d(0.05), nn.Conv2d(512, self.num_classes, kernel_size=1, stride=1, padding=0, bias=True) )
def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1, base_width=64, dilation=1, bn_type=None): super(BasicBlock, self).__init__() if groups != 1 or base_width != 64: raise ValueError( 'BasicBlock only supports groups=1 and base_width=64') if dilation > 1: raise NotImplementedError( "Dilation > 1 not supported in BasicBlock") # Both self.conv1 and self.downsample layers downsample the input when stride != 1 self.conv1 = conv3x3(inplanes, planes, stride) self.bn1 = ModuleHelper.BatchNorm2d(bn_type=bn_type)(planes) self.relu = nn.ReLU(inplace=False) self.relu_in = nn.ReLU(inplace=True) self.conv2 = conv3x3(planes, planes) self.bn2 = ModuleHelper.BatchNorm2d(bn_type=bn_type)(planes) self.downsample = downsample self.stride = stride
def __init__(self, input_num, num1, num2, dilation_rate, drop_out, bn_type): super(_DenseAsppBlock, self).__init__() self.add_module('relu1', nn.ReLU(inplace=False)), self.add_module( 'conv1', nn.Conv2d(in_channels=input_num, out_channels=num1, kernel_size=1)), self.add_module( 'norm2', ModuleHelper.BatchNorm2d(bn_type=bn_type)(num_features=num1)), self.add_module('relu2', nn.ReLU(inplace=False)), self.add_module( 'conv2', nn.Conv2d(in_channels=num1, out_channels=num2, kernel_size=3, dilation=dilation_rate, padding=dilation_rate)), self.add_module( 'norm2', ModuleHelper.BatchNorm2d(bn_type=bn_type)(num_features=input_num)), self.drop_rate = drop_out
def __init__(self, in_channels, channels, kernel_size, stride=(1, 1), padding=(0, 0), dilation=(1, 1), groups=1, bias=True, radix=2, reduction_factor=4, rectify=False, rectify_avg=False, bn_type=None, dropblock_prob=0.0, **kwargs): super(SplAtConv2d, self).__init__() padding = _pair(padding) self.rectify = rectify and (padding[0] > 0 or padding[1] > 0) self.rectify_avg = rectify_avg inter_channels = max(in_channels * radix // reduction_factor, 32) self.radix = radix self.cardinality = groups self.channels = channels self.dropblock_prob = dropblock_prob if self.rectify: from rfconv import RFConv2d self.conv = RFConv2d(in_channels, channels * radix, kernel_size, stride, padding, dilation, groups=groups * radix, bias=bias, average_mode=rectify_avg, **kwargs) else: self.conv = Conv2d(in_channels, channels * radix, kernel_size, stride, padding, dilation, groups=groups * radix, bias=bias, **kwargs) self.use_bn = bn_type is not None self.bn0 = ModuleHelper.BatchNorm2d(bn_type=bn_type)(channels * radix) self.relu = ReLU(inplace=False) self.fc1 = Conv2d(channels, inter_channels, 1, groups=self.cardinality) self.bn1 = ModuleHelper.BatchNorm2d(bn_type=bn_type)(inter_channels) self.fc2 = Conv2d(inter_channels, channels * radix, 1, groups=self.cardinality) if dropblock_prob > 0.0: self.dropblock = DropBlock2D(dropblock_prob, 3) self.rsoftmax = rSoftMax(radix, groups)
def _make_fuse_layers(self, bn_type, bn_momentum=0.1): if self.num_branches == 1: return None num_branches = self.num_branches num_inchannels = self.num_inchannels fuse_layers = [] for i in range(num_branches if self.multi_scale_output else 1): fuse_layer = [] for j in range(num_branches): if j > i: fuse_layer.append( nn.Sequential( nn.Conv2d(num_inchannels[j], num_inchannels[i], 1, 1, 0, bias=False), ModuleHelper.BatchNorm2d(bn_type=bn_type)( num_inchannels[i], momentum=bn_momentum), )) elif j == i: fuse_layer.append(None) else: conv3x3s = [] for k in range(i - j): if k == i - j - 1: num_outchannels_conv3x3 = num_inchannels[i] conv3x3s.append( nn.Sequential( nn.Conv2d(num_inchannels[j], num_outchannels_conv3x3, 3, 2, 1, bias=False), ModuleHelper.BatchNorm2d(bn_type=bn_type)( num_outchannels_conv3x3, momentum=bn_momentum))) else: num_outchannels_conv3x3 = num_inchannels[j] conv3x3s.append( nn.Sequential( nn.Conv2d(num_inchannels[j], num_outchannels_conv3x3, 3, 2, 1, bias=False), ModuleHelper.BatchNorm2d(bn_type=bn_type)( num_outchannels_conv3x3, momentum=bn_momentum), nn.ReLU(inplace=False))) fuse_layer.append(nn.Sequential(*conv3x3s)) fuse_layers.append(nn.ModuleList(fuse_layer)) return nn.ModuleList(fuse_layers)
def __call__(self): arch = self.configer.get('network', 'backbone') from lib.models.backbones.hrnet.hrnet_config import MODEL_CONFIGS if arch == 'hrnet18': arch_net = HighResolutionNet(MODEL_CONFIGS['hrnet18'], bn_type='inplace_abn', bn_momentum=0.1) arch_net = ModuleHelper.load_model(arch_net, pretrained=self.configer.get( 'network', 'pretrained'), all_match=False, network='hrnet') elif arch == 'hrnet32': arch_net = HighResolutionNet(MODEL_CONFIGS['hrnet32'], bn_type='inplace_abn', bn_momentum=0.1) arch_net = ModuleHelper.load_model(arch_net, pretrained=self.configer.get( 'network', 'pretrained'), all_match=False, network='hrnet') elif arch == 'hrnet48': arch_net = HighResolutionNet(MODEL_CONFIGS['hrnet48'], bn_type='inplace_abn', bn_momentum=0.1) arch_net = ModuleHelper.load_model(arch_net, pretrained=self.configer.get( 'network', 'pretrained'), all_match=False, network='hrnet') elif arch == 'hrnet64': arch_net = HighResolutionNet(MODEL_CONFIGS['hrnet64'], bn_type='inplace_abn', bn_momentum=0.1) arch_net = ModuleHelper.load_model(arch_net, pretrained=self.configer.get( 'network', 'pretrained'), all_match=False, network='hrnet') elif arch == 'hrnet2x20': arch_net = HighResolutionNext(MODEL_CONFIGS['hrnet2x20'], bn_type=self.configer.get( 'network', 'bn_type')) arch_net = ModuleHelper.load_model(arch_net, pretrained=self.configer.get( 'network', 'pretrained'), all_match=False, network='hrnet') else: raise Exception('Architecture undefined!') return arch_net
def _resnet(arch, block, layers, pretrained, progress, **kwargs): model = ResNet(block, layers, **kwargs) if pretrained: import os from lib.models.tools.module_helper import ModuleHelper if os.path.exists(pretrained): ModuleHelper.load_model(model, pretrained, all_match=False) else: state_dict = load_state_dict_from_url(pretrained, progress=progress) model.load_state_dict(state_dict, strict=False) return model
def __init__(self, num_classes, dropout=0, bn_type=None, inplane1=512, inplane2=256): super(CE2P_Decoder_Module, self).__init__() self.conv1 = nn.Sequential( nn.Conv2d(inplane1, 256, kernel_size=1, padding=0, dilation=1, bias=False), ModuleHelper.BNReLU(256, bn_type=bn_type), ) self.conv2 = nn.Sequential( nn.Conv2d(inplane2, 48, kernel_size=1, stride=1, padding=0, dilation=1, bias=False), ModuleHelper.BNReLU(48, bn_type=bn_type), ) self.conv3 = nn.Sequential( nn.Conv2d(304, 256, kernel_size=1, padding=0, dilation=1, bias=False), ModuleHelper.BNReLU(256, bn_type=bn_type), nn.Conv2d(256, 256, kernel_size=1, padding=0, dilation=1, bias=False), ModuleHelper.BNReLU(256, bn_type=bn_type), nn.Dropout2d(dropout), ) self.conv4 = nn.Conv2d(256, num_classes, kernel_size=1, padding=0, dilation=1, bias=True)
def __init__(self, inplanes, planes, stride=1, downsample=None, norm_type=None): super(BasicBlock, self).__init__() self.conv1 = conv3x3(inplanes, planes, stride) self.bn1 = ModuleHelper.BatchNorm2d(norm_type=norm_type)(planes) self.relu = nn.ReLU(inplace=True) self.conv2 = conv3x3(planes, planes) self.bn2 = ModuleHelper.BatchNorm2d(norm_type=norm_type)(planes) self.downsample = downsample self.stride = stride
def __init__(self, in_channels, key_channels, value_channels, out_channels=None, bn_type=None): super(SelfAttentionBlock2D, self).__init__() self.in_channels = in_channels self.out_channels = out_channels self.key_channels = key_channels self.value_channels = value_channels if out_channels == None: self.out_channels = in_channels self.f_key = nn.Sequential( nn.Conv2d(self.in_channels, self.key_channels, kernel_size=1, bias=False), ModuleHelper.BNReLU(self.key_channels, bn_type=bn_type), nn.Conv2d(self.key_channels, self.key_channels, kernel_size=1, bias=False), ModuleHelper.BNReLU(self.key_channels, bn_type=bn_type), ) self.f_query = nn.Sequential( nn.Conv2d(self.in_channels, self.key_channels, kernel_size=1, bias=False), ModuleHelper.BNReLU(self.key_channels, bn_type=bn_type), nn.Conv2d(self.key_channels, self.key_channels, kernel_size=1, bias=False), ModuleHelper.BNReLU(self.key_channels, bn_type=bn_type), ) self.f_value = nn.Conv2d(self.in_channels, self.value_channels, kernel_size=1, bias=False) self.W = nn.Sequential( nn.Conv2d(self.value_channels, self.out_channels, kernel_size=1, bias=False), ModuleHelper.BNReLU(self.out_channels, bn_type=bn_type))
def __init__(self, configer): super(SegFix_HRNet, self).__init__() self.configer = configer self.backbone = BackboneSelector(configer).get_backbone() backbone_name = self.configer.get('network', 'backbone') width = int(backbone_name[-2:]) if 'hrnet' in backbone_name: in_channels = width * 15 else: in_channels = width * 31 num_masks = 2 num_directions = DTOffsetConfig.num_classes mid_channels = 256 self.dir_head = nn.Sequential( nn.Conv2d(in_channels, mid_channels, kernel_size=1, stride=1, padding=0, bias=False), ModuleHelper.BNReLU(mid_channels, bn_type=self.configer.get( 'network', 'bn_type')), nn.Conv2d(mid_channels, num_directions, kernel_size=1, stride=1, padding=0, bias=False)) self.mask_head = nn.Sequential( nn.Conv2d(in_channels, mid_channels, kernel_size=1, stride=1, padding=0, bias=False), ModuleHelper.BNReLU(mid_channels, bn_type=self.configer.get( 'network', 'bn_type')), nn.Conv2d(mid_channels, num_masks, kernel_size=1, stride=1, padding=0, bias=False))
def _make_head(self, pre_stage_channels, bn_type, bn_momentum): head_block = Bottleneck head_channels = [32, 64, 128, 256] Log.info("pre_stage_channels: {}".format(pre_stage_channels)) Log.info("head_channels: {}".format(head_channels)) # Increasing the #channels on each resolution # from C, 2C, 4C, 8C to 128, 256, 512, 1024 incre_modules = [] for i, channels in enumerate(pre_stage_channels): incre_module = self._make_layer(head_block, channels, head_channels[i], 1, bn_type=bn_type, bn_momentum=bn_momentum) incre_modules.append(incre_module) incre_modules = nn.ModuleList(incre_modules) # downsampling modules downsamp_modules = [] for i in range(len(pre_stage_channels) - 1): in_channels = head_channels[i] * head_block.expansion out_channels = head_channels[i + 1] * head_block.expansion downsamp_module = nn.Sequential( nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=3, stride=2, padding=1), ModuleHelper.BatchNorm2d(bn_type=bn_type)( out_channels, momentum=bn_momentum), nn.ReLU(inplace=False)) downsamp_modules.append(downsamp_module) downsamp_modules = nn.ModuleList(downsamp_modules) final_layer = nn.Sequential( nn.Conv2d(in_channels=head_channels[3] * head_block.expansion, out_channels=2048, kernel_size=1, stride=1, padding=0), ModuleHelper.BatchNorm2d(bn_type=bn_type)(2048, momentum=bn_momentum), nn.ReLU(inplace=False)) return incre_modules, downsamp_modules, final_layer
def __init__(self, in_channels, key_channels, scale, bn_type=None): super(FastBaseOC_Context_Module, self).__init__() self.object_context_block = ObjectAttentionBlock2D(in_channels, key_channels, scale=1, bn_type=bn_type) self.conv_bn_dropout = nn.Sequential( nn.Conv2d(in_channels, in_channels, kernel_size=1, padding=0), ModuleHelper.BNReLU(in_channels, bn_type=bn_type), )
def __init__(self, configer): self.inplanes = 128 super(DeepLabV3, self).__init__() self.configer = configer self.num_classes = self.configer.get('data', 'num_classes') self.backbone = BackboneSelector(configer).get_backbone() self.head = nn.Sequential( ASPPModule(2048, bn_type=self.configer.get('network', 'bn_type')), nn.Conv2d(512, self.num_classes, kernel_size=1, stride=1, padding=0, bias=True)) self.dsn = nn.Sequential( nn.Conv2d(1024, 512, kernel_size=3, stride=1, padding=1), ModuleHelper.BNReLU(512, bn_type=self.configer.get( 'network', 'bn_type')), nn.Dropout2d(0.1), nn.Conv2d(512, self.num_classes, kernel_size=1, stride=1, padding=0, bias=True))
def __init__(self, configer, flag="", target_class=1): super(DeployClsModel, self).__init__() self.configer = configer self.flag = flag if len(flag) == 0 else "{}_".format(flag) self.backbone = BackboneSelector(configer).get_backbone( backbone_type=configer.get('network.{}backbone'.format(self.flag)), rm_last_stride=configer.get('network', '{}rm_last_stride'.format(self.flag), default=False) ) self.reduction = None fc_dim_out = configer.get('network.{}fc_dim'.format(self.flag), default=None) fc_dim = self.backbone.num_features if fc_dim_out is not None: self.reduction = nn.Conv2d(self.backbone.num_features, fc_dim_out, 1) fc_dim = fc_dim_out self.bn = None if configer.get('network.{}fc_bn'.format(self.flag), default=None): self.bn = nn.BatchNorm2d(fc_dim) if self.configer.get('deploy.extract_score', default=False) or self.configer.get('deploy.extract_cam', default=False): self.linear_lists = nn.ModuleList() for source in range(self.configer.get('data', 'num_data_sources')): linear_list = nn.ModuleList() linear_type = self.configer.get('network', '{}src{}_linear_type'.format(self.flag, source)) for num_classes in self.configer.get('data.src{}_num_classes'.format(source)): linear_list.append(ModuleHelper.Linear(linear_type)(fc_dim, num_classes)) self.linear_lists.append(linear_list)
def __init__(self, bn_type=None, inplane1=512, inplane2=256, outplane=128): super(Decoder_Module, self).__init__() self.conv1 = nn.Sequential( nn.Conv2d(inplane1, 256, kernel_size=1, padding=0, dilation=1, bias=False), ModuleHelper.BNReLU(256, bn_type=bn_type), ) self.conv2 = nn.Sequential( nn.Conv2d(inplane2, 48, kernel_size=1, padding=0, dilation=1, bias=False), ModuleHelper.BNReLU(48, bn_type=bn_type), ) self.conv3 = nn.Sequential( nn.Conv2d(304, outplane, kernel_size=1, padding=0, dilation=1, bias=False), ModuleHelper.BNReLU(outplane, bn_type=bn_type), nn.Conv2d(outplane, outplane, kernel_size=1, padding=0, dilation=1, bias=False), ModuleHelper.BNReLU(outplane, bn_type=bn_type), )
def __init__(self, configer): super(HRNet_W48_ASPOCR, self).__init__() self.configer = configer self.num_classes = self.configer.get('data', 'num_classes') self.backbone = BackboneSelector(configer).get_backbone() # extra added layers in_channels = 720 # 48 + 96 + 192 + 384 from lib.models.modules.spatial_ocr_block import SpatialOCR_ASP_Module self.asp_ocr_head = SpatialOCR_ASP_Module(features=720, hidden_features=256, out_features=256, dilations=(24, 48, 72), num_classes=self.num_classes, bn_type=self.configer.get( 'network', 'bn_type')) self.cls_head = nn.Conv2d(256, self.num_classes, kernel_size=1, stride=1, padding=0, bias=False) self.aux_head = nn.Sequential( nn.Conv2d(in_channels, 512, kernel_size=3, stride=1, padding=1), ModuleHelper.BNReLU(512, bn_type=self.configer.get( 'network', 'bn_type')), nn.Conv2d(512, self.num_classes, kernel_size=1, stride=1, padding=0, bias=False))
def __init__(self, in_channels, key_channels, out_channels, scale=1, dropout=0.1, use_gt=False, use_bg=False, use_oc=True, fetch_attention=False, bn_type=None): super(SpatialOCR_Module, self).__init__() self.use_gt = use_gt self.use_bg = use_bg self.use_oc = use_oc self.fetch_attention = fetch_attention self.object_context_block = ObjectAttentionBlock2D( in_channels, key_channels, scale, use_gt, use_bg, fetch_attention, bn_type) if self.use_bg: if self.use_oc: _in_channels = 3 * in_channels else: _in_channels = 2 * in_channels else: _in_channels = 2 * in_channels self.conv_bn_dropout = nn.Sequential( nn.Conv2d(_in_channels, out_channels, kernel_size=1, padding=0), ModuleHelper.BNReLU(out_channels, bn_type=bn_type), nn.Dropout2d(dropout))
def mobilenetv2(self): model = MobileNetV2() model = ModuleHelper.load_model(model, pretrained=self.configer.get( 'network', 'pretrained'), all_match=False) return model
def __init__(self, configer): self.inplanes = 128 super(ASPOCRNet, self).__init__() self.configer = configer self.num_classes = self.configer.get('data', 'num_classes') self.backbone = BackboneSelector(configer).get_backbone() # extra added layers if "wide_resnet38" in self.configer.get('network', 'backbone'): in_channels = [2048, 4096] else: in_channels = [1024, 2048] # we should increase the dilation rates as the output stride is larger from lib.models.modules.spatial_ocr_block import SpatialOCR_ASP_Module self.asp_ocr_head = SpatialOCR_ASP_Module(features=2048, hidden_features=256, out_features=256, num_classes=self.num_classes, bn_type=self.configer.get('network', 'bn_type')) self.head = nn.Conv2d(256, self.num_classes, kernel_size=1, stride=1, padding=0, bias=True) self.dsn_head = nn.Sequential( nn.Conv2d(in_channels[0], 512, kernel_size=3, stride=1, padding=1), ModuleHelper.BNReLU(512, bn_type=self.configer.get('network', 'bn_type')), nn.Dropout2d(0.1), nn.Conv2d(512, self.num_classes, kernel_size=1, stride=1, padding=0, bias=True) )
def __init__(self, configer, loss_dict=None, flag=""): super(ClsModel, self).__init__() self.configer = configer self.flag = flag if len(flag) == 0 else "{}_".format(flag) self.backbone = slim.__dict__[configer.get('network.{}backbone'.format(self.flag))]( pretrained=configer.get('network.{}pretrained'.format(self.flag)), has_classifier=False ) self.reduction = None fc_dim_out = configer.get('network.{}fc_dim'.format(self.flag), default=None) fc_dim = self.backbone.num_features if fc_dim_out is not None: self.reduction = nn.Conv2d(self.backbone.num_features, fc_dim_out, 1) fc_dim = fc_dim_out self.linear_list = nn.ModuleList() linear_type = configer.get('network', '{}linear_type'.format(self.flag)) for num_classes in configer.get('data.num_classes'): self.linear_list.append(ModuleHelper.Linear(linear_type)(fc_dim, num_classes)) self.embed = None if configer.get('network.{}embed'.format(self.flag), default=True): feat_dim = configer.get('network', '{}feat_dim'.format(self.flag)) self.embed = nn.Sequential( nn.Linear(fc_dim, feat_dim), nn.BatchNorm1d(feat_dim) ) self.bn = nn.BatchNorm1d(fc_dim) nn.init.zeros_(self.bn.bias) self.bn.bias.requires_grad = False self.valid_loss_dict = LOSS_TYPE[configer.get('loss', 'loss_type')] if loss_dict is None else loss_dict
def _make_layer(self, block, planes, blocks, stride=1, norm_type=None, rm_last_stride=False): downsample = None if stride != 1 or self.inplanes != planes * block.expansion: if rm_last_stride: stride = 1 downsample = nn.Sequential( nn.Conv2d(self.inplanes, planes * block.expansion, kernel_size=1, stride=stride, bias=False), ModuleHelper.BatchNorm2d(norm_type=norm_type)(planes * block.expansion), ) layers = [] layers.append( block(self.inplanes, planes, stride, downsample, norm_type=norm_type)) self.inplanes = planes * block.expansion for i in range(1, blocks): layers.append(block(self.inplanes, planes, norm_type=norm_type)) return nn.Sequential(*layers)
def squeezenet(self): """Constructs a ResNet-18 model. Args: pretrained (bool): If True, returns a model pre-trained on Places """ model = SqueezeNet() model = ModuleHelper.load_model(model, pretrained=self.configer.get('network', 'pretrained'), all_match=False) return model
def __init__(self, inplanes, outplanes, kernel_size, stride, padding=1, dilation=1, relu=True, bn_type=None): super(_ConvBatchNormReluBlock, self).__init__() self.relu = relu self.conv = nn.Conv2d(in_channels=inplanes,out_channels=outplanes, kernel_size=kernel_size, stride=stride, padding=padding, dilation = dilation, bias=False) self.bn = ModuleHelper.BatchNorm2d(bn_type=bn_type)(num_features=outplanes) self.relu_f = nn.ReLU()
def __init__(self, inplanes, planes, stride=1, downsample=None, bn_type=None, bn_momentum=0.1): super(BasicBlock, self).__init__() self.conv1 = conv3x3(inplanes, planes, stride) self.bn1 = ModuleHelper.BatchNorm2d(bn_type=bn_type)( planes, momentum=bn_momentum) self.relu = nn.ReLU(inplace=False) self.relu_in = nn.ReLU(inplace=True) self.conv2 = conv3x3(planes, planes) self.bn2 = ModuleHelper.BatchNorm2d(bn_type=bn_type)( planes, momentum=bn_momentum) self.downsample = downsample self.stride = stride