| |
| |
| |
| |
|
|
| 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 import init |
|
|
| __all__ = ['xception'] |
|
|
|
|
| BN = None |
| class SeparableConv2d(nn.Module): |
| def __init__(self,in_channels,out_channels,kernel_size=1,stride=1,padding=0,dilation=1,bias=False): |
| super(SeparableConv2d,self).__init__() |
|
|
| self.conv1 = nn.Conv2d(in_channels,in_channels,kernel_size,stride,padding,dilation,groups=in_channels,bias=bias) |
| self.pointwise = nn.Conv2d(in_channels,out_channels,1,1,0,1,1,bias=bias) |
|
|
| def forward(self,x): |
| x = self.conv1(x) |
| x = self.pointwise(x) |
| return x |
|
|
|
|
| class Block(nn.Module): |
| def __init__(self,in_filters,out_filters,reps,strides=1,start_with_relu=True,grow_first=True): |
| super(Block, self).__init__() |
|
|
| if out_filters != in_filters or strides!=1: |
| self.skip = nn.Conv2d(in_filters,out_filters,1,stride=strides, bias=False) |
| self.skipbn = BN(out_filters) |
| else: |
| self.skip=None |
|
|
| self.relu = nn.ReLU(inplace=True) |
| rep=[] |
|
|
| filters=in_filters |
| if grow_first: |
| rep.append(self.relu) |
| rep.append(SeparableConv2d(in_filters,out_filters,3,stride=1,padding=1,bias=False)) |
| rep.append(BN(out_filters)) |
| filters = out_filters |
|
|
| for i in range(reps-1): |
| rep.append(self.relu) |
| rep.append(SeparableConv2d(filters,filters,3,stride=1,padding=1,bias=False)) |
| rep.append(BN(filters)) |
|
|
| if not grow_first: |
| rep.append(self.relu) |
| rep.append(SeparableConv2d(in_filters,out_filters,3,stride=1,padding=1,bias=False)) |
| rep.append(BN(out_filters)) |
|
|
| if not start_with_relu: |
| rep = rep[1:] |
| else: |
| rep[0] = nn.ReLU(inplace=False) |
|
|
| if strides != 1: |
| rep.append(nn.MaxPool2d(3,strides,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+=skip |
| return x |
|
|
|
|
| class Xception(nn.Module): |
| """ |
| Xception optimized for the ImageNet dataset, as specified in |
| https://arxiv.org/pdf/1610.02357.pdf |
| """ |
| def __init__(self, in_channels = 3, num_classes=1000, bn_group_size=1, |
| bn_group=None, bn_sync_stats=True,feature_visible=False, |
| dropout=0, return_feature_idx=None, bypass_last_bn=False, **kwargs): |
| """ Constructor |
| Args: |
| num_classes: number of classes |
| """ |
| global BN |
|
|
| BN = nn.BatchNorm2d |
|
|
| bypass_bn_weight_list = [] |
| self.inplanes = 64 |
|
|
| super(Xception, self).__init__() |
| self.num_classes = num_classes |
| self.return_feature_idx = return_feature_idx |
| self.feature_visible = feature_visible |
|
|
| self.conv1 = nn.Conv2d(in_channels, 32, 3,2, 0, bias=False) |
| self.bn1 = BN(32) |
| self.relu = nn.ReLU(inplace=True) |
|
|
| self.conv2 = nn.Conv2d(32,64,3,bias=False) |
| self.bn2 = BN(64) |
| |
|
|
| self.block1=Block(64,128,2,2,start_with_relu=False,grow_first=True) |
| self.block2=Block(128,256,2,2,start_with_relu=True,grow_first=True) |
| self.block3=Block(256,728,2,2,start_with_relu=True,grow_first=True) |
|
|
| self.block4=Block(728,728,3,1,start_with_relu=True,grow_first=True) |
| self.block5=Block(728,728,3,1,start_with_relu=True,grow_first=True) |
| self.block6=Block(728,728,3,1,start_with_relu=True,grow_first=True) |
| self.block7=Block(728,728,3,1,start_with_relu=True,grow_first=True) |
|
|
| self.block8=Block(728,728,3,1,start_with_relu=True,grow_first=True) |
| self.block9=Block(728,728,3,1,start_with_relu=True,grow_first=True) |
| self.block10=Block(728,728,3,1,start_with_relu=True,grow_first=True) |
| self.block11=Block(728,728,3,1,start_with_relu=True,grow_first=True) |
|
|
| self.block12=Block(728,1024,2,2,start_with_relu=True,grow_first=False) |
|
|
| self.conv3 = SeparableConv2d(1024,1536,3,1,1) |
| self.bn3 = BN(1536) |
|
|
| |
| self.conv4 = SeparableConv2d(1536,2048,3,1,1) |
| self.bn4 = BN(2048) |
|
|
| self.fc = nn.Linear(2048, num_classes) |
| self.drop = None |
| if dropout > 0: |
| self.drop = nn.Dropout(p=dropout) |
|
|
| 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, torch.nn.SyncBatchNorm) |
| or isinstance(m, nn.BatchNorm2d)): |
| m.weight.data.fill_(1) |
| m.bias.data.zero_() |
|
|
| if bypass_last_bn: |
| for param in bypass_bn_weight_list: |
| param.data.zero_() |
| print('bypass {} bn.weight in BottleneckBlocks'.format(len(bypass_bn_weight_list))) |
|
|
| def att_feature(self, feature): |
| sum_feature = F.relu(torch.sum(feature, dim=1)) |
| sum_feature = sum_feature / (torch.max(sum_feature)+ 1e-9) |
| return sum_feature |
|
|
| def features(self, input): |
| features = [] |
| x = self.conv1(input) |
| x = self.bn1(x) |
| x = self.relu(x) |
|
|
| x = self.conv2(x) |
| x = self.bn2(x) |
| x = self.relu(x) |
| features.append(x) |
|
|
| x_b1 = self.block1(x) |
| features.append(x_b1) |
| x_b2 = self.block2(x_b1) |
| features.append(x_b2) |
| x_b3 = self.block3(x_b2) |
| features.append(x_b3) |
| x_b4 = self.block4(x_b3) |
| features.append(x_b4) |
| x_b5 = self.block5(x_b4) |
| features.append(x_b5) |
| x_b6 = self.block6(x_b5) |
| features.append(x_b6) |
| x_b7 = self.block7(x_b6) |
| features.append(x_b7) |
| x_b8 = self.block8(x_b7) |
| features.append(x_b8) |
| x_b9 = self.block9(x_b8) |
| features.append(x_b9) |
| x_b10 = self.block10(x_b9) |
| features.append(x_b10) |
| x_b11 = self.block11(x_b10) |
| features.append(x_b11) |
| x_b12 = self.block12(x_b11) |
| features.append(x_b12) |
|
|
| x = self.conv3(x_b12) |
| x = self.bn3(x) |
| x = self.relu(x) |
|
|
| x = self.conv4(x) |
| x = self.bn4(x) |
| return x, features |
|
|
| def logits(self, features): |
| x = self.relu(features) |
|
|
| x = F.adaptive_avg_pool2d(x, (1, 1)) |
| x = x.view(x.size(0), -1) |
| if self.drop is not None: |
| x = self.drop(x) |
| out = self.fc(x) |
| return out, x |
|
|
| def forward(self, input): |
| x, features = self.features(input) |
| return x |
| ''' |
| logit, embedding = self.logits(x) |
| features.append(embedding) |
| selected_feature = None |
| if self.return_feature_idx is not None: |
| selected_feature = [features[i] for i in self.return_feature_idx] |
| if self.feature_visible: |
| selected_feature = [self.att_feature(feature) for feature in selected_feature] |
| return logit, selected_feature |
| else: |
| return logit |
| ''' |
|
|
| def get_model_size(model): |
| result = 0 |
| for key,value in model.state_dict().items(): |
| s = 1 |
| for item in value.size(): |
| s *= item |
| result += s |
| print(key) |
| result *= 4 |
| return result |
|
|
| def xception(pretrain_path=None, **kwargs): |
| model = Xception(**kwargs) |
| if pretrain_path != None: |
| state_dict = torch.load(pretrain_path) |
| if 'state_dict' in state_dict.keys(): |
| state_dict = torch.load(pretrain_path) |
|
|
| ''' |
| for name, weights in state_dict.items(): |
| if 'pointwise' in name: |
| print("test") |
| state_dict[name] = weights.unsqueeze(-1).unsqueeze(-1) |
| ''' |
| own_state = model.state_dict() |
| |
| for name, param in state_dict.items(): |
| name = name.replace("module.", "") |
| if name in own_state: |
| if isinstance(param, torch.nn.Parameter): |
| |
| 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())) |
| |
| print("Features Extractor checkpoint loaded.") |
| return model |
|
|
