Avoid calling torch.clamp when neither min and max are given

Author: calvinmccarter-at-lightmatter

download_fixtures.sh

@@ -42,6 +42,12 @@ if [[ $1 == "--all" ]]; 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
+
+  if [[ ! -f mobilenetv2-7.onnx ]]; then
+    echo Downloading mobilenetv2-7
+    curl -LJo mobilenetv2-7.onnx https://github.com/onnx/models/raw/master/vision/classification/mobilenet/model/mobilenetv2-7.onnx\?raw\=true
+  fi
+ 
 fi
 
 echo "All models downloaded."

onnx2pytorch/convert/attribute.py

@@ -124,6 +124,10 @@ def extract_attributes(node):
             kwargs["largest"] = extract_attr_values(attr)
         elif attr.name == "layout":
             kwargs["layout"] = extract_attr_values(attr)
+        elif attr.name == "max":
+            kwargs["max"] = extract_attr_values(attr)
+        elif attr.name == "min":
+            kwargs["min"] = extract_attr_values(attr)
         elif attr.name == "mode":
             kwargs["mode"] = extract_attr_values(attr)
         elif attr.name == "momentum":

onnx2pytorch/convert/operations.py

@@ -93,7 +93,7 @@ def convert_operations(onnx_graph, opset_version, batch_dim=0, enable_pruning=Tr
         elif node.op_type == "Ceil":
             op = OperatorWrapper(torch.ceil)
         elif node.op_type == "Clip":
-            op = OperatorWrapper(torch.clamp)
+            op = Clip(**extract_attributes(node))
         elif node.op_type == "Concat":
             op = partial(torch.cat, **extract_attributes(node))
         elif node.op_type == "Constant":

onnx2pytorch/operations/__init__.py

@@ -2,6 +2,7 @@
 from .batchnorm import BatchNormWrapper
 from .bitshift import BitShift
 from .cast import Cast
+from .clip import Clip
 from .constant import Constant
 from .constantofshape import ConstantOfShape
 from .div import Div
@@ -41,6 +42,7 @@
     "BatchNormWrapper",
     "BitShift",
     "Cast",
+    "Clip",
     "Constant",
     "ConstantOfShape",
     "Div",

onnx2pytorch/operations/clip.py

@@ -0,0 +1,19 @@
+import torch
+from torch import nn
+
+
+class Clip(nn.Module):
+    def __init__(self, min=None, max=None):
+        super().__init__()
+        self.min = min
+        self.max = max
+
+    def forward(self, input, min=None, max=None):
+        if min is None:
+            min = self.min
+        if max is None:
+            max = self.max
+        if min is None and max is None:
+            return input
+        else:
+            return torch.clamp(input, min=min, max=max)

tests/onnx2pytorch/operations/test_clip.py

@@ -0,0 +1,30 @@
+import numpy as np
+import pytest
+import torch
+
+from onnx2pytorch.operations.clip import Clip
+
+
+def test_clip():
+    x_np = np.random.randn(3, 4, 5).astype(np.float32)
+    x = torch.from_numpy(x_np)
+
+    op = Clip(min=-1, max=1)
+    exp_y_np = np.clip(x_np, -1, 1)
+    exp_y = torch.from_numpy(exp_y_np)
+    assert torch.equal(op(x), exp_y)
+
+    op = Clip(min=0)
+    exp_y_np = np.clip(x_np, 0, np.inf)
+    exp_y = torch.from_numpy(exp_y_np)
+    assert torch.equal(op(x), exp_y)
+
+    op = Clip(max=0)
+    exp_y_np = np.clip(x_np, np.NINF, 0)
+    exp_y = torch.from_numpy(exp_y_np)
+    assert torch.equal(op(x), exp_y)
+
+    op = Clip()
+    exp_y_np = np.clip(x_np, np.NINF, np.inf)
+    exp_y = torch.from_numpy(exp_y_np)
+    assert torch.equal(op(x), exp_y)