Feature Xception encoder (#108)

* Add Xception encoder (#102)

Author: qubvel

Diff for: README.md

@@ -7,7 +7,7 @@ The main features of this library are:
 
  - High level API (just two lines to create neural network)
  - 4 models architectures for binary and multi class segmentation (including legendary Unet)
- - 45 available encoders for each architecture
+ - 46 available encoders for each architecture
  - All encoders have pre-trained weights for faster and better convergence
 
 ### Table of content
@@ -111,6 +111,7 @@ preprocess_input = get_preprocessing_fn('resnet18', pretrained='imagenet')
 |efficientnet-b6                 |imagenet                        |40M                             |
 |efficientnet-b7                 |imagenet                        |63M                             |
 |mobilenet_v2                    |imagenet                        |2M                              |
+|xception                        |imagenet                        |22M                             |
 
 ### Models API <a name="api"></a>
 

Diff for: segmentation_models_pytorch/encoders/__init__.py

@@ -10,6 +10,7 @@
 from .inceptionv4 import inceptionv4_encoders
 from .efficientnet import efficient_net_encoders
 from .mobilenet import mobilenet_encoders
+from .xception import xception_encoders
 
 
 from ._preprocessing import preprocess_input
@@ -24,6 +25,7 @@
 encoders.update(inceptionv4_encoders)
 encoders.update(efficient_net_encoders)
 encoders.update(mobilenet_encoders)
+encoders.update(xception_encoders)
 
 
 def get_encoder(name, in_channels=3, depth=5, weights=None):

Diff for: segmentation_models_pytorch/encoders/xception.py

@@ -0,0 +1,105 @@
+import re
+import torch.nn as nn
+
+from pretrainedmodels.models.xception import pretrained_settings
+from pretrainedmodels.models.xception import Xception
+
+from ._base import EncoderMixin
+
+class XceptionEncoder(Xception, EncoderMixin):
+
+    def __init__(self, out_channels, *args, depth=5, **kwargs):
+        super().__init__(*args, **kwargs)
+        
+        self._out_channels = out_channels
+        self._depth = depth
+        self._in_channels = 3
+        
+        #modify padding to maintain output shape
+        self.conv1.padding = 1
+        self.conv2.padding = 1
+        
+        del self.fc
+
+    @staticmethod
+    def _transition(x, transition_block):
+        for module in transition_block:
+            x = module(x)
+            if isinstance(module, nn.ReLU):
+                skip = x
+        return x, skip
+
+    def forward(self, x):
+        features = [x]
+        
+        if self._depth > 0:
+            x = self.conv1(x)
+            x = self.bn1(x)
+            x = self.relu(x)
+
+            x = self.conv2(x)
+            x = self.bn2(x)
+            x0 = self.relu(x)
+            features.append(x0)
+            
+        if self._depth > 1:
+            x1 = self.block1(x0)
+            features.append(x1)
+            
+        if self._depth > 2:
+            x2 = self.block2(x1)
+            features.append(x2)
+        
+        if self._depth > 3:
+            x = self.block3(x2)
+            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)
+            x3 = self.block11(x)
+            features.append(x3)
+            
+        if self._depth > 4:
+            x = self.block12(x)
+
+            x = self.conv3(x)
+            x = self.bn3(x)
+            x = self.relu(x)
+
+            x = self.conv4(x)
+            x4 = self.bn4(x)
+            features.append(x4)
+            
+        return features
+       
+
+    def load_state_dict(self, state_dict):
+        pattern = re.compile(
+            r'^(.*denselayer\d+\.(?:norm|relu|conv))\.((?:[12])\.(?:weight|bias|running_mean|running_var))$')
+        for key in list(state_dict.keys()):
+            res = pattern.match(key)
+            if res:
+                new_key = res.group(1) + res.group(2)
+                state_dict[new_key] = state_dict[key]
+                del state_dict[key]
+
+        # remove linear
+        state_dict.pop('fc.bias')
+        state_dict.pop('fc.weight')
+
+        super().load_state_dict(state_dict)
+
+
+xception_encoders = {
+    'xception': {
+        'encoder': XceptionEncoder,
+        'pretrained_settings': pretrained_settings['xception'],
+        'params': {
+            'out_channels': (3, 64, 128, 256, 728, 2048),
+        }
+    },
+}
+