Spaces:
Running
Running
| import math | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| import torch.utils.model_zoo as model_zoo | |
| from torch.nn.parameter import Parameter | |
| from collections import OrderedDict | |
| class SeparableConv2d(nn.Module): | |
| def __init__(self, inplanes, planes, kernel_size=3, stride=1, padding=0, dilation=1, bias=False): | |
| super(SeparableConv2d, self).__init__() | |
| self.conv1 = nn.Conv2d(inplanes, inplanes, kernel_size, stride, padding, dilation, | |
| groups=inplanes, bias=bias) | |
| self.pointwise = nn.Conv2d(inplanes, planes, 1, 1, 0, 1, 1, bias=bias) | |
| def forward(self, x): | |
| x = self.conv1(x) | |
| x = self.pointwise(x) | |
| return x | |
| def fixed_padding(inputs, kernel_size, rate): | |
| kernel_size_effective = kernel_size + (kernel_size - 1) * (rate - 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_aspp(nn.Module): | |
| def __init__(self, inplanes, planes, kernel_size=3, stride=1, dilation=1, bias=False, padding=0): | |
| super(SeparableConv2d_aspp, self).__init__() | |
| self.depthwise = nn.Conv2d(inplanes, inplanes, kernel_size, stride, padding, dilation, | |
| groups=inplanes, bias=bias) | |
| self.depthwise_bn = nn.BatchNorm2d(inplanes) | |
| self.pointwise = nn.Conv2d(inplanes, planes, 1, 1, 0, 1, 1, bias=bias) | |
| self.pointwise_bn = nn.BatchNorm2d(planes) | |
| self.relu = nn.ReLU() | |
| def forward(self, x): | |
| # x = fixed_padding(x, self.depthwise.kernel_size[0], rate=self.depthwise.dilation[0]) | |
| x = self.depthwise(x) | |
| x = self.depthwise_bn(x) | |
| x = self.relu(x) | |
| x = self.pointwise(x) | |
| x = self.pointwise_bn(x) | |
| x = self.relu(x) | |
| return x | |
| class Decoder_module(nn.Module): | |
| def __init__(self, inplanes, planes, rate=1): | |
| super(Decoder_module, self).__init__() | |
| self.atrous_convolution = SeparableConv2d_aspp(inplanes, planes, 3, stride=1, dilation=rate,padding=1) | |
| def forward(self, x): | |
| x = self.atrous_convolution(x) | |
| return x | |
| class ASPP_module(nn.Module): | |
| def __init__(self, inplanes, planes, rate): | |
| super(ASPP_module, self).__init__() | |
| if rate == 1: | |
| raise RuntimeError() | |
| else: | |
| kernel_size = 3 | |
| padding = rate | |
| self.atrous_convolution = SeparableConv2d_aspp(inplanes, planes, 3, stride=1, dilation=rate, | |
| padding=padding) | |
| def forward(self, x): | |
| x = self.atrous_convolution(x) | |
| return x | |
| class ASPP_module_rate0(nn.Module): | |
| def __init__(self, inplanes, planes, rate=1): | |
| super(ASPP_module_rate0, self).__init__() | |
| if rate == 1: | |
| kernel_size = 1 | |
| padding = 0 | |
| self.atrous_convolution = nn.Conv2d(inplanes, planes, kernel_size=kernel_size, | |
| stride=1, padding=padding, dilation=rate, bias=False) | |
| self.bn = nn.BatchNorm2d(planes, eps=1e-5, affine=True) | |
| self.relu = nn.ReLU() | |
| else: | |
| raise RuntimeError() | |
| def forward(self, x): | |
| x = self.atrous_convolution(x) | |
| x = self.bn(x) | |
| return self.relu(x) | |
| class SeparableConv2d_same(nn.Module): | |
| def __init__(self, inplanes, planes, kernel_size=3, stride=1, dilation=1, bias=False, padding=0): | |
| super(SeparableConv2d_same, self).__init__() | |
| self.depthwise = nn.Conv2d(inplanes, inplanes, kernel_size, stride, padding, dilation, | |
| groups=inplanes, bias=bias) | |
| self.depthwise_bn = nn.BatchNorm2d(inplanes) | |
| self.pointwise = nn.Conv2d(inplanes, planes, 1, 1, 0, 1, 1, bias=bias) | |
| self.pointwise_bn = nn.BatchNorm2d(planes) | |
| def forward(self, x): | |
| x = fixed_padding(x, self.depthwise.kernel_size[0], rate=self.depthwise.dilation[0]) | |
| x = self.depthwise(x) | |
| x = self.depthwise_bn(x) | |
| x = self.pointwise(x) | |
| x = self.pointwise_bn(x) | |
| return x | |
| class Block(nn.Module): | |
| def __init__(self, inplanes, planes, reps, stride=1, dilation=1, 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=2, bias=False) | |
| if is_last: | |
| self.skip = nn.Conv2d(inplanes, planes, 1, stride=1, bias=False) | |
| self.skipbn = nn.BatchNorm2d(planes) | |
| else: | |
| self.skip = None | |
| self.relu = nn.ReLU(inplace=True) | |
| rep = [] | |
| filters = inplanes | |
| if grow_first: | |
| rep.append(self.relu) | |
| rep.append(SeparableConv2d_same(inplanes, planes, 3, stride=1, dilation=dilation)) | |
| # rep.append(nn.BatchNorm2d(planes)) | |
| filters = planes | |
| for i in range(reps - 1): | |
| rep.append(self.relu) | |
| rep.append(SeparableConv2d_same(filters, filters, 3, stride=1, dilation=dilation)) | |
| # rep.append(nn.BatchNorm2d(filters)) | |
| if not grow_first: | |
| rep.append(self.relu) | |
| rep.append(SeparableConv2d_same(inplanes, planes, 3, stride=1, dilation=dilation)) | |
| # rep.append(nn.BatchNorm2d(planes)) | |
| if not start_with_relu: | |
| rep = rep[1:] | |
| if stride != 1: | |
| rep.append(self.relu) | |
| rep.append(SeparableConv2d_same(planes, planes, 3, stride=2,dilation=dilation)) | |
| if is_last: | |
| rep.append(self.relu) | |
| rep.append(SeparableConv2d_same(planes, planes, 3, stride=1,dilation=dilation)) | |
| 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 | |
| # print(x.size(),skip.size()) | |
| x += skip | |
| return x | |
| class Block2(nn.Module): | |
| def __init__(self, inplanes, planes, reps, stride=1, dilation=1, start_with_relu=True, grow_first=True, is_last=False): | |
| super(Block2, self).__init__() | |
| if planes != inplanes or stride != 1: | |
| self.skip = nn.Conv2d(inplanes, planes, 1, stride=stride, bias=False) | |
| self.skipbn = nn.BatchNorm2d(planes) | |
| else: | |
| self.skip = None | |
| self.relu = nn.ReLU(inplace=True) | |
| rep = [] | |
| filters = inplanes | |
| if grow_first: | |
| rep.append(self.relu) | |
| rep.append(SeparableConv2d_same(inplanes, planes, 3, stride=1, dilation=dilation)) | |
| # rep.append(nn.BatchNorm2d(planes)) | |
| filters = planes | |
| for i in range(reps - 1): | |
| rep.append(self.relu) | |
| rep.append(SeparableConv2d_same(filters, filters, 3, stride=1, dilation=dilation)) | |
| # rep.append(nn.BatchNorm2d(filters)) | |
| if not grow_first: | |
| rep.append(self.relu) | |
| rep.append(SeparableConv2d_same(inplanes, planes, 3, stride=1, dilation=dilation)) | |
| # rep.append(nn.BatchNorm2d(planes)) | |
| if not start_with_relu: | |
| rep = rep[1:] | |
| if stride != 1: | |
| self.block2_lastconv = nn.Sequential(*[self.relu,SeparableConv2d_same(planes, planes, 3, stride=2,dilation=dilation)]) | |
| if is_last: | |
| rep.append(SeparableConv2d_same(planes, planes, 3, stride=1)) | |
| self.rep = nn.Sequential(*rep) | |
| def forward(self, inp): | |
| x = self.rep(inp) | |
| low_middle = x.clone() | |
| x1 = x | |
| x1 = self.block2_lastconv(x1) | |
| if self.skip is not None: | |
| skip = self.skip(inp) | |
| skip = self.skipbn(skip) | |
| else: | |
| skip = inp | |
| x1 += skip | |
| return x1,low_middle | |
| class Xception(nn.Module): | |
| """ | |
| Modified Alighed Xception | |
| """ | |
| def __init__(self, inplanes=3, os=16, pretrained=False): | |
| super(Xception, self).__init__() | |
| if os == 16: | |
| entry_block3_stride = 2 | |
| middle_block_rate = 1 | |
| exit_block_rates = (1, 2) | |
| elif os == 8: | |
| entry_block3_stride = 1 | |
| middle_block_rate = 2 | |
| exit_block_rates = (2, 4) | |
| else: | |
| raise NotImplementedError | |
| # Entry flow | |
| self.conv1 = nn.Conv2d(inplanes, 32, 3, stride=2, padding=1, bias=False) | |
| self.bn1 = nn.BatchNorm2d(32) | |
| self.relu = nn.ReLU(inplace=True) | |
| self.conv2 = nn.Conv2d(32, 64, 3, stride=1, padding=1, bias=False) | |
| self.bn2 = nn.BatchNorm2d(64) | |
| self.block1 = Block(64, 128, reps=2, stride=2, start_with_relu=False) | |
| self.block2 = Block2(128, 256, reps=2, stride=2, start_with_relu=True, grow_first=True) | |
| self.block3 = Block(256, 728, reps=2, stride=entry_block3_stride, start_with_relu=True, grow_first=True) | |
| # Middle flow | |
| self.block4 = Block(728, 728, reps=3, stride=1, dilation=middle_block_rate, start_with_relu=True, grow_first=True) | |
| self.block5 = Block(728, 728, reps=3, stride=1, dilation=middle_block_rate, start_with_relu=True, grow_first=True) | |
| self.block6 = Block(728, 728, reps=3, stride=1, dilation=middle_block_rate, start_with_relu=True, grow_first=True) | |
| self.block7 = Block(728, 728, reps=3, stride=1, dilation=middle_block_rate, start_with_relu=True, grow_first=True) | |
| self.block8 = Block(728, 728, reps=3, stride=1, dilation=middle_block_rate, start_with_relu=True, grow_first=True) | |
| self.block9 = Block(728, 728, reps=3, stride=1, dilation=middle_block_rate, start_with_relu=True, grow_first=True) | |
| self.block10 = Block(728, 728, reps=3, stride=1, dilation=middle_block_rate, start_with_relu=True, grow_first=True) | |
| self.block11 = Block(728, 728, reps=3, stride=1, dilation=middle_block_rate, start_with_relu=True, grow_first=True) | |
| self.block12 = Block(728, 728, reps=3, stride=1, dilation=middle_block_rate, start_with_relu=True, grow_first=True) | |
| self.block13 = Block(728, 728, reps=3, stride=1, dilation=middle_block_rate, start_with_relu=True, grow_first=True) | |
| self.block14 = Block(728, 728, reps=3, stride=1, dilation=middle_block_rate, start_with_relu=True, grow_first=True) | |
| self.block15 = Block(728, 728, reps=3, stride=1, dilation=middle_block_rate, start_with_relu=True, grow_first=True) | |
| self.block16 = Block(728, 728, reps=3, stride=1, dilation=middle_block_rate, start_with_relu=True, grow_first=True) | |
| self.block17 = Block(728, 728, reps=3, stride=1, dilation=middle_block_rate, start_with_relu=True, grow_first=True) | |
| self.block18 = Block(728, 728, reps=3, stride=1, dilation=middle_block_rate, start_with_relu=True, grow_first=True) | |
| self.block19 = Block(728, 728, reps=3, stride=1, dilation=middle_block_rate, start_with_relu=True, grow_first=True) | |
| # Exit flow | |
| self.block20 = Block(728, 1024, reps=2, stride=1, dilation=exit_block_rates[0], | |
| start_with_relu=True, grow_first=False, is_last=True) | |
| self.conv3 = SeparableConv2d_aspp(1024, 1536, 3, stride=1, dilation=exit_block_rates[1],padding=exit_block_rates[1]) | |
| # self.bn3 = nn.BatchNorm2d(1536) | |
| self.conv4 = SeparableConv2d_aspp(1536, 1536, 3, stride=1, dilation=exit_block_rates[1],padding=exit_block_rates[1]) | |
| # self.bn4 = nn.BatchNorm2d(1536) | |
| self.conv5 = SeparableConv2d_aspp(1536, 2048, 3, stride=1, dilation=exit_block_rates[1],padding=exit_block_rates[1]) | |
| # self.bn5 = nn.BatchNorm2d(2048) | |
| # Init weights | |
| # self.__init_weight() | |
| # Load pretrained model | |
| if pretrained: | |
| self.__load_xception_pretrained() | |
| def forward(self, x): | |
| # Entry flow | |
| x = self.conv1(x) | |
| x = self.bn1(x) | |
| x = self.relu(x) | |
| # print('conv1 ',x.size()) | |
| x = self.conv2(x) | |
| x = self.bn2(x) | |
| x = self.relu(x) | |
| x = self.block1(x) | |
| # print('block1',x.size()) | |
| # low_level_feat = x | |
| x,low_level_feat = self.block2(x) | |
| # print('block2',x.size()) | |
| x = self.block3(x) | |
| # print('xception block3 ',x.size()) | |
| # 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.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)) | |
| torch.nn.init.kaiming_normal_(m.weight) | |
| elif isinstance(m, nn.BatchNorm2d): | |
| m.weight.data.fill_(1) | |
| m.bias.data.zero_() | |
| def __load_xception_pretrained(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('block12'): | |
| model_dict[k.replace('block12', 'block20')] = v | |
| elif k.startswith('block11'): | |
| 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('conv3'): | |
| model_dict[k] = 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) | |
| class DeepLabv3_plus(nn.Module): | |
| def __init__(self, nInputChannels=3, n_classes=21, os=16, pretrained=False, _print=True): | |
| if _print: | |
| print("Constructing DeepLabv3+ model...") | |
| print("Number of classes: {}".format(n_classes)) | |
| print("Output stride: {}".format(os)) | |
| print("Number of Input Channels: {}".format(nInputChannels)) | |
| super(DeepLabv3_plus, self).__init__() | |
| # Atrous Conv | |
| self.xception_features = Xception(nInputChannels, os, pretrained) | |
| # ASPP | |
| if os == 16: | |
| rates = [1, 6, 12, 18] | |
| elif os == 8: | |
| rates = [1, 12, 24, 36] | |
| raise NotImplementedError | |
| else: | |
| raise NotImplementedError | |
| self.aspp1 = ASPP_module_rate0(2048, 256, rate=rates[0]) | |
| self.aspp2 = ASPP_module(2048, 256, rate=rates[1]) | |
| self.aspp3 = ASPP_module(2048, 256, rate=rates[2]) | |
| self.aspp4 = ASPP_module(2048, 256, rate=rates[3]) | |
| self.relu = nn.ReLU() | |
| self.global_avg_pool = nn.Sequential(nn.AdaptiveAvgPool2d((1, 1)), | |
| nn.Conv2d(2048, 256, 1, stride=1, bias=False), | |
| nn.BatchNorm2d(256), | |
| nn.ReLU() | |
| ) | |
| self.concat_projection_conv1 = nn.Conv2d(1280, 256, 1, bias=False) | |
| self.concat_projection_bn1 = nn.BatchNorm2d(256) | |
| # adopt [1x1, 48] for channel reduction. | |
| self.feature_projection_conv1 = nn.Conv2d(256, 48, 1, bias=False) | |
| self.feature_projection_bn1 = nn.BatchNorm2d(48) | |
| self.decoder = nn.Sequential(Decoder_module(304, 256), | |
| Decoder_module(256, 256) | |
| ) | |
| self.semantic = nn.Conv2d(256, n_classes, kernel_size=1, stride=1) | |
| def forward(self, input): | |
| x, low_level_features = self.xception_features(input) | |
| # print(x.size()) | |
| x1 = self.aspp1(x) | |
| x2 = self.aspp2(x) | |
| x3 = self.aspp3(x) | |
| x4 = self.aspp4(x) | |
| x5 = self.global_avg_pool(x) | |
| x5 = F.upsample(x5, size=x4.size()[2:], mode='bilinear', align_corners=True) | |
| x = torch.cat((x1, x2, x3, x4, x5), dim=1) | |
| x = self.concat_projection_conv1(x) | |
| x = self.concat_projection_bn1(x) | |
| x = self.relu(x) | |
| # print(x.size()) | |
| low_level_features = self.feature_projection_conv1(low_level_features) | |
| low_level_features = self.feature_projection_bn1(low_level_features) | |
| low_level_features = self.relu(low_level_features) | |
| x = F.upsample(x, size=low_level_features.size()[2:], mode='bilinear', align_corners=True) | |
| # print(low_level_features.size()) | |
| # print(x.size()) | |
| x = torch.cat((x, low_level_features), dim=1) | |
| x = self.decoder(x) | |
| x = self.semantic(x) | |
| x = F.upsample(x, size=input.size()[2:], mode='bilinear', align_corners=True) | |
| return x | |
| def freeze_bn(self): | |
| for m in self.xception_features.modules(): | |
| if isinstance(m, nn.BatchNorm2d): | |
| m.eval() | |
| def freeze_totally_bn(self): | |
| for m in self.modules(): | |
| if isinstance(m, nn.BatchNorm2d): | |
| m.eval() | |
| def freeze_aspp_bn(self): | |
| for m in self.aspp1.modules(): | |
| if isinstance(m, nn.BatchNorm2d): | |
| m.eval() | |
| for m in self.aspp2.modules(): | |
| if isinstance(m, nn.BatchNorm2d): | |
| m.eval() | |
| for m in self.aspp3.modules(): | |
| if isinstance(m, nn.BatchNorm2d): | |
| m.eval() | |
| for m in self.aspp4.modules(): | |
| if isinstance(m, nn.BatchNorm2d): | |
| m.eval() | |
| def learnable_parameters(self): | |
| layer_features_BN = [] | |
| layer_features = [] | |
| layer_aspp = [] | |
| layer_projection =[] | |
| layer_decoder = [] | |
| layer_other = [] | |
| model_para = list(self.named_parameters()) | |
| for name,para in model_para: | |
| if 'xception' in name: | |
| if 'bn' in name or 'downsample.1.weight' in name or 'downsample.1.bias' in name: | |
| layer_features_BN.append(para) | |
| else: | |
| layer_features.append(para) | |
| # print (name) | |
| elif 'aspp' in name: | |
| layer_aspp.append(para) | |
| elif 'projection' in name: | |
| layer_projection.append(para) | |
| elif 'decode' in name: | |
| layer_decoder.append(para) | |
| elif 'global' not in name: | |
| layer_other.append(para) | |
| return layer_features_BN,layer_features,layer_aspp,layer_projection,layer_decoder,layer_other | |
| def get_backbone_para(self): | |
| layer_features = [] | |
| other_features = [] | |
| model_para = list(self.named_parameters()) | |
| for name, para in model_para: | |
| if 'xception' in name: | |
| layer_features.append(para) | |
| else: | |
| other_features.append(para) | |
| return layer_features, other_features | |
| def train_fixbn(self, mode=True, freeze_bn=True, freeze_bn_affine=False): | |
| r"""Sets the module in training mode. | |
| This has any effect only on certain modules. See documentations of | |
| particular modules for details of their behaviors in training/evaluation | |
| mode, if they are affected, e.g. :class:`Dropout`, :class:`BatchNorm`, | |
| etc. | |
| Returns: | |
| Module: self | |
| """ | |
| super(DeepLabv3_plus, self).train(mode) | |
| if freeze_bn: | |
| print("Freezing Mean/Var of BatchNorm2D.") | |
| if freeze_bn_affine: | |
| print("Freezing Weight/Bias of BatchNorm2D.") | |
| if freeze_bn: | |
| for m in self.xception_features.modules(): | |
| if isinstance(m, nn.BatchNorm2d): | |
| m.eval() | |
| if freeze_bn_affine: | |
| m.weight.requires_grad = False | |
| m.bias.requires_grad = False | |
| # for m in self.aspp1.modules(): | |
| # if isinstance(m, nn.BatchNorm2d): | |
| # m.eval() | |
| # if freeze_bn_affine: | |
| # m.weight.requires_grad = False | |
| # m.bias.requires_grad = False | |
| # for m in self.aspp2.modules(): | |
| # if isinstance(m, nn.BatchNorm2d): | |
| # m.eval() | |
| # if freeze_bn_affine: | |
| # m.weight.requires_grad = False | |
| # m.bias.requires_grad = False | |
| # for m in self.aspp3.modules(): | |
| # if isinstance(m, nn.BatchNorm2d): | |
| # m.eval() | |
| # if freeze_bn_affine: | |
| # m.weight.requires_grad = False | |
| # m.bias.requires_grad = False | |
| # for m in self.aspp4.modules(): | |
| # if isinstance(m, nn.BatchNorm2d): | |
| # m.eval() | |
| # if freeze_bn_affine: | |
| # m.weight.requires_grad = False | |
| # m.bias.requires_grad = False | |
| # for m in self.global_avg_pool.modules(): | |
| # if isinstance(m, nn.BatchNorm2d): | |
| # m.eval() | |
| # if freeze_bn_affine: | |
| # m.weight.requires_grad = False | |
| # m.bias.requires_grad = False | |
| # for m in self.concat_projection_bn1.modules(): | |
| # if isinstance(m, nn.BatchNorm2d): | |
| # m.eval() | |
| # if freeze_bn_affine: | |
| # m.weight.requires_grad = False | |
| # m.bias.requires_grad = False | |
| # for m in self.feature_projection_bn1.modules(): | |
| # if isinstance(m, nn.BatchNorm2d): | |
| # m.eval() | |
| # if freeze_bn_affine: | |
| # m.weight.requires_grad = False | |
| # m.bias.requires_grad = False | |
| 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, nn.BatchNorm2d): | |
| m.weight.data.fill_(1) | |
| m.bias.data.zero_() | |
| def load_state_dict_new(self, state_dict): | |
| own_state = self.state_dict() | |
| #for name inshop_cos own_state: | |
| # print name | |
| new_state_dict = OrderedDict() | |
| for name, param in state_dict.items(): | |
| name = name.replace('module.','') | |
| new_state_dict[name] = 0 | |
| if name not in own_state: | |
| if 'num_batch' in name: | |
| continue | |
| print ('unexpected key "{}" in state_dict' | |
| .format(name)) | |
| continue | |
| # if isinstance(param, own_state): | |
| if isinstance(param, Parameter): | |
| # backwards compatibility for serialized parameters | |
| param = param.data | |
| try: | |
| own_state[name].copy_(param) | |
| except: | |
| print('While copying the parameter named {}, whose dimensions in the model are' | |
| ' {} and whose dimensions in the checkpoint are {}, ...'.format( | |
| name, own_state[name].size(), param.size())) | |
| continue # i add inshop_cos 2018/02/01 | |
| # raise | |
| # print 'copying %s' %name | |
| # if isinstance(param, own_state): | |
| # backwards compatibility for serialized parameters | |
| own_state[name].copy_(param) | |
| # print 'copying %s' %name | |
| missing = set(own_state.keys()) - set(new_state_dict.keys()) | |
| if len(missing) > 0: | |
| print('missing keys in state_dict: "{}"'.format(missing)) | |
| def get_1x_lr_params(model): | |
| """ | |
| This generator returns all the parameters of the net except for | |
| the last classification layer. Note that for each batchnorm layer, | |
| requires_grad is set to False in deeplab_resnet.py, therefore this function does not return | |
| any batchnorm parameter | |
| """ | |
| b = [model.xception_features] | |
| for i in range(len(b)): | |
| for k in b[i].parameters(): | |
| if k.requires_grad: | |
| yield k | |
| def get_10x_lr_params(model): | |
| """ | |
| This generator returns all the parameters for the last layer of the net, | |
| which does the classification of pixel into classes | |
| """ | |
| b = [model.aspp1, model.aspp2, model.aspp3, model.aspp4, model.conv1, model.conv2, model.last_conv] | |
| for j in range(len(b)): | |
| for k in b[j].parameters(): | |
| if k.requires_grad: | |
| yield k | |
| if __name__ == "__main__": | |
| model = DeepLabv3_plus(nInputChannels=3, n_classes=21, os=16, pretrained=False, _print=True) | |
| model.eval() | |
| image = torch.randn(1, 3, 512, 512)*255 | |
| with torch.no_grad(): | |
| output = model.forward(image) | |
| print(output.size()) | |
| # print(output) | |