def __init__(self, inplanes, planes, stride=1, downsample=None):
super(Resnet_BasicBlock, self).__init__()
self.conv1 = ws.Conv2dws(inplanes, planes, kernel_size=3, stride=1, padding=1)
self.bn1 = sn.SwitchNorm2d(planes)
self.conv2 = ws.Conv2dws(planes, planes, kernel_size=3, stride=1, padding=1)
self.bn2 = sn.SwitchNorm2d(planes)
self.relu = nn.ReLU(inplace=True)
self.downsample = downsample
self.stride = stride
def __init__(self, in_channels, out_channels, kernel_size=1, stride=1, padding=0, dilation=1, groups=1, bias=False):
kernel_size = _pair(kernel_size)
stride = _pair(stride)
padding = _pair(padding)
dilation = _pair(dilation)
super(GatedSpatialConv2d, self).__init__(in_channels, out_channels, kernel_size, stride, padding, dilation,
False, _pair(0), groups, bias, 'zeros')
self.gate_conv = nn.Sequential(
sn.SwitchNorm2d(in_channels+1),
ws.Conv2dws(in_channels+1, in_channels+1, 1),
nn.ReLU(),
ws.Conv2dws(in_channels+1, 1, 1),
sn.SwitchNorm2d(1),
nn.Sigmoid()
)
def ConvSnRelu(in_channels,
out_channels,
using_movavg=True,
using_bn=True,
last_gamma=False,
kernel_size=(3, 3),
stride=1,
padding=1):
conv = ws.Conv2dws(in_channels, out_channels, kernel_size, stride, padding)
norm = sn.SwitchNorm2d(out_channels,
using_movavg=using_movavg,
using_bn=using_bn,
last_gamma=last_gamma)
relu = nn.ReLU()
return nn.Sequential(conv, norm, relu)
def __init__(self, n_channels, n_filters, n_class, using_movavg, using_bn):
super(GSCNN, self).__init__()
self.maxPool = nn.MaxPool2d(kernel_size=2, stride=2, ceil_mode=True)
# down1
self.convd1_1 = ConvSnRelu(n_channels, n_filters, using_movavg,
using_bn)
self.convd1_2 = ConvSnRelu(n_filters, n_filters, using_movavg,
using_bn)
# down2
self.convd2_1 = ConvSnRelu(n_filters, n_filters * 2, using_movavg,
using_bn)
self.convd2_2 = ConvSnRelu(n_filters * 2, n_filters * 2, using_movavg,
using_bn)
# down3
self.convd3_1 = ConvSnRelu(n_filters * 2, n_filters * 4, using_movavg,
using_bn)
self.convd3_2 = ConvSnRelu(n_filters * 4, n_filters * 4, using_movavg,
using_bn)
# down4
self.convd4_1 = ConvSnRelu(n_filters * 4, n_filters * 8, using_movavg,
using_bn)
self.convd4_2 = ConvSnRelu(n_filters * 8, n_filters * 8, using_movavg,
using_bn)
## center
self.conv5_1 = ConvSnRelu(n_filters * 8, n_filters * 16, using_movavg,
using_bn)
self.conv5_2 = ConvSnRelu(n_filters * 16, n_filters * 16, using_movavg,
using_bn)
# down6
#self.convd6_1 = ConvSnRelu(n_filters*16, n_filters*32, using_movavg, using_bn)
#self.convd6_2 = ConvSnRelu(n_filters*32, n_filters*32, using_movavg, using_bn)
## center
#self.conv7_1 = ConvSnRelu(n_filters*32, n_filters*64, using_movavg, using_bn)
#self.conv7_2 = ConvSnRelu(n_filters*64, n_filters*64, using_movavg, using_bn)
self.dsn1 = ws.Conv2dws(64, 1, 1)
self.dsn3 = ws.Conv2dws(256, 1, 1)
self.dsn4 = ws.Conv2dws(512, 1, 1)
self.dsn5 = ws.Conv2dws(1024, 1, 1)
self.res1 = Resnet_BasicBlock(64, 64, stride=1, downsample=None)
self.d1 = ws.Conv2dws(64, 32, 1)
self.res2 = Resnet_BasicBlock(32, 32, stride=1, downsample=None)
self.d2 = ws.Conv2dws(32, 16, 1)
self.res3 = Resnet_BasicBlock(16, 16, stride=1, downsample=None)
self.d3 = ws.Conv2dws(16, 8, 1)
self.fuse = ws.Conv2dws(8, 1, kernel_size=1, padding=0, bias=False)
self.cw = ws.Conv2dws(2, 1, kernel_size=1, padding=0, bias=False)
self.gate1 = GatedSpatialConv2d(32, 32)
self.gate2 = GatedSpatialConv2d(16, 16)
self.gate3 = GatedSpatialConv2d(8, 8)
## ------------------aspp------------------
self.aspp1 = ASPP(
1024, 256, 1, padding=0, dilation=1,
BatchNorm=nn.BatchNorm2d) # 0 1 1024 nn.BatchNorm2d
self.aspp2 = ASPP(1024,
256,
3,
padding=6,
dilation=6,
BatchNorm=nn.BatchNorm2d) # 2
self.aspp3 = ASPP(1024,
256,
3,
padding=12,
dilation=12,
BatchNorm=nn.BatchNorm2d) # 6
self.aspp4 = ASPP(1024,
256,
3,
padding=18,
dilation=18,
BatchNorm=nn.BatchNorm2d) # 12
self.global_avg_pool = nn.Sequential(
nn.AdaptiveAvgPool2d((1, 1)),
nn.Conv2d(1024, 256, 1, stride=1, bias=False), nn.BatchNorm2d(256),
nn.ReLU())
self.edge_conv = nn.Sequential(
nn.AdaptiveAvgPool2d((1, 1)),
nn.Conv2d(1, 256, 1, stride=1, bias=False), nn.BatchNorm2d(256),
nn.ReLU())
self.bot_aspp = ws.Conv2dws(1280 + 256, 256, kernel_size=1, bias=False)
# ConvSnRelu(1536, 256, using_movavg=1, using_bn=1, kernel_size=1, padding=0)
self.bot_fine = ws.Conv2dws(128, 48, kernel_size=1, bias=False)
self.final_seg = nn.Sequential(
ws.Conv2dws(256 + 48, 256, kernel_size=3, padding=1, bias=False),
sn.SwitchNorm2d(256), nn.ReLU(inplace=True),
ws.Conv2dws(256, 256, kernel_size=3, padding=1, bias=False),
sn.SwitchNorm2d(256), nn.ReLU(inplace=True),
ws.Conv2dws(256, n_class, kernel_size=1, bias=False))
self.sigmoid = nn.Sigmoid()