[extractors] removed the network architecture, it is now imported from bob.learn.pytorch

bob/ip/pytorch_extractor/CNN8.py

 import numpy
-import sys
 
 import torch
-import torch.nn as nn
-import torch.nn.functional as F
 from torch.autograd import Variable
 
 import torchvision.transforms as transforms
 
-import bob.core
-logger = bob.core.log.setup("bob.learn.pytorch")
-
-
-CNN8_CONFIG = [64, 64, 'M', 128, 128, 'M', 256, 256, 'M', 512, 512, 'M']
-
-def make_conv_layers(cfg, input_c = 3):
-    layers = []
-    in_channels = input_c
-    for v in cfg:
-        if v == 'M':
-            layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
-        else:
-            conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1)
-            layers += [conv2d, nn.ReLU()]
-            in_channels = v
-    return nn.Sequential(*layers)
-
-class CNN8(nn.Module):
-
-    def __init__(self, num_cls, drop_rate=0.5):
-        super(CNN8, self).__init__()
-        self.num_classes = num_cls
-        self.drop_rate = float(drop_rate)
-        self.conv = make_conv_layers(CNN8_CONFIG)
-        self.avgpool = nn.AvgPool2d(8)
-        self.classifier = nn.Linear(512, self.num_classes)
-
-    def forward(self, x):
-        x = self.conv(x)
-        x = self.avgpool(x)
-        x = x.view(x.size(0), -1)
-        x = F.dropout(x, p = self.drop_rate, training=self.training)
-        return x 
-
-
+from bob.learn.pytorch.architectures import CNN8
 from bob.bio.base.extractor import Extractor
 
 class CNN8Extractor(Extractor):
@@ -55,7 +17,8 @@ class CNN8Extractor(Extractor):
     # model
     self.network = CNN8(num_classes)
     if model_file is None:
-      logger.warning("No model file provided, building network from scratch !")
+      # do nothing (used mainly for unit testing) 
+      pass
     else:
       cp = torch.load(model_file)
       if 'state_dict' in cp:
@@ -85,7 +48,7 @@ class CNN8Extractor(Extractor):
     input_image = self.to_tensor(input_image)
     input_image = self.norm(input_image)
     input_image = input_image.unsqueeze(0)
-    features = self.network.forward(Variable(input_image))
+    _, features = self.network.forward(Variable(input_image))
     features = features.data.numpy().flatten()
 
     # normalize
@@ -93,5 +56,3 @@ class CNN8Extractor(Extractor):
     features /= (std+1e-10)
 
     return features
-
-

bob/ip/pytorch_extractor/CasiaNet.py

 import numpy
 
 import torch
-import torch.nn as nn
-import torch.nn.functional as F
 from torch.autograd import Variable
 
 import torchvision.transforms as transforms
 
-import bob.core
-logger = bob.core.log.setup("bob.learn.pytorch")
-
-CASIA_CONFIG = [32, 64, 'M', 64, 128, 'M', 96, 192, 'M', 128, 256, 'M', 160, 320] 
-
-def make_conv_layers(cfg, input_c = 3):
-    layers = []
-    in_channels = input_c
-    for v in cfg:
-        if v == 'M':
-            layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
-        else:
-            conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1)
-            layers += [conv2d, nn.ReLU()]
-            in_channels = v
-    return nn.Sequential(*layers)
-
-
-class CASIA_NET(nn.Module):
-
-    def __init__(self, num_cls, drop_rate=0.5):
-        super(CASIA_NET, self).__init__()
-        self.num_classes = num_cls
-        self.drop_rate = float(drop_rate)
-        self.conv = make_conv_layers(CASIA_CONFIG)
-        self.avgpool = nn.AvgPool2d(8)
-        self.classifier = nn.Linear(320, self.num_classes)
-
-    def forward(self, x):
-        x = self.conv(x)
-        x = self.avgpool(x)
-        x = x.view(x.size(0), -1)
-        x = F.dropout(x, p = self.drop_rate, training=self.training)
-
-        return x 
-
+from bob.learn.pytorch.architectures import CASIANet
 from bob.bio.base.extractor import Extractor
 
 class CasiaNetExtractor(Extractor):
+  """ The class implementing the feature extraction of CASIA-Net embeddings.
+
+  Attributes
+  ----------
+  network: :py:class:`torch.nn.Module`
+      The network architecture
+  to_tensor: :py:class:`torchvision.transforms`
+      The transform from numpy.array to torch.Tensor
+  norm: :py:class:`torchvision.transforms`
+      The transform to normalize the input 
+
+  """
   
   def __init__(self, model_file=None, num_classes=10575):
-
+    """ Init method
+
+    Parameters
+    ----------
+    model_file: str
+        The path of the trained network to load
+    drop_rate: float
+        The number of classes.
+    
+    """
+        
     Extractor.__init__(self, skip_extractor_training=True)
     
     # model
-    self.network = CASIA_NET(num_classes)
+    self.network = CASIANet(num_classes)
     if model_file is None:
-      logger.warning("No model file provided, building network from scratch !")
+      # do nothing (used mainly for unit testing) 
+      pass
     else:
       cp = torch.load(model_file)
       if 'state_dict' in cp:
@@ -66,16 +52,15 @@ class CasiaNetExtractor(Extractor):
     self.norm = transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
 
   def __call__(self, image):
-    """__call__(image) -> feature
-
-    Extract features
-
-    **Parameters:**
+    """ Extract features from an image
 
+    Parameters
+    ----------
     image : 3D :py:class:`numpy.ndarray` (floats)
       The image to extract the features from. Its size must be 3x128x128
 
-    **Returns:**
+    Returns
+    -------
 
     feature : 2D :py:class:`numpy.ndarray` (floats)
       The extracted features as a 1d array of size 320 
@@ -84,8 +69,6 @@ class CasiaNetExtractor(Extractor):
     input_image = self.to_tensor(input_image)
     input_image = self.norm(input_image)
     input_image = input_image.unsqueeze(0)
-    features = self.network.forward(Variable(input_image))
+    _ , features = self.network.forward(Variable(input_image))
     features = features.data.numpy().flatten()
     return features
-
-