Add support for ArgMax and ArgMin in C2 onnx backend and frontend (pytorch#9050)

Summary:
Pass the end to end test cases in onnx/onnx#1049
Closes pytorch#9050

Reviewed By: hlu1

Differential Revision: D8703757

Pulled By: houseroad

fbshipit-source-id: 63308202e349dfc02d532e87f49495ba1aab085b

Author: houseroad

caffe2/onnx/backend.cc

@@ -335,6 +335,8 @@ Caffe2Backend::get_special_operators() const {
   const static std::
       unordered_map<std::string, Caffe2Backend::SpecialOpConverter>
           kSpecialOperators = {
+              {"ArgMax", &Caffe2Backend::CreateArgMaxMin},
+              {"ArgMin", &Caffe2Backend::CreateArgMaxMin},
               {"Cast", &Caffe2Backend::CreateCast},
               {"Constant", &Caffe2Backend::CreateConstant},
               {"Conv", &Caffe2Backend::CreateConvPoolOpBase},
@@ -363,6 +365,17 @@ Caffe2Backend::get_special_operators() const {
 // Special Operator Converters
 //============================
 
+Caffe2Ops Caffe2Backend::CreateArgMaxMin(
+    OnnxNode* onnx_node,
+    int opset_version) {
+  auto& attributes = onnx_node->attributes;
+  if (!attributes.HasAttribute("axis")) {
+    auto* attr = attributes.AddRewrittenAttribute("axis");
+    attr->set_i(0);
+  }
+  return CommonOnnxNodeToCaffe2Ops(onnx_node, opset_version);
+}
+
 Caffe2Ops Caffe2Backend::CreateCast(OnnxNode* onnx_node, int opset_version) {
   auto c2_op = CommonOnnxNodeToCaffe2Ops(onnx_node, opset_version);
 

caffe2/onnx/backend.h

@@ -160,6 +160,8 @@ class Caffe2Backend {
 
   Caffe2Ops CommonOnnxNodeToCaffe2Ops(OnnxNode* onnx_node, int opset_version);
 
+  Caffe2Ops CreateArgMaxMin(OnnxNode* onnx_node, int opset_version);
+
   Caffe2Ops CreateCast(OnnxNode* onnx_node, int opset_version);
 
   Caffe2Ops CreateConstant(OnnxNode* onnx_node, int opset_version);

caffe2/onnx/onnx_exporter.cc

@@ -222,6 +222,8 @@ const std::unordered_map<std::string, OnnxExporter::SpecialOpConverter>&
 OnnxExporter::get_special_operators() const {
   const static std::unordered_map<std::string, OnnxExporter::SpecialOpConverter>
       kSpecialOperators = {
+          {"ArgMax", &OnnxExporter::CreateArgMaxMinOpNodes},
+          {"ArgMin", &OnnxExporter::CreateArgMaxMinOpNodes},
           {"Add", &OnnxExporter::CreateBinaryElementwiseOpNodes},
           {"Sub", &OnnxExporter::CreateBinaryElementwiseOpNodes},
           {"Mul", &OnnxExporter::CreateBinaryElementwiseOpNodes},
@@ -351,6 +353,25 @@ ConvertedResult OnnxExporter::CommonCaffe2OpToOnnxNodes(
   return result;
 }
 
+ConvertedResult OnnxExporter::CreateArgMaxMinOpNodes(
+    const caffe2::OperatorDef& def,
+    const std::unordered_map<std::string, caffe2::TensorShape>& shapes) {
+  auto result = CommonCaffe2OpToOnnxNodes(def);
+  auto& nodes = result.first;
+
+  CAFFE_ENFORCE_EQ(nodes.size(), 1);
+  auto& node = nodes.back();
+
+  if (!ArgumentHelper::HasArgument(def, "axis")) {
+    const auto& x = def.input(0);
+    const auto& x_shape = shapes.at(x);
+    node.add_attribute()->CopyFrom(
+        MakeAttribute("axis", x_shape.dims().size() - 1));
+  }
+
+  return result;
+}
+
 ConvertedResult OnnxExporter::CreateBinaryElementwiseOpNodes(
     const caffe2::OperatorDef& def,
     const std::unordered_map<std::string, caffe2::TensorShape>& shapes) {

caffe2/onnx/onnx_exporter.h

@@ -52,6 +52,10 @@ class OnnxExporter {
  private:
   ConvertedResult CommonCaffe2OpToOnnxNodes(const caffe2::OperatorDef& def);
 
+  ConvertedResult CreateArgMaxMinOpNodes(
+      const caffe2::OperatorDef& def,
+      const std::unordered_map<std::string, caffe2::TensorShape>& shapes);
+
   ConvertedResult CreateBinaryElementwiseOpNodes(
       const caffe2::OperatorDef& def,
       const std::unordered_map<std::string, caffe2::TensorShape>& shapes);