Merge pull request #1 from ToriML/develop

New models support

Author: Talmaj

README.md

@@ -1,7 +1,8 @@
 # ONNX to PyTorch
-[![GitHub License](https://img.shields.io/badge/License-Apache-2.svg)](https://opensource.org/licenses/Apache-2.0)
+![PyPI - License](https://img.shields.io/pypi/l/onnx2pytorch?color)
 [![CircleCI](https://circleci.com/gh/ToriML/onnx2pytorch.svg?style=shield)](https://app.circleci.com/pipelines/github/ToriML/onnx2pytorch)
 [![Downloads](https://pepy.tech/badge/onnx2pytorch)](https://pepy.tech/project/onnx2pytorch)
+![PyPI](https://img.shields.io/pypi/v/onnx2pytorch)
 
 A library to transform ONNX model to PyTorch. This library enables use of PyTorch 
 backend and all of its great features for manipulation of neural networks.
@@ -19,10 +20,17 @@ pytorch_model = ConvertModel(onnx_model)
 ```
 
 Currently supported and tested models from [onnx_zoo](https://github.com/onnx/models):
-- MobileNet
-- ResNet
-- ShuffleNet
-- Bert
+- [MobileNet](https://github.com/onnx/models/tree/master/vision/classification/mobilenet)
+- [ResNet](https://github.com/onnx/models/tree/master/vision/classification/resnet)
+- [ShuffleNet_V2](https://github.com/onnx/models/tree/master/vision/classification/shufflenet)
+- [BERT-Squad](https://github.com/onnx/models/tree/master/text/machine_comprehension/bert-squad)
+- [EfficientNet-Lite4](https://github.com/onnx/models/tree/master/vision/classification/efficientnet-lite4)
+- [Fast Neural Style Transfer](https://github.com/onnx/models/tree/master/vision/style_transfer/fast_neural_style)
+- [Super Resolution](https://github.com/onnx/models/tree/master/vision/super_resolution/sub_pixel_cnn_2016)
+- [YOLOv4](https://github.com/onnx/models/tree/master/vision/object_detection_segmentation/yolov4)
+  (Not exactly the same, nearest neighbour interpolation in pytorch differs)
+- [U-net](https://pytorch.org/hub/mateuszbuda_brain-segmentation-pytorch_unet/)
+  (Converted from pytorch to onnx and then back)
 
 ## Limitations
 Known current version limitations are:
@@ -48,4 +56,14 @@ Install it into pre-commit hook to always commit nicely formatted code:
 
 ### Testing
 [Pytest](https://docs.pytest.org/en/latest/) and [tox](https://tox.readthedocs.io/en/latest/).  
-```tox```
+```tox```
+#### Test fixtures
+To test the complete conversion of an onnx model download pre-trained models: 
+```./download_fixtures.sh```  
+Use flag `--all` to download more models.
+Add any custom models to `./fixtures` folder to test their conversion.
+
+### Debugging
+Set `ConvertModel(..., debug=True)` to compare each converted
+activation from pytorch with the activation from onnxruntime.  
+This helps identify where in the graph the activations start to differ.

download_fixtures.sh

@@ -1,3 +1,4 @@
+#!/usr/bin/env bash
 mkdir -p fixtures
 cd fixtures
 
@@ -6,22 +7,41 @@ if [[ ! -f mobilenetv2-1.0.onnx ]]; then
   curl -o mobilenetv2-1.0.onnx https://s3.amazonaws.com/onnx-model-zoo/mobilenet/mobilenetv2-1.0/mobilenetv2-1.0.onnx
 fi
 
-#if [[ ! -f resnet18v1.onnx ]]; then
-#  echo Downloading resnet18v1
-#  curl -o resnet18v1.onnx https://s3.amazonaws.com/onnx-model-zoo/resnet/resnet18v1/resnet18v1.onnx
-#fi
-
 if [[ ! -f shufflenet_v2.onnx ]]; then
   echo Downloading shufflenet_v2
   curl -LJo shufflenet_v2.onnx https://github.com/onnx/models/blob/master/vision/classification/shufflenet/model/shufflenet-v2-10.onnx\?raw\=true
 fi
 
-#if [[ ! -f bertsquad-10.onnx ]]; then
-#  echo Downloading bertsquad-10
-#  curl -LJo bertsquad-10.onnx https://github.com/onnx/models/blob/master/text/machine_comprehension/bert-squad/model/bertsquad-10.onnx\?raw\=true
-#fi
+if [[ $1 == "--all" ]]; then
+  if [[ ! -f resnet18v1.onnx ]]; then
+    echo Downloading resnet18v1
+    curl -o resnet18v1.onnx https://s3.amazonaws.com/onnx-model-zoo/resnet/resnet18v1/resnet18v1.onnx
+  fi
+
+  if [[ ! -f bertsquad-10.onnx ]]; then
+    echo Downloading bertsquad-10
+    curl -LJo bertsquad-10.onnx https://github.com/onnx/models/blob/master/text/machine_comprehension/bert-squad/model/bertsquad-10.onnx\?raw\=true
+  fi
+
+  if [[ ! -f yolo_v4.onnx ]]; then
+    echo Downloading yolo_v4
+    curl -LJo yolo_v4.onnx https://github.com/onnx/models/blob/master/vision/object_detection_segmentation/yolov4/model/yolov4.onnx\?raw\=true
+  fi
+
+  if [[ ! -f super_res.onnx ]]; then
+    echo Downloading super_res
+    curl -LJo super_res.onnx https://github.com/onnx/models/blob/master/vision/super_resolution/sub_pixel_cnn_2016/model/super-resolution-10.onnx\?raw\=true
+  fi
+
+  if [[ ! -f fast_neural_style.onnx ]]; then
+    echo Downloading fast_neural_style
+    curl -LJo fast_neural_style.onnx https://github.com/onnx/models/blob/master/vision/style_transfer/fast_neural_style/model/rain-princess-9.onnx\?raw\=true
+  fi
+
+  if [[ ! -f efficientnet-lite4.onnx ]]; then
+    echo Downloading efficientnet-lite4
+    curl -LJo efficientnet-lite4.onnx https://github.com/onnx/models/blob/master/vision/classification/efficientnet-lite4/model/efficientnet-lite4-11.onnx\?raw\=true
+  fi
+fi
 
-#if [[ ! -f yolo_v4.onnx ]]; then
-#  echo Downloading yolo_v4
-#  curl -LJo yolo_v4.onnx https://github.com/onnx/models/blob/master/vision/object_detection_segmentation/yolov4/model/yolov4.onnx\?raw\=true
-#fi
+echo "All models downloaded."

onnx2pytorch/convert/attribute.py

@@ -1,7 +1,12 @@
+import warnings
+
 import onnx
 from onnx import numpy_helper
 
-from onnx2pytorch.utils import to_pytorch_params
+from onnx2pytorch.utils import (
+    extract_padding_params_for_conv_layer,
+    extract_padding_params,
+)
 
 TENSOR_PROTO_MAPPING = dict([i[::-1] for i in onnx.TensorProto.DataType.items()])
 
@@ -54,7 +59,12 @@ def extract_attributes(node):
         elif attr.name == "kernel_shape":
             kwargs["kernel_size"] = extract_attr_values(attr)
         elif attr.name == "pads":
-            kwargs["padding"] = to_pytorch_params(extract_attr_values(attr))
+            params = extract_attr_values(attr)
+            if node.op_type == "Pad":
+                kwargs["padding"] = extract_padding_params(params)
+            else:
+                # Works for Conv, MaxPooling and other layers from convert_layer func
+                kwargs["padding"] = extract_padding_params_for_conv_layer(params)
         elif attr.name == "strides":
             kwargs["stride"] = extract_attr_values(attr)
         elif attr.name == "axis" and node.op_type == "Flatten":
@@ -89,10 +99,28 @@ def extract_attributes(node):
             kwargs["transpose_weight"] = not extract_attr_values(attr)
         elif attr.name == "transA":
             kwargs["transpose_activation"] = bool(extract_attr_values(attr))
+        elif attr.name == "alpha" and node.op_type == "LeakyRelu":
+            kwargs["negative_slope"] = extract_attr_values(attr)
         elif attr.name == "alpha":
             kwargs["weight_multiplier"] = extract_attr_values(attr)
         elif attr.name == "beta":
             kwargs["bias_multiplier"] = extract_attr_values(attr)
+        elif attr.name == "starts":
+            kwargs["starts"] = extract_attr_values(attr)
+        elif attr.name == "ends":
+            kwargs["ends"] = extract_attr_values(attr)
+        elif attr.name == "coordinate_transformation_mode":
+            arg = extract_attr_values(attr)
+            if arg == "align_corners":
+                kwargs["align_corners"] = True
+            else:
+                warnings.warn(
+                    "Pytorch's interpolate uses no coordinate_transformation_mode={}. "
+                    "Result might differ.".format(arg)
+                )
+        elif node.op_type == "Resize":
+            # These parameters are not used, warn in Resize operator
+            kwargs[attr.name] = extract_attr_values(attr)
         elif attr.name == "auto_pad":
             value = extract_attr_values(attr)
             if value == "NOTSET":

onnx2pytorch/convert/debug.py

@@ -0,0 +1,22 @@
+import torch
+import numpy as np
+
+from onnx2pytorch.utils import get_activation_value
+
+
+def debug_model_conversion(onnx_model, inputs, pred_act, node, rtol=1e-3, atol=1e-4):
+    """Compare if the activations of pytorch are the same as from onnxruntime."""
+    if not isinstance(inputs, list):
+        raise TypeError("inputs should be in a list.")
+
+    if not all(isinstance(x, np.ndarray) for x in inputs):
+        inputs = [x.detach().numpy() for x in inputs]
+
+    exp_act = get_activation_value(onnx_model, inputs, list(node.output))
+    if isinstance(pred_act, list):
+        for a, b in zip(exp_act, pred_act):
+            assert torch.allclose(torch.from_numpy(a), b, rtol=rtol, atol=atol)
+    else:
+        a = torch.from_numpy(exp_act[0])
+        b = pred_act
+        assert torch.allclose(a, b, rtol=rtol, atol=atol)

onnx2pytorch/convert/layer.py

@@ -2,7 +2,7 @@
 from torch import nn
 from onnx import numpy_helper
 
-from onnx2pytorch.operations import BatchNormUnsafe
+from onnx2pytorch.operations import BatchNormUnsafe, InstanceNormUnsafe
 from onnx2pytorch.convert.attribute import extract_attributes, extract_attr_values
 
 
@@ -45,6 +45,7 @@ def convert_layer(node, layer_type, params=None):
         )
 
     if params:
+        pad_layer = None
         weight, bias = extract_params(params)
         kwargs["bias"] = bias is not None
         kwargs["in_channels"] = weight.dims[1] * kwargs.get("groups", 1)
@@ -56,9 +57,15 @@ def convert_layer(node, layer_type, params=None):
                 kwargs["in_channels"],
             )
 
+        # if padding is a layer, remove from kwargs and prepend later
+        if isinstance(kwargs["padding"], nn.Module):
+            pad_layer = kwargs.pop("padding")
+
         # initialize layer and load weights
         layer = layer(**kwargs)
         load_params(layer, weight, bias)
+        if pad_layer is not None:
+            layer = nn.Sequential(pad_layer, layer)
     else:
         # initialize operations without parameters (MaxPool, AvgPool, etc.)
         layer = layer(**kwargs)
@@ -80,6 +87,21 @@ def convert_batch_norm_layer(node, params):
     return layer
 
 
+def convert_instance_norm_layer(node, params):
+    kwargs = extract_attributes(node)
+    # Skips input dimension check, not possible before forward pass
+    layer = InstanceNormUnsafe
+
+    kwargs["num_features"] = params[0].dims[0]
+    # initialize layer and load weights
+    layer = layer(**kwargs)
+    key = ["weight", "bias"]
+    for key, value in zip(key, params):
+        getattr(layer, key).data = torch.from_numpy(numpy_helper.to_array(value))
+
+    return layer
+
+
 def convert_linear_layer(node, params):
     """Convert linear layer from onnx node and params."""
     # Default Gemm attributes

onnx2pytorch/convert/model.py

@@ -8,10 +8,13 @@
 from torch.jit import TracerWarning
 from torch.nn.modules.conv import _ConvNd
 from torch.nn.modules.batchnorm import _BatchNorm
+from torch.nn.modules.instancenorm import _InstanceNorm
 from torch.nn.modules.linear import Identity
 
 from onnx2pytorch.operations import Split
+from onnx2pytorch.convert.debug import debug_model_conversion
 from onnx2pytorch.convert.operations import convert_operations
+from onnx2pytorch.utils import get_inputs_names
 
 
 class InitParameters(dict):
@@ -30,7 +33,9 @@ def get(self, item, default):
 
 
 class ConvertModel(nn.Module):
-    def __init__(self, onnx_model: onnx.ModelProto, batch_dim=0, experimental=False):
+    def __init__(
+        self, onnx_model: onnx.ModelProto, batch_dim=0, experimental=False, debug=False
+    ):
         """
         Convert onnx model to pytorch.
 
@@ -53,6 +58,7 @@ def __init__(self, onnx_model: onnx.ModelProto, batch_dim=0, experimental=False)
         self.onnx_model = onnx_model
         self.batch_dim = batch_dim
         self.experimental = experimental
+        self.debug = debug
         self.mapping = {}
         for op_id, op_name, op in convert_operations(onnx_model, batch_dim):
             setattr(self, op_name, op)
@@ -62,6 +68,8 @@ def __init__(self, onnx_model: onnx.ModelProto, batch_dim=0, experimental=False)
             {tensor.name: tensor for tensor in self.onnx_model.graph.initializer}
         )
 
+        self.input_names = get_inputs_names(onnx_model)
+
         if experimental:
             warnings.warn(
                 "Using experimental implementation that allows 'batch_size > 1'."
@@ -74,8 +82,7 @@ def forward(self, *input):
                 "Input with larger batch size than 1 not supported yet."
             )
         # TODO figure out how to store only necessary activations.
-        input_names = [x.name for x in self.onnx_model.graph.input]
-        activations = dict(zip(input_names, input))
+        activations = dict(zip(self.input_names, input))
 
         for node in self.onnx_model.graph.node:
             # Identifying the layer ids and names
@@ -93,7 +100,11 @@ def forward(self, *input):
             # if first layer choose input as in_activations
             # if not in_op_names and len(node.input) == 1:
             #    in_activations = input
-            if isinstance(op, (nn.Linear, _ConvNd, _BatchNorm)):
+            layer_types = (nn.Linear, _ConvNd, _BatchNorm, _InstanceNorm)
+            if isinstance(op, layer_types) or (
+                isinstance(op, nn.Sequential)
+                and any(isinstance(x, layer_types) for x in op.modules())
+            ):
                 in_activations = [
                     activations[in_op_id]
                     for in_op_id in node.input
@@ -122,6 +133,15 @@ def forward(self, *input):
             else:
                 activations[out_op_id] = op(*in_activations)
 
+            if self.debug:
+                # compare if the activations of pytorch are the same as from onnxruntime
+                debug_model_conversion(
+                    self.onnx_model,
+                    [activations[x] for x in self.input_names],
+                    activations[out_op_id],
+                    node,
+                )
+
         # collect all outputs
         outputs = [activations[x.name] for x in self.onnx_model.graph.output]
         if len(outputs) == 1:

onnx2pytorch/convert/operations.py

@@ -2,16 +2,19 @@
 
 import torch
 from torch import nn
+from torch.nn import functional as F
 from onnx import numpy_helper
 
 from onnx2pytorch.convert.attribute import extract_attributes
 from onnx2pytorch.convert.layer import (
     convert_layer,
     convert_linear_layer,
     convert_batch_norm_layer,
+    convert_instance_norm_layer,
 )
 from onnx2pytorch.operations import *
 from onnx2pytorch.operations.base import OperatorWrapper
+from onnx2pytorch.operations import Resize, Upsample
 from onnx2pytorch.utils import value_wrapper
 
 
@@ -41,6 +44,8 @@ def convert_operations(onnx_model, batch_dim=0):
             op = convert_layer(node, "Conv", params)
         elif node.op_type == "Relu":
             op = nn.ReLU(inplace=True)
+        elif node.op_type == "LeakyRelu":
+            op = nn.LeakyReLU(**extract_attributes(node), inplace=True)
         elif node.op_type == "Sigmoid":
             op = nn.Sigmoid()
         elif node.op_type == "MaxPool":
@@ -54,6 +59,8 @@ def convert_operations(onnx_model, batch_dim=0):
             op.feature_dim = batch_dim + 1  # Necessary for transformers
         elif node.op_type == "BatchNormalization":
             op = convert_batch_norm_layer(node, params=params)
+        elif node.op_type == "InstanceNormalization":
+            op = convert_instance_norm_layer(node, params=params)
         elif node.op_type == "Concat":
             op = partial(torch.cat, **extract_attributes(node))
         elif node.op_type == "Constant":
@@ -75,7 +82,7 @@ def convert_operations(onnx_model, batch_dim=0):
         elif node.op_type == "ConstantOfShape":
             op = ConstantOfShape(**extract_attributes(node))
         elif node.op_type == "Slice":
-            op = Slice()
+            op = Slice(**extract_attributes(node))
         elif node.op_type == "Cast":
             op = Cast(**extract_attributes(node))
         elif node.op_type == "Where":
@@ -136,6 +143,10 @@ def convert_operations(onnx_model, batch_dim=0):
             op = convert_layer(node, "ConvTranspose", params)
         elif node.op_type == "Identity":
             op = nn.Identity()
+        elif node.op_type == "Resize":
+            op = Resize(**extract_attributes(node))
+        elif node.op_type == "Upsample":
+            op = Upsample(**extract_attributes(node))
         elif node.op_type == "OneHot":
             op = OneHot(**extract_attributes(node))
         elif node.op_type == "Pad":
@@ -146,6 +157,10 @@ def convert_operations(onnx_model, batch_dim=0):
             op = OperatorWrapper(torch.tanh)
         elif node.op_type == "Erf":
             op = OperatorWrapper(torch.erf)
+        elif node.op_type == "Log":
+            op = OperatorWrapper(torch.log)
+        elif node.op_type == "Exp":
+            op = OperatorWrapper(torch.exp)
         else:
             op = getattr(torch, node.op_type.lower(), None)
             if op is None: