add docs and ResNet architecture

bob/learn/pytorch/architectures/ResNet.py

+import torch.nn as nn
+import math
+import torch.utils.model_zoo as model_zoo
+import torch
+
+__all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101',
+           'resnet152']
+
+
+model_urls = {
+    'resnet18': 'https://s3.amazonaws.com/pytorch/models/resnet18-5c106cde.pth',
+    'resnet34': 'https://s3.amazonaws.com/pytorch/models/resnet34-333f7ec4.pth',
+    'resnet50': 'https://s3.amazonaws.com/pytorch/models/resnet50-19c8e357.pth',
+    'resnet101': 'https://s3.amazonaws.com/pytorch/models/resnet101-5d3b4d8f.pth',
+    'resnet152': 'https://s3.amazonaws.com/pytorch/models/resnet152-b121ed2d.pth',
+}
+
+
+def conv3x3(in_planes, out_planes, stride=1):
+    "3x3 convolution with padding"
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+                     padding=1, bias=False)
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None):
+        super(BasicBlock, self).__init__()
+        self.conv1 = conv3x3(inplanes, planes, stride)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = conv3x3(planes, planes)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None):
+        super(Bottleneck, self).__init__()
+        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
+                               padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
+        self.bn3 = nn.BatchNorm2d(planes * 4)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+
+class ResNet(nn.Module):
+    def __init__(self, block, layers, num_classes=1000, tp='cls', bn_dim=128):
+        self.net_type = tp
+
+        self.inplanes = 64
+        super(ResNet, self).__init__()
+        self.conv1 = nn.Conv2d(1, 64, kernel_size=5, stride=2, padding=2,
+                               bias=False)
+        self.bn1 = nn.BatchNorm2d(64)
+        self.relu = nn.ReLU(inplace=True)
+        self.maxpool = nn.MaxPool2d(kernel_size=(3, 1), stride=(2, 1), padding=0)
+        self.layer1 = self._make_layer(block, 64, layers[0], stride=2)
+        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
+        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
+        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
+        self.conv4 = nn.Conv2d(256, 256, kernel_size=(1, 9), stride=1, padding=0,
+                               bias=False)
+        self.bn4 = nn.BatchNorm2d(256)
+        self.conv5 = nn.Conv2d(256, 512, kernel_size=(1, 9), stride=1, padding=0,
+                               bias=False)
+        self.bn5 = nn.BatchNorm2d(512)
+        self.fc = nn.Linear(1024, num_classes)
+        self.fc1 = nn.Linear(bn_dim, num_classes)
+        self.em = nn.Linear(1024, bn_dim)
+        self.drp = nn.Dropout(p=0.2)
+
+        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, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+
+    def _make_layer(self, block, planes, blocks, stride=1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(self.inplanes, planes * block.expansion,
+                          kernel_size=1, stride=stride, bias=False),
+                nn.BatchNorm2d(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample))
+        self.inplanes = planes * block.expansion
+        for i in range(1, blocks):
+            layers.append(block(self.inplanes, planes))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        if self.net_type == 'prefc':
+            return self.forward_cnn(x)
+        if self.net_type == 'cls':
+            return self.forward_cls(x)
+        if self.net_type == 'cls2':
+            return self.forward_cls2(x)
+        if self.net_type == 'emb':
+            return self.forward_emb(x)
+        if self.net_type == 'emb2':
+            return self.forward_emb2(x)
+
+    def forward_cnn(self, x):
+        # x = x.unsqueeze(1)
+        try:
+            x = self.conv1(x)
+        except ValueError:
+            print("Error with ", x.size())
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.conv4(x)
+        x = self.bn4(x)
+        x = self.relu(x)
+        x = self.conv5(x)
+        x = self.bn5(x)
+        x = self.relu(x)
+        t = torch.std(x, 2)
+        m = torch.mean(x, 2)
+        x = torch.cat((m, t), 1)
+        x = x.view(x.size(0), -1)
+        return x
+
+    def forward_cls(self, x):
+        x = self.forward_cnn(x)
+        x = self.fc(x)
+        return x
+
+    def forward_cls2(self, x):
+        x = self.forward_cnn(x)
+        x = self.drp(x)
+        x = self.em(x)
+        x = self.drp(x)
+        x = self.fc1(x)
+        return x
+
+    def forward_emb(self, x):
+        x = self.forward_cnn(x)
+        x = self.em(x)
+        x = torch.div(x, torch.norm(x, 2, 1).unsqueeze(1).expand_as(x))        
+        return x
+    
+
+    def forward_emb2(self, x):
+        x = self.forward_cnn(x)
+        x = torch.div(x, torch.norm(x, 2, 1).unsqueeze(1).expand_as(x))        
+        return x
+
+    def forward_emb_test(self, x, step=1):
+        x = self.conv1(x)
+        if step == 1:
+            return x
+        x = self.bn1(x)
+        if step == 2:
+            return x
+        x = self.relu(x)
+        x = self.maxpool(x)
+        if step == 3:
+            return x
+        x = self.layer1(x)
+        if step == 4:
+            return x
+        x = self.layer2(x)
+        if step == 5:
+            return x
+        x = self.layer3(x)
+        if step == 6:
+            return x
+        x = self.conv4(x)
+        if step == 7:
+            return x
+        x = self.bn4(x)
+        x = self.relu(x)
+        x = self.conv5(x)
+        x = self.bn5(x)
+        x = self.relu(x)
+        if step == 8:
+            return x
+        t = torch.std(x, 2)
+        m = torch.mean(x, 2)
+        x = torch.cat((m, t), 1)
+        x = x.view(x.size(0), -1)
+        return x
+
+
+def resnet18(pretrained=False):
+    """Constructs a ResNet-18 model.
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
+    model = ResNet(BasicBlock, [2, 2, 2, 2])
+    if pretrained:
+        model.load_state_dict(model_zoo.load_url(model_urls['resnet18']))
+    return model
+
+
+def resnet34(pretrained=False, num_classes=1251, tp='cls', bn_dim=128):
+    """Constructs a ResNet-34 model.
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
+    model = ResNet(BasicBlock, [3, 4, 6, 3], num_classes=num_classes, tp=tp, bn_dim=bn_dim)
+    if pretrained:
+        model.load_state_dict(model_zoo.load_url(model_urls['resnet34']))
+    return model
+
+
+def resnet50(pretrained=False, num_classes=1000):
+    """Constructs a ResNet-50 model.
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
+    model = ResNet(Bottleneck, [3, 4, 14, 3], num_classes=num_classes)
+    if pretrained:
+        model.load_state_dict(model_zoo.load_url(model_urls['resnet50']))
+    return model
+
+
+def resnet101(pretrained=False):
+    """Constructs a ResNet-101 model.
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
+    model = ResNet(Bottleneck, [3, 4, 23, 3])
+    if pretrained:
+        model.load_state_dict(model_zoo.load_url(model_urls['resnet101']))
+    return model
+
+
+def resnet152(pretrained=False):
+    """Constructs a ResNet-152 model.
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
+    model = ResNet(Bottleneck, [3, 8, 36, 3])
+    if pretrained:
+        model.load_state_dict(model_zoo.load_url(model_urls['resnet152']))
+    return model

bob/learn/pytorch/extractor/audio/DltResNet.py

+import numpy
+from bob.bio.base.extractor import Extractor
+
+class DltResNetExtractor(Extractor):
+  """ The Dummy class for passing the extracted embedding.
+
+  """
+  pass

bob/learn/pytorch/extractor/audio/__init__.py

+from .DltResNet import DltResNetExtractor
+
+# gets sphinx autodoc done right - don't remove it
+def __appropriate__(*args):
+    """Says object was actually declared here, and not in the import module.
+    Fixing sphinx warnings of not being able to find classes, when path is
+    shortened. Parameters:
+
+      *args: An iterable of objects to modify
+
+    Resolves `Sphinx referencing issues
+    <https://github.com/sphinx-doc/sphinx/issues/3048>`
+    """
+
+    for obj in args:
+        obj.__module__ = __name__
+
+__appropriate__(
+    DltResNetExtractor,
+)
+
+# gets sphinx autodoc done right - don't remove it
+__all__ = [_ for _ in dir() if not _.startswith('_')]

bob/learn/pytorch/preprocessor/audio/DltResNet.py

@@ -8,6 +8,10 @@ import bob.io.base
 import torch.backends.cudnn as cudnn
 from .spectrogram_utils import spectrogram_computation, read
 import pickle
+import pkg_resources
+from bob.extension import rc
+from bob.extension.download import download_file
+import os
 
 class DltResNetExtractor(Preprocessor):
   """ The class implementing the feature extraction of DltResNet embeddings.
