def __init__(self,
num_classes,
enable_auxiliary_loss=True,
align_corners=False,
pretrained=None):
super().__init__()
self.learning_to_downsample = LearningToDownsample(32, 48, 64)
self.global_feature_extractor = GlobalFeatureExtractor(
in_channels=64,
block_channels=[64, 96, 128],
out_channels=128,
expansion=6,
num_blocks=[3, 3, 3],
align_corners=True)
self.feature_fusion = FeatureFusionModule(64, 128, 128, align_corners)
self.classifier = Classifier(128, num_classes)
if enable_auxiliary_loss:
self.auxlayer = layers.AuxLayer(64, 32, num_classes)
self.enable_auxiliary_loss = enable_auxiliary_loss
self.align_corners = align_corners
self.pretrained = pretrained
self.init_weight()
def __init__(self,
num_classes,
backbone_indices,
backbone_channels,
pp_out_channels,
bin_sizes,
enable_auxiliary_loss=True):
super().__init__()
self.backbone_indices = backbone_indices
self.psp_module = layers.PPModule(in_channels=backbone_channels[1],
out_channels=pp_out_channels,
bin_sizes=bin_sizes)
self.dropout = nn.Dropout(p=0.1) # dropout_prob
self.conv = nn.Conv2D(in_channels=pp_out_channels,
out_channels=num_classes,
kernel_size=1)
if enable_auxiliary_loss:
self.auxlayer = layers.AuxLayer(
in_channels=backbone_channels[0],
inter_channels=backbone_channels[0] // 4,
out_channels=num_classes)
self.enable_auxiliary_loss = enable_auxiliary_loss
def __init__(self,
num_classes,
backbone,
backbone_indices=[0, 1, 2, 3],
enable_auxiliary_loss=True,
align_corners=True,
jpu=True,
pretrained=None):
super().__init__()
self.nclass = num_classes
self.aux = enable_auxiliary_loss
self.jpu = jpu
self.backbone = backbone
self.backbone_indices = backbone_indices
self.align_corners = align_corners
self.jpu = layers.JPU([512, 1024, 2048], width=512) if jpu else None
self.head = GIHead(in_channels=2048, nclass=num_classes)
if self.aux:
self.auxlayer = layers.AuxLayer(1024,
1024 // 4,
num_classes,
bias_attr=False)
self.pretrained = pretrained
self.init_weight()
def __init__(self,
num_classes,
backbone,
backbone_indices=[1, 2, 3],
num_codes=32,
mid_channels=512,
use_se_loss=True,
add_lateral=False,
pretrained=None):
super().__init__()
self.add_lateral = add_lateral
self.num_codes = num_codes
self.backbone = backbone
self.backbone_indices = backbone_indices
in_channels = [
self.backbone.feat_channels[index] for index in backbone_indices
]
self.bottleneck = layers.ConvBNReLU(
in_channels[-1],
mid_channels,
3,
padding=1,
)
if self.add_lateral:
self.lateral_convs = nn.LayerList()
for in_ch in in_channels[:-1]:
self.lateral_convs.append(
layers.ConvBNReLU(
in_ch,
mid_channels,
1,
))
self.fusion = layers.ConvBNReLU(
len(in_channels) * mid_channels,
mid_channels,
3,
padding=1,
)
self.enc_module = EncModule(mid_channels, num_codes)
self.head = nn.Conv2D(mid_channels, num_classes, 1)
self.fcn_head = layers.AuxLayer(self.backbone.feat_channels[2],
mid_channels, num_classes)
self.use_se_loss = use_se_loss
if use_se_loss:
self.se_layer = nn.Linear(mid_channels, num_classes)
self.pretrained = pretrained
self.init_weight()
def __init__(self,
num_classes,
backbone_indices,
backbone_channels,
key_value_channels,
inter_channels,
psp_size,
enable_auxiliary_loss=True):
super().__init__()
low_in_channels = backbone_channels[0]
high_in_channels = backbone_channels[1]
self.fusion = AFNB(
low_in_channels=low_in_channels,
high_in_channels=high_in_channels,
out_channels=high_in_channels,
key_channels=key_value_channels,
value_channels=key_value_channels,
dropout_prob=0.05,
repeat_sizes=([1]),
psp_size=psp_size)
self.context = nn.Sequential(
layers.ConvBNReLU(
in_channels=high_in_channels,
out_channels=inter_channels,
kernel_size=3,
padding=1),
APNB(
in_channels=inter_channels,
out_channels=inter_channels,
key_channels=key_value_channels,
value_channels=key_value_channels,
dropout_prob=0.05,
repeat_sizes=([1]),
psp_size=psp_size))
self.cls = nn.Conv2D(
in_channels=inter_channels, out_channels=num_classes, kernel_size=1)
self.auxlayer = layers.AuxLayer(
in_channels=low_in_channels,
inter_channels=low_in_channels // 2,
out_channels=num_classes,
dropout_prob=0.05)
self.backbone_indices = backbone_indices
self.enable_auxiliary_loss = enable_auxiliary_loss
def __init__(self,
inplane,
num_class,
fpn_inplanes,
fpn_dim=256,
enable_auxiliary_loss=False):
super(SFNetHead, self).__init__()
self.ppm = layers.PPModule(in_channels=inplane,
out_channels=fpn_dim,
bin_sizes=(1, 2, 3, 6),
dim_reduction=True,
align_corners=True)
self.enable_auxiliary_loss = enable_auxiliary_loss
self.fpn_in = []
for fpn_inplane in fpn_inplanes[:-1]:
self.fpn_in.append(
nn.Sequential(nn.Conv2D(fpn_inplane, fpn_dim, 1),
layers.SyncBatchNorm(fpn_dim), nn.ReLU()))
self.fpn_in = nn.LayerList(self.fpn_in)
self.fpn_out = []
self.fpn_out_align = []
self.dsn = []
for i in range(len(fpn_inplanes) - 1):
self.fpn_out.append(
nn.Sequential(
layers.ConvBNReLU(fpn_dim, fpn_dim, 3, bias_attr=False)))
self.fpn_out_align.append(
AlignedModule(inplane=fpn_dim, outplane=fpn_dim // 2))
if self.enable_auxiliary_loss:
self.dsn.append(
nn.Sequential(layers.AuxLayer(fpn_dim, fpn_dim,
num_class)))
self.fpn_out = nn.LayerList(self.fpn_out)
self.fpn_out_align = nn.LayerList(self.fpn_out_align)
if self.enable_auxiliary_loss:
self.dsn = nn.LayerList(self.dsn)
self.conv_last = nn.Sequential(
layers.ConvBNReLU(len(fpn_inplanes) * fpn_dim,
fpn_dim,
3,
bias_attr=False),
nn.Conv2D(fpn_dim, num_class, kernel_size=1))
def __init__(self,
num_classes,
backbone_indices,
backbone_channels,
gc_channels,
ratio,
enable_auxiliary_loss=True):
super().__init__()
in_channels = backbone_channels[1]
self.conv_bn_relu1 = layers.ConvBNReLU(in_channels=in_channels,
out_channels=gc_channels,
kernel_size=3,
padding=1)
self.gc_block = GlobalContextBlock(gc_channels=gc_channels,
in_channels=gc_channels,
ratio=ratio)
self.conv_bn_relu2 = layers.ConvBNReLU(in_channels=gc_channels,
out_channels=gc_channels,
kernel_size=3,
padding=1)
self.conv_bn_relu3 = layers.ConvBNReLU(in_channels=in_channels +
gc_channels,
out_channels=gc_channels,
kernel_size=3,
padding=1)
self.dropout = nn.Dropout(p=0.1)
self.conv = nn.Conv2D(in_channels=gc_channels,
out_channels=num_classes,
kernel_size=1)
if enable_auxiliary_loss:
self.auxlayer = layers.AuxLayer(
in_channels=backbone_channels[0],
inter_channels=backbone_channels[0] // 4,
out_channels=num_classes)
self.backbone_indices = backbone_indices
self.enable_auxiliary_loss = enable_auxiliary_loss
def __init__(self, num_classes, enable_auxiliary_loss=True,
pretrained=None):
super().__init__()
self.learning_to_downsample = LearningToDownsample(32, 48, 64)
self.global_feature_extractor = GlobalFeatureExtractor(
64, [64, 96, 128], 128, 6, [3, 3, 3])
self.feature_fusion = FeatureFusionModule(64, 128, 128)
self.classifier = Classifier(128, num_classes)
if enable_auxiliary_loss:
self.auxlayer = layers.AuxLayer(64, 32, num_classes)
self.enable_auxiliary_loss = enable_auxiliary_loss
self.pretrained = pretrained
self.init_weight()
def __init__(self,
num_classes,
backbone,
num_codes=32,
mid_channels=512,
use_jpu=True,
aux_loss=True,
use_se_loss=True,
add_lateral=False,
pretrained=None):
super().__init__()
self.add_lateral = add_lateral
self.num_codes = num_codes
self.backbone = backbone
self.use_jpu = use_jpu
in_channels = self.backbone.feat_channels
if use_jpu:
self.jpu_layer = layers.JPU(in_channels, mid_channels)
in_channels[-1] = mid_channels * 4
self.bottleneck = layers.ConvBNReLU(
in_channels[-1],
mid_channels,
1,
padding=0,
bias_attr=False,
)
else:
self.bottleneck = layers.ConvBNReLU(
in_channels[-1],
mid_channels,
3,
padding=1,
bias_attr=False,
)
if self.add_lateral:
self.lateral_convs = nn.LayerList([
layers.ConvBNReLU(in_channels[0],
mid_channels,
1,
bias_attr=False),
layers.ConvBNReLU(in_channels[1],
mid_channels,
1,
bias_attr=False),
])
self.fusion = layers.ConvBNReLU(
3 * mid_channels,
mid_channels,
3,
padding=1,
bias_attr=False,
)
self.enc_module = EncModule(mid_channels, num_codes)
self.cls_seg = nn.Conv2D(mid_channels, num_classes, 1)
self.aux_loss = aux_loss
if self.aux_loss:
self.fcn_head = layers.AuxLayer(in_channels[-2], mid_channels,
num_classes)
self.use_se_loss = use_se_loss
if use_se_loss:
self.se_layer = nn.Linear(mid_channels, num_classes)
self.pretrained = pretrained
self.init_weight()
