Merge pull request #9 from ToriML/develop

Version 0.3.0

Author: Talmaj

onnx2pytorch/convert/attribute.py

@@ -101,6 +101,8 @@ def extract_attributes(node):
             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" and node.op_type == "Elu":
+            kwargs["alpha"] = extract_attr_values(attr)
         elif attr.name == "alpha":
             kwargs["weight_multiplier"] = extract_attr_values(attr)
         elif attr.name == "beta":

onnx2pytorch/convert/layer.py

@@ -44,8 +44,8 @@ def convert_layer(node, layer_type, params=None):
             "Unexpected length of kernel_size dimension: {}".format(kernel_size_length)
         )
 
+    pad_layer = 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)
@@ -58,18 +58,23 @@ def convert_layer(node, layer_type, params=None):
             )
 
         # if padding is a layer, remove from kwargs and prepend later
-        if isinstance(kwargs["padding"], nn.Module):
+        if "padding" in kwargs and 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.)
+
+        # if padding is a layer, remove from kwargs and prepend later
+        if "padding" in kwargs and isinstance(kwargs["padding"], nn.Module):
+            pad_layer = kwargs.pop("padding")
         layer = layer(**kwargs)
 
+    if pad_layer is not None:
+        layer = nn.Sequential(pad_layer, layer)
+
     return layer
 
 
@@ -80,8 +85,8 @@ def convert_batch_norm_layer(node, params):
     kwargs["num_features"] = params[0].dims[0]
     # initialize layer and load weights
     layer = layer(**kwargs)
-    key = ["weight", "bias", "running_mean", "running_var"]
-    for key, value in zip(key, params):
+    keys = ["weight", "bias", "running_mean", "running_var"]
+    for key, value in zip(keys, params):
         getattr(layer, key).data = torch.from_numpy(numpy_helper.to_array(value))
 
     return layer
@@ -95,8 +100,8 @@ def convert_instance_norm_layer(node, params):
     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):
+    keys = ["weight", "bias"]
+    for key, value in zip(keys, params):
         getattr(layer, key).data = torch.from_numpy(numpy_helper.to_array(value))
 
     return layer

onnx2pytorch/convert/operations.py

@@ -20,7 +20,7 @@
 
 def convert_operations(onnx_model, batch_dim=0):
     """
-    Convert onnx model operations. Yields onnx's operator_id, opeartor_name and
+    Convert onnx model operations. Yields onnx's operator_id, operator_name and
     converted pytorch operator.
 
     Parameters
@@ -35,6 +35,7 @@ def convert_operations(onnx_model, batch_dim=0):
     iterator: (op_id, op_name, op)
     """
     weights = {tensor.name: tensor for tensor in onnx_model.graph.initializer}
+    opset_version = onnx_model.opset_import[0].version
 
     for i, node in enumerate(onnx_model.graph.node):
         # extract only useful inputs
@@ -46,6 +47,8 @@ def convert_operations(onnx_model, batch_dim=0):
             op = nn.ReLU(inplace=True)
         elif node.op_type == "LeakyRelu":
             op = nn.LeakyReLU(**extract_attributes(node), inplace=True)
+        elif node.op_type == "Elu":
+            op = nn.ELU(**extract_attributes(node), inplace=True)
         elif node.op_type == "Sigmoid":
             op = nn.Sigmoid()
         elif node.op_type == "MaxPool":
@@ -73,14 +76,18 @@ def convert_operations(onnx_model, batch_dim=0):
             op = Reshape(shape)
         elif node.op_type == "Shape":
             op = Shape()
+        elif node.op_type == "Expand":
+            op = Expand()
         elif node.op_type == "Gather":
             op = Gather(**extract_attributes(node))
         elif node.op_type == "Squeeze":
-            op = Squeeze(**extract_attributes(node))
+            op = Squeeze(opset_version=opset_version, **extract_attributes(node))
         elif node.op_type == "Unsqueeze":
-            op = partial(torch.unsqueeze, **extract_attributes(node))
+            op = Unsqueeze(opset_version=opset_version, **extract_attributes(node))
         elif node.op_type == "ConstantOfShape":
             op = ConstantOfShape(**extract_attributes(node))
+        elif node.op_type == "Range":
+            op = Range()
         elif node.op_type == "Slice":
             op = Slice(**extract_attributes(node))
         elif node.op_type == "Cast":
@@ -161,6 +168,14 @@ def convert_operations(onnx_model, batch_dim=0):
             op = OperatorWrapper(torch.log)
         elif node.op_type == "Exp":
             op = OperatorWrapper(torch.exp)
