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