Implemented parser for pytorch models.

Author: rbodo

snntoolbox/bin/utils.py

@@ -65,6 +65,7 @@ def run_pipeline(config, queue=None):
 
     normset, testset = get_dataset(config)
 
+    results = None
     parsed_model = None
     if config.getboolean('tools', 'parse') and not is_stop(queue):
 
@@ -80,9 +81,10 @@ def run_pipeline(config, queue=None):
         if config.getboolean('tools', 'evaluate_ann') and not is_stop(queue):
             print("Evaluating input model on {} samples...".format(
                 num_to_test))
-            model_lib.evaluate(input_model['val_fn'],
-                               config.getint('simulation', 'batch_size'),
-                               num_to_test, **testset)
+            acc = model_lib.evaluate(input_model['val_fn'],
+                                     config.getint('simulation', 'batch_size'),
+                                     num_to_test, **testset)
+            results = [acc]
 
         # ____________________________ PARSE ________________________________ #
 
@@ -100,8 +102,10 @@ def run_pipeline(config, queue=None):
         if config.getboolean('tools', 'evaluate_ann') and not is_stop(queue):
             print("Evaluating parsed model on {} samples...".format(
                 num_to_test))
-            model_parser.evaluate(config.getint(
-                'simulation', 'batch_size'), num_to_test, **testset)
+            score = model_parser.evaluate(
+                config.getint('simulation', 'batch_size'),
+                num_to_test, **testset)
+            results = [score[1]]
 
         # Write parsed model to disk
         parsed_model.save(str(
@@ -153,7 +157,7 @@ def run(snn, **test_set):
         if queue:
             queue.put(results)
 
-        return results
+    return results
 
 
 def is_stop(queue):

snntoolbox/config_defaults

@@ -116,7 +116,7 @@ plotproperties = {
 # validity of config.
 
 [restrictions]
-model_libs = {'keras', 'lasagne', 'caffe'}
+model_libs = {'keras', 'lasagne', 'caffe', 'pytorch'}
 dataset_formats = {'npz', 'jpg', 'aedat'}
 frame_gen_method = {'signed_sum', 'rectified_sum'}
 maxpool_types = {'fir_max', 'exp_max', 'avg_max'}

snntoolbox/parsing/model_libs/pytorch_input_lib.py

@@ -0,0 +1,120 @@
+# -*- coding: utf-8 -*-
+"""Keras model parser.
+
+@author: rbodo
+"""
+
+import os
+import numpy as np
+
+import keras
+import torch
+import onnx
+from onnx2keras import onnx_to_keras
+import onnxruntime
+
+from snntoolbox.parsing.model_libs import keras_input_lib
+from snntoolbox.utils.utils import import_script
+
+
+def to_numpy(tensor):
+    return tensor.detach().cpu().numpy() if tensor.requires_grad \
+        else tensor.cpu().numpy()
+
+
+class ModelParser(keras_input_lib.ModelParser):
+
+    def try_insert_flatten(self, layer, idx, name_map):
+        return False
+        
+
+def load(path, filename):
+    """Load network from file.
+
+    Parameters
+    ----------
+
+    path: str
+        Path to directory where to load model from.
+
+    filename: str
+        Name of file to load model from.
+
+    Returns
+    -------
+
+    : dict[str, Union[keras.models.Sequential, function]]
+        A dictionary of objects that constitute the input model. It must
+        contain the following two keys:
+
+        - 'model': keras.models.Sequential
+            Keras model instance of the network.
+        - 'val_fn': function
+            Function that allows evaluating the original model.
+    """
+
+    filepath = str(os.path.join(path, filename))
+
+    # Load the Pytorch model.
+    mod = import_script(path, filename)
+    model_pytorch = mod.Model()
+    model_pytorch.load_state_dict(torch.load(filepath + '.pkl'))
+
+    # Switch from train to eval mode to ensure Dropout / BatchNorm is handled
+    # correctly.
+    model_pytorch.eval()
+
+    # Run on dummy input with correct shape to trace the Pytorch model.
+    input_shape = [1] + list(model_pytorch.input_shape)
+    input_numpy = np.random.random_sample(input_shape).astype(np.float32)
+    input_torch = torch.from_numpy(input_numpy).float()
+    output_torch = model_pytorch(input_torch)
+    output_numpy = to_numpy(output_torch)
+
+    # Export as onnx model, and then reload.
+    input_names = ['input_0']
+    output_names = ['output_{}'.format(i) for i in range(len(output_torch))]
+    dynamic_axes = {'input_0': {0: 'batch_size'}}
+    dynamic_axes.update({name: {0: 'batch_size'} for name in output_names})
+    torch.onnx.export(model_pytorch, input_torch, filepath + '.onnx',
+                      input_names=input_names,
+                      output_names=output_names,
+                      dynamic_axes=dynamic_axes)
+    model_onnx = onnx.load(filepath + '.onnx')
+    # onnx.checker.check_model(model_onnx)  # Crashes with segmentation fault.
+
+    # Compute ONNX Runtime output prediction.
+    ort_session = onnxruntime.InferenceSession(filepath + '.onnx')
+    input_onnx = {ort_session.get_inputs()[0].name: input_numpy}
+    output_onnx = ort_session.run(None, input_onnx)
+
+    # Compare ONNX Runtime and PyTorch results.
+    err_msg = "Pytorch model could not be ported to ONNX. Output difference: "
+    np.testing.assert_allclose(output_numpy, output_onnx[0],
+                               rtol=1e-03, atol=1e-05, err_msg=err_msg)
+    print("Pytorch model was successfully ported to ONNX.")
+
+    # Port ONNX model to Keras.
+    model_keras = onnx_to_keras(model_onnx, input_names, [input_shape[1:]])
+
+    # Save the keras model.
+    keras.models.save_model(model_keras, filepath + '.h5')
+
+    # Loading the model here is a workaround for version conflicts with
+    # TF > 2.0.1 and keras > 2.2.5. Should be able to remove this later.
+    model_keras = keras.models.load_model(filepath + '.h5')
+    model_keras.compile('sgd', 'categorical_crossentropy',
+                        ['accuracy', keras.metrics.top_k_categorical_accuracy])
+
+    # Compute Keras output and compare against ONNX.
+    output_keras = model_keras.predict(input_numpy)
+    err_msg = "ONNX model could not be ported to Keras. Output difference: "
+    np.testing.assert_allclose(output_numpy, output_keras,
+                               rtol=1e-03, atol=1e-05, err_msg=err_msg)
+    print("ONNX model was successfully ported to Keras.")
+
+    return {'model': model_keras, 'val_fn': model_keras.evaluate}
+
+
+def evaluate(*args, **kwargs):
+    return keras_input_lib.evaluate(*args, **kwargs)

tests/conftest.py

@@ -227,6 +227,79 @@ def _model_3(_dataset):
     return model
 
 
+@pytest.fixture(scope='session')
+def _model_4(_dataset):
+
+    if not is_module_installed('torch'):
+        return
+
+    import torch
+    import torch.nn as nn
+    from tests.parsing.models import pytorch
+
+    x_train, y_train, x_test, y_test = _dataset
+
+    # Pytorch doesn't support one-hot labels.
+    y_train = np.argmax(y_train, 1)
+    y_test = np.argmax(y_test, 1)
+
+    # Pytorch needs channel dimension first.
+    if keras.backend.image_data_format() == 'channels_last':
+        x_train = np.moveaxis(x_train, 3, 1)
+        x_test = np.moveaxis(x_test, 3, 1)
+
+    class PytorchDataset(torch.utils.data.Dataset):
+        def __init__(self, data, target, transform=None):
+            self.data = torch.from_numpy(data).float()
+            self.target = torch.from_numpy(target).long()
+            self.transform = transform
+
+        def __getitem__(self, index):
+            x = self.data[index]
+
+            if self.transform:
+                x = self.transform(x)
+
+            return x, self.target[index]
+
+        def __len__(self):
+            return len(self.data)
+
+    trainset = torch.utils.data.DataLoader(PytorchDataset(x_train, y_train),
+                                           batch_size=64)
+    testset = torch.utils.data.DataLoader(PytorchDataset(x_test, y_test),
+                                          batch_size=64)
+
+    model = pytorch.Model()
+
+    criterion = nn.CrossEntropyLoss()
+    optimizer = torch.optim.Adam(model.parameters())
+
+    acc = 0
+    for epoch in range(3):
+        for i, (xx, y) in enumerate(trainset):
+            optimizer.zero_grad()
+            outputs = model(xx)
+            loss = criterion(outputs, y)
+            loss.backward()
+            optimizer.step()
+
+        total = 0
+        correct = 0
+        with torch.no_grad():
+            for xx, y in testset:
+                outputs = model(xx)
+                _, predicted = torch.max(outputs.data, 1)
+                total += y.size(0)
+                correct += (predicted == y).sum().item()
+        acc = correct / total
+
+    print("Test accuracy: {:.2%}".format(acc))
+    # assert acc > 0.96, "Test accuracy after training not high enough."
+
+    return model
+
+
 spinnaker_conditions = (is_module_installed('keras_rewiring') and
                         is_module_installed('pynn_object_serialisation') and
                         (is_module_installed('pyNN.spiNNaker') or
@@ -243,6 +316,12 @@ def _model_3(_dataset):
 brian2_skip_if_dependency_missing = pytest.mark.skipif(
     not brian2_conditions, reason="Brian2 dependency missing.")
 
+pytorch_conditions = (is_module_installed('torch') and
+                      is_module_installed('onnx') and
+                      is_module_installed('onnx2keras'))
+pytorch_skip_if_dependency_missing = pytest.mark.skipif(
+    not pytorch_conditions, reason='Pytorch dependencies missing')
+
 
 def get_examples():
     path = os.path.abspath(os.path.join(os.path.dirname(__file__), '..',

tests/core/test_models.py

@@ -112,12 +112,11 @@ def test_inisim(self, _model_2, _config):
                                 os.path.join(path_wd, model_name + '.h5'))
 
         updates = {
-            'tools': {'evaluate_ann': False}, 'simulation': {
-                'simulator': 'INI',
+            'tools': {'evaluate_ann': False},
+            'simulation': {
                 'duration': 100,
                 'num_to_test': 100,
-                'batch_size': 50,
-                'keras_backend': 'tensorflow'},
+                'batch_size': 50},
             'output': {
                 'log_vars': {'activations_n_b_l', 'spiketrains_n_b_l_t'}}}
 

tests/parsing/models/__init__.py

@@ -0,0 +1,45 @@
+import torch
+import torch.nn as nn
+
+
+class Model(nn.Module):
+    def __init__(self):
+        super(Model, self).__init__()
+
+        self.input_shape = (1, 28, 28)
+
+        layers_trunk = [
+            nn.Conv2d(1, 16, kernel_size=5, stride=2),
+            # BatchNorm doesn't work with Keras==2.3.1 because for some reason
+            # they put the batch-norm axis in a list.
+            # nn.BatchNorm2d(16),
+            nn.ReLU(),
+            nn.AvgPool2d(kernel_size=2, stride=2)]
+        layers_branch1 = [
+            nn.Conv2d(16, 32, kernel_size=3, padding=1),
+            nn.ReLU()]
+        layers_branch2 = [
+            nn.Conv2d(16, 8, kernel_size=1),
+            nn.ReLU()]
+        layers_head = [
+            nn.Conv2d(40, 8, kernel_size=1),
+            nn.ReLU()]
+        layers_classifier = [
+            nn.Dropout(1e-5),
+            nn.Linear(288, 10),
+            nn.Softmax(1)]
+        self.trunk = nn.Sequential(*layers_trunk)
+        self.branch1 = nn.Sequential(*layers_branch1)
+        self.branch2 = nn.Sequential(*layers_branch2)
+        self.head = nn.Sequential(*layers_head)
+        self.classifier = nn.Sequential(*layers_classifier)
+
+    def forward(self, x):
+        x = self.trunk(x)
+        x1 = self.branch1(x)
+        x2 = self.branch2(x)
+        x = torch.cat([x1, x2], 1)
+        x = self.head(x)
+        x = x.view(-1, 288)  # Flatten
+        x = self.classifier(x)
+        return x