def __init__(self,
in_channel,
out_channel,
stride=1,
dilation=1,
downsample=None,
residual=True,
norm='bn'):
super(Bottleneck, self).__init__()
# 1x1
self.conv_1_1x1 = nn.Conv2d(in_channel, out_channel, 1, bias=False)
self.bn_1 = get_norm(norm, channels=out_channel)
# 3x3
self.conv_2_3x3 = nn.Conv2d(out_channel,
out_channel,
3,
stride=stride,
padding=dilation,
dilation=dilation,
bias=False)
self.bn_2 = get_norm(norm, channels=out_channel)
# 1x1
self.conv_3_1x1 = nn.Conv2d(out_channel,
4 * out_channel,
1,
bias=False)
self.bn_3 = get_norm(norm, channels=4 * out_channel)
self.downsample = downsample
self.residual = residual
self.relu = nn.ReLU(inplace=True)
def __init__(self,
in_channels,
out_channels,
se_loss,
dim_codes,
norm='bn'):
super(EncCore, self).__init__()
self.top = nn.Sequential(
nn.Conv2d(in_channels,
512,
kernel_size=3,
stride=1,
padding=1,
bias=False), get_norm(norm, channels=512),
nn.ReLU(inplace=True))
self.encoding = nn.Sequential(
nn.Conv2d(512, 512, kernel_size=1, bias=False),
get_norm(norm, channels=512), nn.ReLU(inplace=True),
Encoding(D=512, K=dim_codes),
get_norm(
'{}1d'.format(norm) if norm == 'bn' or norm == 'sn' else 'gn',
channels=dim_codes), nn.ReLU(inplace=True), Mean(dim=1))
self.fc = nn.Sequential(nn.Linear(512, 512), nn.Sigmoid())
self.se_loss = se_loss
if se_loss:
self.se_layer = nn.Linear(512, out_channels)
self.tail = nn.Sequential(nn.Dropout2d(0.1, False),
nn.Conv2d(512, out_channels, kernel_size=1))
def __init__(self,
in_channel,
out_channel,
stride=1,
dilation=1,
downsample=None,
residual=True,
norm='bn'):
super(BasicBlock, self).__init__()
self.conv_1_3x3 = nn.Conv2d(in_channel,
out_channel,
kernel_size=3,
stride=stride,
padding=dilation,
dilation=dilation,
bias=False)
self.bn_1 = get_norm(norm, channels=out_channel)
self.conv_2_3x3 = nn.Conv2d(out_channel,
out_channel,
kernel_size=3,
stride=stride,
padding=dilation,
dilation=dilation,
bias=False)
self.bn_2 = get_norm(norm, channels=out_channel)
self.downsample = downsample
self.residual = residual
self.relu = nn.ReLU(inplace=True)
def __init__(self, in_channels, out_channels, norm='bn', up_method='conv'):
super(Up, self).__init__()
self.in_channels = in_channels
self.out_channels = out_channels
if up_method == 'upsample':
self.up = nn.Upsample(scale_factor=2,
mode='bilinear',
align_corners=True)
else:
self.up = nn.ConvTranspose2d(in_channels,
in_channels,
kernel_size=2,
stride=2)
self.conv_up = nn.Sequential(
nn.Conv2d(in_channels * 2, out_channels, kernel_size=3, padding=1),
get_norm(name=norm, channels=out_channels), nn.ReLU(inplace=True),
nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
get_norm(name=norm, channels=out_channels), nn.ReLU(inplace=True))
self.conv = nn.Sequential(
nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
get_norm(name=norm, channels=out_channels), nn.ReLU(inplace=True),
nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
get_norm(name=norm, channels=out_channels), nn.ReLU(inplace=True))
def _make_layers(self,
block,
in_channel,
out_channel,
nbr_blocks,
stride=1,
dilation=1,
norm='bn'):
downsample = None
if stride != 1 or in_channel != out_channel * block.expansion:
downsample = nn.Sequential(
nn.Conv2d(in_channel,
out_channel * block.expansion,
kernel_size=1,
stride=stride,
bias=False),
get_norm(norm, channels=out_channel * block.expansion))
layers = []
layers.append(
block(in_channel,
out_channel,
stride=stride,
downsample=downsample,
dilation=1 if dilation == 1 else dilation // 2,
norm=norm))
in_channel = out_channel * block.expansion
for _ in range(1, nbr_blocks):
layers.append(
block(in_channel, out_channel, dilation=dilation, norm=norm))
return nn.Sequential(*layers)
def _make_layers(self, cfg, batch_norm=True, dilation=True, norm='bn'):
layers = []
in_channels = 3
multi_grid = [1, 2, 4, 8]
i = 0
for v in cfg:
if v == 'M':
layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
i = 0
else:
if not dilation:
conv2d = nn.Conv2d(in_channels,
v,
kernel_size=3,
padding=1)
else:
conv2d = nn.Conv2d(in_channels,
v,
kernel_size=3,
padding=multi_grid[i],
dilation=multi_grid[i])
i += 1
if batch_norm:
layers += [
conv2d,
get_norm(norm, channels=v),
nn.ReLU(inplace=True)
]
else:
layers += [conv2d, nn.ReLU(inplace=True)]
in_channels = v
return nn.Sequential(*layers)
def __init__(self,
nbr_classes,
backbone='xception',
deep_supervision=True,
os=16,
norm='bn',
**kwargs):
super(DeepLabV3Plus, self).__init__()
self.nbr_classes = nbr_classes
self.up_method = {'mode': 'bilinear', 'align_corners': True}
self.backbone = get_backbone(backbone, norm=norm, **kwargs)
self.core = DeepLabV3PlusCore(in_channels=2048,
out_channels=nbr_classes,
backbone=self.backbone,
up_method=self.up_method,
os=os)
if deep_supervision:
self.aux_branch = nn.Sequential(
nn.Conv2d(1024,
256,
kernel_size=3,
stride=1,
padding=1,
bias=False), get_norm(norm, channels=256),
nn.ReLU(False), nn.Dropout2d(0.1, False),
nn.Conv2d(256, nbr_classes, kernel_size=1, stride=1))
self.deep_supervision = deep_supervision
def __init__(self,
nbr_classes,
deep_supervision=True,
backbone='resnet50',
se_loss=True,
norm='bn',
**kwargs):
super(EncNet, self).__init__()
self.up_method = {'mode': 'bilinear', 'align_corners': True}
self.nbr_classes = nbr_classes
self.backbone = get_backbone(backbone, norm=norm, **kwargs)
self.core = EncCore(in_channels=2048,
out_channels=nbr_classes,
se_loss=se_loss,
dim_codes=32,
norm=norm)
self.deep_supervision = deep_supervision
if self.deep_supervision:
self.aux_branch = nn.Sequential(
nn.Conv2d(1024,
256,
kernel_size=3,
stride=1,
padding=1,
bias=False), get_norm(norm, channels=256),
nn.ReLU(False), nn.Dropout2d(0.1, False),
nn.Conv2d(256, nbr_classes, kernel_size=1, stride=1))
def __init__(self, in_channels, out_channels, up_method, norm='bn'):
super(ASPPPooling, self).__init__()
self.up_method = up_method
self.global_pooling = nn.Sequential(nn.AdaptiveAvgPool2d(1),
nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False),
get_norm(norm, channels=out_channels),
nn.ReLU(True))
def __init__(self,
nbr_classes,
backbone='vgg16',
norm='bn',
deep_supervision=True,
**kwargs):
super(UNet, self).__init__()
self.nbr_classes = nbr_classes
self.deep_supervision = deep_supervision
self.up_method = {'mode': 'bilinear', 'align_corners': True}
up_method = 'conv'
self.backbone = get_backbone(backbone, **kwargs)
self.core = UNetCore(out_channels=nbr_classes,
norm=norm,
up_method=up_method,
skip_dims=self.backbone.skip_dims)
if deep_supervision:
self.aux_branch = nn.Sequential(
nn.Conv2d(self.backbone.aux_dim,
256,
kernel_size=3,
stride=1,
padding=1,
bias=False), get_norm(norm, channels=256),
nn.ReLU(False), nn.Dropout2d(0.1, False),
nn.Conv2d(256, nbr_classes, kernel_size=1, stride=1))
def fcn_conv(in_channels, out_channels, norm):
return nn.Sequential(
nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1, bias=False),
get_norm(norm, channels=in_channels),
nn.ReLU(False),
nn.Dropout2d(0.1, False),
nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1)
)
def __init__(self, in_channels, out_channels, up_method, norm='bn'):
super(PSPCore, self).__init__()
branch_channels = in_channels // 4
self.branch_1 = nn.Sequential(
nn.AdaptiveAvgPool2d(1),
nn.Conv2d(in_channels, branch_channels, 1, bias=False),
get_norm(norm, channels=branch_channels),
nn.ReLU(True)
)
self.branch_2 = nn.Sequential(
nn.AdaptiveAvgPool2d(2),
nn.Conv2d(in_channels, branch_channels, 1, bias=False),
get_norm(norm, channels=branch_channels),
nn.ReLU(True)
)
self.branch_3 = nn.Sequential(
nn.AdaptiveAvgPool2d(3),
nn.Conv2d(in_channels, branch_channels, 1, bias=False),
get_norm(norm, channels=branch_channels),
nn.ReLU(True)
)
self.branch_6 = nn.Sequential(
nn.AdaptiveAvgPool2d(6),
nn.Conv2d(in_channels, branch_channels, 1, bias=False),
get_norm(norm, channels=branch_channels),
nn.ReLU(True)
)
self.tail = nn.Sequential(
nn.Conv2d(in_channels*2, branch_channels, kernel_size=3, stride=1, padding=1, bias=False),
get_norm(norm, channels=branch_channels),
nn.ReLU(True),
nn.Dropout2d(0.1, False),
nn.Conv2d(branch_channels, out_channels, kernel_size=1, stride=1)
)
self.up_method = up_method
def __init__(self, in_channels, out_channels, up_method, os=16, norm='bn'):
super(DeepLabV3Core, self).__init__()
rate = 16 // os
inter_channels = in_channels // os
self.up_method = up_method
self.aspp = ASPP(in_channels, inter_channels, self.up_method, rate=rate)
self.tail = nn.Sequential(
nn.Conv2d(inter_channels, inter_channels, kernel_size=3, padding=1, bias=False),
get_norm(norm, channels=inter_channels),
nn.ReLU(True),
nn.Dropout2d(0.1, False),
nn.Conv2d(inter_channels, out_channels, 1)
)
def __init__(self,
in_channels,
out_channels,
backbone,
up_method,
os=16,
norm='bn'):
super(DeepLabV3PlusCore, self).__init__()
rate = 16 // os
inter_channels = in_channels // os
self.up_method = up_method
self.aspp = ASPP(in_channels,
inter_channels,
self.up_method,
rate=rate)
self.backbone = backbone
self.outer_branch = nn.Sequential(
nn.Conv2d(256, 48, 1, 1, padding=0, bias=False),
get_norm(norm, channels=48), nn.ReLU(inplace=True))
self.os = os
self.merge_layers = nn.Sequential(
nn.Conv2d(inter_channels + 48,
inter_channels,
3,
1,
padding=1,
bias=False), get_norm(norm, channels=inter_channels),
nn.ReLU(inplace=True),
nn.Conv2d(inter_channels,
inter_channels,
3,
1,
padding=1,
bias=False), get_norm(norm, channels=inter_channels),
nn.ReLU(inplace=True), nn.Dropout2d(0.1, False))
self.cls = nn.Conv2d(inter_channels, out_channels, 1)
def __init__(self, in_channels, out_channels, up_method, rate=1, norm='bn'):
super(ASPP, self).__init__()
self.branch_1 = nn.Sequential(
nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0, bias=False),
get_norm(norm, channels=out_channels),
nn.ReLU(True)
)
self.branch_2 = nn.Sequential(
nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, dilation=6 * rate, padding=6 * rate,
bias=False),
get_norm(norm, channels=out_channels),
nn.ReLU(True)
)
self.branch_3 = nn.Sequential(
nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, dilation=12 * rate, padding=12 * rate,
bias=False),
get_norm(norm, channels=out_channels),
nn.ReLU(True)
)
self.branch_4 = nn.Sequential(
nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, dilation=18 * rate, padding=18 * rate,
bias=False),
get_norm(norm, channels=out_channels),
nn.ReLU(True)
)
self.up_method = up_method
self.pooling_branch = ASPPPooling(in_channels, out_channels, self.up_method)
self.merge_branch = nn.Sequential(
nn.Conv2d(5 * out_channels, out_channels, 1, bias=False),
get_norm(norm, channels=out_channels),
nn.ReLU(True),
nn.Dropout2d(0.1, False)
)
def __init__(self,
nbr_classes=1000,
is_bottleneck=True,
nbr_layers=50,
sk_conn=False,
norm='bn'):
super(ResNet, self).__init__()
global net_structures
if nbr_layers not in net_structures:
raise RuntimeError(
"nbr_layers can only be 50, 101 or 152, but got {}".format(
nbr_layers))
net_structure = net_structures[nbr_layers]
self.sk_conn = sk_conn
self.aux_dim = 1024
self.conv1 = nn.Sequential(conv3x3(3, 64, stride=2),
get_norm(norm, channels=64),
nn.ReLU(inplace=True))
self.conv2 = nn.Sequential(conv3x3(64, 64), get_norm(norm,
channels=64),
nn.ReLU(inplace=True))
self.conv3 = nn.Sequential(conv3x3(64, 128),
get_norm(norm, channels=128),
nn.ReLU(inplace=True))
self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
residual_block = Bottleneck if is_bottleneck else BasicBlock
self.block_1 = self._make_layers(residual_block,
128,
64,
net_structure[0],
norm=norm)
self.block_2 = self._make_layers(residual_block,
64 * residual_block.expansion,
128,
net_structure[1],
stride=2,
norm=norm)
self.block_3 = self._make_layers(residual_block,
128 * residual_block.expansion,
256,
net_structure[2],
dilation=2,
norm=norm)
self.block_4 = self._make_layers(residual_block,
256 * residual_block.expansion,
512,
net_structure[3],
dilation=4,
norm=norm)
self.avg_pool = nn.AdaptiveAvgPool2d(1)
self.fc = nn.Linear(512 * residual_block.expansion, nbr_classes)
if self.sk_conn:
self.skip_dims = [2048, 1024, 256, 128]
self.ratio_mapping = {2: -1, 4: -2, 8: -3}
for m in self.modules():
if isinstance(m, nn.Conv2d):
n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
m.weight.data.normal_(0, math.sqrt(2. / n))
elif isinstance(m, nn.BatchNorm2d):
m.weight.data.fill_(1)
m.bias.data.zero_()
elif isinstance(m, nn.GroupNorm):
m.weight.data.fill_(1)
m.bias.data.zero_()
print("init resnet")
