Merge pull request #1 from qubvel/feature-linknet

Linknet model

Author: qubvel

segmentation_models_pytorch/__init__.py

@@ -1,3 +1,4 @@
 from .unet import Unet
+from .linknet import Linknet
 
 from . import encoders

segmentation_models_pytorch/linknet/__init__.py

@@ -0,0 +1 @@
+from .model import Linknet

segmentation_models_pytorch/linknet/decoder.py

@@ -0,0 +1,75 @@
+import torch.nn as nn
+
+from ..common.blocks import Conv2dReLU
+from ..base.model import Model
+
+
+class TransposeX2(nn.Module):
+
+    def __init__(self, in_channels, out_channels, use_batchnorm=True, **batchnorm_params):
+        super().__init__()
+        layers = [
+            nn.ConvTranspose2d(in_channels, out_channels, kernel_size=4, stride=2, padding=1),
+            nn.ReLU(inplace=True),
+        ]
+        if use_batchnorm:
+            layers.insert(1, nn.BatchNorm2d(out_channels, **batchnorm_params))
+
+        self.block = nn.Sequential(*layers)
+
+    def forward(self, x):
+        return self.block(x)
+
+
+class DecoderBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, use_batchnorm=True):
+        super().__init__()
+
+        self.block = nn.Sequential(
+            Conv2dReLU(in_channels, in_channels // 4, kernel_size=1, use_batchnorm=use_batchnorm),
+            TransposeX2(in_channels // 4, in_channels // 4, use_batchnorm=use_batchnorm),
+            Conv2dReLU(in_channels // 4, out_channels, kernel_size=1, use_batchnorm=use_batchnorm),
+        )
+
+    def forward(self, x):
+        x, skip = x
+        x = self.block(x)
+        if skip is not None:
+            x += skip
+        return x
+
+
+class LinknetDecoder(Model):
+
+    def __init__(
+            self,
+            encoder_channels,
+            prefinal_channels=32,
+            final_channels=1,
+            use_batchnorm=True,
+    ):
+        super().__init__()
+
+        in_channels = encoder_channels
+
+        self.layer1 = DecoderBlock(in_channels[0], in_channels[1], use_batchnorm=use_batchnorm)
+        self.layer2 = DecoderBlock(in_channels[1], in_channels[2], use_batchnorm=use_batchnorm)
+        self.layer3 = DecoderBlock(in_channels[2], in_channels[3], use_batchnorm=use_batchnorm)
+        self.layer4 = DecoderBlock(in_channels[3], in_channels[4], use_batchnorm=use_batchnorm)
+        self.layer5 = DecoderBlock(in_channels[4], prefinal_channels, use_batchnorm=use_batchnorm)
+        self.final_conv = nn.Conv2d(prefinal_channels, final_channels, kernel_size=(1, 1))
+
+        self.initialize()
+
+    def forward(self, x):
+        encoder_head = x[0]
+        skips = x[1:]
+
+        x = self.layer1([encoder_head, skips[0]])
+        x = self.layer2([x, skips[1]])
+        x = self.layer3([x, skips[2]])
+        x = self.layer4([x, skips[3]])
+        x = self.layer5([x, None])
+        x = self.final_conv(x)
+
+        return x

segmentation_models_pytorch/linknet/model.py

@@ -0,0 +1,31 @@
+from .decoder import LinknetDecoder
+from ..base import EncoderDecoder
+from ..encoders import get_encoder
+
+
+class Linknet(EncoderDecoder):
+
+    def __init__(
+            self,
+            encoder_name='resnet34',
+            encoder_weights='imagenet',
+            decoder_use_batchnorm=True,
+            classes=1,
+            activation='sigmoid',
+    ):
+
+        encoder = get_encoder(
+            encoder_name,
+            encoder_weights=encoder_weights
+        )
+
+        decoder = LinknetDecoder(
+            encoder_channels=encoder.out_shapes,
+            prefinal_channels=32,
+            final_channels=classes,
+            use_batchnorm=decoder_use_batchnorm,
+        )
+
+        super().__init__(encoder, decoder, activation)
+
+        self.name = 'link-{}'.format(encoder_name)