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, 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))
class PSPNet(nn.Sequential):
def __init__(self, configer):
super(PSPNet, self).__init__()
self.configer = configer
self.num_classes = self.configer.get('data', 'num_classes')
self.backbone = BackboneSelector(configer).get_backbone()
num_features = self.backbone.get_num_features()
self.dsn = nn.Sequential(
_ConvBatchNormReluBlock(num_features // 2, num_features // 4, 3, 1,
bn_type=self.configer.get('network', 'bn_type')),
nn.Dropout2d(0.1),
nn.Conv2d(num_features // 4, self.num_classes, 1, 1, 0)
)
self.ppm = PPMBilinearDeepsup(fc_dim=num_features, bn_type=self.configer.get('network', 'bn_type'))
self.cls = nn.Sequential(
nn.Conv2d(num_features + 4 * 512, 512, kernel_size=3, padding=1, bias=False),
ModuleHelper.BNReLU(512, bn_type=self.configer.get('network', 'bn_type')),
nn.Dropout2d(0.1),
nn.Conv2d(512, self.num_classes, kernel_size=1)
)
def forward(self, x_):
x = self.backbone(x_)
aux_x = self.dsn(x[-2])
x = self.ppm(x[-1])
x = self.cls(x)
aux_x = F.interpolate(aux_x, size=(x_.size(2), x_.size(3)), mode="bilinear", align_corners=True)
x = F.interpolate(x, size=(x_.size(2), x_.size(3)), mode="bilinear", align_corners=True)
return aux_x, x
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)
)
class PSPNet(nn.Sequential):
def __init__(self, configer):
super(PSPNet, self).__init__()
self.configer = configer
self.num_classes = self.configer.get('data', 'num_classes')
self.backbone = BackboneSelector(configer).get_backbone()
num_features = self.backbone.get_num_features()
self.low_features = nn.Sequential(
self.backbone.conv1, self.backbone.bn1, self.backbone.relu,
self.backbone.maxpool,
self.backbone.layer1,
)
self.high_features1 = nn.Sequential(self.backbone.layer2, self.backbone.layer3)
self.high_features2 = nn.Sequential(self.backbone.layer4)
self.decoder = PPMBilinearDeepsup(num_class=self.num_classes, fc_dim=num_features,
bn_type=self.configer.get('network', 'bn_type'))
def forward(self, x_):
low = self.low_features(x_)
aux = self.high_features1(low)
x = self.high_features2(aux)
x, aux = self.decoder([x, aux])
x = F.interpolate(x, size=(x_.size(2), x_.size(3)), mode="bilinear", align_corners=False)
return x, aux
def __init__(self, configer):
super(DeepLabV3Contrast, self).__init__()
self.configer = configer
self.num_classes = self.configer.get('data', 'num_classes')
self.backbone = BackboneSelector(configer).get_backbone()
self.proj_dim = self.configer.get('contrast', 'proj_dim')
# extra added layers
if "wide_resnet38" in self.configer.get('network', 'backbone'):
in_channels = [2048, 4096]
else:
in_channels = [1024, 2048]
self.proj_head = ProjectionHead(dim_in=in_channels[1],
proj_dim=self.proj_dim)
self.decoder = DeepLabHead(num_classes=self.num_classes,
bn_type=self.configer.get(
'network', 'bn_type'))
for modules in [self.proj_head, self.decoder]:
for m in modules.modules():
if isinstance(m, nn.Conv2d):
nn.init.kaiming_normal_(m.weight.data)
if m.bias is not None:
m.bias.data.zero_()
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, 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):
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, configer):
super(PSPNet, self).__init__()
self.configer = configer
self.num_classes = self.configer.get('data', 'num_classes')
self.backbone = BackboneSelector(configer).get_backbone()
num_features = self.backbone.get_num_features()
self.low_features = nn.Sequential(
self.backbone.conv1, self.backbone.bn1, self.backbone.relu,
self.backbone.maxpool,
self.backbone.layer1,
)
self.high_features1 = nn.Sequential(self.backbone.layer2, self.backbone.layer3)
self.high_features2 = nn.Sequential(self.backbone.layer4)
self.decoder = PPMBilinearDeepsup(num_class=self.num_classes, fc_dim=num_features,
bn_type=self.configer.get('network', 'bn_type'))
def __init__(self, configer):
super(PSPNet, self).__init__()
self.configer = configer
self.num_classes = self.configer.get('data', 'num_classes')
self.backbone = BackboneSelector(configer).get_backbone()
num_features = self.backbone.get_num_features()
self.dsn = nn.Sequential(
_ConvBatchNormReluBlock(num_features // 2, num_features // 4, 3, 1,
bn_type=self.configer.get('network', 'bn_type')),
nn.Dropout2d(0.1),
nn.Conv2d(num_features // 4, self.num_classes, 1, 1, 0)
)
self.ppm = PPMBilinearDeepsup(fc_dim=num_features, bn_type=self.configer.get('network', 'bn_type'))
self.cls = nn.Sequential(
nn.Conv2d(num_features + 4 * 512, 512, kernel_size=3, padding=1, bias=False),
ModuleHelper.BNReLU(512, bn_type=self.configer.get('network', 'bn_type')),
nn.Dropout2d(0.1),
nn.Conv2d(512, self.num_classes, kernel_size=1)
)
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 __init__(self, configer):
super(HRNet_W48_OCR, self).__init__()
self.configer = configer
self.num_classes = self.configer.get('data', 'num_classes')
self.backbone = BackboneSelector(configer).get_backbone()
in_channels = 720
self.conv3x3 = 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')),
)
from lib.models.modules.spatial_ocr_block import SpatialGather_Module
self.ocr_gather_head = SpatialGather_Module(self.num_classes)
from lib.models.modules.spatial_ocr_block import SpatialOCR_Module
self.ocr_distri_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.cls_head = nn.Conv2d(512,
self.num_classes,
kernel_size=1,
stride=1,
padding=0,
bias=True)
self.aux_head = nn.Sequential(
nn.Conv2d(in_channels,
in_channels,
kernel_size=3,
stride=1,
padding=1),
ModuleHelper.BNReLU(in_channels,
bn_type=self.configer.get(
'network', 'bn_type')),
nn.Conv2d(in_channels,
self.num_classes,
kernel_size=1,
stride=1,
padding=0,
bias=True))
def __init__(self, configer):
self.inplanes = 128
super(FcnNet_wo_dsn, 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]
elif "mobilenetv2" in self.configer.get('network', 'backbone'):
in_channels = [160, 320]
else:
in_channels = [1024, 2048]
self.cls_head = 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')), nn.Dropout2d(0.10),
nn.Conv2d(512,
self.num_classes,
kernel_size=1,
stride=1,
padding=0,
bias=True))
if "mobilenetv2" in self.configer.get('network', 'backbone'):
self.cls_head = nn.Sequential(
nn.Conv2d(in_channels[1],
256,
kernel_size=3,
stride=1,
padding=1),
ModuleHelper.BNReLU(256,
bn_type=self.configer.get(
'network', 'bn_type')),
nn.Dropout2d(0.10),
nn.Conv2d(256,
self.num_classes,
kernel_size=1,
stride=1,
padding=0,
bias=False))
def __init__(self, configer):
self.inplanes = 128
super(AspOCNet, 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]
from lib.models.modules.asp_oc_block import ASP_OC_Module
self.context = 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")),
ASP_OC_Module(512,
256,
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, configer):
self.inplanes = 128
super(FastAspOCNet, self).__init__()
self.configer = configer
self.num_classes = self.configer.get('data', 'num_classes')
self.backbone = BackboneSelector(configer).get_backbone()
# extra added layers
self.down = nn.Sequential(
nn.Conv2d(2048, 512, kernel_size=3, stride=1, padding=1),
ModuleHelper.BNReLU(512, bn_type=self.configer.get('network', 'bn_type')),
)
# self.object_head = ObjectContext_Module(num_classes)
from lib.models.modules.fast_asp_oc_block import Fast_ASP_OC_Module
self.fast_oc_head = Fast_ASP_OC_Module(512, 512, num_classes=self.num_classes,
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(1024, 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, configer):
self.inplanes = 128
super(AspOCNet, self).__init__()
self.configer = configer
self.num_classes = self.configer.get('data', 'num_classes')
self.backbone = BackboneSelector(configer).get_backbone()
# extra added layers
from lib.models.modules.asp_oc_block import ASP_OC_Module
self.context = nn.Sequential(
nn.Conv2d(2048, 512, kernel_size=3, stride=1, padding=1),
ModuleHelper.BNReLU(512,
bn_type=self.configer.get(
'network', 'bn_type')),
ASP_OC_Module(512,
512,
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(1024, 512, kernel_size=3, stride=1, padding=1),
ModuleHelper.BNReLU(512,
bn_type=self.configer.get(
'network', 'bn_type')), nn.Dropout2d(0.10),
nn.Conv2d(512,
self.num_classes,
kernel_size=1,
stride=1,
padding=0,
bias=True))
self.context.apply(self.weights_init)
self.cls.apply(self.weights_init)
self.dsn.apply(self.weights_init)
def __init__(self, configer):
super(HRNet_W48, 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
self.cls_head = nn.Sequential(
nn.Conv2d(in_channels,
in_channels,
kernel_size=3,
stride=1,
padding=1),
ModuleHelper.BNReLU(in_channels,
bn_type=self.configer.get(
'network', 'bn_type')), nn.Dropout2d(0.10),
nn.Conv2d(in_channels,
self.num_classes,
kernel_size=1,
stride=1,
padding=0,
bias=False))
def __init__(self, configer):
self.inplanes = 128
super(PyramidOCNet, self).__init__()
self.configer = configer
self.num_classes = self.configer.get('data', 'num_classes')
self.backbone = BackboneSelector(configer).get_backbone()
self.layer5 = nn.Sequential(
nn.Conv2d(2048, 512, kernel_size=3, stride=1, padding=1),
ModuleHelper.BNReLU(512,
bn_type=self.configer.get(
'network', 'bn_type')))
# extra added layers
from lib.models.modules.pyramid_oc_block import Pyramid_OC_Module
self.context = Pyramid_OC_Module(in_channels=512,
out_channels=512,
dropout=0.05,
sizes=([1, 2, 3, 6]),
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(1024, 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, configer):
self.inplanes = 128
super(ISANet, self).__init__()
self.configer = configer
self.num_classes = self.configer.get('data', 'num_classes')
self.backbone = BackboneSelector(configer).get_backbone()
# extra added layers
bn_type = self.configer.get('network', 'bn_type')
factors = self.configer.get('network', 'factors')
from lib.models.modules.isa_block import ISA_Module
self.isa_head = nn.Sequential(
nn.Conv2d(2048, 512, kernel_size=3, stride=1, padding=1, bias=False),
ModuleHelper.BNReLU(512, bn_type=bn_type),
ISA_Module(in_channels=512, key_channels=256, value_channels=512,
out_channels=512, down_factors=factors, dropout=0.05, bn_type=bn_type),
)
self.cls_head = nn.Conv2d(512, self.num_classes, kernel_size=1, stride=1, padding=0, bias=True)
self.dsn_head = nn.Sequential(
nn.Conv2d(1024, 512, kernel_size=3, stride=1, padding=1, bias=False),
ModuleHelper.BNReLU(512, bn_type=bn_type),
nn.Dropout2d(0.05),
nn.Conv2d(512, self.num_classes, kernel_size=1, stride=1, padding=0, bias=True),
)
def __init__(self, configer):
super(DenseASPP, self).__init__()
self.configer = configer
dropout0 = 0.1
dropout1 = 0.1
self.backbone = BackboneSelector(configer).get_backbone()
num_features = self.backbone.get_num_features()
self.trans = _Transition(num_input_features=self.num_features,
num_output_features=self.num_features // 2,
bn_type=self.configer.get(
'network', 'bn_type'))
self.num_features = self.num_features // 2
self.ASPP_3 = _DenseAsppBlock(input_num=num_features,
num1=256,
num2=64,
dilation_rate=3,
drop_out=dropout0,
bn_type=self.configer.get(
'network', 'bn_type'))
self.ASPP_6 = _DenseAsppBlock(input_num=num_features + 64 * 1,
num1=256,
num2=64,
dilation_rate=6,
drop_out=dropout0,
bn_type=self.configer.get(
'network', 'bn_type'))
self.ASPP_12 = _DenseAsppBlock(input_num=num_features + 64 * 2,
num1=256,
num2=64,
dilation_rate=12,
drop_out=dropout0,
bn_type=self.configer.get(
'network', 'bn_type'))
self.ASPP_18 = _DenseAsppBlock(input_num=num_features + 64 * 3,
num1=256,
num2=64,
dilation_rate=18,
drop_out=dropout0,
bn_type=self.configer.get(
'network', 'bn_type'))
self.ASPP_24 = _DenseAsppBlock(input_num=num_features + 64 * 4,
num1=256,
num2=64,
dilation_rate=24,
drop_out=dropout0,
bn_type=self.configer.get(
'network', 'bn_type'))
num_features = num_features + 5 * 64
self.classification = nn.Sequential(
nn.Dropout2d(p=dropout1),
nn.Conv2d(in_channels=num_features,
out_channels=self.configer.get('network',
'out_channels'),
kernel_size=1,
padding=0))
for m in self.modules():
if isinstance(m, nn.Conv2d):
nn.init.xavier_uniform_(m.weight.data)
elif isinstance(
m,
ModuleHelper.BatchNorm2d(
bn_type=self.configer.get('network', 'bn_type'))):
m.weight.data.fill_(1)
m.bias.data.zero_()
class DenseASPP(nn.Module):
"""
* output_scale can only set as 8 or 16
"""
def __init__(self, configer):
super(DenseASPP, self).__init__()
self.configer = configer
dropout0 = 0.1
dropout1 = 0.1
self.backbone = BackboneSelector(configer).get_backbone()
num_features = self.backbone.get_num_features()
self.trans = _Transition(num_input_features=self.num_features,
num_output_features=self.num_features // 2,
bn_type=self.configer.get(
'network', 'bn_type'))
self.num_features = self.num_features // 2
self.ASPP_3 = _DenseAsppBlock(input_num=num_features,
num1=256,
num2=64,
dilation_rate=3,
drop_out=dropout0,
bn_type=self.configer.get(
'network', 'bn_type'))
self.ASPP_6 = _DenseAsppBlock(input_num=num_features + 64 * 1,
num1=256,
num2=64,
dilation_rate=6,
drop_out=dropout0,
bn_type=self.configer.get(
'network', 'bn_type'))
self.ASPP_12 = _DenseAsppBlock(input_num=num_features + 64 * 2,
num1=256,
num2=64,
dilation_rate=12,
drop_out=dropout0,
bn_type=self.configer.get(
'network', 'bn_type'))
self.ASPP_18 = _DenseAsppBlock(input_num=num_features + 64 * 3,
num1=256,
num2=64,
dilation_rate=18,
drop_out=dropout0,
bn_type=self.configer.get(
'network', 'bn_type'))
self.ASPP_24 = _DenseAsppBlock(input_num=num_features + 64 * 4,
num1=256,
num2=64,
dilation_rate=24,
drop_out=dropout0,
bn_type=self.configer.get(
'network', 'bn_type'))
num_features = num_features + 5 * 64
self.classification = nn.Sequential(
nn.Dropout2d(p=dropout1),
nn.Conv2d(in_channels=num_features,
out_channels=self.configer.get('network',
'out_channels'),
kernel_size=1,
padding=0))
for m in self.modules():
if isinstance(m, nn.Conv2d):
nn.init.xavier_uniform_(m.weight.data)
elif isinstance(
m,
ModuleHelper.BatchNorm2d(
bn_type=self.configer.get('network', 'bn_type'))):
m.weight.data.fill_(1)
m.bias.data.zero_()
def forward(self, x):
feature = self.backbone(x)
aspp3 = self.ASPP_3(feature)
feature = torch.cat((aspp3, feature), dim=1)
aspp6 = self.ASPP_6(feature)
feature = torch.cat((aspp6, feature), dim=1)
aspp12 = self.ASPP_12(feature)
feature = torch.cat((aspp12, feature), dim=1)
aspp18 = self.ASPP_18(feature)
feature = torch.cat((aspp18, feature), dim=1)
aspp24 = self.ASPP_24(feature)
feature = torch.cat((aspp24, feature), dim=1)
cls = self.classification(feature)
out = F.interpolate(cls,
scale_factor=8,
mode='bilinear',
align_corners=True)
return out
def __init__(self, configer):
self.inplanes = 128
super(CE2P_IdealOCRNet, self).__init__()
self.configer = configer
self.num_classes = self.configer.get('data', 'num_classes')
self.backbone = BackboneSelector(configer).get_backbone()
from lib.models.modules.edge_block import Edge_Module
from lib.models.modules.decoder_block import Decoder_Module
if "wide_resnet38" in self.configer.get('network', 'backbone'):
in_channels = [2048, 4096]
self.edgelayer = Edge_Module(256,
2,
bn_type=self.configer.get(
'network', 'bn_type'),
factor=2)
self.decoder = Decoder_Module(self.num_classes,
dropout=0.1,
bn_type=self.configer.get(
'network', 'bn_type'),
inplane1=512,
inplane2=512)
else:
in_channels = [1024, 2048]
self.edgelayer = Edge_Module(256,
2,
bn_type=self.configer.get(
'network', 'bn_type'),
factor=1)
self.decoder = Decoder_Module(self.num_classes,
dropout=0.1,
bn_type=self.configer.get(
'network', 'bn_type'),
inplane1=512,
inplane2=256)
# extra added layers
from lib.models.modules.spatial_ocr_block import SpatialGather_Module, SpatialOCR_Module
self.spatial_context_head = SpatialGather_Module(self.num_classes,
use_gt=True)
self.spatial_ocr_head = SpatialOCR_Module(in_channels=2048,
key_channels=256,
out_channels=512,
scale=1,
dropout=0,
use_gt=True,
bn_type=self.configer.get(
'network', 'bn_type'))
self.cls = nn.Sequential(
nn.Conv2d(1024,
256,
kernel_size=1,
padding=0,
dilation=1,
bias=False),
ModuleHelper.BNReLU(256,
bn_type=self.configer.get(
'network', 'bn_type')),
nn.Conv2d(256,
self.num_classes,
kernel_size=1,
padding=0,
dilation=1,
bias=True))
self.dsn = nn.Sequential(
nn.Conv2d(in_channels[0],
512,
kernel_size=3,
stride=1,
padding=1,
bias=False),
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 = BackboneSelector(self.configer).get_backbone(
backbone_type=self.configer.get('network.{}backbone'.format(
self.flag)),
pretrained_model=self.configer.get('network.{}pretrained'.format(
self.flag)),
rm_last_stride=self.configer.get('network.{}rm_last_stride'.format(
self.flag),
default=False))
self.reduction = None
fc_dim_out = self.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 self.configer.get('network.{}fc_bn'.format(self.flag),
default=True):
self.bn = nn.BatchNorm1d(fc_dim)
nn.init.zeros_(self.bn.bias)
self.bn.bias.requires_grad = False
self.relu = None
if self.configer.get('network.{}fc_relu'.format(self.flag),
default=True):
self.relu = nn.ReLU()
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)
self.global_linear = None
if self.configer.get('data.global_num_classes',
default=None) is not None:
global_linear_type = self.configer.get(
'network', '{}global_linear_type'.format(self.flag))
self.global_linear = ModuleHelper.Linear(global_linear_type)(
fc_dim,
self.configer.get('data.global_num_classes', default=None))
self.embed_after_norm = self.configer.get('network.embed_after_norm',
default=True)
self.embed = None
if self.configer.get('network.{}embed'.format(self.flag),
default=True):
feat_dim = self.configer.get('network',
'{}feat_dim'.format(self.flag))
embed = []
embed.append(nn.Linear(fc_dim, feat_dim))
if self.configer.get('network.{}embed_norm_type'.format(
self.flag)) == 'L2':
embed.append(LpNormalize(p=2, dim=1))
elif self.configer.get('network.{}embed_norm_type'.format(
self.flag)) == 'BN':
embed.append(nn.BatchNorm1d(feat_dim))
self.embed = nn.Sequential(*embed)
self.valid_loss_dict = LOSS_TYPE[self.configer.get(
'loss', 'loss_type')] if loss_dict is None else loss_dict