nomnigraph - easy - use new test utils in converter_nomnigraph_test (#15751)

Summary:
Use new test utils in converter_nomnigraph_test , and add utils to set device option name, external inputs, outputs.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/15751

Differential Revision: D13586228

Pulled By: duc0

fbshipit-source-id: ff809dd7bf9f30641ce2a6fef7e2810f005521c2

diff --git a/caffe2/core/test_utils.cc b/caffe2/core/test_utils.cc
index 3da7397..f81288c 100644 (file)
--- a/caffe2/core/test_utils.cc
+++ b/caffe2/core/test_utils.cc
@@ -44,7 +44,7 @@ void assertTensorListEquals(
     const std::vector<std::string>& tensorNames,
     const Workspace& workspace1,
     const Workspace& workspace2) {
-  for (const string& tensorName : tensorNames) {
+  for (const std::string& tensorName : tensorNames) {
     CAFFE_ENFORCE(workspace1.HasBlob(tensorName));
     CAFFE_ENFORCE(workspace2.HasBlob(tensorName));
     auto& tensor1 = getTensor(workspace1, tensorName);
@@ -68,8 +68,8 @@ caffe2::Tensor* createTensor(
 
 caffe2::OperatorDef* createOperator(
     const std::string& type,
-    const std::vector<string>& inputs,
-    const std::vector<string>& outputs,
+    const std::vector<std::string>& inputs,
+    const std::vector<std::string>& outputs,
     caffe2::NetDef* net) {
   auto* op = net->add_op();
   op->set_type(type);
@@ -84,11 +84,33 @@ caffe2::OperatorDef* createOperator(
 
 NetMutator& NetMutator::newOp(
     const std::string& type,
-    const std::vector<string>& inputs,
-    const std::vector<string>& outputs) {
+    const std::vector<std::string>& inputs,
+    const std::vector<std::string>& outputs) {
   lastCreatedOp_ = createOperator(type, inputs, outputs, net_);
   return *this;
 }
 
+NetMutator& NetMutator::externalInputs(
+    const std::vector<std::string>& externalInputs) {
+  for (auto& blob : externalInputs) {
+    net_->add_external_input(blob);
+  }
+  return *this;
+}
+
+NetMutator& NetMutator::externalOutputs(
+    const std::vector<std::string>& externalOutputs) {
+  for (auto& blob : externalOutputs) {
+    net_->add_external_output(blob);
+  }
+  return *this;
+}
+
+NetMutator& NetMutator::setDeviceOptionName(const std::string& name) {
+  CAFFE_ENFORCE(lastCreatedOp_ != nullptr);
+  lastCreatedOp_->mutable_device_option()->set_node_name(name);
+  return *this;
+}
+
 } // namespace testing
 } // namespace caffe2

diff --git a/caffe2/core/test_utils.h b/caffe2/core/test_utils.h
index 53094d1..fcf069d 100644 (file)
--- a/caffe2/core/test_utils.h
+++ b/caffe2/core/test_utils.h
@@ -6,6 +6,7 @@
 #include "caffe2/utils/proto_utils.h"
 
 #include <cmath>
+#include <string>
 #include <vector>
 
 // Utilities that make it easier to write caffe2 C++ unit tests.
@@ -67,15 +68,15 @@ caffe2::Tensor* createTensor(
 // Create a new operator in the net.
 caffe2::OperatorDef* createOperator(
     const std::string& type,
-    const std::vector<string>& inputs,
-    const std::vector<string>& outputs,
+    const std::vector<std::string>& inputs,
+    const std::vector<std::string>& outputs,
     caffe2::NetDef* net);
 
 // Fill data from a vector to a tensor.
 template <typename T>
 void fillTensor(
-    const vector<int64_t>& shape,
-    const vector<T>& data,
+    const std::vector<int64_t>& shape,
+    const std::vector<T>& data,
     TensorCPU* tensor) {
   tensor->Resize(shape);
   CAFFE_ENFORCE_EQ(data.size(), tensor->numel());
@@ -88,9 +89,9 @@ void fillTensor(
 // Create a tensor and fill data.
 template <typename T>
 caffe2::Tensor* createTensorAndFill(
-    const string& name,
-    const vector<int64_t>& shape,
-    const vector<T>& data,
+    const std::string& name,
+    const std::vector<int64_t>& shape,
+    const std::vector<T>& data,
     Workspace* workspace) {
   auto* tensor = createTensor(name, workspace);
   fillTensor<T>(shape, data, tensor);
@@ -113,8 +114,8 @@ void constantFillTensor(
 // Create a tensor and fill a constant.
 template <typename T>
 caffe2::Tensor* createTensorAndConstantFill(
-    const string& name,
-    const vector<int64_t>& shape,
+    const std::string& name,
+    const std::vector<int64_t>& shape,
     const T& data,
     Workspace* workspace) {
   auto* tensor = createTensor(name, workspace);
@@ -123,23 +124,30 @@ caffe2::Tensor* createTensorAndConstantFill(
 }
 
 // Concise util class to mutate a net in a chaining fashion.
-class NetMutator {
+class CAFFE2_API NetMutator {
  public:
   explicit NetMutator(caffe2::NetDef* net) : net_(net) {}
 
   NetMutator& newOp(
       const std::string& type,
-      const std::vector<string>& inputs,
-      const std::vector<string>& outputs);
+      const std::vector<std::string>& inputs,
+      const std::vector<std::string>& outputs);
+
+  NetMutator& externalInputs(const std::vector<std::string>& externalInputs);
+
+  NetMutator& externalOutputs(const std::vector<std::string>& externalOutputs);
 
   // Add argument to the last created op.
   template <typename T>
-  NetMutator& addArgument(const string& name, const T& value) {
+  NetMutator& addArgument(const std::string& name, const T& value) {
     CAFFE_ENFORCE(lastCreatedOp_ != nullptr);
     AddArgument(name, value, lastCreatedOp_);
     return *this;
   }
 
+  // Set device name for the last created op.
+  NetMutator& setDeviceOptionName(const std::string& name);
+
  private:
   caffe2::NetDef* net_;
   caffe2::OperatorDef* lastCreatedOp_;
@@ -154,9 +162,9 @@ class WorkspaceMutator {
   // New tensor filled by a data vector.
   template <typename T>
   WorkspaceMutator& newTensor(
-      const string& name,
-      const vector<int64_t>& shape,
-      const vector<T>& data) {
+      const std::string& name,
+      const std::vector<int64_t>& shape,
+      const std::vector<T>& data) {
     createTensorAndFill<T>(name, shape, data, workspace_);
     return *this;
   }
@@ -164,7 +172,7 @@ class WorkspaceMutator {
   // New tensor filled by a constant.
   template <typename T>
   WorkspaceMutator& newTensorConst(
-      const string& name,
+      const std::string& name,
       const std::vector<int64_t>& shape,
       const T& data) {
     createTensorAndConstantFill<T>(name, shape, data, workspace_);

diff --git a/caffe2/opt/converter_nomigraph_test.cc b/caffe2/opt/converter_nomigraph_test.cc
index ec87365..26735bf 100644 (file)
--- a/caffe2/opt/converter_nomigraph_test.cc
+++ b/caffe2/opt/converter_nomigraph_test.cc
@@ -1,34 +1,24 @@
+#include "caffe2/core/test_utils.h"
 #include "caffe2/opt/converter.h"
 
 #include <gtest/gtest.h>
 
-#define ADD_ARG(_op, _name, _type, _val)                                       \
-{                                                                            \
-  caffe2::Argument *arg = _op->add_arg();                                    \
-  arg->set_name(_name);                                                      \
-  arg->set_##_type(_val);                                                    \
-}
-
 TEST(Converter, Basic) {
+  using namespace caffe2::testing;
   caffe2::NetDef net;
   for (auto i = 0; i < 10; ++i) {
     if (rand() % 2) {
-      caffe2::OperatorDef *def = net.add_op();
-      def->set_type("Conv");
-      def->add_input("X");
-      def->add_input("W" + c10::to_string(i)); // different weights
-      ADD_ARG(def, "kernel", i, 3);
-      ADD_ARG(def, "stride", i, 1);
-      ADD_ARG(def, "pad", i, 0);
-      ADD_ARG(def, "order", s, "NCHW");
-      def->add_output("X");
-      def->mutable_device_option()->set_node_name("conv_runner");
+      NetMutator(&net)
+          .newOp("Conv", {"X", "W" + c10::to_string(i)}, {"X"})
+          .addArgument("kernel", 3)
+          .addArgument("stride", 1)
+          .addArgument("pad", 0)
+          .addArgument("order", string("NCHW"))
+          .setDeviceOptionName("conv_runner");
     } else {
-      caffe2::OperatorDef *def = net.add_op();
-      def->set_type("Relu");
-      def->add_input("X");
-      def->add_output("X");
-      def->mutable_device_option()->set_node_name("relu_runner");
+      NetMutator(&net)
+          .newOp("Relu", {"X"}, {"X"})
+          .setDeviceOptionName("relu_runner");
     }
   }
   auto nn = caffe2::convertToNNModule(net);
@@ -36,44 +26,24 @@ TEST(Converter, Basic) {
 }
 
 TEST(Converter, UnknownType) {
+  using namespace caffe2::testing;
   caffe2::NetDef net;
-
-  caffe2::OperatorDef *def = net.add_op();
-  def->set_type("NeverSeen");
-  def->add_input("X");
-  def->add_output("X");
-  def->mutable_device_option()->set_node_name(
-      "device_" + c10::to_string(rand() % 2));
+  NetMutator(&net)
+      .newOp("NeverSeen", {"X"}, {"X"})
+      .setDeviceOptionName("device_" + c10::to_string(rand() % 2));
   auto nn = caffe2::convertToNNModule(net);
   auto new_netdef = caffe2::convertToCaffe2Proto(nn);
 }
 
 caffe2::NetDef fakeNet() {
+  using namespace caffe2::testing;
   caffe2::NetDef net;
-
-  {
-    caffe2::OperatorDef* def = net.add_op();
-    def->set_type("Fake");
-    def->add_input("X");
-    def->add_output("Y");
-  }
-  {
-    caffe2::OperatorDef* def = net.add_op();
-    def->set_type("Fake");
-    def->add_input("Y");
-    def->add_output("Z");
-  }
-  {
-    caffe2::OperatorDef* def = net.add_op();
-    def->set_type("Fake");
-    def->add_input("Z");
-    def->add_input("X");
-    def->add_output("W");
-  }
-  net.add_external_input("X");
-  net.add_external_output("Y");
-  net.add_external_output("W");
-
+  NetMutator(&net)
+      .newOp("Fake", {"X"}, {"Y"})
+      .newOp("Fake", {"Y"}, {"Z"})
+      .newOp("Fake", {"Z", "X"}, {"W"})
+      .externalInputs({"X"})
+      .externalOutputs({"Y", "W"});
   return net;
 }