Spaces:
Running
on
Zero
Running
on
Zero
| 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()) | |