[caffe2][nomnigraph] Add optimize function to opt:: namespace that ta…

…kes in a level and optimizes the graph/workspace accordingly. Adding it to predictor and speed_benchmark arguments (#7558)
sungin-h · May 15, 2018 · 4b6c884 · 4b6c884
1 parent 469c6c8
commit 4b6c884
Show file tree

Hide file tree

Showing 8 changed files with 85 additions and 9 deletions.
diff --git a/binaries/speed_benchmark.cc b/binaries/speed_benchmark.cc
@@ -19,6 +19,7 @@
 #include "caffe2/core/init.h"
 #include "caffe2/core/logging.h"
 #include "caffe2/core/operator.h"
+#include "caffe2/opt/optimizer.h"
 #include "caffe2/proto/caffe2.pb.h"
 #include "caffe2/utils/proto_utils.h"
 #include "caffe2/utils/string_utils.h"
@@ -63,6 +64,7 @@ CAFFE2_DEFINE_string(
     "folder must already exist in the file system.");
 CAFFE2_DEFINE_int(warmup, 0, "The number of iterations to warm up.");
 CAFFE2_DEFINE_int(iter, 10, "The number of iterations to run.");
+CAFFE2_DEFINE_int(opt, 0, "The level of optimization to run automatically.");
 CAFFE2_DEFINE_bool(
     run_individual,
     false,
@@ -170,6 +172,10 @@ int main(int argc, char** argv) {
           ->set_s(caffe2::FLAGS_algo);
     }
   }
+  if (caffe2::FLAGS_opt) {
+    net_def = caffe2::opt::optimize(net_def, workspace.get(), caffe2::FLAGS_opt);
+  }
+
   caffe2::NetBase* net = workspace->CreateNet(net_def);
   CHECK_NOTNULL(net);
   CAFFE_ENFORCE(net->Run());

diff --git a/caffe2/core/predictor.cc b/caffe2/core/predictor.cc
@@ -1,4 +1,5 @@
 #include "caffe2/core/predictor.h"
+#include "caffe2/opt/optimizer.h"
 
 #include <unordered_set>
 
@@ -81,12 +82,18 @@ Predictor::Predictor(
     const NetDef& init_net,
     const NetDef& run_net,
     Workspace* parent,
-    bool run_init)
+    bool run_init,
+    int optimization)
     : run_net_(run_net), ws_(parent) {
+
   if (run_init) {
     CAFFE_ENFORCE(ws_.RunNetOnce(init_net));
   }
 
+  if (optimization) {
+    run_net_ = opt::optimize(run_net_, &ws_, optimization);
+  }
+
   // real model inputs can be fed later in run* functions
   const auto& initialized_vec = ws_.Blobs();
   const std::unordered_set<std::string> initialized{initialized_vec.begin(),
@@ -97,6 +104,7 @@ Predictor::Predictor(
       blob->template GetMutable<TensorCPU>();
     }
   }
+
   CAFFE_ENFORCE(ws_.CreateNet(run_net));
 }
 

diff --git a/caffe2/core/predictor.h b/caffe2/core/predictor.h
@@ -26,7 +26,8 @@ class Predictor {
       const NetDef& init_net,
       const NetDef& run_net,
       Workspace* parent = nullptr,
-      bool run_init = true);
+      bool run_init = true,
+      int optimization = 1);
 
   ~Predictor() {}
 

diff --git a/caffe2/opt/converter.h b/caffe2/opt/converter.h
@@ -1,5 +1,5 @@
-#ifndef NOM_CONVERTERS_CAFFE2_H
-#define NOM_CONVERTERS_CAFFE2_H
+#ifndef CAFFE2_OPT_CONVERTER_H
+#define CAFFE2_OPT_CONVERTER_H
 
 #include "nomnigraph/Graph/Graph.h"
 #include "nomnigraph/Representations/ControlFlow.h"
@@ -55,4 +55,4 @@ std::unique_ptr<nom::repr::NeuralNetOperator> convertToOperatorDef(caffe2::Opera
 } // namespace caffe2 
 
 
-#endif // NOM_CONVERTERS_CAFFE2_H
+#endif // CAFFE2_OPT_CONVERTER_H
diff --git a/caffe2/opt/fusion.h b/caffe2/opt/fusion.h
@@ -17,9 +17,7 @@
 #ifndef CAFFE2_OPT_FUSION_H_
 #define CAFFE2_OPT_FUSION_H_
 
-#include "caffe2/core/common.h"
 #include "caffe2/core/workspace.h"
-#include "caffe2/proto/caffe2.pb.h"
 #include "nomnigraph/Representations/NeuralNet.h"
 
 namespace caffe2 {

diff --git a/caffe2/opt/mobile.h b/caffe2/opt/mobile.h
@@ -1,8 +1,6 @@
 #ifndef CAFFE2_OPT_MOBILE_H_
 #define CAFFE2_OPT_MOBILE_H_
 
-#include "caffe2/core/common.h"
-#include "caffe2/proto/caffe2.pb.h"
 #include "nomnigraph/Representations/NeuralNet.h"
 
 namespace caffe2 {

diff --git a/caffe2/opt/optimizer.cc b/caffe2/opt/optimizer.cc
@@ -0,0 +1,48 @@
+#include "caffe2/opt/optimizer.h"
+
+#include "caffe2/opt/converter.h"
+#include "caffe2/opt/mobile.h"
+#include "caffe2/opt/fusion.h"
+
+namespace caffe2 {
+namespace opt {
+
+void workspaceOptimizations(nom::repr::NNModule* nn, Workspace* ws, int level) {
+  switch (level) {
+    case 1:
+      opt::fuseConvBN(nn, ws);
+    case 0:
+    default:
+      break;
+  }
+}
+
+void graphOptimzations(nom::repr::NNModule* nn, int level) {
+  switch (level) {
+    case 1:
+#ifdef USE_NNPACK 
+      opt::addNNPACK(nn, false);
+      opt::fuseNNPACKConvRelu(nn);
+#endif
+    case 0:
+    default:
+      break;
+  }
+}
+
+NetDef optimize(NetDef net, Workspace* ws, int level) {
+  auto nn = convertToNNModule(net);
+  graphOptimzations(&nn, level);
+  workspaceOptimizations(&nn, ws, level);
+  return convertToCaffe2Proto(nn, net);
+}
+
+NetDef optimize(NetDef net, int level) {
+  auto nn = convertToNNModule(net);
+  graphOptimzations(&nn, level);
+  return convertToCaffe2Proto(nn, net);
+}
+
+} // namespace opt
+} // namespace caffe2
+
diff --git a/caffe2/opt/optimizer.h b/caffe2/opt/optimizer.h
@@ -0,0 +1,17 @@
+#ifndef CAFFE2_OPT_OPTIMIZER_H
+#define CAFFE2_OPT_OPTIMIZER_H
+
+#include "caffe2/core/common.h"
+#include "caffe2/proto/caffe2.pb.h"
+#include "caffe2/core/workspace.h"
+
+namespace caffe2 {
+namespace opt {
+
+NetDef optimize(NetDef net, Workspace* ws, int level = 1);
+NetDef optimize(NetDef net, int level = 1);
+
+} // namespace opt
+} // namespace caffe2
+
+#endif // CAFFE2_OPT_OPTIMIZER_H