Merge pull request #16 from MathGaron/feature/modules

add common modules

Author: MathGaron

examples/classification/cat_dog_net.py

@@ -1,43 +1,33 @@
 import torch.nn.functional as F
 import torch.nn as nn
-from pytorch_toolbox.network.network_base import NetworkBase
+from pytorch_toolbox.network_base import NetworkBase
+from pytorch_toolbox.modules.conv2d_module import ConvBlock
+from pytorch_toolbox.modules.fc_module import FCBlock
 
 
 class CatDogNet(NetworkBase):
     def __init__(self):
         super(CatDogNet, self).__init__()
-        self.conv1 = nn.Conv2d(3, 24, 5)
-        self.conv1_bn = nn.BatchNorm2d(24)
-        self.dropout1 = nn.Dropout2d(0.25)
-        self.conv2 = nn.Conv2d(24, 48, 3)
-        self.conv2_bn = nn.BatchNorm2d(48)
-        self.dropout2 = nn.Dropout2d(0.25)
-        self.conv3 = nn.Conv2d(48, 48, 3)
-        self.conv3_bn = nn.BatchNorm2d(48)
-        self.dropout3 = nn.Dropout2d(0.25)
-        self.conv4 = nn.Conv2d(48, 96, 3)
-        self.conv4_bn = nn.BatchNorm2d(96)
+        self.conv1 = ConvBlock(3, 24, 5, dropout=True, batchnorm=True, maxpool=True, activation=F.elu)
+        self.conv2 = ConvBlock(24, 48, 3, dropout=True, batchnorm=True, maxpool=True, activation=F.elu)
+        self.conv3 = ConvBlock(48, 48, 3, dropout=True, batchnorm=True, maxpool=True, activation=F.elu)
+        self.conv4 = ConvBlock(48, 48, 3, dropout=True, batchnorm=True, maxpool=True, activation=F.elu)
 
         self.view_size = 96 * 6 * 6
-        self.fc1 = nn.Linear(self.view_size, 250)
-        self.fc_bn1 = nn.BatchNorm1d(250)
-        self.fc2 = nn.Linear(250, 2)
 
-        self.dropout1 = nn.Dropout()
-        self.dropout2 = nn.Dropout()
+        self.fc1 = FCBlock(self.view_size, 250, dropout=True, batchnorm=True, activation=F.elu)
+        self.fc2 = FCBlock(250, 2, dropout=False, batchnorm=False, activation=None)
 
         self.criterion = nn.NLLLoss()
 
     def forward(self, x):
-        x = self.dropout1(F.max_pool2d(F.elu(self.conv1(x)), 2))
-        x = self.dropout2(F.max_pool2d(F.elu(self.conv2_bn(self.conv2(x))), 2))
-        x = self.dropout3(F.max_pool2d(F.elu(self.conv3_bn(self.conv3(x))), 2))
-        x = F.max_pool2d(F.elu(self.conv4_bn(self.conv4(x))), 2)
-        x = x.view(-1, self.view_size)
-        x = self.dropout1(x)
-        x = F.elu(self.fc_bn1(self.fc1(x)))
-        x = self.dropout2(x)
-        x = F.log_softmax(self.fc2(x))
+        x = self.conv1(x)
+        x = self.conv2(x)
+        x = self.conv3(x)
+        x = self.conv4(x)
+        x = self.fc1(x)
+        x = self.fc2(x)
+        x = F.log_softmax(x)
         return x
 
     def loss(self, predictions, targets):

pytorch_toolbox/modules/__init__.py

