reducemax

Author: calvinmccarter-at-lightmatter

onnx2pytorch/convert/operations.py

@@ -199,9 +199,7 @@ def convert_operations(onnx_graph, opset_version, batch_dim=0, enable_pruning=Tr
         elif node.op_type == "Reciprocal":
             op = OperatorWrapper(torch.reciprocal)
         elif node.op_type == "ReduceMax":
-            kwargs = dict(keepdim=True)
-            kwargs.update(extract_attributes(node))
-            op = partial(torch.max, **kwargs)
+            op = ReduceMax(**extract_attributes(node))
         elif node.op_type == "ReduceMean":
             kwargs = dict(keepdim=True)
             kwargs.update(extract_attributes(node))

onnx2pytorch/operations/__init__.py

@@ -21,6 +21,7 @@
 from .pad import Pad
 from .prelu import PRelu
 from .range import Range
+from .reducemax import ReduceMax
 from .reducesum import ReduceSum
 from .reshape import Reshape
 from .resize import Resize, Upsample
@@ -61,6 +62,7 @@
     "Pad",
     "PRelu",
     "Range",
+    "ReduceMax",
     "ReduceSum",
     "Reshape",
     "Resize",

onnx2pytorch/operations/reducemax.py

@@ -0,0 +1,15 @@
+import torch
+from torch import nn
+
+
+class ReduceMax(nn.Module):
+    def __init__(self, dim=None, keepdim=True):
+        self.dim = dim
+        self.keepdim = keepdim
+        super().__init__()
+
+    def forward(self, data: torch.Tensor):
+        dim = self.dim
+        if dim is None:
+            dim = tuple(range(data.ndim))
+        return torch.amax(data, dim=dim, keepdim=self.keepdim)