diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..a9545d8 --- /dev/null +++ b/.gitignore @@ -0,0 +1,3 @@ +*.pyc +*.pth +/__pycache__ diff --git a/README.md b/README.md index 8619899..a3660ef 100644 --- a/README.md +++ b/README.md @@ -1,3 +1,2 @@ -# -BUAA_FRB_- - -这里是由阳神带领的冯如小队,完毕! \ No newline at end of file +# FRB_lttprivate +save code of myself diff --git a/deeplab.py b/deeplab.py index 50b1e90..2c58938 100644 --- a/deeplab.py +++ b/deeplab.py @@ -61,6 +61,14 @@ def get_10x_lr_params(self): yield p +def getinput(): + root_path = "../data20200111/" + img_path = "img/" + fea_path = "feature/" + label_path = "label/" + + + if __name__ == "__main__": model = DeepLab(backbone='mobilenet', output_stride=16) model.eval() diff --git a/modeling/__init__.py b/modeling/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/modeling/aspp.py b/modeling/aspp.py new file mode 100644 index 0000000..5a97879 --- /dev/null +++ b/modeling/aspp.py @@ -0,0 +1,95 @@ +import math +import torch +import torch.nn as nn +import torch.nn.functional as F +from modeling.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) \ No newline at end of file diff --git a/modeling/backbone/.__init__.py.swp b/modeling/backbone/.__init__.py.swp new file mode 100644 index 0000000..e418868 Binary files /dev/null and b/modeling/backbone/.__init__.py.swp differ diff --git a/modeling/backbone/.mobilenet.py.swp b/modeling/backbone/.mobilenet.py.swp new file mode 100644 index 0000000..1bb895e Binary files /dev/null and b/modeling/backbone/.mobilenet.py.swp differ diff --git a/modeling/backbone/__init__.py b/modeling/backbone/__init__.py new file mode 100644 index 0000000..5864c11 --- /dev/null +++ b/modeling/backbone/__init__.py @@ -0,0 +1,13 @@ +from modeling.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 diff --git a/modeling/backbone/drn.py b/modeling/backbone/drn.py new file mode 100644 index 0000000..d07b577 --- /dev/null +++ b/modeling/backbone/drn.py @@ -0,0 +1,402 @@ +import torch.nn as nn +import math +import torch.utils.model_zoo as model_zoo +from modeling.sync_batchnorm.batchnorm import SynchronizedBatchNorm2d + +webroot = 'https://tigress-web.princeton.edu/~fy/drn/models/' + +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()) diff --git a/modeling/backbone/mobilenet.py b/modeling/backbone/mobilenet.py new file mode 100644 index 0000000..95410e9 --- /dev/null +++ b/modeling/backbone/mobilenet.py @@ -0,0 +1,156 @@ +import torch +import torch.nn.functional as F +import torch.nn as nn +import math +from modeling.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), #in_ch, out_ch, kernel, stride, pad + 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): + #ltt add + pretrained = False + super(MobileNetV2, self).__init__() + block = InvertedResidual + input_channel = 32 #????P的32, 网络的第一层由conv_bn 的channel决定 + 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(1, input_channel, 2, BatchNorm)] #input_channel成为输出层 + print("first layer success") + 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 #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() + print("finish init") + 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): + print("before low features") + x = torch.tensor(x, dtype=torch.float32) + low_level_feat = self.low_level_features(x) + print("after low feature") + 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()) diff --git a/modeling/backbone/resnet.py b/modeling/backbone/resnet.py new file mode 100644 index 0000000..8eeac69 --- /dev/null +++ b/modeling/backbone/resnet.py @@ -0,0 +1,162 @@ +import math +import torch.nn as nn +import torch.utils.model_zoo as model_zoo +from modeling.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): + 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: + 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()) \ No newline at end of file diff --git a/modeling/backbone/xception.py b/modeling/backbone/xception.py new file mode 100644 index 0000000..4752367 --- /dev/null +++ b/modeling/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 modeling.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 model_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()) \ No newline at end of file diff --git a/modeling/decoder.py b/modeling/decoder.py new file mode 100644 index 0000000..5ed41d0 --- /dev/null +++ b/modeling/decoder.py @@ -0,0 +1,57 @@ +import math +import torch +import torch.nn as nn +import torch.nn.functional as F +from modeling.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)) + 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) \ No newline at end of file diff --git a/modeling/sync_batchnorm/__init__.py b/modeling/sync_batchnorm/__init__.py new file mode 100644 index 0000000..540118f --- /dev/null +++ b/modeling/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 \ No newline at end of file diff --git a/modeling/sync_batchnorm/batchnorm.py b/modeling/sync_batchnorm/batchnorm.py new file mode 100644 index 0000000..aa9dd37 --- /dev/null +++ b/modeling/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) \ No newline at end of file diff --git a/modeling/sync_batchnorm/comm.py b/modeling/sync_batchnorm/comm.py new file mode 100644 index 0000000..6005430 --- /dev/null +++ b/modeling/sync_batchnorm/comm.py @@ -0,0 +1,133 @@ +# -*- 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 sys +if sys.version > '3': + import queue +else: + import Queue as 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) diff --git a/modeling/sync_batchnorm/replicate.py b/modeling/sync_batchnorm/replicate.py new file mode 100644 index 0000000..3734266 --- /dev/null +++ b/modeling/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 \ No newline at end of file diff --git a/modeling/sync_batchnorm/unittest.py b/modeling/sync_batchnorm/unittest.py new file mode 100644 index 0000000..f826560 --- /dev/null +++ b/modeling/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()) + ) diff --git a/train.py b/train.py new file mode 100644 index 0000000..518e5df --- /dev/null +++ b/train.py @@ -0,0 +1,48 @@ +from torchvision import transforms, utils +from torch.utils.data import Dataset, DataLoader +import matplotlib.pyplot as plt +from scipy.io import loadmat +from PIL import Image +from deeplab import DeepLab + + +DATADIR = "../data20200111/" + +class MyDataset(Dataset): + def __init__(self, imgtxt, labeltxt, transform=None, target_transform=None): + fh = open(imgtxt, 'r').readlines() + fhlabel = open(labeltxt, 'r').readlines() + imgs = [] + for i in range(len(fh)): + print(i) + lineimg = fh[i].strip('\n') + linelabel = fhlabel[i].strip('\n') + imgs.append((lineimg, linelabel)) + self.imgs = imgs + self.transform = transform + self.target_transform = target_transform + + + def __getitem__(self, index): + fn, label = self.imgs[index] + img = loadmat(DATADIR+fn)['data'] + label = loadmat(DATADIR+label)['data'] + if self.transform is not None: + img = self.transform(img) + label = self.transform(label) + return img,label + + def __len__(self): + return len(self.imgs) + +train_data=MyDataset(imgtxt='../data20200111/img_list.txt', labeltxt='../data20200111/label_list.txt', transform=transforms.ToTensor()) +data_loader = DataLoader(train_data, batch_size=16,shuffle=True) +for epoch in range(10): + print("epoch, ", epoch) + for i, data in enumerate(data_loader): + inputs, labels = data + model = DeepLab(backbone='mobilenet', output_stride=16) + model.eval() + print(inputs.shape) + output = model(inputs) + print(output.size) diff --git a/visualMat.py b/visualMat.py new file mode 100644 index 0000000..55f09f6 --- /dev/null +++ b/visualMat.py @@ -0,0 +1,32 @@ +import os +import matplotlib.pyplot as plt +from scipy.io import loadmat +import argparse + +def readMatListandSave(path, name): + for obj in open(path+name).readlines(): + obj = obj.strip("\n") + mat = loadmat(path+obj)['data'] + plt.imshow(mat) + plt.savefig(path+obj.split('.')[0]+".png") + print("save ",path+obj.split('.')[0]+".png") + print("finish") + +def showPic(path, name): + mat = loadmat(path+name)['data'] + plt.imshow(mat) + plt.show() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description='input mat name') # 首先创建一个ArgumentParser对象 + parser.add_argument('path', type=str, help='please input mat or txt path') + parser.add_argument('name', type=str, help='please input file name') + + args = parser.parse_args() + path = args.path + name = args.name + if name.split('.')[-1] == "txt": #接受list文件可视化并保存 + readMatListandSave(path, name) + else: + showPic(path, name)