@@ -0,0 +1,45 @@
+import torch.nn as nn
+import torch
+
+
+class ConvBlock(nn.Module):
+    """
+    Generic convolution block
+
+    """
+    def __init__(self, in_channels, out_channels, kernel_size, stride=1,
+                 padding=0, dropout=False, dropout_rate=0.25,
+                 batchnorm=True, maxpool=True, maxpool_size=2, activation=None):
+        """
+
+        :param in_channels:
+        :param out_channels:
+        :param kernel_size:
+        :param stride:
+        :param padding:
+        :param dropout:
+        :param dropout_rate:
+        :param batchnorm:
+        :param maxpool:
+        :param maxpool_size:
+        :param activation:      Function from function api
+        """
+        super(ConvBlock, self).__init__()
+        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride=stride, padding=padding)
+        self.batch_norm = nn.BatchNorm2d(out_channels) if batchnorm else None
+        self.dropout = nn.Dropout2d(dropout_rate) if dropout else None
+        self.maxpool = nn.MaxPool2d(maxpool_size) if maxpool else None
+        self.activation = activation
+
+    def forward(self, x):
+
+        x = self.conv(x)
+        if self.batch_norm:
+            x = self.batch_norm(x)
+        if self.activation:
+            x = self.activation(x)
+        if self.maxpool:
+            x = self.maxpool(x)
+        if self.dropout:
+            x = self.dropout(x)
+        return x

pytorch_toolbox/modules/conv2d_module.py

@@ -0,0 +1,35 @@
+import torch.nn as nn
+import torch
+
+
+class FCBlock(nn.Module):
+    """
+        generic Fully connected block
+    """
+    def __init__(self, in_features, out_features, dropout=True, dropout_rate=0.5,
+                 batchnorm=True, activation=None):
+        """
+
+        :param in_features:
+        :param out_features:
+        :param dropout:
+        :param dropout_rate:
+        :param batchnorm:
+        :param activation:      Function from function api
+        """
+        super(FCBlock, self).__init__()
+        self.fc = nn.Linear(in_features, out_features)
+        self.batch_norm = nn.BatchNorm1d(out_features) if batchnorm else None
+        self.dropout = nn.Dropout(dropout_rate) if dropout else None
+        self.activation = activation
+
+    def forward(self, x):
+
+        x = self.fc(x)
+        if self.batch_norm:
+            x = self.batch_norm(x)
+        if self.activation:
+            x = self.activation(x)
+        if self.dropout:
+            x = self.dropout(x)
+        return x

pytorch_toolbox/modules/fc_module.py

@@ -0,0 +1,25 @@
+import torch.nn as nn
+import torch
+
+
+class Fire(nn.Module):
+    """
+    From SqueezeNet : https://github.com/pytorch/vision/blob/master/torchvision/models/squeezenet.py
+
+    """
+    def __init__(self, inplanes, squeeze_planes, expand1x1_planes, expand3x3_planes):
+        super(Fire, self).__init__()
+        self.inplanes = inplanes
+        self.squeeze = nn.Conv2d(inplanes, squeeze_planes, kernel_size=1)
+        self.squeeze_activation = nn.ELU(inplace=True)
+        self.expand1x1 = nn.Conv2d(squeeze_planes, expand1x1_planes, kernel_size=1)
+        self.expand1x1_activation = nn.ELU(inplace=True)
+        self.expand3x3 = nn.Conv2d(squeeze_planes, expand3x3_planes, kernel_size=3, padding=1)
+        self.expand3x3_activation = nn.ELU(inplace=True)
+
+    def forward(self, x):
+        x = self.squeeze_activation(self.squeeze(x))
+        return torch.cat([
+            self.expand1x1_activation(self.expand1x1(x)),
+            self.expand3x3_activation(self.expand3x3(x))
+        ], 1)

pytorch_toolbox/modules/fire_module.py

@@ -0,0 +1,22 @@
+from torch import nn
+
+
+class SELayer(nn.Module):
+    """
+    From Squeeze-and-Excitation Networks : https://github.com/moskomule/senet.pytorch/blob/master/se_module.py
+    """
+    def __init__(self, channel, reduction=16):
+        super(SELayer, self).__init__()
+        self.avg_pool = nn.AdaptiveAvgPool2d(1)
+        self.fc = nn.Sequential(
+                nn.Linear(channel, reduction),
+                nn.ReLU(inplace=True),
+                nn.Linear(reduction, channel),
+                nn.Sigmoid()
+        )
+
+    def forward(self, x):
+        b, c, _, _ = x.size()
+        y = self.avg_pool(x).view(b, c)
+        y = self.fc(y).view(b, c, 1, 1)
+        return x * y