Repository: onnx/onnx
Pull request: add eliminate nop monotone argmax pass
fixes #1429

Problem statement:
Optimization Pass: Monotonic growing function is useless before argmax calculation.
It's standard to use softmax to generate a probability distribution at the end of a model. If we want to get the predicted class we end up with `net(x).softmax(dim=1).argmax(dim=1)`. 

In reality though because softmax is a monotonic growing operation we can replace the expression with `net(x).argmax(dim=1)`

Hints from issue discussion:

Commit message: Merge branch 'master' into argmax_monotone

Patch to review:
diff --git a/onnx/common/interned_strings.h b/onnx/common/interned_strings.h
index 542a02e204b..ce6ed29ea82 100644
--- a/onnx/common/interned_strings.h
+++ b/onnx/common/interned_strings.h
@@ -2,170 +2,167 @@
 // Adventurous users should note that the APIs will probably change.
 
 #pragma once
-#include <vector>
 #include <stdint.h>
 #include <string>
 #include <unordered_map>
+#include <vector>
 
 namespace ONNX_NAMESPACE {
 
 #define FORALL_BUILTIN_SYMBOLS(_) \
-_(PythonOp) \
-_(CppOp) \
-_(Param) \
-_(Select) \
-_(Return) \
-_(Eval) \
-_(add) \
-_(Add) \
-_(Div) \
-_(Mul) \
-_(Neg) \
-_(Sub) \
-_(Pow) \
-_(Sigmoid) \
-_(Softmax) \
-_(LogSoftmax) \
-_(Tanh) \
-_(mul) \
-_(neg) \
-_(sigmoid) \
-_(tanh) \
-_(Constant) \
-_(cat) \
-_(Slice) \
-_(Squeeze) \
-_(Undefined) \
-_(FusionGroup) \
-_(Gemm) \
-_(Tile) \
-_(SubConstant) \
-_(Scale) \
-_(Transpose) \
-_(Pad) \
-_(Reshape) \
-_(split) \
-_(chunk) \
-_(Offset) \
-_(value) \
-_(Subgraph) \
-_(BatchNormalization) \
-_(Conv) \
-_(ConvTranspose) \
-_(is_test) \
-_(epsilon) \
-_(expand) \
-_(Expand) \
-_(order) \
-_(momentum) \
-_(consumed_inputs) \
-_(kernels) \
-_(kernel_shape) \
-_(kernel) \
-_(scale) \
-_(strides) \
-_(stride) \
-_(pads) \
-_(pad) \
-_(beta) \
-_(alpha) \
-_(dilations) \
-_(dilation) \
-_(broadcast) \
-_(axis) \
-_(size) \
-_(dim) \
-_(perm) \
-_(shape) \
-_(axes) \
-_(group) \
-_(inplace) \
-_(transA) \
-_(transB) \
-_(other) \
-_(__and__) \
-_(__lshift__) \
-_(__or__) \
-_(__rshift__) \
-_(__xor__) \
-_(abs) \
-_(acos) \
-_(asin) \
-_(atan) \
-_(atan2) \
-_(ceil) \
-_(clamp) \
-_(cos) \
-_(cosh) \
-_(div) \
-_(eq) \
-_(equal) \
-_(exp) \
-_(expm1) \
-_(floor) \
-_(fmod) \
-_(frac) \
-_(ge) \
-_(gt) \
-_(le) \
-_(lerp) \
-_(lgamma) \
-_(log) \
-_(Log) \
-_(log1p) \
-_(lt) \
-_(max) \
-_(min) \
-_(ne) \
-_(ones) \
-_(pow) \
-_(reciprocal) \
-_(remainder) \
-_(round) \
-_(rsqrt) \
-_(sin) \
-_(sinh) \
-_(sqrt) \
-_(sub) \
-_(tan) \
-_(trunc) \
-_(zeros) \
-_(exponent) \
-_(device) \
-_(Identity) \
-_(Loop) \
-_(If) \
-_(body) \
-_(then_branch) \
-_(else_branch) \
-_(Captured) \
-_(__control_inputs) \
-_(count_include_pad) \
-_(storage_order) \
-_(Unsqueeze)
+  _(PythonOp)                     \
+  _(CppOp)                        \
+  _(Param)                        \
+  _(Select)                       \
+  _(Return)                       \
+  _(Eval)                         \
+  _(add)                          \
+  _(Add)                          \
+  _(Div)                          \
+  _(Mul)                          \
+  _(Neg)                          \
+  _(Sub)                          \
+  _(Pow)                          \
+  _(Sigmoid)                      \
+  _(ArgMax)                       \
+  _(Softmax)                      \
+  _(LogSoftmax)                   \
+  _(Tanh)                         \
+  _(mul)                          \
+  _(neg)                          \
+  _(sigmoid)                      \
+  _(tanh)                         \
+  _(Constant)                     \
+  _(cat)                          \
+  _(Slice)                        \
+  _(Squeeze)                      \
+  _(Undefined)                    \
+  _(FusionGroup)                  \
+  _(Gemm)                         \
+  _(Tile)                         \
+  _(SubConstant)                  \
+  _(Scale)                        \
+  _(Transpose)                    \
+  _(Pad)                          \
+  _(Reshape)                      \
+  _(split)                        \
+  _(chunk)                        \
+  _(Offset)                       \
+  _(value)                        \
+  _(Subgraph)                     \
+  _(BatchNormalization)           \
+  _(Conv)                         \
+  _(ConvTranspose)                \
+  _(is_test)                      \
+  _(epsilon)                      \
+  _(expand)                       \
+  _(Expand)                       \
+  _(order)                        \
+  _(momentum)                     \
+  _(consumed_inputs)              \
+  _(kernels)                      \
+  _(kernel_shape)                 \
+  _(kernel)                       \
+  _(scale)                        \
+  _(strides)                      \
+  _(stride)                       \
+  _(pads)                         \
+  _(pad)                          \
+  _(beta)                         \
+  _(alpha)                        \
+  _(dilations)                    \
+  _(dilation)                     \
+  _(broadcast)                    \
+  _(axis)                         \
+  _(size)                         \
+  _(dim)                          \
+  _(perm)                         \
+  _(shape)                        \
+  _(axes)                         \
+  _(group)                        \
+  _(inplace)                      \
+  _(transA)                       \
+  _(transB)                       \
+  _(other)                        \
+  _(__and__)                      \
+  _(__lshift__)                   \
+  _(__or__)                       \
+  _(__rshift__)                   \
+  _(__xor__)                      \
+  _(abs)                          \
+  _(acos)                         \
+  _(asin)                         \
+  _(atan)                         \
+  _(atan2)                        \
+  _(ceil)                         \
+  _(clamp)                        \
+  _(cos)                          \
+  _(cosh)                         \
+  _(div)                          \
+  _(eq)                           \
+  _(equal)                        \
+  _(Exp)                          \
+  _(expm1)                        \
+  _(floor)                        \
+  _(fmod)                         \
+  _(frac)                         \
+  _(ge)                           \
+  _(gt)                           \
+  _(le)                           \
+  _(lerp)                         \
+  _(lgamma)                       \
+  _(Log)                          \
+  _(log1p)                        \
+  _(lt)                           \
+  _(max)                          \
+  _(min)                          \
+  _(ne)                           \
+  _(ones)                         \
+  _(pow)                          \
+  _(reciprocal)                   \
+  _(remainder)                    \
+  _(round)                        \
+  _(rsqrt)                        \
+  _(sin)                          \
+  _(sinh)                         \
+  _(Sqrt)                         \
+  _(sub)                          \
+  _(tan)                          \
+  _(trunc)                        \
+  _(zeros)                        \
+  _(exponent)                     \
+  _(device)                       \
+  _(Identity)                     \
+  _(Loop)                         \
+  _(If)                           \
+  _(body)                         \
+  _(then_branch)                  \
+  _(else_branch)                  \
+  _(Captured)                     \
+  _(__control_inputs)             \
+  _(count_include_pad)            \
+  _(storage_order)                \
+  _(Unsqueeze)
 
 enum BuiltinSymbol {
-  #define DEFINE_SYMBOL(s) \
-    k##s,
+#define DEFINE_SYMBOL(s) k##s,
   FORALL_BUILTIN_SYMBOLS(DEFINE_SYMBOL)
-  #undef DEFINE_SYMBOL
-  kLastSymbol, //where we start counting for new symbols
+#undef DEFINE_SYMBOL
+      kLastSymbol, // where we start counting for new symbols
 };
 
-
 struct Symbol {
   Symbol() {}
-  /*implicit*/ Symbol(BuiltinSymbol value)
-  : value(value) {}
-  explicit Symbol(const std::string & s);
-  explicit Symbol(uint32_t value)
-  : value(value) {}
+  /*implicit*/ Symbol(BuiltinSymbol value) : value(value) {}
+  explicit Symbol(const std::string& s);
+  explicit Symbol(uint32_t value) : value(value) {}
 
   operator uint32_t() const {
     return value;
   }
-  const char * toString() const;
-private:
+  const char* toString() const;
+
+ private:
   uint32_t value;
 };
 
@@ -180,7 +177,7 @@ static inline bool operator==(Symbol lhs, BuiltinSymbol rhs) {
   return static_cast<uint32_t>(lhs) == static_cast<uint32_t>(rhs);
 }
 
-inline Symbol operator "" _sym(const char * s, size_t) {
+inline Symbol operator"" _sym(const char* s, size_t) {
   return Symbol(s);
 }
 
@@ -188,10 +185,10 @@ inline Symbol operator "" _sym(const char * s, size_t) {
 
 // make symbol behave like an integer in hash tables
 namespace std {
-  template<>
-  struct hash<ONNX_NAMESPACE::Symbol> {
-    std::size_t operator()(ONNX_NAMESPACE::Symbol s) const {
-      return std::hash<uint32_t>()(static_cast<uint32_t>(s));
-    }
-  };
-} // namespace std
+template <>
+struct hash<ONNX_NAMESPACE::Symbol> {
+  std::size_t operator()(ONNX_NAMESPACE::Symbol s) const {
+    return std::hash<uint32_t>()(static_cast<uint32_t>(s));
+  }
+};
+} // namespace std
\ No newline at end of file
diff --git a/onnx/examples/optimize_onnx.ipynb b/onnx/examples/optimize_onnx.ipynb
index df14455fb1b..44a15ff0d40 100644
--- a/onnx/examples/optimize_onnx.ipynb
+++ b/onnx/examples/optimize_onnx.ipynb
@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": 2,
    "metadata": {},
    "outputs": [
     {
@@ -103,7 +103,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": 3,
    "metadata": {},
    "outputs": [
     {
@@ -113,6 +113,7 @@
       "Available optimization passes:\n",
       "eliminate_deadend\n",
       "eliminate_identity\n",
+      "eliminate_nop_monotone_argmax\n",
       "eliminate_nop_pad\n",
       "eliminate_nop_transpose\n",
       "eliminate_unused_initializer\n",
diff --git a/onnx/optimizer/pass_registry.h b/onnx/optimizer/pass_registry.h
index 459534e24fe..5128bc1f992 100644
--- a/onnx/optimizer/pass_registry.h
+++ b/onnx/optimizer/pass_registry.h
@@ -8,6 +8,7 @@
 #include "onnx/common/stl_backports.h"
 #include "onnx/optimizer/passes/eliminate_deadend.h"
 #include "onnx/optimizer/passes/eliminate_identity.h"
+#include "onnx/optimizer/passes/eliminate_nop_monotone_argmax.h"
 #include "onnx/optimizer/passes/eliminate_nop_pad.h"
 #include "onnx/optimizer/passes/eliminate_nop_transpose.h"
 #include "onnx/optimizer/passes/eliminate_unused_initializer.h"
@@ -38,6 +39,7 @@ struct GlobalPassRegistry {
     registerPass<NopEmptyPass>();
     registerPass<EliminateDeadEnd>();
     registerPass<EliminateIdentity>();
+    registerPass<EliminateNopMonotoneArgmax>();
     registerPass<EliminateNopPad>();
     registerPass<EliminateNopTranspose>();
     registerPass<EliminateUnusedInitializer>();
diff --git a/onnx/optimizer/passes/eliminate_nop_monotone_argmax.h b/onnx/optimizer/passes/eliminate_nop_monotone_argmax.h
new file mode 100644
index 00000000000..06680729d42
--- /dev/null
+++ b/onnx/optimizer/passes/eliminate_nop_monotone_argmax.h
@@ -0,0 +1,56 @@
+// ATTENTION: The code in this file is highly EXPERIMENTAL.
+// Adventurous users should note that the APIs will probably change.
+#pragma once
+
+#include "onnx/optimizer/pass.h"
+
+namespace ONNX_NAMESPACE {
+namespace optimization {
+
+const std::unordered_set<NodeKind> monotone_node_no_axis_kind{kLog,
+                                                              kExp,
+                                                              kSqrt};
+
+const std::unordered_set<NodeKind> monotone_node_axis_kind{kSoftmax,
+                                                           kLogSoftmax};
+
+struct EliminateNopMonotoneArgmax final : public PredicateBasedPass {
+  explicit EliminateNopMonotoneArgmax()
+      : PredicateBasedPass(
+            PassType::Nop,
+            PassEfficiency::Partial,
+            PassOptimizationType::Compute) {}
+
+  std::string getPassName() const override {
+    return "eliminate_nop_monotone_argmax";
+  }
+
+  static inline bool satisfies_monotone_condition(int64_t axis, Node* node) {
+    if (monotone_node_no_axis_kind.find(node->kind()) !=
+        monotone_node_no_axis_kind.end()) {
+      return true;
+    }
+    if (monotone_node_axis_kind.find(node->kind()) !=
+        monotone_node_axis_kind.end()) {
+      return axis == node->i(kaxis);
+    }
+    return false;
+  }
+
+  bool patternMatchPredicate(Node* node) override {
+    return node->kind() == kArgMax && node->inputs().size() == 1 &&
+        satisfies_monotone_condition(node->i(kaxis), node->input()->node());
+  }
+  bool runTransform(Node* node, Graph& graph, NodeDestroyType& destroy_current)
+      override {
+    Node* monotone_node = node->input()->node();
+    if (monotone_node->output()->uses().size() == 1) {
+      monotone_node->output()->replaceAllUsesWith(monotone_node->input());
+      monotone_node->destroy();
+      return true;
+    }
+    return false;
+  }
+};
+} // namespace optimization
+} // namespace ONNX_NAMESPACE
\ No newline at end of file
diff --git a/onnx/test/optimizer_test.py b/onnx/test/optimizer_test.py
index 14f3d5ecb8f..7fa7aecd891 100644
--- a/onnx/test/optimizer_test.py
+++ b/onnx/test/optimizer_test.py
@@ -926,6 +926,140 @@ def test_deadend_elimination_simple(self):  # type: () -> None
     def test_deadend_elimination_simple_fixed(self):  # type: () -> None
         self._internal_test_deadend_elimination(True)
 
+    def test_eliminate_nop_monotone_argmax_basic_no_node_axis(self):  # type: () -> None
+        for node_name in ["Log", "Exp", "Sqrt"]:
+            for axis in range(3):
+                node = helper.make_node(node_name, ["X"], ["Y"])
+                argmax = helper.make_node("ArgMax", ["Y"], ["Z"], axis=axis)
+                graph = helper.make_graph(
+                    [node, argmax],
+                    "test",
+                    [helper.make_tensor_value_info(
+                        "X", TensorProto.FLOAT, (5, 7, 11))],
+                    [helper.make_tensor_value_info("Z", TensorProto.FLOAT, (5, 7, 11))])
+                optimized_model = self._optimized(
+                    graph, ["eliminate_nop_monotone_argmax"])
+                assert len(optimized_model.graph.output) == 1
+                assert len(optimized_model.graph.node) == 1
+                assert optimized_model.graph.output[0].type.tensor_type.elem_type == TensorProto.FLOAT
+                assert optimized_model.graph.node[0].op_type == "ArgMax"
+                assert optimized_model.graph.node[0].attribute[0].name == "axis"
+                assert optimized_model.graph.node[0].attribute[0].i == axis
+
+    def test_eliminate_nop_monotone_argmax_basic_with_node_axis(self):  # type: () -> None
+        for node_name in ["Softmax", "LogSoftmax"]:
+            for axis_n in range(3):
+                for axis_max in range(3):
+                    node = helper.make_node(node_name, ["X"], ["Y"], axis=axis_n)
+                    argmax = helper.make_node("ArgMax", ["Y"], ["Z"], axis=axis_max)
+                    graph = helper.make_graph(
+                        [node, argmax],
+                        "test",
+                        [helper.make_tensor_value_info(
+                            "X", TensorProto.FLOAT, (5, 7, 11))],
+                        [helper.make_tensor_value_info("Z", TensorProto.FLOAT, (5, 7, 11))])
+                    optimized_model = self._optimized(
+                        graph, ["eliminate_nop_monotone_argmax"])
+                    if axis_max == axis_n:
+                        assert len(optimized_model.graph.output) == 1
+                        assert len(optimized_model.graph.node) == 1
+                        assert optimized_model.graph.output[0].type.tensor_type.elem_type == TensorProto.FLOAT
+                        assert optimized_model.graph.node[0].op_type == "ArgMax"
+                        assert optimized_model.graph.node[0].attribute[0].name == "axis"
+                        assert optimized_model.graph.node[0].attribute[0].i == axis_max
+                    else:
+                        assert optimized_model.graph == graph
+
+    def test_eliminate_nop_monotone_argmax_multiple_out(self):  # type: () -> None
+        for node_name in ["Log", "Exp", "Sqrt"]:
+            for axis in range(3):
+                node = helper.make_node(node_name, ["X"], ["Y"])
+                node2 = helper.make_node(node_name, ["Y"], ["Z1"])
+                argmax = helper.make_node("ArgMax", ["Y"], ["Z"], axis=axis)
+                graph = helper.make_graph(
+                    [node, node2, argmax],
+                    "test",
+                    [helper.make_tensor_value_info(
+                        "X", TensorProto.FLOAT, (5, 7, 11))],
+                    [helper.make_tensor_value_info("Z", TensorProto.FLOAT, (5, 7, 11)),
+                     helper.make_tensor_value_info("Z1", TensorProto.FLOAT, (5, 7, 11))])
+                optimized_model = self._optimized(
+                    graph, ["eliminate_nop_monotone_argmax"])
+                assert optimized_model.graph == graph
+
+    def test_eliminate_nop_monotone_argmax_consecutive(self):  # type: () -> None
+        def _assertion(graph, optimized_model, axis_aligned, true_axis):  # type: (GraphProto, ModelProto, bool, int) -> None
+            if axis_aligned:
+                assert len(optimized_model.graph.output) == 1
+                assert len(optimized_model.graph.node) == 1
+                assert optimized_model.graph.output[0].type.tensor_type.elem_type == TensorProto.FLOAT
+                assert optimized_model.graph.node[0].op_type == "ArgMax"
+                assert optimized_model.graph.node[0].attribute[0].name == "axis"
+                assert optimized_model.graph.node[0].attribute[0].i == true_axis
+            else:
+                assert optimized_model.graph == graph
+        # no axis X no axis test
+        for node_name_0 in ["Log", "Exp", "Sqrt"]:
+            for node_name_1 in ["Log", "Exp", "Sqrt"]:
+                for axis in range(3):
+                    node = helper.make_node(node_name_0, ["X"], ["Y"])
+                    node2 = helper.make_node(node_name_1, ["Y"], ["Y1"])
+                    argmax = helper.make_node("ArgMax", ["Y1"], ["Z"], axis=axis)
+                    graph = helper.make_graph(
+                        [node, node2, argmax],
+                        "test",
+                        [helper.make_tensor_value_info(
+                            "X", TensorProto.FLOAT, (5, 7, 11))],
+                        [helper.make_tensor_value_info("Z", TensorProto.FLOAT, (5, 7, 11))])
+                    optimized_model = self._optimized(
+                        graph, ["eliminate_nop_monotone_argmax"], True)
+                    _assertion(graph, optimized_model, True, axis)
+        # no axis X axis test
+        for node_name_0 in ["Log", "Exp", "Sqrt"]:
+            for node_name_1 in ["Softmax", "LogSoftmax"]:
+                for axis_0 in range(3):
+                    for axis_1 in range(3):
+                        node = helper.make_node(node_name_0, ["X"], ["Y"])
+                        node2 = helper.make_node(node_name_1, ["Y"], ["Y1"], axis=axis_0)
+                        argmax = helper.make_node("ArgMax", ["Y1"], ["Z"], axis=axis_1)
+                        graph = helper.make_graph(
+                            [node, node2, argmax],
+                            "test",
+                            [helper.make_tensor_value_info(
+                                "X", TensorProto.FLOAT, (5, 7, 11))],
+                            [helper.make_tensor_value_info("Z", TensorProto.FLOAT, (5, 7, 11))])
+                        optimized_model = self._optimized(
+                            graph, ["eliminate_nop_monotone_argmax"], True)
+                        _assertion(graph, optimized_model, axis_0 == axis_1, axis_1)
+        # axis X axis test
+        for node_name_0 in ["Softmax", "LogSoftmax"]:
+            for node_name_1 in ["Softmax", "LogSoftmax"]:
+                for axis_0 in range(3):
+                    for axis_1 in range(3):
+                        for axis_2 in range(3):
+                            node = helper.make_node(node_name_0, ["X"], ["Y"], axis=axis_0)
+                            node2 = helper.make_node(node_name_1, ["Y"], ["Y1"], axis=axis_1)
+                            argmax = helper.make_node("ArgMax", ["Y1"], ["Z"], axis=axis_2)
+                            graph = helper.make_graph(
+                                [node, node2, argmax],
+                                "test",
+                                [helper.make_tensor_value_info(
+                                    "X", TensorProto.FLOAT, (5, 7, 11))],
+                                [helper.make_tensor_value_info("Z", TensorProto.FLOAT, (5, 7, 11))])
+                            optimized_model = self._optimized(
+                                graph, ["eliminate_nop_monotone_argmax"], True)
+                            if axis_0 == axis_1:  # we can reduce both of the monotonic ops
+                                _assertion(graph, optimized_model, axis_1 == axis_2, axis_2)
+                            elif axis_1 == axis_2:  # we can reduce one of the monotonic ops
+                                assert len(optimized_model.graph.output) == 1
+                                assert len(optimized_model.graph.node) == 2
+                                assert optimized_model.graph.output[0].type.tensor_type.elem_type == TensorProto.FLOAT
+                                assert optimized_model.graph.node[-1].op_type == "ArgMax"
+                                assert optimized_model.graph.node[-1].attribute[0].name == "axis"
+                                assert optimized_model.graph.node[-1].attribute[0].i == axis_2
+                            else:  # we can't reduce anything
+                                assert optimized_model.graph == graph
+
 
 if __name__ == '__main__':
     unittest.main()


Code review comments:
diff_hunk: @@ -0,0 +1,56 @@
+// ATTENTION: The code in this file is highly EXPERIMENTAL.
+// Adventurous users should note that the APIs will probably change.
+#pragma once
+
+#include "onnx/optimizer/pass.h"
+
+namespace ONNX_NAMESPACE {
+namespace optimization {
+
+const std::unordered_set<NodeKind> monotone_node_no_axis_kind{kLog,
+                                                              kExp,
+                                                              kSqrt};
+
+const std::unordered_set<NodeKind> monotone_node_axis_kind{kSoftmax,
+                                                           kLogSoftmax};
+
+struct EliminateNopMonotoneArgmax final : public PredicateBasedPass {
+  explicit EliminateNopMonotoneArgmax()
+      : PredicateBasedPass(
+            PassType::Nop,
+            PassEfficiency::Partial,
+            PassOptimizationType::Compute) {}
+
+  std::string getPassName() const override {
+    return "eliminate_nop_monotone_argmax";
+  }
+
+  static inline bool satisfies_monotone_condition(int64_t axis, Node* node) {
+    if (monotone_node_no_axis_kind.find(node->kind()) !=
+        monotone_node_no_axis_kind.end()) {
+      return true;
+    }
+    if (monotone_node_axis_kind.find(node->kind()) !=
+        monotone_node_axis_kind.end()) {
+      return axis == node->i(kaxis);
path: onnx/optimizer/passes/eliminate_nop_monotone_argmax.h
text: @user1:
if axis is not specified, default value is used. Please check and handle such condition in the code.