+        elif node.op_type == "Reciprocal":
+            op = OperatorWrapper(torch.reciprocal)
+        elif node.op_type == "And":
+            op = OperatorWrapper(torch.logical_and)
+        elif node.op_type == "Or":
+            op = OperatorWrapper(torch.logical_or)
+        elif node.op_type == "Not":
+            op = OperatorWrapper(torch.logical_not)
         else:
             op = getattr(torch, node.op_type.lower(), None)
             if op is None:

onnx2pytorch/operations/__init__.py

@@ -1,5 +1,6 @@
 from .add import Add
 from .batchnorm import BatchNormUnsafe
+from .expand import Expand
 from .instancenorm import InstanceNormUnsafe
 from .cast import Cast
 from .constant import ConstantOfShape
@@ -8,16 +9,19 @@
 from .onehot import OneHot
 from .pad import Pad
 from .pooling import GlobalAveragePool
+from .range import Range
 from .reshape import Reshape
 from .shape import Shape
 from .slice import Slice
 from .split import Split
 from .squeeze import Squeeze
 from .resize import Resize, Upsample
+from .unsqueeze import Unsqueeze
 
 __all__ = [
     "Add",
     "BatchNormUnsafe",
+    "Expand",
     "InstanceNormUnsafe",
     "Cast",
     "ConstantOfShape",
@@ -26,11 +30,13 @@
     "OneHot",
     "Pad",
     "GlobalAveragePool",
+    "Range",
     "Reshape",
     "Shape",
     "Slice",
     "Split",
     "Squeeze",
     "Resize",
+    "Unsqueeze",
     "Upsample",
 ]

onnx2pytorch/operations/expand.py

@@ -0,0 +1,8 @@
+import torch
+from torch import nn
+
+
+class Expand(nn.Module):
+    def forward(self, input: torch.Tensor, shape: torch.Tensor):
+        # return input.expand(torch.Size(shape))
+        return input * torch.ones(torch.Size(shape), dtype=input.dtype)

onnx2pytorch/operations/range.py

@@ -0,0 +1,7 @@
+import torch
+from torch import nn
+
+
+class Range(nn.Module):
+    def forward(self, start: torch.Tensor, limit: torch.Tensor, delta: torch.Tensor):
+        return torch.arange(start=start, end=limit, step=delta)

onnx2pytorch/operations/slice.py

@@ -28,6 +28,8 @@ def forward(
         if steps is None:
             steps = tuple(1 for _ in axes)
 
+        axes = [input.ndim + x if x < 0 else x for x in axes]
+
         selection = [slice(None) for _ in range(max(axes) + 1)]
         for i, axis in enumerate(axes):
             selection[axis] = slice(starts[i], ends[i], steps[i])

onnx2pytorch/operations/squeeze.py

@@ -5,16 +5,22 @@
 
 
 class Squeeze(Operator):
-    def __init__(self, dim=None):
+    def __init__(self, opset_version, dim=None):
+        self.opset_version = opset_version
         self.dim = dim
         super().__init__()
 
-    def forward(self, input):
-        if self.dim is None:
+    def forward(self, input: torch.Tensor, axes: torch.Tensor = None):
+        if self.opset_version < 13:
+            dims = self.dim
+        else:
+            dims = axes
+
+        if dims is None:
             return torch.squeeze(input)
-        elif isinstance(self.dim, int):
-            return torch.squeeze(input, dim=self.dim)
+        elif isinstance(dims, int):
+            return torch.squeeze(input, dim=dims)
         else:
-            for dim in sorted(self.dim, reverse=True):
+            for dim in sorted(dims, reverse=True):
                 input = torch.squeeze(input, dim=dim)
             return input

onnx2pytorch/operations/unsqueeze.py

@@ -0,0 +1,25 @@
+import torch
+from torch import nn
+
+from onnx2pytorch.operations.base import Operator
+
+
+class Unsqueeze(Operator):
+    def __init__(self, opset_version, dim=None):
+        self.opset_version = opset_version
+        self.dim = dim
+        super().__init__()
+
+    def forward(self, data: torch.Tensor, axes: torch.Tensor = None):
+        if self.opset_version < 13:
+            dims = self.dim
+        else:
+            dims = torch.Size(axes)
+        if dims is None:
+            raise ValueError("Unsqueeze expects axes")
+        elif isinstance(dims, int):
+            return torch.unsqueeze(data, dim=dims)
+        else:
+            for dim in sorted(dims, reverse=True):
+                data = torch.unsqueeze(data, dim=dim)
+            return data

setup.py

@@ -12,7 +12,7 @@
 
 setup(
     name="onnx2pytorch",
-    version="0.2.0",
+    version="0.3.0",
     description="Library to transform onnx model to pytorch.",
     license="apache-2.0",
     author="Talmaj Marinc",