init + inference.py патч

data/MBD/model/__init__.py

@@ -0,0 +1,50 @@
+import torchvision.models as models
+from model.densenetccnl import *
+from model.unetnc import *
+from model.gienet import *
+
+
+def get_model(name, n_classes=1, filters=64,version=None,in_channels=3, is_batchnorm=True, norm='batch', model_path=None, use_sigmoid=True, layers=3,img_size=512):
+    model = _get_model_instance(name)
+
+
+    if name == 'dnetccnl':
+        model = model(img_size=128, in_channels=in_channels, out_channels=n_classes, filters=32)
+    elif name == 'dnetccnl512':
+        model = model(img_size=img_size, in_channels=in_channels, out_channels=n_classes, filters=32)
+    elif name == 'unetnc':
+        model = model(input_nc=in_channels, output_nc=n_classes, num_downs=7)
+    elif name == 'gie':
+        model = model(input_nc=in_channels, output_nc=n_classes, num_downs=7)
+    elif name == 'giecbam':
+        model = model(input_nc=in_channels, output_nc=n_classes, num_downs=7)
+    elif name == 'gie2head':
+        model = model(input_nc=in_channels, output_nc=n_classes, num_downs=7)
+    elif name == 'giemask':
+        model = model(input_nc=in_channels, output_nc=n_classes, num_downs=7)
+    elif name == 'giemask2':
+        model = model(input_nc=in_channels, output_nc=n_classes, num_downs=7)
+    elif name == 'giedilated':
+        model = model(input_nc=in_channels, output_nc=n_classes, num_downs=7)
+    elif name == 'bmp':
+        model = model(input_nc=in_channels, output_nc=n_classes, num_downs=7)
+    elif name == 'displacement':
+        model = model(n_classes=2, num_filter=32, BatchNorm='GN', in_channels=5)
+    return model
+
+def _get_model_instance(name):
+    try:
+        return {
+            'dnetccnl': dnetccnl,
+            'dnetccnl512': dnetccnl512,
+            'unetnc': UnetGenerator,
+            'gie':GieGenerator,
+            'giecbam':GiecbamGenerator,
+            'giedilated':DilatedSingleUnet,
+            'gie2head':Gie2headGenerator,
+            'giemask':GiemaskGenerator,
+            'giemask2':Giemask2Generator,
+            'bmp':BmpGenerator,
+        }[name]
+    except:
+        print('Model {} not available'.format(name))

data/MBD/model/cbam.py

@@ -0,0 +1,95 @@
+import torch
+import math
+import torch.nn as nn
+import torch.nn.functional as F
+
+class BasicConv(nn.Module):
+    def __init__(self, in_planes, out_planes, kernel_size, stride=1, padding=0, dilation=1, groups=1, relu=True, bn=True, bias=False):
+        super(BasicConv, self).__init__()
+        self.out_channels = out_planes
+        self.conv = nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias)
+        self.bn = nn.BatchNorm2d(out_planes,eps=1e-5, momentum=0.01, affine=True) if bn else None
+        self.relu = nn.ReLU() if relu else None
+
+    def forward(self, x):
+        x = self.conv(x)
+        if self.bn is not None:
+            x = self.bn(x)
+        if self.relu is not None:
+            x = self.relu(x)
+        return x
+
+class Flatten(nn.Module):
+    def forward(self, x):
+        return x.view(x.size(0), -1)
+
+class ChannelGate(nn.Module):
+    def __init__(self, gate_channels, reduction_ratio=16, pool_types=['avg', 'max']):
+        super(ChannelGate, self).__init__()
+        self.gate_channels = gate_channels
+        self.mlp = nn.Sequential(
+            Flatten(),
+            nn.Linear(gate_channels, gate_channels // reduction_ratio),
+            nn.ReLU(),
+            nn.Linear(gate_channels // reduction_ratio, gate_channels)
+            )
+        self.pool_types = pool_types
+    def forward(self, x):
+        channel_att_sum = None
+        for pool_type in self.pool_types:
+            if pool_type=='avg':
+                avg_pool = F.avg_pool2d( x, (x.size(2), x.size(3)), stride=(x.size(2), x.size(3)))
+                channel_att_raw = self.mlp( avg_pool )
+            elif pool_type=='max':
+                max_pool = F.max_pool2d( x, (x.size(2), x.size(3)), stride=(x.size(2), x.size(3)))
+                channel_att_raw = self.mlp( max_pool )
+            elif pool_type=='lp':
+                lp_pool = F.lp_pool2d( x, 2, (x.size(2), x.size(3)), stride=(x.size(2), x.size(3)))
+                channel_att_raw = self.mlp( lp_pool )
+            elif pool_type=='lse':
+                # LSE pool only
+                lse_pool = logsumexp_2d(x)
+                channel_att_raw = self.mlp( lse_pool )
+
+            if channel_att_sum is None:
+                channel_att_sum = channel_att_raw
+            else:
+                channel_att_sum = channel_att_sum + channel_att_raw
+
+        scale = F.sigmoid( channel_att_sum ).unsqueeze(2).unsqueeze(3).expand_as(x)
+        return x * scale
+
+def logsumexp_2d(tensor):
+    tensor_flatten = tensor.view(tensor.size(0), tensor.size(1), -1)
+    s, _ = torch.max(tensor_flatten, dim=2, keepdim=True)
+    outputs = s + (tensor_flatten - s).exp().sum(dim=2, keepdim=True).log()
+    return outputs
+
+class ChannelPool(nn.Module):
+    def forward(self, x):
+        return torch.cat( (torch.max(x,1)[0].unsqueeze(1), torch.mean(x,1).unsqueeze(1)), dim=1 )
+
+class SpatialGate(nn.Module):
+    def __init__(self):
+        super(SpatialGate, self).__init__()
+        kernel_size = 7
+        self.compress = ChannelPool()
+        self.spatial = BasicConv(2, 1, kernel_size, stride=1, padding=(kernel_size-1) // 2, relu=False)
+    def forward(self, x):
+        x_compress = self.compress(x)
+        x_out = self.spatial(x_compress)
+        scale = F.sigmoid(x_out) # broadcasting
+        return x * scale
+
+class CBAM(nn.Module):
+    def __init__(self, gate_channels, reduction_ratio=16, pool_types=['avg', 'max'], no_spatial=False):
+        super(CBAM, self).__init__()
+        self.ChannelGate = ChannelGate(gate_channels, reduction_ratio, pool_types)
+        self.no_spatial=no_spatial
+        if not no_spatial:
+            self.SpatialGate = SpatialGate()
+    def forward(self, x):
+        x_out = self.ChannelGate(x)
+        if not self.no_spatial:
+            x_out = self.SpatialGate(x_out)
+        return x_out

data/MBD/model/deep_lab_model/aspp.py

@@ -0,0 +1,95 @@
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from model.deep_lab_model.sync_batchnorm.batchnorm import SynchronizedBatchNorm2d
+
+class _ASPPModule(nn.Module):
+    def __init__(self, inplanes, planes, kernel_size, padding, dilation, BatchNorm):
+        super(_ASPPModule, self).__init__()
+        self.atrous_conv = nn.Conv2d(inplanes, planes, kernel_size=kernel_size,
+                                            stride=1, padding=padding, dilation=dilation, bias=False)
+        self.bn = BatchNorm(planes)
+        self.relu = nn.ReLU()
+
+        self._init_weight()
+
+    def forward(self, x):
+        x = self.atrous_conv(x)
+        x = self.bn(x)
+
+        return self.relu(x)
+
+    def _init_weight(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                torch.nn.init.kaiming_normal_(m.weight)
+            elif isinstance(m, SynchronizedBatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+
+class ASPP(nn.Module):
+    def __init__(self, backbone, output_stride, BatchNorm):
+        super(ASPP, self).__init__()
+        if backbone == 'drn':
+            inplanes = 512
+        elif backbone == 'mobilenet':
+            inplanes = 320
+        else:
+            inplanes = 2048
+        if output_stride == 16:
+            dilations = [1, 6, 12, 18]
+        elif output_stride == 8:
+            dilations = [1, 12, 24, 36]
+        else:
+            raise NotImplementedError
+
+        self.aspp1 = _ASPPModule(inplanes, 256, 1, padding=0, dilation=dilations[0], BatchNorm=BatchNorm)
+        self.aspp2 = _ASPPModule(inplanes, 256, 3, padding=dilations[1], dilation=dilations[1], BatchNorm=BatchNorm)
+        self.aspp3 = _ASPPModule(inplanes, 256, 3, padding=dilations[2], dilation=dilations[2], BatchNorm=BatchNorm)
+        self.aspp4 = _ASPPModule(inplanes, 256, 3, padding=dilations[3], dilation=dilations[3], BatchNorm=BatchNorm)
+
+        self.global_avg_pool = nn.Sequential(nn.AdaptiveAvgPool2d((1, 1)),
+                                             nn.Conv2d(inplanes, 256, 1, stride=1, bias=False),
+                                             BatchNorm(256),
+                                             nn.ReLU())
+        self.conv1 = nn.Conv2d(1280, 256, 1, bias=False)
+        self.bn1 = BatchNorm(256)
+        self.relu = nn.ReLU()
+        self.dropout = nn.Dropout(0.5)
+        self._init_weight()
+
+    def forward(self, x):
+        x1 = self.aspp1(x)
+        x2 = self.aspp2(x)
+        x3 = self.aspp3(x)
+        x4 = self.aspp4(x)
+        x5 = self.global_avg_pool(x)
+        x5 = F.interpolate(x5, size=x4.size()[2:], mode='bilinear', align_corners=True)
+        x = torch.cat((x1, x2, x3, x4, x5), dim=1)
+
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+
+        return self.dropout(x)
+
+    def _init_weight(self):
+        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))
+                torch.nn.init.kaiming_normal_(m.weight)
+            elif isinstance(m, SynchronizedBatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+
+
+def build_aspp(backbone, output_stride, BatchNorm):
+    return ASPP(backbone, output_stride, BatchNorm)

data/MBD/model/deep_lab_model/backbone/__init__.py

@@ -0,0 +1,13 @@
+from model.deep_lab_model.backbone import resnet, xception, drn, mobilenet
+
+def build_backbone(backbone, output_stride, BatchNorm):
+    if backbone == 'resnet':
+        return resnet.ResNet101(output_stride, BatchNorm)
+    elif backbone == 'xception':
+        return xception.AlignedXception(output_stride, BatchNorm)
+    elif backbone == 'drn':
+        return drn.drn_d_54(BatchNorm)
+    elif backbone == 'mobilenet':
+        return mobilenet.MobileNetV2(output_stride, BatchNorm)
+    else:
+        raise NotImplementedError

data/MBD/model/deep_lab_model/backbone/drn.py

@@ -0,0 +1,402 @@
+import torch.nn as nn
+import math
+import torch.utils.model_zoo as model_zoo
+from model.deep_lab_model.sync_batchnorm.batchnorm import SynchronizedBatchNorm2d
+
+webroot = 'http://dl.yf.io/drn/'
+
+model_urls = {
+    'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
+    'drn-c-26': webroot + 'drn_c_26-ddedf421.pth',
+    'drn-c-42': webroot + 'drn_c_42-9d336e8c.pth',
+    'drn-c-58': webroot + 'drn_c_58-0a53a92c.pth',
+    'drn-d-22': webroot + 'drn_d_22-4bd2f8ea.pth',
+    'drn-d-38': webroot + 'drn_d_38-eebb45f0.pth',
+    'drn-d-54': webroot + 'drn_d_54-0e0534ff.pth',
+    'drn-d-105': webroot + 'drn_d_105-12b40979.pth'
+}
+
+
+def conv3x3(in_planes, out_planes, stride=1, padding=1, dilation=1):
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+                     padding=padding, bias=False, dilation=dilation)
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None,
+                 dilation=(1, 1), residual=True, BatchNorm=None):
+        super(BasicBlock, self).__init__()
+        self.conv1 = conv3x3(inplanes, planes, stride,
+                             padding=dilation[0], dilation=dilation[0])
+        self.bn1 = BatchNorm(planes)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = conv3x3(planes, planes,
+                             padding=dilation[1], dilation=dilation[1])
+        self.bn2 = BatchNorm(planes)
+        self.downsample = downsample
+        self.stride = stride
+        self.residual = residual
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+        if self.residual:
+            out += residual
+        out = self.relu(out)
+
+        return out
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None,
+                 dilation=(1, 1), residual=True, BatchNorm=None):
+        super(Bottleneck, self).__init__()
+        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
+        self.bn1 = BatchNorm(planes)
+        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
+                               padding=dilation[1], bias=False,
+                               dilation=dilation[1])
+        self.bn2 = BatchNorm(planes)
+        self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
+        self.bn3 = BatchNorm(planes * 4)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+
+class DRN(nn.Module):
+
+    def __init__(self, block, layers, arch='D',
+                 channels=(16, 32, 64, 128, 256, 512, 512, 512),
+                 BatchNorm=None):
+        super(DRN, self).__init__()
+        self.inplanes = channels[0]
+        self.out_dim = channels[-1]
+        self.arch = arch
+
+        if arch == 'C':
+            self.conv1 = nn.Conv2d(3, channels[0], kernel_size=7, stride=1,
+                                   padding=3, bias=False)
+            self.bn1 = BatchNorm(channels[0])
+            self.relu = nn.ReLU(inplace=True)
+
+            self.layer1 = self._make_layer(
+                BasicBlock, channels[0], layers[0], stride=1, BatchNorm=BatchNorm)
+            self.layer2 = self._make_layer(
+                BasicBlock, channels[1], layers[1], stride=2, BatchNorm=BatchNorm)
+
+        elif arch == 'D':
+            self.layer0 = nn.Sequential(
+                nn.Conv2d(3, channels[0], kernel_size=7, stride=1, padding=3,
+                          bias=False),
+                BatchNorm(channels[0]),
+                nn.ReLU(inplace=True)
+            )
+
+            self.layer1 = self._make_conv_layers(
+                channels[0], layers[0], stride=1, BatchNorm=BatchNorm)
+            self.layer2 = self._make_conv_layers(
+                channels[1], layers[1], stride=2, BatchNorm=BatchNorm)
+
+        self.layer3 = self._make_layer(block, channels[2], layers[2], stride=2, BatchNorm=BatchNorm)
+        self.layer4 = self._make_layer(block, channels[3], layers[3], stride=2, BatchNorm=BatchNorm)
+        self.layer5 = self._make_layer(block, channels[4], layers[4],
+                                       dilation=2, new_level=False, BatchNorm=BatchNorm)
+        self.layer6 = None if layers[5] == 0 else \
+            self._make_layer(block, channels[5], layers[5], dilation=4,
+                             new_level=False, BatchNorm=BatchNorm)
+
+        if arch == 'C':
+            self.layer7 = None if layers[6] == 0 else \
+                self._make_layer(BasicBlock, channels[6], layers[6], dilation=2,
+                                 new_level=False, residual=False, BatchNorm=BatchNorm)
+            self.layer8 = None if layers[7] == 0 else \
+                self._make_layer(BasicBlock, channels[7], layers[7], dilation=1,
+                                 new_level=False, residual=False, BatchNorm=BatchNorm)
+        elif arch == 'D':
+            self.layer7 = None if layers[6] == 0 else \
+                self._make_conv_layers(channels[6], layers[6], dilation=2, BatchNorm=BatchNorm)
+            self.layer8 = None if layers[7] == 0 else \
+                self._make_conv_layers(channels[7], layers[7], dilation=1, BatchNorm=BatchNorm)
+
+        self._init_weight()
+
+    def _init_weight(self):
+        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, SynchronizedBatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+
+
+    def _make_layer(self, block, planes, blocks, stride=1, dilation=1,
+                    new_level=True, residual=True, BatchNorm=None):
+        assert dilation == 1 or dilation % 2 == 0
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(self.inplanes, planes * block.expansion,
+                          kernel_size=1, stride=stride, bias=False),
+                BatchNorm(planes * block.expansion),
+            )
+
+        layers = list()
+        layers.append(block(
+            self.inplanes, planes, stride, downsample,
+            dilation=(1, 1) if dilation == 1 else (
+                dilation // 2 if new_level else dilation, dilation),
+            residual=residual, BatchNorm=BatchNorm))
+        self.inplanes = planes * block.expansion
+        for i in range(1, blocks):
+            layers.append(block(self.inplanes, planes, residual=residual,
+                                dilation=(dilation, dilation), BatchNorm=BatchNorm))
+
+        return nn.Sequential(*layers)
+
+    def _make_conv_layers(self, channels, convs, stride=1, dilation=1, BatchNorm=None):
+        modules = []
+        for i in range(convs):
+            modules.extend([
+                nn.Conv2d(self.inplanes, channels, kernel_size=3,
+                          stride=stride if i == 0 else 1,
+                          padding=dilation, bias=False, dilation=dilation),
+                BatchNorm(channels),
+                nn.ReLU(inplace=True)])
+            self.inplanes = channels
+        return nn.Sequential(*modules)
+
+    def forward(self, x):
+        if self.arch == 'C':
+            x = self.conv1(x)
+            x = self.bn1(x)
+            x = self.relu(x)
+        elif self.arch == 'D':
+            x = self.layer0(x)
+
+        x = self.layer1(x)
+        x = self.layer2(x)
+
+        x = self.layer3(x)
+        low_level_feat = x
+
+        x = self.layer4(x)
+        x = self.layer5(x)
+
+        if self.layer6 is not None:
+            x = self.layer6(x)
+
+        if self.layer7 is not None:
+            x = self.layer7(x)
+
+        if self.layer8 is not None:
+            x = self.layer8(x)
+
+        return x, low_level_feat
+
+
+class DRN_A(nn.Module):
+
+    def __init__(self, block, layers, BatchNorm=None):
+        self.inplanes = 64
+        super(DRN_A, self).__init__()
+        self.out_dim = 512 * block.expansion
+        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3,
+                               bias=False)
+        self.bn1 = BatchNorm(64)
+        self.relu = nn.ReLU(inplace=True)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        self.layer1 = self._make_layer(block, 64, layers[0], BatchNorm=BatchNorm)
+        self.layer2 = self._make_layer(block, 128, layers[1], stride=2, BatchNorm=BatchNorm)
+        self.layer3 = self._make_layer(block, 256, layers[2], stride=1,
+                                       dilation=2, BatchNorm=BatchNorm)
+        self.layer4 = self._make_layer(block, 512, layers[3], stride=1,
+                                       dilation=4, BatchNorm=BatchNorm)
+
+        self._init_weight()
+
+    def _init_weight(self):
+        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, SynchronizedBatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+
+    def _make_layer(self, block, planes, blocks, stride=1, dilation=1, BatchNorm=None):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(self.inplanes, planes * block.expansion,
+                          kernel_size=1, stride=stride, bias=False),
+                BatchNorm(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample, BatchNorm=BatchNorm))
+        self.inplanes = planes * block.expansion
+        for i in range(1, blocks):
+            layers.append(block(self.inplanes, planes,
+                                dilation=(dilation, dilation, ), BatchNorm=BatchNorm))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
+
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+
+        return x
+
+def drn_a_50(BatchNorm, pretrained=True):
+    model = DRN_A(Bottleneck, [3, 4, 6, 3], BatchNorm=BatchNorm)
+    if pretrained:
+        model.load_state_dict(model_zoo.load_url(model_urls['resnet50']))
+    return model
+
+
+def drn_c_26(BatchNorm, pretrained=True):
+    model = DRN(BasicBlock, [1, 1, 2, 2, 2, 2, 1, 1], arch='C', BatchNorm=BatchNorm)
+    if pretrained:
+        pretrained = model_zoo.load_url(model_urls['drn-c-26'])
+        del pretrained['fc.weight']
+        del pretrained['fc.bias']
+        model.load_state_dict(pretrained)
+    return model
+
+
+def drn_c_42(BatchNorm, pretrained=True):
+    model = DRN(BasicBlock, [1, 1, 3, 4, 6, 3, 1, 1], arch='C', BatchNorm=BatchNorm)
+    if pretrained:
+        pretrained = model_zoo.load_url(model_urls['drn-c-42'])
+        del pretrained['fc.weight']
+        del pretrained['fc.bias']
+        model.load_state_dict(pretrained)
+    return model
+
+
+def drn_c_58(BatchNorm, pretrained=True):
+    model = DRN(Bottleneck, [1, 1, 3, 4, 6, 3, 1, 1], arch='C', BatchNorm=BatchNorm)
+    if pretrained:
+        pretrained = model_zoo.load_url(model_urls['drn-c-58'])
+        del pretrained['fc.weight']
+        del pretrained['fc.bias']
+        model.load_state_dict(pretrained)
+    return model
+
+
+def drn_d_22(BatchNorm, pretrained=True):
+    model = DRN(BasicBlock, [1, 1, 2, 2, 2, 2, 1, 1], arch='D', BatchNorm=BatchNorm)
+    if pretrained:
+        pretrained = model_zoo.load_url(model_urls['drn-d-22'])
+        del pretrained['fc.weight']
+        del pretrained['fc.bias']
+        model.load_state_dict(pretrained)
+    return model
+
+
+def drn_d_24(BatchNorm, pretrained=True):
+    model = DRN(BasicBlock, [1, 1, 2, 2, 2, 2, 2, 2], arch='D', BatchNorm=BatchNorm)
+    if pretrained:
+        pretrained = model_zoo.load_url(model_urls['drn-d-24'])
+        del pretrained['fc.weight']
+        del pretrained['fc.bias']
+        model.load_state_dict(pretrained)
+    return model
+
+
+def drn_d_38(BatchNorm, pretrained=True):
+    model = DRN(BasicBlock, [1, 1, 3, 4, 6, 3, 1, 1], arch='D', BatchNorm=BatchNorm)
+    if pretrained:
+        pretrained = model_zoo.load_url(model_urls['drn-d-38'])
+        del pretrained['fc.weight']
+        del pretrained['fc.bias']
+        model.load_state_dict(pretrained)
+    return model
+
+
+def drn_d_40(BatchNorm, pretrained=True):
+    model = DRN(BasicBlock, [1, 1, 3, 4, 6, 3, 2, 2], arch='D', BatchNorm=BatchNorm)
+    if pretrained:
+        pretrained = model_zoo.load_url(model_urls['drn-d-40'])
+        del pretrained['fc.weight']
+        del pretrained['fc.bias']
+        model.load_state_dict(pretrained)
+    return model
+
+
+def drn_d_54(BatchNorm, pretrained=True):
+    model = DRN(Bottleneck, [1, 1, 3, 4, 6, 3, 1, 1], arch='D', BatchNorm=BatchNorm)
+    if pretrained:
+        pretrained = model_zoo.load_url(model_urls['drn-d-54'])
+        del pretrained['fc.weight']
+        del pretrained['fc.bias']
+        model.load_state_dict(pretrained)
+    return model
+
+
+def drn_d_105(BatchNorm, pretrained=True):
+    model = DRN(Bottleneck, [1, 1, 3, 4, 23, 3, 1, 1], arch='D', BatchNorm=BatchNorm)
+    if pretrained:
+        pretrained = model_zoo.load_url(model_urls['drn-d-105'])
+        del pretrained['fc.weight']
+        del pretrained['fc.bias']
+        model.load_state_dict(pretrained)
+    return model
+
+if __name__ == "__main__":
+    import torch
+    model = drn_a_50(BatchNorm=nn.BatchNorm2d, pretrained=True)
+    input = torch.rand(1, 3, 512, 512)
+    output, low_level_feat = model(input)
+    print(output.size())
+    print(low_level_feat.size())

