Enable Conv fusion optimizations in optimizeForIdeep (pytorch#9255)

Summary: Enable fusion for IDEEP in optimizeForIdeep including Conv+ReLU, Conv+Sum, Conv+Sum+ReLU, Conv+BN Pull Request resolved: pytorch#9255 Reviewed By: bddppq Differential Revision: D8809030 Pulled By: yinghai fbshipit-source-id: af30bad3b96cb965bd26a4dfa810370faec4bb88
goodlux · Aug 15, 2018 · e1d10d5 · e1d10d5
1 parent 421b18c
commit e1d10d5
Show file tree

Hide file tree

Showing 7 changed files with 548 additions and 53 deletions.
diff --git a/caffe2/opt/converter.cc b/caffe2/opt/converter.cc
@@ -42,6 +42,14 @@ std::vector<int> getDilations(std::map<std::string, caffe2::Argument> argMap) {
   return dilations;
 }
 
+int getGroup(std::map<std::string, caffe2::Argument>& argMap) {
+  if (argMap.count("group")) {
+    CAFFE_ENFORCE(argMap["group"].has_i() && "Invalid group argument");
+    return static_cast<int>(argMap["group"].i());
+  }
+  return 1;
+}
+
 } // namespace
 
 namespace caffe2 {
@@ -115,6 +123,8 @@ class ConvConverter : public Converter {
     c->setStrides(getStrides(argMap));
     c->setPads(getPads(argMap));
     c->setDilations(getDilations(argMap));
+    c->setGroup(getGroup(argMap));
+
     return nnOp;
   }
   // Does not override default converter to OperatorDef