Check mappings ONNX -> Caffe2 bear the same argument names (#6317)

* Check mappings ONNX -> Caffe2 bear the same argument names

When adding an extra arg to an input ONNX op, if it's not supported in Caffe2, the exporter would just silently pass it to NetDef and ignore it in the implementation. It's pretty error-prone. Caffe2 also has an OpSchema description and we can enforce that all arguments explicitly appear in schema or listed explicitly in Caffe2.

See also https://github.com/caffe2/caffe2/pull/2478

Add test for C2 argument checking

* Some operators do not log arguments, which prevents argument checks.
Invite users to file an issue to fix the schema.
diff --git a/caffe2/onnx/backend.cc b/caffe2/onnx/backend.cc
index 6877ca2..282c090 100644
--- a/caffe2/onnx/backend.cc
+++ b/caffe2/onnx/backend.cc
@@ -901,17 +901,55 @@
   return OnnxNodeToCaffe2Ops(init_model, pred_model, &onnx_node, opset_version);
 }
 
+void Caffe2Backend::CheckOpSchemaArguments(
+    const caffe2::OpSchema& schema,
+    const caffe2::OperatorDef& op) {
+  const auto& schema_args = schema.args();
+  if (schema_args.size() > 0){
+    std::vector<std::string> argnames;
+    std::transform(
+        schema_args.begin(),
+        schema_args.end(),
+        std::back_inserter(argnames),
+        [](caffe2::OpSchema::Argument elem) { return elem.name(); });
+
+    for (const auto& arg : op.arg()) {
+      if (std::count(argnames.begin(), argnames.end(), arg.name()) == 0) {
+        CAFFE_THROW(
+            "Don't know how to map unexpected argument ",
+            arg.name(),
+            " (from operator ",
+            op.type(), ")");
+      }
+    }
+  } else {
+    // A number of C2 operators do not declare proper arguments. Let's log the error
+    VLOG(2) << "Operator " << op.type() << " does not declare arguments in its schema. Please file a Caffe2 issue.";
+  }
+}
+
 Caffe2Ops Caffe2Backend::OnnxNodeToCaffe2Ops(
     const ModelProto& init_model,
     const ModelProto& pred_model,
     OnnxNode* onnx_node,
     int opset_version) {
+  Caffe2Ops res;
   if (get_special_operators().count(onnx_node->node.op_type())) {
-    return (this->*get_special_operators().at(onnx_node->node.op_type()))(
+    res = (this->*get_special_operators().at(onnx_node->node.op_type()))(
         onnx_node, opset_version);
   } else {
-    return CommonOnnxNodeToCaffe2Ops(onnx_node, opset_version);
+    res = CommonOnnxNodeToCaffe2Ops(onnx_node, opset_version);
   }
+
+  for (const auto& result_op: res.ops){
+    const auto* schema = OpSchemaRegistry::Schema(result_op.type());
+    if (schema) {
+      CheckOpSchemaArguments(*schema, result_op);
+    } else {
+      CAFFE_THROW("Caffe2 has no such operator, could not find schema for ", result_op.type());
+    }
+  }
+  return res;
 }
 
 void Caffe2Backend::OnnxToCaffe2(
diff --git a/caffe2/onnx/backend.h b/caffe2/onnx/backend.h
index 2a83fb2..c53a61b 100644
--- a/caffe2/onnx/backend.h
+++ b/caffe2/onnx/backend.h
@@ -125,6 +125,8 @@
       bool include_initializers,
       const std::vector<Caffe2Ops>& extras);
 
+  void CheckOpSchemaArguments(const caffe2::OpSchema& schema, const caffe2::OperatorDef& op);
+
   Caffe2Ops OnnxNodeToCaffe2Ops(
       const ModelProto& init_model,
       const ModelProto& pred_model,
diff --git a/caffe2/python/onnx/tests/c2_ref_test.py b/caffe2/python/onnx/tests/c2_ref_test.py
index 5d395e9..f23b2f8 100644
--- a/caffe2/python/onnx/tests/c2_ref_test.py
+++ b/caffe2/python/onnx/tests/c2_ref_test.py
@@ -27,6 +27,8 @@
 from caffe2.python.onnx.helper import dummy_name
 from caffe2.python.onnx.tests.test_utils import TestCase
 
+import caffe2.python._import_c_extension as C
+
 
 class TestCaffe2Basic(TestCase):
     def test_dummy_name(self):
@@ -34,6 +36,22 @@
         n2 = dummy_name()
         assert n1 != n2, "Got same names in different calls: {}".format(n1)
 
+    def test_check_arguments(self):
+        X = np.random.randn(3, 2).astype(np.float32)
+        Y = np.random.randn(3, 2).astype(np.float32)
+        Z = np.zeros((3, 2)).astype(np.float32)
+
+        b2 = C.Caffe2Backend()
+
+        node_def = make_node("Add", inputs = ["X", "Y"], outputs = ["Z"], broadcast = 0)
+        output = b2.convert_node(node_def.SerializeToString(), 6)
+
+        bad_node_def = make_node("Add", inputs = ["X", "Y"], outputs = ["Z"], foo = 42, bar = 56)
+        with self.assertRaisesRegexp(
+            RuntimeError,
+            ".*?Don't know how to map unexpected argument (foo|bar) \(from operator .*?\).*$"):
+            output = b2.convert_node(bad_node_def.SerializeToString(), 6)
+
     def test_relu_graph(self):
         X = np.random.randn(3, 2).astype(np.float32)
         Y_ref = np.clip(X, 0, np.inf)