@@ -19,7 +23,7 @@ class DltResNetExtractor(Preprocessor):
 
   """
   
-  def __init__(self, model_file=None, num_classes=1211, bn_dim=128, cuda_flag=0):
+  def __init__(self, model_file=None, num_classes=1211, bn_dim=128, cuda_flag=0, test=False):
     """ Init method
 
     Parameters
@@ -32,6 +36,8 @@ class DltResNetExtractor(Preprocessor):
         The embedding dimension (Default: 128).
     cuda_flag: int32
         Use gpu for extracting the embeddings (Default: 0).
+    test: bool
+        Called for unit test (Default: False).
     
     """
         
@@ -44,8 +50,22 @@ class DltResNetExtractor(Preprocessor):
     self.cuda_flag = cuda_flag
 
     if model_file is None:
-      # do nothing (used mainly for unit testing) 
-      pass
+      if not test:
+        model_path = pkg_resources.resource_filename(__name__, 'data/drn34.pth.tar')
+        if rc['bob.learn.pytorch.extractor.audio.drn_modelpath'] != None:
+          model_path = rc['bob.learn.pytorch.extractor.audio.drn_modelpath']
+        if not os.path.exists(model_path):
+          os.makedirs(os.path.dirname(model_path), exist_ok=True)
+          url = 'https://www.idiap.ch/software/bob/data/bob/bob.learn.pytorch/master/models_emb128_2drp_dlt_aug_relu__voxceleb___rn34-best.cls.pth.tar'
+          download_file(url, model_path)
+        model_file = model_path
+
+      else:
+        # do nothing (used mainly for unit testing) 
+        pass
+    
+    if test:
+      self.network = torch.nn.DataParallel(self.network)
     else:
       if cuda_flag == 1:
         self.network = torch.nn.DataParallel(self.network).cuda()

bob/learn/pytorch/preprocessor/audio/ResNet.py

+import numpy
+import torch
+from torch.autograd import Variable
+import torchvision.transforms as transforms
+from bob.learn.pytorch.architectures.ResNet import resnet34
+from bob.bio.base.preprocessor import Preprocessor
+import bob.io.base
+import torch.backends.cudnn as cudnn
+from .spectrogram_utils import spectrogram_computation, read
+import pickle
+import pkg_resources
+from bob.extension import rc
+from bob.extension.download import download_file
+import os
+
+class ResNetExtractor(Preprocessor):
+  """ The class implementing the feature extraction of ResNet embeddings.
+
+  Attributes
+  ----------
+  network: :py:class:`torch.nn.Module`
+      The network architecture
+
+  """
+  
+  def __init__(self, model_file=None, num_classes=1211, bn_dim=128, cuda_flag=0, test=False):
+    """ Init method
+
+    Parameters
+    ----------
+    model_file: str
+        The path of the trained network to load
+    num_classes: float
+        The number of classes (Default: 1211).
+    bn_dim: int32
+        The embedding dimension (Default: 128).
+    cuda_flag: int32
+        Use gpu for extracting the embeddings (Default: 0).
+    test: bool
+        Called for unit test (Default: False).
+    
+    """
+        
+    Preprocessor.__init__(self, min_preprocessed_file_size=bn_dim)
+    
+    # model
+    model_type = "emb"
+    
+    self.network = resnet34(num_classes=num_classes, tp=model_type, bn_dim=bn_dim)
+    self.cuda_flag = cuda_flag
+
+    if model_file is None:
+      if not test:
+        model_path = pkg_resources.resource_filename(__name__, 'data/rn34.pth.tar')
+        if rc['bob.learn.pytorch.extractor.audio.rn_modelpath'] != None:
+          model_path = rc['bob.learn.pytorch.extractor.audio.rn_modelpath']
+        if not os.path.exists(model_path):
+          os.makedirs(os.path.dirname(model_path), exist_ok=True)
+          url = 'https://www.idiap.ch/software/bob/data/bob/bob.learn.pytorch/master/models_emb128_2drp_aug_relu__voxceleb___rn34-best.cls.pth.tar'
+          download_file(url, model_path)
+        model_file = model_path
+
+      else:
+        # do nothing (used mainly for unit testing) 
+        pass
+    
+    if test:
+      self.network = torch.nn.DataParallel(self.network)
+    else:
+      if cuda_flag == 1:
+        self.network = torch.nn.DataParallel(self.network).cuda()
+        cudnn.benchmark = True
+        checkpoint = torch.load(model_file)
+      else:
+        self.network = torch.nn.DataParallel(self.network)
+        checkpoint = torch.load(model_file, map_location=lambda storage, loc: storage)
+      if 'state_dict' in checkpoint:
+        self.network.load_state_dict(checkpoint['state_dict'], strict=False)
+        self.network.net_type = model_type
+        self.network.module.net_type = self.network.net_type
+        self.network.eval()
+
+
+  def extract_embeddings(self, model, specs, L, batch_size=1, cuda_flag=0, min_length=30):
+    """ Extract features from spectrogram
+
+    Parameters
+    ----------
+    model : :py:class:`torch.nn.Module`
+      The trained pytorch model.
+    specs : :py:class:`numpy.ndarray` (floats)
+      The spectrogram of audio file.
+    L : int32
+      Maximum chunk size for embedding extraction. 0 for using all the data (Default: 0).
+    batch_size : int32
+      Batch size for extracting the embedding (Default: 1).
+    cuda_flag : int32
+      Use gpu for extracting the embeddings
+    min_length : int32
+      Minimum chunk size for embedding extraction (Default: 30).
+
+    Returns
+    -------
+    feature : 2D :py:class:`numpy.ndarray` (floats)
+      The extracted features as a 1d array of size 128 
+    
+    """
+
+    feats = []
+    batch = []
+
+    if L == 0 or L > specs.shape[0]:
+        L = specs.shape[0]
+
+    embedding_count = specs.shape[0] // L
+    if embedding_count == 0 :
+        embedding_count = 1
+    elif specs.shape[0] % L > min_length:
+        embedding_count += 1
+
+    for i in range(embedding_count):
+        if (i+1)* L < specs.shape[0]:
+            k = specs[i*L:(i+1)*L,:]
+        else:
+            k = specs[i*L:,:]
+        
+        X = (k - k.mean(axis=0)) / (k.std(axis=0) + 1e-7)
+        batch.append(X)
+        if len(batch) == batch_size or i == embedding_count - 2 or i == embedding_count - 1: 
+          with torch.no_grad():
+            in_batch = numpy.array(batch)
+            if cuda_flag == 1:
+                input_var = torch.autograd.Variable(torch.Tensor(in_batch).unsqueeze(1).cuda(async=True))
+            else:
+                input_var = torch.autograd.Variable(torch.Tensor(in_batch).unsqueeze(1))
+            fX = model(input_var).data.cpu().numpy()
+            feats.append(fX)
+            batch = []
+  
+    return numpy.vstack(feats)
+
+  def __call__(self, audio, annotations=None):
+    """ Extract features from an audio file
+
+    Parameters
+    ----------
+    audio : :py:class:`numpy.ndarray` (floats)
+      The audio file to extract the features from. 
+    annotations : None
+      Apply annotations if needed
+
+    Returns
+    -------
+    feature : 2D :py:class:`numpy.ndarray` (floats)
+      The extracted features as a 1d array of size 128 
+    
+    """
+    data = audio[1]
+    if data.dtype =='int16':
+      data = numpy.cast['float'](data)
+    if numpy.max(numpy.abs(data)) < 1:
+      data = data * 2**15
+    rate = audio[0]
+    input_feat = spectrogram_computation(rate, data)
+    features = self.extract_embeddings(self.network, input_feat, 0, batch_size=1, cuda_flag=self.cuda_flag)
+    return features
+
+
+  def write_data(self, data, data_file, compression=0):
+    """Writes the given *preprocessed* data to a file with the given name.
+    """
+    f = bob.io.base.HDF5File(data_file, 'w')
+    f.set("feats", data.astype('float'), compression=compression)
+
+
+  def read_data(self, data_file):
+    f= bob.io.base.HDF5File(data_file)
+    feats = f.read("feats")
+    return feats
\ No newline at end of file

bob/learn/pytorch/preprocessor/audio/__init__.py

 from .DltResNet import DltResNetExtractor
+from .ResNet import ResNetExtractor
 
 # gets sphinx autodoc done right - don't remove it
 def __appropriate__(*args):
@@ -17,6 +18,7 @@ def __appropriate__(*args):
 
 __appropriate__(
     DltResNetExtractor,
+    ResNetExtractor,
 )
 
 # gets sphinx autodoc done right - don't remove it

bob/learn/pytorch/test/test_audio_extractor.py

@@ -17,10 +17,24 @@ def test_drn():
         output is an embedding of dimension 128
     """
     from bob.learn.pytorch.preprocessor.audio import DltResNetExtractor