data/MBD/model/deep_lab_model/backbone/mobilenet.py

@@ -0,0 +1,151 @@
+import torch
+import torch.nn.functional as F
+import torch.nn as nn
+import math
+from model.deep_lab_model.sync_batchnorm.batchnorm import SynchronizedBatchNorm2d
+import torch.utils.model_zoo as model_zoo
+
+def conv_bn(inp, oup, stride, BatchNorm):
+    return nn.Sequential(
+        nn.Conv2d(inp, oup, 3, stride, 1, bias=False),
+        BatchNorm(oup),
+        nn.ReLU6(inplace=True)
+    )
+
+
+def fixed_padding(inputs, kernel_size, dilation):
+    kernel_size_effective = kernel_size + (kernel_size - 1) * (dilation - 1)
+    pad_total = kernel_size_effective - 1
+    pad_beg = pad_total // 2
+    pad_end = pad_total - pad_beg
+    padded_inputs = F.pad(inputs, (pad_beg, pad_end, pad_beg, pad_end))
+    return padded_inputs
+
+
+class InvertedResidual(nn.Module):
+    def __init__(self, inp, oup, stride, dilation, expand_ratio, BatchNorm):
+        super(InvertedResidual, self).__init__()
+        self.stride = stride
+        assert stride in [1, 2]
+
+        hidden_dim = round(inp * expand_ratio)
+        self.use_res_connect = self.stride == 1 and inp == oup
+        self.kernel_size = 3
+        self.dilation = dilation
+
+        if expand_ratio == 1:
+            self.conv = nn.Sequential(
+                # dw
+                nn.Conv2d(hidden_dim, hidden_dim, 3, stride, 0, dilation, groups=hidden_dim, bias=False),
+                BatchNorm(hidden_dim),
+                nn.ReLU6(inplace=True),
+                # pw-linear
+                nn.Conv2d(hidden_dim, oup, 1, 1, 0, 1, 1, bias=False),
+                BatchNorm(oup),
+            )
+        else:
+            self.conv = nn.Sequential(
+                # pw
+                nn.Conv2d(inp, hidden_dim, 1, 1, 0, 1, bias=False),
+                BatchNorm(hidden_dim),
+                nn.ReLU6(inplace=True),
+                # dw
+                nn.Conv2d(hidden_dim, hidden_dim, 3, stride, 0, dilation, groups=hidden_dim, bias=False),
+                BatchNorm(hidden_dim),
+                nn.ReLU6(inplace=True),
+                # pw-linear
+                nn.Conv2d(hidden_dim, oup, 1, 1, 0, 1, bias=False),
+                BatchNorm(oup),
+            )
+
+    def forward(self, x):
+        x_pad = fixed_padding(x, self.kernel_size, dilation=self.dilation)
+        if self.use_res_connect:
+            x = x + self.conv(x_pad)
+        else:
+            x = self.conv(x_pad)
+        return x
+
+
+class MobileNetV2(nn.Module):
+    def __init__(self, output_stride=8, BatchNorm=None, width_mult=1., pretrained=True):
+        super(MobileNetV2, self).__init__()
+        block = InvertedResidual
+        input_channel = 32
+        current_stride = 1
+        rate = 1
+        interverted_residual_setting = [
+            # t, c, n, s
+            [1, 16, 1, 1],
+            [6, 24, 2, 2],
+            [6, 32, 3, 2],
+            [6, 64, 4, 2],
+            [6, 96, 3, 1],
+            [6, 160, 3, 2],
+            [6, 320, 1, 1],
+        ]
+
+        # building first layer
+        input_channel = int(input_channel * width_mult)
+        self.features = [conv_bn(3, input_channel, 2, BatchNorm)]
+        current_stride *= 2
+        # building inverted residual blocks
+        for t, c, n, s in interverted_residual_setting:
+            if current_stride == output_stride:
+                stride = 1
+                dilation = rate
+                rate *= s
+            else:
+                stride = s
+                dilation = 1
+                current_stride *= s
+            output_channel = int(c * width_mult)
+            for i in range(n):
+                if i == 0:
+                    self.features.append(block(input_channel, output_channel, stride, dilation, t, BatchNorm))
+                else:
+                    self.features.append(block(input_channel, output_channel, 1, dilation, t, BatchNorm))
+                input_channel = output_channel
+        self.features = nn.Sequential(*self.features)
+        self._initialize_weights()
+
+        if pretrained:
+            self._load_pretrained_model()
+
+        self.low_level_features = self.features[0:4]
+        self.high_level_features = self.features[4:]
+
+    def forward(self, x):
+        low_level_feat = self.low_level_features(x)
+        x = self.high_level_features(low_level_feat)
+        return x, low_level_feat
+
+    def _load_pretrained_model(self):
+        pretrain_dict = model_zoo.load_url('http://jeff95.me/models/mobilenet_v2-6a65762b.pth')
+        model_dict = {}
+        state_dict = self.state_dict()
+        for k, v in pretrain_dict.items():
+            if k in state_dict:
+                model_dict[k] = v
+        state_dict.update(model_dict)
+        self.load_state_dict(state_dict)
+
+    def _initialize_weights(self):
+        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))
+                torch.nn.init.kaiming_normal_(m.weight)
+            elif isinstance(m, SynchronizedBatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+
+if __name__ == "__main__":
+    input = torch.rand(1, 3, 512, 512)
+    model = MobileNetV2(output_stride=16, BatchNorm=nn.BatchNorm2d)
+    output, low_level_feat = model(input)
+    print(output.size())
+    print(low_level_feat.size())

data/MBD/model/deep_lab_model/backbone/resnet.py

@@ -0,0 +1,170 @@
+import math
+import torch.nn as nn
+import torch.utils.model_zoo as model_zoo
+from model.deep_lab_model.sync_batchnorm.batchnorm import SynchronizedBatchNorm2d
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1, dilation=1, downsample=None, BatchNorm=None):
+        super(Bottleneck, self).__init__()
+        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
+        self.bn1 = BatchNorm(planes)
+        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
+                               dilation=dilation, padding=dilation, bias=False)
+        self.bn2 = BatchNorm(planes)
+        self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
+        self.bn3 = BatchNorm(planes * 4)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+        self.dilation = dilation
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+class ResNet(nn.Module):
+
+    def __init__(self, block, layers, output_stride, BatchNorm, pretrained=True):
+        self.inplanes = 64
+        super(ResNet, self).__init__()
+        blocks = [1, 2, 4]
+        if output_stride == 16:
+            strides = [1, 2, 2, 1]
+            dilations = [1, 1, 1, 2]
+        elif output_stride == 8:
+            strides = [1, 2, 1, 1]
+            dilations = [1, 1, 2, 4]
+        else:
+            raise NotImplementedError
+
+        # Modules
+        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3,
+                                bias=False)
+        self.bn1 = BatchNorm(64)
+        self.relu = nn.ReLU(inplace=True)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+
+        self.layer1 = self._make_layer(block, 64, layers[0], stride=strides[0], dilation=dilations[0], BatchNorm=BatchNorm)
+        self.layer2 = self._make_layer(block, 128, layers[1], stride=strides[1], dilation=dilations[1], BatchNorm=BatchNorm)
+        self.layer3 = self._make_layer(block, 256, layers[2], stride=strides[2], dilation=dilations[2], BatchNorm=BatchNorm)
+        self.layer4 = self._make_MG_unit(block, 512, blocks=blocks, stride=strides[3], dilation=dilations[3], BatchNorm=BatchNorm)
+        # self.layer4 = self._make_layer(block, 512, layers[3], stride=strides[3], dilation=dilations[3], BatchNorm=BatchNorm)
+        self._init_weight()
+
+        # if pretrained:
+        #     self._load_pretrained_model()
+
+    def _make_layer(self, block, planes, blocks, stride=1, dilation=1, BatchNorm=None):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(self.inplanes, planes * block.expansion,
+                          kernel_size=1, stride=stride, bias=False),
+                BatchNorm(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, dilation, downsample, BatchNorm))
+        self.inplanes = planes * block.expansion
+        for i in range(1, blocks):
+            layers.append(block(self.inplanes, planes, dilation=dilation, BatchNorm=BatchNorm))
+
+        return nn.Sequential(*layers)
+
+    def _make_MG_unit(self, block, planes, blocks, stride=1, dilation=1, BatchNorm=None):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(self.inplanes, planes * block.expansion,
+                          kernel_size=1, stride=stride, bias=False),
+                BatchNorm(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, dilation=blocks[0]*dilation,
+                            downsample=downsample, BatchNorm=BatchNorm))
+        self.inplanes = planes * block.expansion
+        for i in range(1, len(blocks)):
+            layers.append(block(self.inplanes, planes, stride=1,
+                                dilation=blocks[i]*dilation, BatchNorm=BatchNorm))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, input):
+        x = self.conv1(input)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
+
+        x = self.layer1(x)
+        low_level_feat = x
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+        return x, low_level_feat
+
+    def _init_weight(self):
+        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, SynchronizedBatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+
+    def _load_pretrained_model(self):
+
+        import urllib.request
+        import ssl
+        ssl._create_default_https_context = ssl._create_unverified_context
+        response = urllib.request.urlopen('https://download.pytorch.org/models/resnet101-5d3b4d8f.pth')
+
+        pretrain_dict = model_zoo.load_url('https://download.pytorch.org/models/resnet101-5d3b4d8f.pth')
+        model_dict = {}
+        state_dict = self.state_dict()
+        for k, v in pretrain_dict.items():
+            if k in state_dict:
+                # if 'conv1' in k:
+                #     continue 
+                model_dict[k] = v
+        state_dict.update(model_dict)
+        self.load_state_dict(state_dict)
+
+def ResNet101(output_stride, BatchNorm, pretrained=True):
+    """Constructs a ResNet-101 model.
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
+    model = ResNet(Bottleneck, [3, 4, 23, 3], output_stride, BatchNorm, pretrained=pretrained)
+    return model
+
+if __name__ == "__main__":
+    import torch
+    model = ResNet101(BatchNorm=nn.BatchNorm2d, pretrained=True, output_stride=8)
+    input = torch.rand(1, 3, 512, 512)
+    output, low_level_feat = model(input)
+    print(output.size())
+    print(low_level_feat.size())

data/MBD/model/deep_lab_model/backbone/xception.py

@@ -0,0 +1,288 @@
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.model_zoo as model_zoo
+from model.deep_lab_model.sync_batchnorm.batchnorm import SynchronizedBatchNorm2d
+
+def fixed_padding(inputs, kernel_size, dilation):
+    kernel_size_effective = kernel_size + (kernel_size - 1) * (dilation - 1)
+    pad_total = kernel_size_effective - 1
+    pad_beg = pad_total // 2
+    pad_end = pad_total - pad_beg
+    padded_inputs = F.pad(inputs, (pad_beg, pad_end, pad_beg, pad_end))
+    return padded_inputs
+
+
+class SeparableConv2d(nn.Module):
+    def __init__(self, inplanes, planes, kernel_size=3, stride=1, dilation=1, bias=False, BatchNorm=None):
+        super(SeparableConv2d, self).__init__()
+
+        self.conv1 = nn.Conv2d(inplanes, inplanes, kernel_size, stride, 0, dilation,
+                               groups=inplanes, bias=bias)
+        self.bn = BatchNorm(inplanes)
+        self.pointwise = nn.Conv2d(inplanes, planes, 1, 1, 0, 1, 1, bias=bias)
+
+    def forward(self, x):
+        x = fixed_padding(x, self.conv1.kernel_size[0], dilation=self.conv1.dilation[0])
+        x = self.conv1(x)
+        x = self.bn(x)
+        x = self.pointwise(x)
+        return x
+
+
+class Block(nn.Module):
+    def __init__(self, inplanes, planes, reps, stride=1, dilation=1, BatchNorm=None,
+                 start_with_relu=True, grow_first=True, is_last=False):
+        super(Block, self).__init__()
+
+        if planes != inplanes or stride != 1:
+            self.skip = nn.Conv2d(inplanes, planes, 1, stride=stride, bias=False)
+            self.skipbn = BatchNorm(planes)
+        else:
+            self.skip = None
+
+        self.relu = nn.ReLU(inplace=True)
+        rep = []
+
+        filters = inplanes
+        if grow_first:
+            rep.append(self.relu)
+            rep.append(SeparableConv2d(inplanes, planes, 3, 1, dilation, BatchNorm=BatchNorm))
+            rep.append(BatchNorm(planes))
+            filters = planes
+
+        for i in range(reps - 1):
+            rep.append(self.relu)
+            rep.append(SeparableConv2d(filters, filters, 3, 1, dilation, BatchNorm=BatchNorm))
+            rep.append(BatchNorm(filters))
+
+        if not grow_first:
+            rep.append(self.relu)
+            rep.append(SeparableConv2d(inplanes, planes, 3, 1, dilation, BatchNorm=BatchNorm))
+            rep.append(BatchNorm(planes))
+
+        if stride != 1:
+            rep.append(self.relu)
+            rep.append(SeparableConv2d(planes, planes, 3, 2, BatchNorm=BatchNorm))
+            rep.append(BatchNorm(planes))
+
+        if stride == 1 and is_last:
+            rep.append(self.relu)
+            rep.append(SeparableConv2d(planes, planes, 3, 1, BatchNorm=BatchNorm))
+            rep.append(BatchNorm(planes))
+
+        if not start_with_relu:
+            rep = rep[1:]
+
+        self.rep = nn.Sequential(*rep)
+
+    def forward(self, inp):
+        x = self.rep(inp)
+
+        if self.skip is not None:
+            skip = self.skip(inp)
+            skip = self.skipbn(skip)
+        else:
+            skip = inp
+
+        x = x + skip
+
+        return x
+
+
+class AlignedXception(nn.Module):
+    """
+    Modified Alighed Xception
+    """
+    def __init__(self, output_stride, BatchNorm,
+                 pretrained=True):
+        super(AlignedXception, self).__init__()
+
+        if output_stride == 16:
+            entry_block3_stride = 2
+            middle_block_dilation = 1
+            exit_block_dilations = (1, 2)
+        elif output_stride == 8:
+            entry_block3_stride = 1
+            middle_block_dilation = 2
+            exit_block_dilations = (2, 4)
+        else:
+            raise NotImplementedError
+
+
+        # Entry flow
+        self.conv1 = nn.Conv2d(3, 32, 3, stride=2, padding=1, bias=False)
+        self.bn1 = BatchNorm(32)
+        self.relu = nn.ReLU(inplace=True)
+
+        self.conv2 = nn.Conv2d(32, 64, 3, stride=1, padding=1, bias=False)
+        self.bn2 = BatchNorm(64)
+
+        self.block1 = Block(64, 128, reps=2, stride=2, BatchNorm=BatchNorm, start_with_relu=False)
+        self.block2 = Block(128, 256, reps=2, stride=2, BatchNorm=BatchNorm, start_with_relu=False,
+                            grow_first=True)
+        self.block3 = Block(256, 728, reps=2, stride=entry_block3_stride, BatchNorm=BatchNorm,
+                            start_with_relu=True, grow_first=True, is_last=True)
+
+        # Middle flow
+        self.block4  = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
+                             BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
+        self.block5  = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
+                             BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
+        self.block6  = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
+                             BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
+        self.block7  = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
+                             BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
+        self.block8  = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
+                             BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
+        self.block9  = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
+                             BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
+        self.block10 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
+                             BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
+        self.block11 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
+                             BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
+        self.block12 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
+                             BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
+        self.block13 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
+                             BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
+        self.block14 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
+                             BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
+        self.block15 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
+                             BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
+        self.block16 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
+                             BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
+        self.block17 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
+                             BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
+        self.block18 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
+                             BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
+        self.block19 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
+                             BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
+
+        # Exit flow
+        self.block20 = Block(728, 1024, reps=2, stride=1, dilation=exit_block_dilations[0],
+                             BatchNorm=BatchNorm, start_with_relu=True, grow_first=False, is_last=True)
+
+        self.conv3 = SeparableConv2d(1024, 1536, 3, stride=1, dilation=exit_block_dilations[1], BatchNorm=BatchNorm)
+        self.bn3 = BatchNorm(1536)
+
+        self.conv4 = SeparableConv2d(1536, 1536, 3, stride=1, dilation=exit_block_dilations[1], BatchNorm=BatchNorm)
+        self.bn4 = BatchNorm(1536)
+
+        self.conv5 = SeparableConv2d(1536, 2048, 3, stride=1, dilation=exit_block_dilations[1], BatchNorm=BatchNorm)
+        self.bn5 = BatchNorm(2048)
+
+        # Init weights
+        self._init_weight()
+
+        # Load pretrained model
+        if pretrained:
+            self._load_pretrained_model()
+
+    def forward(self, x):
+        # Entry flow
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+
+        x = self.conv2(x)
+        x = self.bn2(x)
+        x = self.relu(x)
+
+        x = self.block1(x)
+        # add relu here
+        x = self.relu(x)
+        low_level_feat = x
+        x = self.block2(x)
+        x = self.block3(x)
+
+        # Middle flow
+        x = self.block4(x)
+        x = self.block5(x)
+        x = self.block6(x)
+        x = self.block7(x)
+        x = self.block8(x)
+        x = self.block9(x)
+        x = self.block10(x)
+        x = self.block11(x)
+        x = self.block12(x)
+        x = self.block13(x)
+        x = self.block14(x)
+        x = self.block15(x)
+        x = self.block16(x)
+        x = self.block17(x)
+        x = self.block18(x)
+        x = self.block19(x)
+
+        # Exit flow
+        x = self.block20(x)
+        x = self.relu(x)
+        x = self.conv3(x)
+        x = self.bn3(x)
+        x = self.relu(x)
+
+        x = self.conv4(x)
+        x = self.bn4(x)
+        x = self.relu(x)
+
+        x = self.conv5(x)
+        x = self.bn5(x)
+        x = self.relu(x)
+
+        return x, low_level_feat
+
+    def _init_weight(self):
+        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, SynchronizedBatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+
+
+    def _load_pretrained_model(self):
+        pretrain_dict = model_zoo.load_url('http://data.lip6.fr/cadene/pretrainedmodels/xception-b5690688.pth')
+        model_dict = {}
+        state_dict = self.state_dict()
+
+        for k, v in pretrain_dict.items():
+            if k in state_dict:
+                if 'pointwise' in k:
+                    v = v.unsqueeze(-1).unsqueeze(-1)
+                if k.startswith('block11'):
+                    model_dict[k] = v
+                    model_dict[k.replace('block11', 'block12')] = v
+                    model_dict[k.replace('block11', 'block13')] = v
+                    model_dict[k.replace('block11', 'block14')] = v
+                    model_dict[k.replace('block11', 'block15')] = v
+                    model_dict[k.replace('block11', 'block16')] = v
+                    model_dict[k.replace('block11', 'block17')] = v
+                    model_dict[k.replace('block11', 'block18')] = v
+                    model_dict[k.replace('block11', 'block19')] = v
+                elif k.startswith('block12'):
+                    model_dict[k.replace('block12', 'block20')] = v
+                elif k.startswith('bn3'):
+                    model_dict[k] = v
+                    model_dict[k.replace('bn3', 'bn4')] = v
+                elif k.startswith('conv4'):
+                    model_dict[k.replace('conv4', 'conv5')] = v
+                elif k.startswith('bn4'):
+                    model_dict[k.replace('bn4', 'bn5')] = v
+                else:
+                    model_dict[k] = v
+        state_dict.update(model_dict)
+        self.load_state_dict(state_dict)
+
+
+
+if __name__ == "__main__":
+    import torch
+    model = AlignedXception(BatchNorm=nn.BatchNorm2d, pretrained=True, output_stride=16)
+    input = torch.rand(1, 3, 512, 512)
+    output, low_level_feat = model(input)
+    print(output.size())
+    print(low_level_feat.size())

data/MBD/model/deep_lab_model/decoder.py

@@ -0,0 +1,59 @@
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from model.deep_lab_model.sync_batchnorm.batchnorm import SynchronizedBatchNorm2d
+
+class Decoder(nn.Module):
+    def __init__(self, num_classes, backbone, BatchNorm):
+        super(Decoder, self).__init__()
+        if backbone == 'resnet' or backbone == 'drn':
+            low_level_inplanes = 256
+        elif backbone == 'xception':
+            low_level_inplanes = 128
+        elif backbone == 'mobilenet':
+            low_level_inplanes = 24
+        else:
+            raise NotImplementedError
+
+        self.conv1 = nn.Conv2d(low_level_inplanes, 48, 1, bias=False)
+        self.bn1 = BatchNorm(48)
+        self.relu = nn.ReLU()
+        self.last_conv = nn.Sequential(nn.Conv2d(304, 256, kernel_size=3, stride=1, padding=1, bias=False),
+                                       BatchNorm(256),
+                                       nn.ReLU(),
+                                       nn.Dropout(0.5),
+                                       nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1, bias=False),
+                                       BatchNorm(256),
+                                       nn.ReLU(),
+                                       nn.Dropout(0.1),
+                                       nn.Conv2d(256, num_classes, kernel_size=1, stride=1),
+                                       nn.Sigmoid()
+                                       )
+        self._init_weight()
+
+
+    def forward(self, x, low_level_feat):
+        low_level_feat = self.conv1(low_level_feat)
+        low_level_feat = self.bn1(low_level_feat)
+        low_level_feat = self.relu(low_level_feat)
+
+        x = F.interpolate(x, size=low_level_feat.size()[2:], mode='bilinear', align_corners=True)
+        x = torch.cat((x, low_level_feat), dim=1)
+        x = self.last_conv(x)
+
+        return x
+
+    def _init_weight(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                torch.nn.init.kaiming_normal_(m.weight)
+            elif isinstance(m, SynchronizedBatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+
+def build_decoder(num_classes, backbone, BatchNorm):
+    return Decoder(num_classes, backbone, BatchNorm)

data/MBD/model/deep_lab_model/deeplab.py

@@ -0,0 +1,81 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from model.deep_lab_model.sync_batchnorm.batchnorm import SynchronizedBatchNorm2d
+from model.deep_lab_model.aspp import build_aspp
+from model.deep_lab_model.decoder import build_decoder
+from model.deep_lab_model.backbone import build_backbone
+
+class DeepLab(nn.Module):
+    def __init__(self, backbone='resnet', output_stride=16, num_classes=21,
+                 sync_bn=True, freeze_bn=False):
+        super(DeepLab, self).__init__()
+        if backbone == 'drn':
+            output_stride = 8
+
+        if sync_bn == True:
+            BatchNorm = SynchronizedBatchNorm2d
+        else:
+            BatchNorm = nn.BatchNorm2d
+
+        self.backbone = build_backbone(backbone, output_stride, BatchNorm)
+        self.aspp = build_aspp(backbone, output_stride, BatchNorm)
+        self.decoder = build_decoder(num_classes, backbone, BatchNorm)
+
+        self.freeze_bn = freeze_bn
+
+    def forward(self, input):
+        x, low_level_feat = self.backbone(input)
+        x = self.aspp(x)
+        x = self.decoder(x, low_level_feat)
+        x = F.interpolate(x, size=input.size()[2:], mode='bilinear', align_corners=True)
+
+        return x
+
+    def freeze_bn(self):
+        for m in self.modules():
+            if isinstance(m, SynchronizedBatchNorm2d):
+                m.eval()
+            elif isinstance(m, nn.BatchNorm2d):
+                m.eval()
+
+    def get_1x_lr_params(self):
+        modules = [self.backbone]
+        for i in range(len(modules)):
+            for m in modules[i].named_modules():
+                if self.freeze_bn:
+                    if isinstance(m[1], nn.Conv2d):
+                        for p in m[1].parameters():
+                            if p.requires_grad:
+                                yield p
+                else:
+                    if isinstance(m[1], nn.Conv2d) or isinstance(m[1], SynchronizedBatchNorm2d) \
+                            or isinstance(m[1], nn.BatchNorm2d):
+                        for p in m[1].parameters():
+                            if p.requires_grad:
+                                yield p
+
+    def get_10x_lr_params(self):
+        modules = [self.aspp, self.decoder]
+        for i in range(len(modules)):
+            for m in modules[i].named_modules():
+                if self.freeze_bn:
+                    if isinstance(m[1], nn.Conv2d):
+                        for p in m[1].parameters():
+                            if p.requires_grad:
+                                yield p
+                else:
+                    if isinstance(m[1], nn.Conv2d) or isinstance(m[1], SynchronizedBatchNorm2d) \
+                            or isinstance(m[1], nn.BatchNorm2d):
+                        for p in m[1].parameters():
+                            if p.requires_grad:
+                                yield p
+
+if __name__ == "__main__":
+    model = DeepLab(backbone='mobilenet', output_stride=16)
+    model.eval()
+    input = torch.rand(1, 3, 513, 513)
+    output = model(input)
+    print(output.size())
+
+

data/MBD/model/deep_lab_model/sync_batchnorm/__init__.py

@@ -0,0 +1,12 @@
+# -*- coding: utf-8 -*-
+# File   : __init__.py
+# Author : Jiayuan Mao
+# Email  : maojiayuan@gmail.com
+# Date   : 27/01/2018
+#
+# This file is part of Synchronized-BatchNorm-PyTorch.
+# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch
+# Distributed under MIT License.
+
+from .batchnorm import SynchronizedBatchNorm1d, SynchronizedBatchNorm2d, SynchronizedBatchNorm3d
+from .replicate import DataParallelWithCallback, patch_replication_callback

data/MBD/model/deep_lab_model/sync_batchnorm/batchnorm.py

@@ -0,0 +1,282 @@
+# -*- coding: utf-8 -*-
+# File   : batchnorm.py
+# Author : Jiayuan Mao
+# Email  : maojiayuan@gmail.com
+# Date   : 27/01/2018
+#
+# This file is part of Synchronized-BatchNorm-PyTorch.
+# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch
+# Distributed under MIT License.
+
+import collections
+
+import torch
+import torch.nn.functional as F
+
+from torch.nn.modules.batchnorm import _BatchNorm
+from torch.nn.parallel._functions import ReduceAddCoalesced, Broadcast
+
+from .comm import SyncMaster
+
+__all__ = ['SynchronizedBatchNorm1d', 'SynchronizedBatchNorm2d', 'SynchronizedBatchNorm3d']
+
+
+def _sum_ft(tensor):
+    """sum over the first and last dimention"""
+    return tensor.sum(dim=0).sum(dim=-1)
+
+
+def _unsqueeze_ft(tensor):
+    """add new dementions at the front and the tail"""
+    return tensor.unsqueeze(0).unsqueeze(-1)
+
+
+_ChildMessage = collections.namedtuple('_ChildMessage', ['sum', 'ssum', 'sum_size'])
+_MasterMessage = collections.namedtuple('_MasterMessage', ['sum', 'inv_std'])
+
+
+class _SynchronizedBatchNorm(_BatchNorm):
+    def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True):
+        super(_SynchronizedBatchNorm, self).__init__(num_features, eps=eps, momentum=momentum, affine=affine)
+
+        self._sync_master = SyncMaster(self._data_parallel_master)
+
+        self._is_parallel = False
+        self._parallel_id = None
+        self._slave_pipe = None
+
+    def forward(self, input):
+        # If it is not parallel computation or is in evaluation mode, use PyTorch's implementation.
+        if not (self._is_parallel and self.training):
+            return F.batch_norm(
+                input, self.running_mean, self.running_var, self.weight, self.bias,
+                self.training, self.momentum, self.eps)
+
+        # Resize the input to (B, C, -1).
+        input_shape = input.size()
+        input = input.view(input.size(0), self.num_features, -1)
+
+        # Compute the sum and square-sum.
+        sum_size = input.size(0) * input.size(2)
+        input_sum = _sum_ft(input)
+        input_ssum = _sum_ft(input ** 2)
+
+        # Reduce-and-broadcast the statistics.
+        if self._parallel_id == 0:
+            mean, inv_std = self._sync_master.run_master(_ChildMessage(input_sum, input_ssum, sum_size))
+        else:
+            mean, inv_std = self._slave_pipe.run_slave(_ChildMessage(input_sum, input_ssum, sum_size))
+
+        # Compute the output.
+        if self.affine:
+            # MJY:: Fuse the multiplication for speed.
+            output = (input - _unsqueeze_ft(mean)) * _unsqueeze_ft(inv_std * self.weight) + _unsqueeze_ft(self.bias)
+        else:
+            output = (input - _unsqueeze_ft(mean)) * _unsqueeze_ft(inv_std)
+
+        # Reshape it.
+        return output.view(input_shape)
+
+    def __data_parallel_replicate__(self, ctx, copy_id):
+        self._is_parallel = True
+        self._parallel_id = copy_id
+
+        # parallel_id == 0 means master device.
+        if self._parallel_id == 0:
+            ctx.sync_master = self._sync_master
+        else:
+            self._slave_pipe = ctx.sync_master.register_slave(copy_id)
+
+    def _data_parallel_master(self, intermediates):
+        """Reduce the sum and square-sum, compute the statistics, and broadcast it."""
+
+        # Always using same "device order" makes the ReduceAdd operation faster.
+        # Thanks to:: Tete Xiao (http://tetexiao.com/)
+        intermediates = sorted(intermediates, key=lambda i: i[1].sum.get_device())
+
+        to_reduce = [i[1][:2] for i in intermediates]
+        to_reduce = [j for i in to_reduce for j in i]  # flatten
+        target_gpus = [i[1].sum.get_device() for i in intermediates]
+
+        sum_size = sum([i[1].sum_size for i in intermediates])
+        sum_, ssum = ReduceAddCoalesced.apply(target_gpus[0], 2, *to_reduce)
+        mean, inv_std = self._compute_mean_std(sum_, ssum, sum_size)
+
+        broadcasted = Broadcast.apply(target_gpus, mean, inv_std)
+
+        outputs = []
+        for i, rec in enumerate(intermediates):
+            outputs.append((rec[0], _MasterMessage(*broadcasted[i * 2:i * 2 + 2])))
+
+        return outputs
+
+    def _compute_mean_std(self, sum_, ssum, size):
+        """Compute the mean and standard-deviation with sum and square-sum. This method
+        also maintains the moving average on the master device."""
+        assert size > 1, 'BatchNorm computes unbiased standard-deviation, which requires size > 1.'
+        mean = sum_ / size
+        sumvar = ssum - sum_ * mean
+        unbias_var = sumvar / (size - 1)
+        bias_var = sumvar / size
+
+        self.running_mean = (1 - self.momentum) * self.running_mean + self.momentum * mean.data
+        self.running_var = (1 - self.momentum) * self.running_var + self.momentum * unbias_var.data
+
+        return mean, bias_var.clamp(self.eps) ** -0.5
+
+
+class SynchronizedBatchNorm1d(_SynchronizedBatchNorm):
+    r"""Applies Synchronized Batch Normalization over a 2d or 3d input that is seen as a
+    mini-batch.
+    .. math::
+        y = \frac{x - mean[x]}{ \sqrt{Var[x] + \epsilon}} * gamma + beta
+    This module differs from the built-in PyTorch BatchNorm1d as the mean and
+    standard-deviation are reduced across all devices during training.
+    For example, when one uses `nn.DataParallel` to wrap the network during
+    training, PyTorch's implementation normalize the tensor on each device using
+    the statistics only on that device, which accelerated the computation and
+    is also easy to implement, but the statistics might be inaccurate.
+    Instead, in this synchronized version, the statistics will be computed
+    over all training samples distributed on multiple devices.
+
+    Note that, for one-GPU or CPU-only case, this module behaves exactly same
+    as the built-in PyTorch implementation.
+    The mean and standard-deviation are calculated per-dimension over
+    the mini-batches and gamma and beta are learnable parameter vectors
+    of size C (where C is the input size).
+    During training, this layer keeps a running estimate of its computed mean
+    and variance. The running sum is kept with a default momentum of 0.1.
+    During evaluation, this running mean/variance is used for normalization.
+    Because the BatchNorm is done over the `C` dimension, computing statistics
+    on `(N, L)` slices, it's common terminology to call this Temporal BatchNorm
+    Args:
+        num_features: num_features from an expected input of size
+            `batch_size x num_features [x width]`
+        eps: a value added to the denominator for numerical stability.
+            Default: 1e-5
+        momentum: the value used for the running_mean and running_var
+            computation. Default: 0.1
+        affine: a boolean value that when set to ``True``, gives the layer learnable
+            affine parameters. Default: ``True``
+    Shape:
+        - Input: :math:`(N, C)` or :math:`(N, C, L)`
+        - Output: :math:`(N, C)` or :math:`(N, C, L)` (same shape as input)
+    Examples:
+        >>> # With Learnable Parameters
+        >>> m = SynchronizedBatchNorm1d(100)
+        >>> # Without Learnable Parameters
+        >>> m = SynchronizedBatchNorm1d(100, affine=False)
+        >>> input = torch.autograd.Variable(torch.randn(20, 100))
+        >>> output = m(input)
+    """
+
+    def _check_input_dim(self, input):
+        if input.dim() != 2 and input.dim() != 3:
+            raise ValueError('expected 2D or 3D input (got {}D input)'
+                             .format(input.dim()))
+        super(SynchronizedBatchNorm1d, self)._check_input_dim(input)
+
+
+class SynchronizedBatchNorm2d(_SynchronizedBatchNorm):
+    r"""Applies Batch Normalization over a 4d input that is seen as a mini-batch
+    of 3d inputs
+    .. math::
+        y = \frac{x - mean[x]}{ \sqrt{Var[x] + \epsilon}} * gamma + beta
+    This module differs from the built-in PyTorch BatchNorm2d as the mean and
+    standard-deviation are reduced across all devices during training.
+    For example, when one uses `nn.DataParallel` to wrap the network during
+    training, PyTorch's implementation normalize the tensor on each device using
+    the statistics only on that device, which accelerated the computation and
+    is also easy to implement, but the statistics might be inaccurate.
+    Instead, in this synchronized version, the statistics will be computed
+    over all training samples distributed on multiple devices.
+
+    Note that, for one-GPU or CPU-only case, this module behaves exactly same
+    as the built-in PyTorch implementation.
+    The mean and standard-deviation are calculated per-dimension over
+    the mini-batches and gamma and beta are learnable parameter vectors
+    of size C (where C is the input size).
+    During training, this layer keeps a running estimate of its computed mean
+    and variance. The running sum is kept with a default momentum of 0.1.
+    During evaluation, this running mean/variance is used for normalization.
+    Because the BatchNorm is done over the `C` dimension, computing statistics
+    on `(N, H, W)` slices, it's common terminology to call this Spatial BatchNorm
+    Args:
+        num_features: num_features from an expected input of
+            size batch_size x num_features x height x width
+        eps: a value added to the denominator for numerical stability.
+            Default: 1e-5
+        momentum: the value used for the running_mean and running_var
+            computation. Default: 0.1
+        affine: a boolean value that when set to ``True``, gives the layer learnable
+            affine parameters. Default: ``True``
+    Shape:
+        - Input: :math:`(N, C, H, W)`
+        - Output: :math:`(N, C, H, W)` (same shape as input)
+    Examples:
+        >>> # With Learnable Parameters
+        >>> m = SynchronizedBatchNorm2d(100)
+        >>> # Without Learnable Parameters
+        >>> m = SynchronizedBatchNorm2d(100, affine=False)
+        >>> input = torch.autograd.Variable(torch.randn(20, 100, 35, 45))
+        >>> output = m(input)
+    """
+
+    def _check_input_dim(self, input):
+        if input.dim() != 4:
+            raise ValueError('expected 4D input (got {}D input)'
+                             .format(input.dim()))
+        super(SynchronizedBatchNorm2d, self)._check_input_dim(input)
+
+
+class SynchronizedBatchNorm3d(_SynchronizedBatchNorm):
+    r"""Applies Batch Normalization over a 5d input that is seen as a mini-batch
+    of 4d inputs
+    .. math::
+        y = \frac{x - mean[x]}{ \sqrt{Var[x] + \epsilon}} * gamma + beta
+    This module differs from the built-in PyTorch BatchNorm3d as the mean and
+    standard-deviation are reduced across all devices during training.
+    For example, when one uses `nn.DataParallel` to wrap the network during
+    training, PyTorch's implementation normalize the tensor on each device using
+    the statistics only on that device, which accelerated the computation and
+    is also easy to implement, but the statistics might be inaccurate.
+    Instead, in this synchronized version, the statistics will be computed
+    over all training samples distributed on multiple devices.
+
+    Note that, for one-GPU or CPU-only case, this module behaves exactly same
+    as the built-in PyTorch implementation.
+    The mean and standard-deviation are calculated per-dimension over
+    the mini-batches and gamma and beta are learnable parameter vectors
+    of size C (where C is the input size).
+    During training, this layer keeps a running estimate of its computed mean
+    and variance. The running sum is kept with a default momentum of 0.1.
+    During evaluation, this running mean/variance is used for normalization.
+    Because the BatchNorm is done over the `C` dimension, computing statistics
+    on `(N, D, H, W)` slices, it's common terminology to call this Volumetric BatchNorm
+    or Spatio-temporal BatchNorm
+    Args:
+        num_features: num_features from an expected input of
+            size batch_size x num_features x depth x height x width
+        eps: a value added to the denominator for numerical stability.
+            Default: 1e-5
+        momentum: the value used for the running_mean and running_var
+            computation. Default: 0.1
+        affine: a boolean value that when set to ``True``, gives the layer learnable
+            affine parameters. Default: ``True``
+    Shape:
+        - Input: :math:`(N, C, D, H, W)`
+        - Output: :math:`(N, C, D, H, W)` (same shape as input)
+    Examples:
+        >>> # With Learnable Parameters
+        >>> m = SynchronizedBatchNorm3d(100)
+        >>> # Without Learnable Parameters
+        >>> m = SynchronizedBatchNorm3d(100, affine=False)
+        >>> input = torch.autograd.Variable(torch.randn(20, 100, 35, 45, 10))
+        >>> output = m(input)
+    """
+
+    def _check_input_dim(self, input):
+        if input.dim() != 5:
+            raise ValueError('expected 5D input (got {}D input)'
+                             .format(input.dim()))
+        super(SynchronizedBatchNorm3d, self)._check_input_dim(input)

data/MBD/model/deep_lab_model/sync_batchnorm/comm.py

@@ -0,0 +1,129 @@
+# -*- coding: utf-8 -*-
+# File   : comm.py
+# Author : Jiayuan Mao
+# Email  : maojiayuan@gmail.com
+# Date   : 27/01/2018
+#
+# This file is part of Synchronized-BatchNorm-PyTorch.
+# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch
+# Distributed under MIT License.
+
+import queue
+import collections
+import threading
+
+__all__ = ['FutureResult', 'SlavePipe', 'SyncMaster']
+
+
+class FutureResult(object):
+    """A thread-safe future implementation. Used only as one-to-one pipe."""
+
+    def __init__(self):
+        self._result = None
+        self._lock = threading.Lock()
+        self._cond = threading.Condition(self._lock)
+
+    def put(self, result):
+        with self._lock:
+            assert self._result is None, 'Previous result has\'t been fetched.'
+            self._result = result
+            self._cond.notify()
+
+    def get(self):
+        with self._lock:
+            if self._result is None:
+                self._cond.wait()
+
+            res = self._result
+            self._result = None
+            return res
+
+
+_MasterRegistry = collections.namedtuple('MasterRegistry', ['result'])
+_SlavePipeBase = collections.namedtuple('_SlavePipeBase', ['identifier', 'queue', 'result'])
+
+
+class SlavePipe(_SlavePipeBase):
+    """Pipe for master-slave communication."""
+
+    def run_slave(self, msg):
+        self.queue.put((self.identifier, msg))
+        ret = self.result.get()
+        self.queue.put(True)
+        return ret
+
+
+class SyncMaster(object):
+    """An abstract `SyncMaster` object.
+    - During the replication, as the data parallel will trigger an callback of each module, all slave devices should
+    call `register(id)` and obtain an `SlavePipe` to communicate with the master.
+    - During the forward pass, master device invokes `run_master`, all messages from slave devices will be collected,
+    and passed to a registered callback.
+    - After receiving the messages, the master device should gather the information and determine to message passed
+    back to each slave devices.
+    """
+
+    def __init__(self, master_callback):
+        """
+        Args:
+            master_callback: a callback to be invoked after having collected messages from slave devices.
+        """
+        self._master_callback = master_callback
+        self._queue = queue.Queue()
+        self._registry = collections.OrderedDict()
+        self._activated = False
+
+    def __getstate__(self):
+        return {'master_callback': self._master_callback}
+
+    def __setstate__(self, state):
+        self.__init__(state['master_callback'])
+
+    def register_slave(self, identifier):
+        """
+        Register an slave device.
+        Args:
+            identifier: an identifier, usually is the device id.
+        Returns: a `SlavePipe` object which can be used to communicate with the master device.
+        """
+        if self._activated:
+            assert self._queue.empty(), 'Queue is not clean before next initialization.'
+            self._activated = False
+            self._registry.clear()
+        future = FutureResult()
+        self._registry[identifier] = _MasterRegistry(future)
+        return SlavePipe(identifier, self._queue, future)
+
+    def run_master(self, master_msg):
+        """
+        Main entry for the master device in each forward pass.
+        The messages were first collected from each devices (including the master device), and then
+        an callback will be invoked to compute the message to be sent back to each devices
+        (including the master device).
+        Args:
+            master_msg: the message that the master want to send to itself. This will be placed as the first
+            message when calling `master_callback`. For detailed usage, see `_SynchronizedBatchNorm` for an example.
+        Returns: the message to be sent back to the master device.
+        """
+        self._activated = True
+
+        intermediates = [(0, master_msg)]
+        for i in range(self.nr_slaves):
+            intermediates.append(self._queue.get())
+
+        results = self._master_callback(intermediates)
+        assert results[0][0] == 0, 'The first result should belongs to the master.'
+
+        for i, res in results:
+            if i == 0:
+                continue
+            self._registry[i].result.put(res)
+
+        for i in range(self.nr_slaves):
+            assert self._queue.get() is True
+
+        return results[0][1]
+
+    @property
+    def nr_slaves(self):
+        return len(self._registry)

data/MBD/model/deep_lab_model/sync_batchnorm/replicate.py

@@ -0,0 +1,88 @@
+# -*- coding: utf-8 -*-
+# File   : replicate.py
+# Author : Jiayuan Mao
+# Email  : maojiayuan@gmail.com
+# Date   : 27/01/2018
+#
+# This file is part of Synchronized-BatchNorm-PyTorch.
+# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch
+# Distributed under MIT License.
+
+import functools
+
+from torch.nn.parallel.data_parallel import DataParallel
+
+__all__ = [
+    'CallbackContext',
+    'execute_replication_callbacks',
+    'DataParallelWithCallback',
+    'patch_replication_callback'
+]
+
+
+class CallbackContext(object):
+    pass
+
+
+def execute_replication_callbacks(modules):
+    """
+    Execute an replication callback `__data_parallel_replicate__` on each module created by original replication.
+    The callback will be invoked with arguments `__data_parallel_replicate__(ctx, copy_id)`
+    Note that, as all modules are isomorphism, we assign each sub-module with a context
+    (shared among multiple copies of this module on different devices).
+    Through this context, different copies can share some information.
+    We guarantee that the callback on the master copy (the first copy) will be called ahead of calling the callback
+    of any slave copies.
+    """
+    master_copy = modules[0]
+    nr_modules = len(list(master_copy.modules()))
+    ctxs = [CallbackContext() for _ in range(nr_modules)]
+
+    for i, module in enumerate(modules):
+        for j, m in enumerate(module.modules()):
+            if hasattr(m, '__data_parallel_replicate__'):
+                m.__data_parallel_replicate__(ctxs[j], i)
+
+
+class DataParallelWithCallback(DataParallel):
+    """
+    Data Parallel with a replication callback.
+    An replication callback `__data_parallel_replicate__` of each module will be invoked after being created by
+    original `replicate` function.
+    The callback will be invoked with arguments `__data_parallel_replicate__(ctx, copy_id)`
+    Examples:
+        > sync_bn = SynchronizedBatchNorm1d(10, eps=1e-5, affine=False)
+        > sync_bn = DataParallelWithCallback(sync_bn, device_ids=[0, 1])
+        # sync_bn.__data_parallel_replicate__ will be invoked.
+    """
+
+    def replicate(self, module, device_ids):
+        modules = super(DataParallelWithCallback, self).replicate(module, device_ids)
+        execute_replication_callbacks(modules)
+        return modules
+
+
+def patch_replication_callback(data_parallel):
+    """
+    Monkey-patch an existing `DataParallel` object. Add the replication callback.
+    Useful when you have customized `DataParallel` implementation.
+    Examples:
+        > sync_bn = SynchronizedBatchNorm1d(10, eps=1e-5, affine=False)
+        > sync_bn = DataParallel(sync_bn, device_ids=[0, 1])
+        > patch_replication_callback(sync_bn)
+        # this is equivalent to
+        > sync_bn = SynchronizedBatchNorm1d(10, eps=1e-5, affine=False)
+        > sync_bn = DataParallelWithCallback(sync_bn, device_ids=[0, 1])
+    """
+
+    assert isinstance(data_parallel, DataParallel)
+
+    old_replicate = data_parallel.replicate
+
+    @functools.wraps(old_replicate)
+    def new_replicate(module, device_ids):
+        modules = old_replicate(module, device_ids)
+        execute_replication_callbacks(modules)
+        return modules
+
+    data_parallel.replicate = new_replicate

data/MBD/model/deep_lab_model/sync_batchnorm/unittest.py

@@ -0,0 +1,29 @@
+# -*- coding: utf-8 -*-
+# File   : unittest.py
+# Author : Jiayuan Mao
+# Email  : maojiayuan@gmail.com
+# Date   : 27/01/2018
+#
+# This file is part of Synchronized-BatchNorm-PyTorch.
+# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch
+# Distributed under MIT License.
+
+import unittest
+
+import numpy as np
+from torch.autograd import Variable
+
+
+def as_numpy(v):
+    if isinstance(v, Variable):
+        v = v.data
+    return v.cpu().numpy()
+
+
+class TorchTestCase(unittest.TestCase):
+    def assertTensorClose(self, a, b, atol=1e-3, rtol=1e-3):
+        npa, npb = as_numpy(a), as_numpy(b)
+        self.assertTrue(
+                np.allclose(npa, npb, atol=atol),
+                'Tensor close check failed\n{}\n{}\nadiff={}, rdiff={}'.format(a, b, np.abs(npa - npb).max(), np.abs((npa - npb) / np.fmax(npa, 1e-5)).max())
+        )

data/MBD/model/densenetccnl.py

@@ -0,0 +1,382 @@
+# Densenet decoder encoder with intermediate fully connected layers and dropout
+
+import torch
+import torch.backends.cudnn as cudnn
+import torch.nn as nn
+import torch.nn.functional as F
+import functools
+from torch.autograd import gradcheck
+from torch.autograd import Function
+from torch.autograd import Variable
+from torch.autograd import gradcheck
+from torch.autograd import Function
+import numpy as np
+
+
+def add_coordConv_channels(t):
+    n,c,h,w=t.size()
+    xx_channel=np.ones((h, w))
+    xx_range=np.array(range(h))
+    xx_range=np.expand_dims(xx_range,-1)
+    xx_coord=xx_channel*xx_range
+    yy_coord=xx_coord.transpose()
+
+    xx_coord=xx_coord/(h-1)
+    yy_coord=yy_coord/(h-1)
+    xx_coord=xx_coord*2 - 1
+    yy_coord=yy_coord*2 - 1
+    xx_coord=torch.from_numpy(xx_coord).float()
+    yy_coord=torch.from_numpy(yy_coord).float()
+
+    if t.is_cuda:
+    	xx_coord=xx_coord.cuda()
+    	yy_coord=yy_coord.cuda()
+
+    xx_coord=xx_coord.unsqueeze(0).unsqueeze(0).repeat(n,1,1,1)
+    yy_coord=yy_coord.unsqueeze(0).unsqueeze(0).repeat(n,1,1,1)
+
+    t_cc=torch.cat((t,xx_coord,yy_coord),dim=1)
+
+    return t_cc
+
+
+
+class DenseBlockEncoder(nn.Module):
+    def __init__(self, n_channels, n_convs, activation=nn.ReLU, args=[False]):
+        super(DenseBlockEncoder, self).__init__()
+        assert(n_convs > 0)
+
+        self.n_channels = n_channels
+        self.n_convs    = n_convs
+        self.layers     = nn.ModuleList()
+        for i in range(n_convs):
+            self.layers.append(nn.Sequential(
+                    nn.BatchNorm2d(n_channels),
+                    activation(*args),
+                    nn.Conv2d(n_channels, n_channels, 3, stride=1, padding=1, bias=False),))
+
+    def forward(self, inputs):
+        outputs = []
+
+        for i, layer in enumerate(self.layers):
+            if i > 0:
+                next_output = 0
+                for no in outputs:
+                    next_output = next_output + no
+                outputs.append(next_output)
+            else:
+                outputs.append(layer(inputs))
+        return outputs[-1]
+
+# Dense block in encoder.
+class DenseBlockDecoder(nn.Module):
+    def __init__(self, n_channels, n_convs, activation=nn.ReLU, args=[False]):
+        super(DenseBlockDecoder, self).__init__()
+        assert(n_convs > 0)
+
+        self.n_channels = n_channels
+        self.n_convs    = n_convs
+        self.layers = nn.ModuleList()
+        for i in range(n_convs):
+            self.layers.append(nn.Sequential(
+                    nn.BatchNorm2d(n_channels),
+                    activation(*args),
+                    nn.ConvTranspose2d(n_channels, n_channels, 3, stride=1, padding=1, bias=False),))
+
+    def forward(self, inputs):
+        outputs = []
+
+        for i, layer in enumerate(self.layers):
+            if i > 0:
+                next_output = 0
+                for no in outputs:
+                    next_output = next_output + no
+                outputs.append(next_output)
+            else:
+                outputs.append(layer(inputs))
+        return outputs[-1]
+
+class DenseTransitionBlockEncoder(nn.Module):
+    def __init__(self, n_channels_in, n_channels_out, mp, activation=nn.ReLU, args=[False]):
+        super(DenseTransitionBlockEncoder, self).__init__()
+        self.n_channels_in  = n_channels_in
+        self.n_channels_out = n_channels_out
+        self.mp             = mp
+        self.main           = nn.Sequential(
+                nn.BatchNorm2d(n_channels_in),
+                activation(*args),
+                nn.Conv2d(n_channels_in, n_channels_out, 1, stride=1, padding=0, bias=False),
+                nn.MaxPool2d(mp),
+        )
+    def forward(self, inputs):
+        # print(inputs.shape,'222222222222222',self.main(inputs).shape)
+        return self.main(inputs)
+
+
+class DenseTransitionBlockDecoder(nn.Module):
+    def __init__(self, n_channels_in, n_channels_out, activation=nn.ReLU, args=[False]):
+        super(DenseTransitionBlockDecoder, self).__init__()
+        self.n_channels_in  = n_channels_in
+        self.n_channels_out = n_channels_out
+        self.main           = nn.Sequential(
+                nn.BatchNorm2d(n_channels_in),
+                activation(*args),
+                nn.ConvTranspose2d(n_channels_in, n_channels_out, 4, stride=2, padding=1, bias=False),
+        )
+    def forward(self, inputs):
+        # print(inputs.shape,'333333333333',self.main(inputs).shape)
+        return self.main(inputs)
+
+## Dense encoders and decoders for image of size 128 128
+class waspDenseEncoder128(nn.Module):
+    def __init__(self, nc=1, ndf = 32, ndim = 128, activation=nn.LeakyReLU, args=[0.2, False], f_activation=nn.Tanh, f_args=[]):
+        super(waspDenseEncoder128, self).__init__()
+        self.ndim = ndim
+
+        self.main = nn.Sequential(
+                # input is (nc) x 128 x 128
+                nn.BatchNorm2d(nc),
+                nn.ReLU(True),
+                nn.Conv2d(nc, ndf, 4, stride=2, padding=1),
+
+                # state size. (ndf) x 64 x 64
+                DenseBlockEncoder(ndf, 6),
+                DenseTransitionBlockEncoder(ndf, ndf*2, 2, activation=activation, args=args),
+
+                # state size. (ndf*2) x 32 x 32
+                DenseBlockEncoder(ndf*2, 12),
+                DenseTransitionBlockEncoder(ndf*2, ndf*4, 2, activation=activation, args=args),
+
+                # state size. (ndf*4) x 16 x 16
+                DenseBlockEncoder(ndf*4, 16),
+                DenseTransitionBlockEncoder(ndf*4, ndf*8, 2, activation=activation, args=args),
+
+                # state size. (ndf*4) x 8 x 8
+                DenseBlockEncoder(ndf*8, 16),
+                DenseTransitionBlockEncoder(ndf*8, ndf*8, 2, activation=activation, args=args),
+
+                # state size. (ndf*8) x 4 x 4
+                DenseBlockEncoder(ndf*8, 16),
+                DenseTransitionBlockEncoder(ndf*8, ndim, 4, activation=activation, args=args),
+                f_activation(*f_args),
+        )
+
+    def forward(self, input):
+        input=add_coordConv_channels(input)
+        output = self.main(input).view(-1,self.ndim)
+        #print(output.size())
+        return output
+
+class waspDenseDecoder128(nn.Module):
+    def __init__(self, nz=128, nc=1, ngf=32, lb=0, ub=1, activation=nn.ReLU, args=[False], f_activation=nn.Hardtanh, f_args=[]):
+        super(waspDenseDecoder128, self).__init__()
+        self.main   = nn.Sequential(
+            # input is Z, going into convolution
+            nn.BatchNorm2d(nz),
+            activation(*args),
+            nn.ConvTranspose2d(nz, ngf * 8, 4, 1, 0, bias=False),
+
+            # state size. (ngf*8) x 4 x 4
+            DenseBlockDecoder(ngf*8, 16),
+            DenseTransitionBlockDecoder(ngf*8, ngf*8),
+
+            # state size. (ngf*4) x 8 x 8
+            DenseBlockDecoder(ngf*8, 16),
+            DenseTransitionBlockDecoder(ngf*8, ngf*4),
+
+            # state size. (ngf*2) x 16 x 16
+            DenseBlockDecoder(ngf*4, 12),
+            DenseTransitionBlockDecoder(ngf*4, ngf*2),
+
+            # state size. (ngf) x 32 x 32
+            DenseBlockDecoder(ngf*2, 6),
+            DenseTransitionBlockDecoder(ngf*2, ngf),
+
+            # state size. (ngf) x 64 x 64
+            DenseBlockDecoder(ngf, 6),
+            DenseTransitionBlockDecoder(ngf, ngf),
+
+            # state size (ngf) x 128 x 128
+            nn.BatchNorm2d(ngf),
+            activation(*args),
+            nn.ConvTranspose2d(ngf, nc, 3, stride=1, padding=1, bias=False),
+            f_activation(*f_args),
+        )
+        # self.smooth=nn.Sequential(
+        #     nn.Conv2d(nc, nc, 1, stride=1, padding=0, bias=False),
+        #     f_activation(*f_args),
+        # )
+    def forward(self, inputs):
+        # return self.smooth(self.main(inputs))
+        return self.main(inputs)
+
+
+
+## Dense encoders and decoders for image of size 512 512
+class waspDenseEncoder512(nn.Module):
+    def __init__(self, nc=1, ndf = 32, ndim = 128, activation=nn.LeakyReLU, args=[0.2, False], f_activation=nn.Tanh, f_args=[]):
+        super(waspDenseEncoder512, self).__init__()
+        self.ndim = ndim
+
+        self.main = nn.Sequential(
+                # input is (nc) x 128 x 128  > *4
+                nn.BatchNorm2d(nc),
+                nn.ReLU(True),
+                nn.Conv2d(nc, ndf, 4, stride=2, padding=1),
+
+                # state size. (ndf) x 64 x 64  > *4
+                DenseBlockEncoder(ndf, 6),
+                DenseTransitionBlockEncoder(ndf, ndf*2, 2, activation=activation, args=args),
+
+                # state size. (ndf*2) x 32 x 32 > *4
+                DenseBlockEncoder(ndf*2, 12),
+                DenseTransitionBlockEncoder(ndf*2, ndf*4, 2, activation=activation, args=args),
+
+                # state size. (ndf*4) x 16 x 16  > *4
+                DenseBlockEncoder(ndf*4, 16),
+                DenseTransitionBlockEncoder(ndf*4, ndf*8, 2, activation=activation, args=args),
+
+                # state size. (ndf*8) x 8 x 8   *4
+                DenseBlockEncoder(ndf*8, 16),
+                DenseTransitionBlockEncoder(ndf*8, ndf*8, 2, activation=activation, args=args),
+
+                # state size. (ndf*8) x 4 x 4 > *4
+                DenseBlockEncoder(ndf*8, 16),
+                DenseTransitionBlockEncoder(ndf*8, ndf*8, 4, activation=activation, args=args),
+                f_activation(*f_args),
+
+                # state size. (ndf*8) x 2 x 2 > *4
+                DenseBlockEncoder(ndf*8, 16),
+                DenseTransitionBlockEncoder(ndf*8, ndim, 4, activation=activation, args=args),
+                f_activation(*f_args),
+        )
+
+    def forward(self, input):
+        input=add_coordConv_channels(input)
+        output = self.main(input).view(-1,self.ndim)
+        # output = self.main(input).view(8,-1)
+        # print(input.shape,'---------------------')
+        #print(output.size())
+        return output
+
+class waspDenseDecoder512(nn.Module):
+    def __init__(self, nz=128, nc=1, ngf=32, lb=0, ub=1, activation=nn.ReLU, args=[False], f_activation=nn.Tanh, f_args=[]):
+        super(waspDenseDecoder512, self).__init__()
+        self.main   = nn.Sequential(
+            # input is Z, going into convolution
+            nn.BatchNorm2d(nz),
+            activation(*args),
+            nn.ConvTranspose2d(nz, ngf * 8, 4, 1, 0, bias=False),
+
+            # state size. (ngf*8) x 4 x 4
+            DenseBlockDecoder(ngf*8, 16),
+            DenseTransitionBlockDecoder(ngf*8, ngf*8),
+
+            # state size. (ngf*8) x 8 x 8
+            DenseBlockDecoder(ngf*8, 16),
+            DenseTransitionBlockDecoder(ngf*8, ngf*8),
+
+            # state size. (ngf*4) x 16 x 16
+            DenseBlockDecoder(ngf*8, 16),
+            DenseTransitionBlockDecoder(ngf*8, ngf*4),
+
+            # state size. (ngf*2) x 32 x 32
+            DenseBlockDecoder(ngf*4, 12),
+            DenseTransitionBlockDecoder(ngf*4, ngf*2),
+
+            # state size. (ngf) x 64 x 64
+            DenseBlockDecoder(ngf*2, 6),
+            DenseTransitionBlockDecoder(ngf*2, ngf),
+
+            # state size. (ngf) x 128 x 128
+            DenseBlockDecoder(ngf, 6),
+            DenseTransitionBlockDecoder(ngf, ngf),
+
+            # state size. (ngf) x 256 x 256
+            DenseBlockDecoder(ngf, 6),
+            DenseTransitionBlockDecoder(ngf, ngf),
+
+            # state size (ngf) x 512 x 512
+            nn.BatchNorm2d(ngf),
+            activation(*args),
+            nn.ConvTranspose2d(ngf, nc, 3, stride=1, padding=1, bias=False),
+            f_activation(*f_args),
+        )
+        # self.smooth=nn.Sequential(
+        #     nn.Conv2d(nc, nc, 1, stride=1, padding=0, bias=False),
+        #     f_activation(*f_args),
+        # )
+    def forward(self, inputs):
+        # return self.smooth(self.main(inputs))
+        return self.main(inputs)
+
+
+class dnetccnl(nn.Module):
+    #in_channels -> nc      | encoder first layer
+    #filters -> ndf    | encoder first layer
+    #img_size(h,w) -> ndim
+    #out_channels  -> optical flow (x,y)
+
+    def __init__(self, img_size=448, in_channels=3, out_channels=2, filters=32,fc_units=100):
+        super(dnetccnl, self).__init__()
+        self.nc=in_channels
+        self.nf=filters
+        self.ndim=img_size
+        self.oc=out_channels
+        self.fcu=fc_units
+
+        self.encoder=waspDenseEncoder128(nc=self.nc+2,ndf=self.nf,ndim=self.ndim)
+        self.decoder=waspDenseDecoder128(nz=self.ndim,nc=self.oc,ngf=self.nf)
+        # self.fc_layers= nn.Sequential(nn.Linear(self.ndim, self.fcu),
+        #                               nn.ReLU(True),
+        #                               nn.Dropout(0.25),
+        #                               nn.Linear(self.fcu,self.ndim),
+        #                               nn.ReLU(True),
+        #                               nn.Dropout(0.25),
+        #                               )
+
+    def forward(self, inputs):
+
+        encoded=self.encoder(inputs)
+        encoded=encoded.unsqueeze(-1).unsqueeze(-1)
+        decoded=self.decoder(encoded)
+        # print torch.max(decoded)
+        # print torch.min(decoded)
+        # print(decoded.shape,'11111111111111111',encoded.shape)
+
+        return decoded
+
+class dnetccnl512(nn.Module):
+    #in_channels -> nc      | encoder first layer
+    #filters -> ndf    | encoder first layer
+    #img_size(h,w) -> ndim
+    #out_channels  -> optical flow (x,y)
+
+    def __init__(self, img_size=448, in_channels=3, out_channels=2, filters=32,fc_units=100):
+        super(dnetccnl512, self).__init__()
+        self.nc=in_channels
+        self.nf=filters
+        self.ndim=img_size
+        self.oc=out_channels
+        self.fcu=fc_units
+
+        self.encoder=waspDenseEncoder512(nc=self.nc+2,ndf=self.nf,ndim=self.ndim)
+        self.decoder=waspDenseDecoder512(nz=self.ndim,nc=self.oc,ngf=self.nf)
+        # self.fc_layers= nn.Sequential(nn.Linear(self.ndim, self.fcu),
+        #                               nn.ReLU(True),
+        #                               nn.Dropout(0.25),
+        #                               nn.Linear(self.fcu,self.ndim),
+        #                               nn.ReLU(True),
+        #                               nn.Dropout(0.25),
+        #                               )
+
+    def forward(self, inputs):
+
+        encoded=self.encoder(inputs)
+        encoded=encoded.unsqueeze(-1).unsqueeze(-1)
+        decoded=self.decoder(encoded)
+        # print torch.max(decoded)
+        # print torch.min(decoded)
+        # print(decoded.shape,'11111111111111111',encoded.shape)
+
+        return decoded

data/MBD/model/gienet.py

@@ -0,0 +1,742 @@
+from math import log
+import torch
+import torch.nn as nn
+from torch.nn import init
+import functools
+from model.cbam import CBAM
+# Defines the Unet generator.
+# |num_downs|: number of downsamplings in UNet. For example,
+# if |num_downs| == 7, image of size 128x128 will become of size 1x1
+# at the bottleneck
+class SingleConv(nn.Module):
+    """(convolution => [BN] => ReLU) * 2"""
+
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        self.double_conv = nn.Sequential(
+            nn.ReflectionPad2d(1),
+            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=0,stride=1),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True),
+            # nn.ReflectionPad2d(1),
+            # nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=0,stride=1),
+            # nn.BatchNorm2d(out_channels),
+            # nn.ReLU(inplace=True)
+        )
+
+    def forward(self, x):
+        return self.double_conv(x)
+class Down_single(nn.Module):
+    """Downscaling with maxpool then double conv"""
+
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        self.maxpool_conv = nn.Sequential(
+            nn.MaxPool2d(2),
+            SingleConv(in_channels, out_channels)
+        )
+
+    def forward(self, x):
+        return self.maxpool_conv(x)
+class Up_single(nn.Module):
+    """Upscaling then double conv"""
+    def __init__(self, in_channels, out_channels, bilinear=True):
+        super().__init__()
+        self.up = nn.Upsample(scale_factor=2, mode='nearest')
+        self.conv = SingleConv(in_channels, out_channels)
+        self.deconv = nn.ConvTranspose2d(in_channels, out_channels,kernel_size=4, stride=2,padding=1, bias=True)
+    def forward(self, x1, x2):
+        x1 = self.deconv(x1)
+        # input is BCHW
+        x = torch.cat([x2, x1], dim=1)
+        return self.conv(x)
+class DoubleConv(nn.Module):
+    """(convolution => [BN] => ReLU) * 2"""
+
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        self.double_conv = nn.Sequential(
+            nn.ReflectionPad2d(1),
+            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=0,stride=1),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True),
+            nn.ReflectionPad2d(1),
+            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=0,stride=1),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True)
+        )
+
+    def forward(self, x):
+        return self.double_conv(x)
+class Down(nn.Module):
+    """Downscaling with maxpool then double conv"""
+
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        self.maxpool_conv = nn.Sequential(
+            nn.MaxPool2d(2),
+            DoubleConv(in_channels, out_channels)
+        )
+
+    def forward(self, x):
+        return self.maxpool_conv(x)
+class Up(nn.Module):
+    """Upscaling then double conv"""
+    def __init__(self, in_channels, out_channels, bilinear=True):
+        super().__init__()
+        self.up = nn.Upsample(scale_factor=2, mode='nearest')
+        self.conv = DoubleConv(in_channels, out_channels)
+        self.deconv = nn.ConvTranspose2d(in_channels, out_channels,kernel_size=4, stride=2,padding=1, bias=True)
+    def forward(self, x1, x2):
+        x1 = self.deconv(x1)
+        # input is BCHW
+        x = torch.cat([x2, x1], dim=1)
+        return self.conv(x)
+
+class OutConv(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super(OutConv, self).__init__()
+        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)
+        self.tanh = nn.Tanh()
+        self.hardtanh = nn.Hardtanh()
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, x1):
+        x = self.conv(x1)
+        # x = self.sigmoid(x)
+        # x = self.hardtanh(x)
+        # x = (x+1)/2
+        return x
+class GiemaskGenerator(nn.Module):
+    """Create a Unet-based generator"""
+
+    def __init__(self, input_nc, output_nc, num_downs, ngf=64, biline=True, norm_layer=nn.BatchNorm2d, use_dropout=False):
+        """Construct a Unet generator
+        Parameters:
+            input_nc (int)  -- the number of channels in input images
+            output_nc (int) -- the number of channels in output images
+            num_downs (int) -- the number of downsamplings in UNet. For example, # if |num_downs| == 7,
+                                image of size 128x128 will become of size 1x1 # at the bottleneck
+            ngf (int)       -- the number of filters in the last conv layer
+            norm_layer      -- normalization layer
+
+        We construct the U-Net from the innermost layer to the outermost layer.
+        It is a recursive process.
+        """
+        super(GiemaskGenerator, self).__init__()
+        self.init_channel =32
+        self.inc = DoubleConv(3,self.init_channel)
+        self.down1 = Down(self.init_channel, self.init_channel*2)
+        self.down2 = Down(self.init_channel*2, self.init_channel*4)
+        self.down3 = Down(self.init_channel*4, self.init_channel*8)
+        self.down4 = Down(self.init_channel*8, self.init_channel*16)
+        self.down5 = Down(self.init_channel*16, self.init_channel*32)
+
+        self.up1 = Up(self.init_channel*32, self.init_channel*16)
+        self.up2 = Up(self.init_channel*16, self.init_channel*8)
+        self.up3 = Up(self.init_channel*8, self.init_channel*4)
+        self.up4 = Up(self.init_channel*4,self.init_channel*2)
+        self.up5 = Up(self.init_channel*2, self.init_channel)
+        self.outc = OutConv(self.init_channel, 1)
+        self.up1_1 = Up_single(self.init_channel*32, self.init_channel*16)
+        self.up2_1 = Up_single(self.init_channel*16, self.init_channel*8)
+        self.up3_1 = Up_single(self.init_channel*8, self.init_channel*4)
+        self.up4_1 = Up_single(self.init_channel*4,self.init_channel*2)
+        self.up5_1 = Up_single(self.init_channel*2, self.init_channel)
+        self.outc_1 = OutConv(self.init_channel, 1)
+  #      self.dropout = nn.Dropout(p=0.5)
+    def forward(self, input):
+        x1 = self.inc(input)
+        x2 = self.down1(x1)
+        x3 = self.down2(x2)
+        x4 = self.down3(x3)
+        x5 = self.down4(x4)
+        x6 = self.down5(x5)
+
+
+        x_1 = self.up1_1(x6, x5)
+        x_1 = self.up2_1(x_1, x4)
+        x_1 = self.up3_1(x_1, x3)
+        x_1 = self.up4_1(x_1, x2)
+        x_1 = self.up5_1(x_1, x1)
+        mask = self.outc_1(x_1)
+
+        x = self.up1(x6, x5)
+#        x = self.dropout(x)
+        x = self.up2(x, x4)
+#        x = self.dropout(x)
+        x = self.up3(x, x3)
+#        x = self.dropout(x)
+        x = self.up4(x, x2)
+#        x = self.dropout(x)
+        x = self.up5(x, x1)
+#        x = self.dropout(x)
+        depth = self.outc(x)
+        return depth,mask
+    """Create a Unet-based generator"""
+class Giemask2Generator(nn.Module):
+    """Create a Unet-based generator"""
+
+    def __init__(self, input_nc, output_nc, num_downs, ngf=64, biline=True, norm_layer=nn.BatchNorm2d, use_dropout=False):
+        """Construct a Unet generator
+        Parameters:
+            input_nc (int)  -- the number of channels in input images
+            output_nc (int) -- the number of channels in output images
+            num_downs (int) -- the number of downsamplings in UNet. For example, # if |num_downs| == 7,
+                                image of size 128x128 will become of size 1x1 # at the bottleneck
+            ngf (int)       -- the number of filters in the last conv layer
+            norm_layer      -- normalization layer
+
+        We construct the U-Net from the innermost layer to the outermost layer.
+        It is a recursive process.
+        """
+        super(Giemask2Generator, self).__init__()
+        self.init_channel =32
+        self.inc = DoubleConv(3,self.init_channel)
+        self.down1 = Down(self.init_channel, self.init_channel*2)
+        self.down2 = Down(self.init_channel*2, self.init_channel*4)
+        self.down3 = Down(self.init_channel*4, self.init_channel*8)
+        self.down4 = Down(self.init_channel*8, self.init_channel*16)
+        self.down5 = Down(self.init_channel*16, self.init_channel*32)
+
+        self.up1 = Up(self.init_channel*32, self.init_channel*16)
+        self.up2 = Up(self.init_channel*16, self.init_channel*8)
+        self.up3 = Up(self.init_channel*8, self.init_channel*4)
+        self.up4 = Up(self.init_channel*4,self.init_channel*2)
+        self.up5 = Up(self.init_channel*2, self.init_channel)
+        self.outc = OutConv(self.init_channel, 1)
+        self.up1_1 = Up_single(self.init_channel*32, self.init_channel*16)
+        self.up2_1 = Up_single(self.init_channel*16, self.init_channel*8)
+        self.up3_1 = Up_single(self.init_channel*8, self.init_channel*4)
+        self.up4_1 = Up_single(self.init_channel*4,self.init_channel*2)
+        self.up5_1 = Up_single(self.init_channel*2, self.init_channel)
+        self.outc_1 = OutConv(self.init_channel, 1)
+        self.outc_2 = OutConv(self.init_channel, 1)
+  #      self.dropout = nn.Dropout(p=0.5)
+    def forward(self, input):
+        x1 = self.inc(input)
+        x2 = self.down1(x1)
+        x3 = self.down2(x2)
+        x4 = self.down3(x3)
+        x5 = self.down4(x4)
+        x6 = self.down5(x5)
+
+
+        x_1 = self.up1_1(x6, x5)
+        x_1 = self.up2_1(x_1, x4)
+        x_1 = self.up3_1(x_1, x3)
+        x_1 = self.up4_1(x_1, x2)
+        x_1 = self.up5_1(x_1, x1)
+        mask = self.outc_1(x_1)
+        edge = self.outc_2(x_1)
+
+        x = self.up1(x6, x5)
+#        x = self.dropout(x)
+        x = self.up2(x, x4)
+#        x = self.dropout(x)
+        x = self.up3(x, x3)
+#        x = self.dropout(x)
+        x = self.up4(x, x2)
+#        x = self.dropout(x)
+        x = self.up5(x, x1)
+#        x = self.dropout(x)
+        depth = self.outc(x)
+        return depth,mask,edge
+    """Create a Unet-based generator"""
+class GieGenerator(nn.Module):
+    def __init__(self, input_nc, output_nc, num_downs, ngf=64, biline=True, norm_layer=nn.BatchNorm2d, use_dropout=False):
+        """Construct a Unet generator
+        Parameters:
+            input_nc (int)  -- the number of channels in input images
+            output_nc (int) -- the number of channels in output images
+            num_downs (int) -- the number of downsamplings in UNet. For example, # if |num_downs| == 7,
+                                image of size 128x128 will become of size 1x1 # at the bottleneck
+            ngf (int)       -- the number of filters in the last conv layer
+            norm_layer      -- normalization layer
+
+        We construct the U-Net from the innermost layer to the outermost layer.
+        It is a recursive process.
+        """
+        super(GieGenerator, self).__init__()
+        self.init_channel =32
+        self.inc = DoubleConv(input_nc,self.init_channel)
+        self.down1 = Down(self.init_channel, self.init_channel*2)
+        self.down2 = Down(self.init_channel*2, self.init_channel*4)
+        self.down3 = Down(self.init_channel*4, self.init_channel*8)
+        self.down4 = Down(self.init_channel*8, self.init_channel*16)
+        self.down5 = Down(self.init_channel*16, self.init_channel*32)
+
+        self.up1 = Up(self.init_channel*32, self.init_channel*16)
+        self.up2 = Up(self.init_channel*16, self.init_channel*8)
+        self.up3 = Up(self.init_channel*8, self.init_channel*4)
+        self.up4 = Up(self.init_channel*4,self.init_channel*2)
+        self.up5 = Up(self.init_channel*2, self.init_channel)
+        self.outc = OutConv(self.init_channel, 2)
+  #      self.dropout = nn.Dropout(p=0.5)
+    def forward(self, input):
+        x1 = self.inc(input)
+        x2 = self.down1(x1)
+        x3 = self.down2(x2)
+        x4 = self.down3(x3)
+        x5 = self.down4(x4)
+        x6 = self.down5(x5)
+
+        x = self.up1(x6, x5)
+#        x = self.dropout(x)
+        x = self.up2(x, x4)
+#        x = self.dropout(x)
+        x = self.up3(x, x3)
+#        x = self.dropout(x)
+        x = self.up4(x, x2)
+#        x = self.dropout(x)
+        x = self.up5(x, x1)
+#        x = self.dropout(x)
+        logits1 = self.outc(x)
+        return logits1
+
+
+class GiecbamGenerator(nn.Module):
+    def __init__(self, input_nc, output_nc, num_downs, ngf=64, biline=True, norm_layer=nn.BatchNorm2d, use_dropout=False):
+        """Construct a Unet generator
+        Parameters:
+            input_nc (int)  -- the number of channels in input images
+            output_nc (int) -- the number of channels in output images
+            num_downs (int) -- the number of downsamplings in UNet. For example, # if |num_downs| == 7,
+                                image of size 128x128 will become of size 1x1 # at the bottleneck
+            ngf (int)       -- the number of filters in the last conv layer
+            norm_layer      -- normalization layer
+
+        We construct the U-Net from the innermost layer to the outermost layer.
+        It is a recursive process.
+        """
+        super(GiecbamGenerator, self).__init__()
+        self.init_channel =32
+        self.inc = DoubleConv(input_nc,self.init_channel)
+        self.down1 = Down(self.init_channel, self.init_channel*2)
+        self.down2 = Down(self.init_channel*2, self.init_channel*4)
+        self.down3 = Down(self.init_channel*4, self.init_channel*8)
+        self.down4 = Down(self.init_channel*8, self.init_channel*16)
+        self.down5 = Down(self.init_channel*16, self.init_channel*32)
+        self.cbam = CBAM(gate_channels=self.init_channel*32)
+        self.up1 = Up(self.init_channel*32, self.init_channel*16)
+        self.up2 = Up(self.init_channel*16, self.init_channel*8)
+        self.up3 = Up(self.init_channel*8, self.init_channel*4)
+        self.up4 = Up(self.init_channel*4,self.init_channel*2)
+        self.up5 = Up(self.init_channel*2, self.init_channel)
+        self.outc = OutConv(self.init_channel, 2)
+        self.dropout = nn.Dropout(p=0.1)
+    def forward(self, input):
+        x1 = self.inc(input)
+        x2 = self.down1(x1)
+        x3 = self.down2(x2)
+        x4 = self.down3(x3)
+        x5 = self.down4(x4)
+        x6 = self.down5(x5)
+        x6 = self.cbam(x6)
+        x = self.up1(x6, x5)
+        x = self.up2(x, x4)
+        x = self.up3(x, x3)
+        x = self.up4(x, x2)
+        x = self.up5(x, x1)
+        x = self.dropout(x)
+        logits1 = self.outc(x)
+        return logits1
+
+
+
+
+class Gie2headGenerator(nn.Module):
+    def __init__(self, input_nc, output_nc, num_downs, ngf=64, biline=True, norm_layer=nn.BatchNorm2d, use_dropout=False):
+        """Construct a Unet generator
+        Parameters:
+            input_nc (int)  -- the number of channels in input images
+            output_nc (int) -- the number of channels in output images
+            num_downs (int) -- the number of downsamplings in UNet. For example, # if |num_downs| == 7,
+                                image of size 128x128 will become of size 1x1 # at the bottleneck
+            ngf (int)       -- the number of filters in the last conv layer
+            norm_layer      -- normalization layer
+
+        We construct the U-Net from the innermost layer to the outermost layer.
+        It is a recursive process.
+        """
+        super(Gie2headGenerator, self).__init__()
+        self.init_channel =32
+        self.inc = DoubleConv(input_nc,self.init_channel)
+        self.down1 = Down(self.init_channel, self.init_channel*2)
+        self.down2 = Down(self.init_channel*2, self.init_channel*4)
+        self.down3 = Down(self.init_channel*4, self.init_channel*8)
+        self.down4 = Down(self.init_channel*8, self.init_channel*16)
+        self.down5 = Down(self.init_channel*16, self.init_channel*32)
+
+        self.up1_1 = Up(self.init_channel*32, self.init_channel*16)
+        self.up2_1 = Up(self.init_channel*16, self.init_channel*8)
+        self.up3_1 = Up(self.init_channel*8, self.init_channel*4)
+        self.up4_1 = Up(self.init_channel*4,self.init_channel*2)
+        self.up5_1 = Up(self.init_channel*2, self.init_channel)
+        self.outc_1 = OutConv(self.init_channel, 1)
+
+        self.up1_2 = Up(self.init_channel*32, self.init_channel*16)
+        self.up2_2 = Up(self.init_channel*16, self.init_channel*8)
+        self.up3_2 = Up(self.init_channel*8, self.init_channel*4)
+        self.up4_2 = Up(self.init_channel*4,self.init_channel*2)
+        self.up5_2 = Up(self.init_channel*2, self.init_channel)
+        self.outc_2 = OutConv(self.init_channel, 1)
+
+    def forward(self, input):
+        x1 = self.inc(input)
+        x2 = self.down1(x1)
+        x3 = self.down2(x2)
+        x4 = self.down3(x3)
+        x5 = self.down4(x4)
+        x6 = self.down5(x5)
+
+        x_1 = self.up1_1(x6, x5)
+        x_1 = self.up2_1(x_1, x4)
+        x_1 = self.up3_1(x_1, x3)
+        x_1 = self.up4_1(x_1, x2)
+        x_1 = self.up5_1(x_1, x1)
+        logits_1 = self.outc_1(x_1)
+
+        x_2 = self.up1_2(x6, x5)
+        x_2 = self.up2_2(x_2, x4)
+        x_2 = self.up3_2(x_2, x3)
+        x_2 = self.up4_2(x_2, x2)
+        x_2 = self.up5_2(x_2, x1)
+        logits_2 = self.outc_2(x_2)
+
+        logits = torch.cat((logits_1,logits_2),1)
+
+        return logits
+
+
+
+class BmpGenerator(nn.Module):
+    def __init__(self, input_nc, output_nc, num_downs, ngf=64, biline=True, norm_layer=nn.BatchNorm2d, use_dropout=False):
+        """Construct a Unet generator
+        Parameters:
+            input_nc (int)  -- the number of channels in input images
+            output_nc (int) -- the number of channels in output images
+            num_downs (int) -- the number of downsamplings in UNet. For example, # if |num_downs| == 7,
+                                image of size 128x128 will become of size 1x1 # at the bottleneck
+            ngf (int)       -- the number of filters in the last conv layer
+            norm_layer      -- normalization layer
+
+        We construct the U-Net from the innermost layer to the outermost layer.
+        It is a recursive process.
+        """
+        super(BmpGenerator, self).__init__()
+        self.init_channel =32
+        self.output_nc = output_nc
+        self.inc = DoubleConv(input_nc,self.init_channel)
+        self.down1 = Down(self.init_channel, self.init_channel*2)
+        self.down2 = Down(self.init_channel*2, self.init_channel*4)
+        self.down3 = Down(self.init_channel*4, self.init_channel*8)
+        self.down4 = Down(self.init_channel*8, self.init_channel*16)
+        self.down5 = Down(self.init_channel*16, self.init_channel*32)
+
+        self.up1 = Up(self.init_channel*32, self.init_channel*16)
+        self.up2 = Up(self.init_channel*16, self.init_channel*8)
+        self.up3 = Up(self.init_channel*8, self.init_channel*4)
+        self.up4 = Up(self.init_channel*4,self.init_channel*2)
+        self.up5 = Up(self.init_channel*2, self.init_channel)
+        self.outc = OutConv(self.init_channel, self.output_nc)
+  #      self.dropout = nn.Dropout(p=0.5)
+    def forward(self, input):
+        x1 = self.inc(input)
+        x2 = self.down1(x1)
+        x3 = self.down2(x2)
+        x4 = self.down3(x3)
+        x5 = self.down4(x4)
+        x6 = self.down5(x5)
+
+        x = self.up1(x6, x5)
+#        x = self.dropout(x)
+        x = self.up2(x, x4)
+#        x = self.dropout(x)
+        x = self.up3(x, x3)
+#        x = self.dropout(x)
+        x = self.up4(x, x2)
+#        x = self.dropout(x)
+        x = self.up5(x, x1)
+#        x = self.dropout(x)
+        logits1 = self.outc(x)
+        return logits1
+class Bmp2Generator(nn.Module):
+    """Create a Unet-based generator"""
+
+    def __init__(self, input_nc, output_nc, num_downs, ngf=64, biline=True, norm_layer=nn.BatchNorm2d, use_dropout=False):
+        """Construct a Unet generator
+        Parameters:
+            input_nc (int)  -- the number of channels in input images
+            output_nc (int) -- the number of channels in output images
+            num_downs (int) -- the number of downsamplings in UNet. For example, # if |num_downs| == 7,
+                                image of size 128x128 will become of size 1x1 # at the bottleneck
+            ngf (int)       -- the number of filters in the last conv layer
+            norm_layer      -- normalization layer
+
+        We construct the U-Net from the innermost layer to the outermost layer.
+        It is a recursive process.
+        """
+        super(Bmp2Generator, self).__init__()
+        #gienet
+        self.init_channel =32
+        self.inc = DoubleConv(3,self.init_channel)
+        self.down1 = Down(self.init_channel, self.init_channel*2)
+        self.down2 = Down(self.init_channel*2, self.init_channel*4)
+        self.down3 = Down(self.init_channel*4, self.init_channel*8)
+        self.down4 = Down(self.init_channel*8, self.init_channel*16)
+        self.down5 = Down(self.init_channel*16, self.init_channel*32)
+
+        self.up1 = Up(self.init_channel*32, self.init_channel*16)
+        self.up2 = Up(self.init_channel*16, self.init_channel*8)
+        self.up3 = Up(self.init_channel*8, self.init_channel*4)
+        self.up4 = Up(self.init_channel*4,self.init_channel*2)
+        self.up5 = Up(self.init_channel*2, self.init_channel)
+        self.outc = OutConv(self.init_channel, 1)
+        self.up1_1 = Up_single(self.init_channel*32, self.init_channel*16)
+        self.up2_1 = Up_single(self.init_channel*16, self.init_channel*8)
+        self.up3_1 = Up_single(self.init_channel*8, self.init_channel*4)
+        self.up4_1 = Up_single(self.init_channel*4,self.init_channel*2)
+        self.up5_1 = Up_single(self.init_channel*2, self.init_channel)
+        self.outc_1 = OutConv(self.init_channel, 1)
+        self.outc_2 = OutConv(self.init_channel, 1)
+
+        #bpm net
+        self.inc_b = DoubleConv(4,self.init_channel)
+        self.down1_b = Down(self.init_channel, self.init_channel*2)
+        self.down2_b = Down(self.init_channel*2, self.init_channel*4)
+        self.down3_b = Down(self.init_channel*4, self.init_channel*8)
+        self.down4_b = Down(self.init_channel*8, self.init_channel*16)
+        self.down5_b = Down(self.init_channel*16, self.init_channel*32)
+
+        self.up1_b = Up(self.init_channel*32, self.init_channel*16)
+        self.up2_b = Up(self.init_channel*16, self.init_channel*8)
+        self.up3_b = Up(self.init_channel*8, self.init_channel*4)
+        self.up4_b = Up(self.init_channel*4,self.init_channel*2)
+        self.up5_b = Up(self.init_channel*2, self.init_channel)
+        self.outc_b = OutConv(self.init_channel, 2)
+  #      self.dropout = nn.Dropout(p=0.5)
+    def forward(self, input):
+        #gienet
+        x1 = self.inc(input)
+        x2 = self.down1(x1)
+        x3 = self.down2(x2)
+        x4 = self.down3(x3)
+        x5 = self.down4(x4)
+        x6 = self.down5(x5)
+
+        x_1 = self.up1_1(x6, x5)
+        x_1 = self.up2_1(x_1, x4)
+        x_1 = self.up3_1(x_1, x3)
+        x_1 = self.up4_1(x_1, x2)
+        x_1 = self.up5_1(x_1, x1)
+        mask = self.outc_1(x_1)
+        edge = self.outc_2(x_1)
+
+        x = self.up1(x6, x5)
+        x = self.up2(x, x4)
+        x = self.up3(x, x3)
+        x = self.up4(x, x2)
+        x = self.up5(x, x1)
+        depth = self.outc(x)
+        
+        #bmpnet
+        mask[mask>0.5]=1.
+        mask[mask<=0.5]=0.
+        image_cat_depth = torch.cat((input*mask,depth*mask),dim=1)
+        x1_b = self.inc_b(image_cat_depth)
+        x2_b = self.down1_b(x1_b)
+        x3_b = self.down2_b(x2_b)
+        x4_b = self.down3_b(x3_b)
+        x5_b = self.down4_b(x4_b)
+        x6_b = self.down5_b(x5_b)
+        x_b = self.up1_b(x6_b, x5_b)
+        x_b = self.up2_b(x_b, x4_b)
+        x_b = self.up3_b(x_b, x3_b)
+        x_b = self.up4_b(x_b, x2_b)
+        x_b = self.up5_b(x_b, x1_b)
+        bm = self.outc_b(x_b)
+        # return depth,mask,edge,bm
+        return bm
+class UnetGenerator(nn.Module):
+    def __init__(self, input_nc, output_nc, num_downs, ngf=64,
+                 norm_layer=nn.BatchNorm2d, use_dropout=False):
+        super(UnetGenerator, self).__init__()
+
+        # construct unet structure
+        unet_block = UnetSkipConnectionBlock(ngf * 8, ngf * 8, input_nc=None, submodule=None, norm_layer=norm_layer, innermost=True)
+        for i in range(num_downs - 5):
+            unet_block = UnetSkipConnectionBlock(ngf * 8, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer, use_dropout=use_dropout)
+        unet_block = UnetSkipConnectionBlock(ngf * 4, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
+        unet_block = UnetSkipConnectionBlock(ngf * 2, ngf * 4, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
+        unet_block = UnetSkipConnectionBlock(ngf, ngf * 2, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
+        unet_block = UnetSkipConnectionBlock(output_nc, ngf, input_nc=input_nc, submodule=unet_block, outermost=True, norm_layer=norm_layer)
+
+        self.model = unet_block
+
+    def forward(self, input):
+        return self.model(input)
+
+#class GieGenerator(nn.Module):
+#    def __init__(self, input_nc, output_nc, num_downs, ngf=64,
+#                 norm_layer=nn.BatchNorm2d, use_dropout=False):
+#        super(GieGenerator, self).__init__()
+#
+#        # construct unet structure
+#        unet_block = UnetSkipConnectionBlock(ngf * 8, ngf * 8, input_nc=None, submodule=None, norm_layer=norm_layer, innermost=True)
+#        for i in range(num_downs - 5):
+#            unet_block = UnetSkipConnectionBlock(ngf * 8, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer, use_dropout=use_dropout)
+#        unet_block = UnetSkipConnectionBlock(ngf * 4, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
+#        unet_block = UnetSkipConnectionBlock(ngf * 2, ngf * 4, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
+#        unet_block = UnetSkipConnectionBlock(ngf, ngf * 2, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
+#        unet_block = UnetSkipConnectionBlock(output_nc, ngf, input_nc=input_nc, submodule=unet_block, outermost=True, norm_layer=norm_layer)
+#
+#        self.model = unet_block
+#
+#    def forward(self, input):
+#        return self.model(input)
+
+# Defines the submodule with skip connection.
+# X -------------------identity---------------------- X
+#   |-- downsampling -- |submodule| -- upsampling --|
+class UnetSkipConnectionBlock(nn.Module):
+    def __init__(self, outer_nc, inner_nc, input_nc=None,
+                 submodule=None, outermost=False, innermost=False, norm_layer=nn.BatchNorm2d, use_dropout=False):
+        super(UnetSkipConnectionBlock, self).__init__()
+        self.outermost = outermost
+        if type(norm_layer) == functools.partial:
+            use_bias = norm_layer.func == nn.InstanceNorm2d
+        else:
+            use_bias = norm_layer == nn.InstanceNorm2d
+        if input_nc is None:
+            input_nc = outer_nc
+        downconv = nn.Conv2d(input_nc, inner_nc, kernel_size=4,
+                             stride=2, padding=1, bias=use_bias)
+        downrelu = nn.LeakyReLU(0.2, True)
+        downnorm = norm_layer(inner_nc)
+        uprelu = nn.ReLU(True)
+        upnorm = norm_layer(outer_nc)
+
+        if outermost:
+            upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc,
+                                        kernel_size=4, stride=2,
+                                        padding=1)
+            down = [downconv]
+            up = [uprelu, upconv, nn.Tanh()]
+            model = down + [submodule] + up
+        elif innermost:
+           # resize = nn.Upsample(scale_factor=2)
+           # conv = nn.Conv2d(inner_nc,outer_nc,kernel_size=4,stride=2,padding=1,bias=use_bias)
+            upconv = nn.ConvTranspose2d(inner_nc, outer_nc,
+                                        kernel_size=4, stride=2,
+                                        padding=1, bias=use_bias)
+            down = [downrelu, downconv]
+            up = [uprelu, upconv, upnorm]
+            #up = [uprelu, resize, conv, upnorm]
+            model = down + up
+        else:
+            upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc,
+                                        kernel_size=4, stride=2,
+                                        padding=1, bias=use_bias)
+            down = [downrelu, downconv, downnorm]
+            up = [uprelu, upconv, upnorm]
+
+            if use_dropout:
+                model = down + [submodule] + up + [nn.Dropout(0.5)]
+            else:
+                model = down + [submodule] + up
+
+        self.model = nn.Sequential(*model)
+
+    def forward(self, x):
+        if self.outermost:
+            return self.model(x)
+        else:
+            return torch.cat([x, self.model(x)], 1)
+
+
+
+##===================================================================================================
+class DilatedDoubleConv(nn.Module):
+    """(convolution => [BN] => ReLU) * 2"""
+
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        self.double_conv = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=4,stride=1,dilation=4),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=4,stride=1,dilation=4),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True)
+        )
+
+    def forward(self, x):
+        return self.double_conv(x)
+
+class DilatedDown(nn.Module):
+    """Downscaling with maxpool then double conv"""
+
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        self.maxpool_conv = nn.Sequential(
+            nn.MaxPool2d(2),
+            DilatedDoubleConv(in_channels, out_channels)
+        )
+
+    def forward(self, x):
+        return self.maxpool_conv(x)
+
+class DilatedUp(nn.Module):
+    """Upscaling then double conv"""
+    def __init__(self, in_channels, out_channels, bilinear=True):
+        super().__init__()
+        self.up = nn.Upsample(scale_factor=2, mode='nearest')
+        self.conv = DilatedDoubleConv(in_channels, out_channels)
+
+        self.conv1 = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=4,stride=1,dilation=4),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True),
+        )
+        # self.deconv = nn.ConvTranspose2d(in_channels, out_channels,kernel_size=4, stride=2,padding=1, bias=True)
+    def forward(self, x1, x2):
+        x1 = self.up(x1)
+        x1 = self.conv1(x1)
+        # x1 = self.deconv(x1)
+        # input is BCHW
+        x = torch.cat([x2, x1], dim=1)
+        return self.conv(x)
+class DilatedSingleUnet(nn.Module):
+    def __init__(self, input_nc, output_nc, num_downs, ngf=64, biline=True, norm_layer=nn.BatchNorm2d, use_dropout=False):
+        super(DilatedSingleUnet, self).__init__()
+        self.init_channel = 32
+        self.inc = DilatedDoubleConv(input_nc,self.init_channel)
+        self.down1 = DilatedDown(self.init_channel, self.init_channel*2)
+        self.down2 = DilatedDown(self.init_channel*2, self.init_channel*4)
+        self.down3 = DilatedDown(self.init_channel*4, self.init_channel*8)
+        self.down4 = DilatedDown(self.init_channel*8, self.init_channel*16)
+        self.down5 = DilatedDown(self.init_channel*16, self.init_channel*32)
+        self.cbam = CBAM(gate_channels=self.init_channel*32)
+
+        self.up1 = DilatedUp(self.init_channel*32, self.init_channel*16)
+        self.up2 = DilatedUp(self.init_channel*16, self.init_channel*8)
+        self.up3 = DilatedUp(self.init_channel*8, self.init_channel*4)
+        self.up4 = DilatedUp(self.init_channel*4,self.init_channel*2)
+        self.up5 = DilatedUp(self.init_channel*2, self.init_channel)
+        self.outc = OutConv(self.init_channel, output_nc)
+    def forward(self, input):
+        x1 = self.inc(input)
+        x2 = self.down1(x1)
+        x3 = self.down2(x2)
+        x4 = self.down3(x3)
+        x5 = self.down4(x4)
+        x6 = self.down5(x5)
+        x6 = self.cbam(x6)
+        x = self.up1(x6, x5)
+        x = self.up2(x, x4)
+        x = self.up3(x, x3)
+        x = self.up4(x, x2)
+        x = self.up5(x, x1)
+        logits1 = self.outc(x)
+        return logits1

data/MBD/model/unetnc.py

@@ -0,0 +1,86 @@
+import torch
+import torch.nn as nn
+from torch.nn import init
+import functools
+
+# Defines the Unet generator.
+# |num_downs|: number of downsamplings in UNet. For example,
+# if |num_downs| == 7, image of size 128x128 will become of size 1x1
+# at the bottleneck
+class UnetGenerator(nn.Module):
+    def __init__(self, input_nc, output_nc, num_downs, ngf=64,
+                 norm_layer=nn.BatchNorm2d, use_dropout=False):
+        super(UnetGenerator, self).__init__()
+
+        # construct unet structure
+        unet_block = UnetSkipConnectionBlock(ngf * 8, ngf * 8, input_nc=None, submodule=None, norm_layer=norm_layer, innermost=True)
+        for i in range(num_downs - 5):
+            unet_block = UnetSkipConnectionBlock(ngf * 8, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer, use_dropout=use_dropout)
+        unet_block = UnetSkipConnectionBlock(ngf * 4, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
+        unet_block = UnetSkipConnectionBlock(ngf * 2, ngf * 4, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
+        unet_block = UnetSkipConnectionBlock(ngf, ngf * 2, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
+        unet_block = UnetSkipConnectionBlock(output_nc, ngf, input_nc=input_nc, submodule=unet_block, outermost=True, norm_layer=norm_layer)
+
+        self.model = unet_block
+
+    def forward(self, input):
+        return self.model(input)
+
+
+    def forward(self, input):
+        return self.model(input)
+
+# Defines the submodule with skip connection.
+# X -------------------identity---------------------- X
+#   |-- downsampling -- |submodule| -- upsampling --|
+class UnetSkipConnectionBlock(nn.Module):
+    def __init__(self, outer_nc, inner_nc, input_nc=None,
+                 submodule=None, outermost=False, innermost=False, norm_layer=nn.BatchNorm2d, use_dropout=False):
+        super(UnetSkipConnectionBlock, self).__init__()
+        self.outermost = outermost
+        if type(norm_layer) == functools.partial:
+            use_bias = norm_layer.func == nn.InstanceNorm2d
+        else:
+            use_bias = norm_layer == nn.InstanceNorm2d
+        if input_nc is None:
+            input_nc = outer_nc
+        downconv = nn.Conv2d(input_nc, inner_nc, kernel_size=4,
+                             stride=2, padding=1, bias=use_bias)
+        downrelu = nn.LeakyReLU(0.2, True)
+        downnorm = norm_layer(inner_nc)
+        uprelu = nn.ReLU(True)
+        upnorm = norm_layer(outer_nc)
+
+        if outermost:
+            upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc,
+                                        kernel_size=4, stride=2,
+                                        padding=1)
+            down = [downconv]
+            up = [uprelu, upconv, nn.Tanh()]
+            model = down + [submodule] + up
+        elif innermost:
+            upconv = nn.ConvTranspose2d(inner_nc, outer_nc,
+                                        kernel_size=4, stride=2,
+                                        padding=1, bias=use_bias)
+            down = [downrelu, downconv]
+            up = [uprelu, upconv, upnorm]
+            model = down + up
+        else:
+            upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc,
+                                        kernel_size=4, stride=2,
+                                        padding=1, bias=use_bias)
+            down = [downrelu, downconv, downnorm]
+            up = [uprelu, upconv, upnorm]
+
+            if use_dropout:
+                model = down + [submodule] + up + [nn.Dropout(0.5)]
+            else:
+                model = down + [submodule] + up
+
+        self.model = nn.Sequential(*model)
+
+    def forward(self, x):
+        if self.outermost:
+            return self.model(x)
+        else:
+            return torch.cat([x, self.model(x)], 1)