-    extractor = DltResNetExtractor()
+    extractor = DltResNetExtractor(test=True)
     # this architecture expects Tx257 
     data = numpy.random.rand(512, 257).astype("float32")
     output = extractor.extract_embeddings(extractor.network, data, 0)
-    assert output.shape[1] == 128
+    assert output.size == 128, output.shape
 
 
+def test_rn():
+    """ test for the ResNet architecture
+
+        This architecture takes Tx257 audio chunk as input,
+        where T is the number of frames and 257 is unique numbers from
+        spectrogram with 512 FFT resolution.
+        output is an embedding of dimension 128
+    """
+    from bob.learn.pytorch.preprocessor.audio import ResNetExtractor
+    extractor = ResNetExtractor(test=True)
+    # this architecture expects Tx257 
+    data = numpy.random.rand(512, 257).astype("float32")
+    output = extractor.extract_embeddings(extractor.network, data, 0)
+    assert output.size == 128, output.shape
\ No newline at end of file

doc/guide_audio_extractor.rst

+==========
+User guide
+==========
+
+This package has been done with the intent of using it in conjunction with
+both ``bob.learn.pytorch`` to run speaker verification experiments.
+
+In particular, the goal here is to extract features from an input audio using CNN models 
+trained with PyTorch.
+
+For this purpose, you can specify your feature extractor in configuration
+file to be used together with the ``verifiy.py`` script from ``bob.bio.base``.
+
+
+DltResNet Model
+---------------
+
+
+:ref:`bob.bio.base <bob.bio.base>` wrapper DltResNet model.
+
+
+.. note::
+
+   The models will automatically download to the data folder of this package and save it in 
+   ``[env-path]./bob/learn/pytorch/data/drn34.pth.tar``.
+   If you want to set another path for this model do::
+   
+   $ bob config set bob.learn.pytorch.extractor.audio.drn_modelpath /path/to/mymodel
+
+
+ResNet Model
+---------------
+
+
+:ref:`bob.bio.base <bob.bio.base>` wrapper ResNet model.
+
+
+.. note::
+
+   The models will automatically download to the data folder of this package and save it in 
+   ``[env-path]./bob/learn/pytorch/data/rn34.pth.tar``.
+   If you want to set another path for this model do::
+   
+   $ bob config set bob.learn.pytorch.extractor.audio.rn_modelpath /path/to/mymodel
+
+
+A concrete example
+------------------
+
+Imagine that you have the DltResNet model and you would 
+like to use the embedding layer as a feature to encode identity. 
+
+Your `preprocessor` in Bob pipe-lines should be defined this way in the configuration file:
+
+.. code:: python
+
+  from bob.learn.pytorch.preprocessor.audio import DltResNetExtractor
+
+  _model = 'path/to/your/model.pth'
+  _num_classes = 1211
+  preprocessor = ResNetDltExtractor(_model, _num_classes)
+
+Note that the number of classes is irrelevant here, but is required to build the 
+network (before loading it). In this set we just need a dummy `extractor` in Bob pipeline which can be defined in the configuration file in this way:
+
+.. code:: python
+
+  from bob.learn.pytorch.extractor.audio import DltResNetExtractor as DummyExtractor
+
+  extractor = DummyExtractor()
+
+
+You can easily implement your own extractor based on your own network too. Just have
+a look at the code in ``bob/learn/pytorch/preprocessor/audio``.
+