[Importer] Support arg "order" in Caffe2 Conv, MaxPool and AveragePoo…

…l ops.
pytorch · Oct 12, 2018 · cb2b12a · cb2b12a
1 parent e0d3754
commit cb2b12a
Show file tree

Hide file tree

Showing 8 changed files with 374 additions and 11 deletions.
diff --git a/lib/Importer/Caffe2ModelLoader.cpp b/lib/Importer/Caffe2ModelLoader.cpp
@@ -153,6 +153,7 @@ void Caffe2ModelLoader::loadOperator(const caffe2::OperatorDef &op) {
     std::vector<unsigned_t> pads = getPads(dict);
     std::vector<unsigned_t> kernels = getSizeHW(dict, "kernel", 0);
     unsigned_t group = dict.count("group") ? loadInt(dict["group"]) : 1;
+    std::string order = dict.count("order") ? loadStr(dict["order"]) : "NCHW";
 
     auto in = getNodeValueOrCreateVariableByName(op.input(0));
     Tensor *w = getTensorByName(op.input(1));
@@ -186,23 +187,30 @@ void Caffe2ModelLoader::loadOperator(const caffe2::OperatorDef &op) {
     }
     auto *bias = G_.getParent()->createConstant("conv.bias", biasTensor);
 
-    // Caffe passes the input as NCHW, and we expect the input to be NHWC.
-    auto *tr = G_.createTranspose(opName, in, NCHW2NHWC);
+    // We expect the input to be NHWC.
+    Node *tr;
+    if (order == "NCHW") {
+      tr = G_.createTranspose(opName, in, NCHW2NHWC);
+    } else {
+      tr = in;
+    }
 
     // Calculate the size and allocate the output buffer.
-    ShapeNHWC idim = ShapeNHWC(tr->getResult().dims());
+    ShapeNHWC idim = ShapeNHWC(tr->getType(0)->dims());
     auto outSz =
         calculateConvPoolOutputDims(idim.h, idim.w, kernels, strides, pads);
     std::array<size_t, 4> outDims = {
         {idim.n, outSz.first, outSz.second, depth}};
     auto outTy = G_.getParent()->uniqueType(ElemKind::FloatTy, outDims);
 
-    auto *node = G_.createConv(opName, tr, filter, bias, outTy, kernels,
+    Node *node = G_.createConv(opName, tr, filter, bias, outTy, kernels,
                                strides, pads, group);
 
-    // Transpose the output back.
-    auto *N = G_.createTranspose(opName, node, NHWC2NCHW);
-    addNodeAsOutput(op, N);
+    if (order == "NCHW") {
+      // Transpose the output back.
+      node = G_.createTranspose(opName, node, NHWC2NCHW);
+    }
+    addNodeAsOutput(op, node);
     return;
   }
 
@@ -212,8 +220,14 @@ void Caffe2ModelLoader::loadOperator(const caffe2::OperatorDef &op) {
     std::vector<unsigned_t> strides = getSizeHW(dict, "stride", 1);
     std::vector<unsigned_t> kernels = getSizeHW(dict, "kernel", 0);
     std::vector<unsigned_t> pads = getPads(dict);
-
-    auto *tr = G_.createTranspose(opName, in, NCHW2NHWC);
+    std::string order = dict.count("order") ? loadStr(dict["order"]) : "NCHW";
+    // We expect the input to be NHWC.
+    Node *tr;
+    if (order == "NCHW") {
+      tr = G_.createTranspose(opName, in, NCHW2NHWC);
+    } else {
+      tr = in;
+    }
 
     // If 'global_pooling' is set then the operation will pool over the size of
     // the input by doing: kernels = {height, width}.
@@ -229,8 +243,11 @@ void Caffe2ModelLoader::loadOperator(const caffe2::OperatorDef &op) {
     } else {
       node = G_.createAvgPool(opName, tr, kernels, strides, pads);
     }
-    auto *N = G_.createTranspose(opName, node, NHWC2NCHW);
-    addNodeAsOutput(op, N);
+    if (order == "NCHW") {
+      // Transpose the output back.
+      node = G_.createTranspose(opName, node, NHWC2NCHW);
+    }
+    addNodeAsOutput(op, node);
     return;
   }
 

diff --git a/tests/models/caffe2Models/avgpool_nhwc_predict_net.pbtxt b/tests/models/caffe2Models/avgpool_nhwc_predict_net.pbtxt
@@ -0,0 +1,23 @@
+name: "averagePool_test"
+op {
+  input: "inputs"
+  output: "output"
+  type: "AveragePool"
+  arg {
+    name: "order"
+    s: "NHWC"
+  }
+  arg {
+    name: "kernel"
+    i: 2
+  }
+  arg {
+    name: "stride"
+    i: 1
+  }
+  arg {
+    name: "pad"
+    i: 1
+  }
+}
+external_output: "output"
diff --git a/tests/models/caffe2Models/avgpool_predict_net.pbtxt b/tests/models/caffe2Models/avgpool_predict_net.pbtxt
@@ -0,0 +1,19 @@
+name: "averagePool_test"
+op {
+  input: "inputs"
+  output: "output"
+  type: "AveragePool"
+  arg {
+    name: "kernel"
+    i: 2
+  }
+  arg {
+    name: "stride"
+    i: 1
+  }
+  arg {
+    name: "pad"
+    i: 1
+  }
+}
+external_output: "output"
diff --git a/tests/models/caffe2Models/conv_nhwc_init_net.pbtxt b/tests/models/caffe2Models/conv_nhwc_init_net.pbtxt
@@ -0,0 +1,32 @@
+name: "init"
+op {
+  output: "conv_w"
+  type: "GivenTensorFill"
+  arg {
+    name: "shape"
+    ints: 1
+    ints: 1
+    ints: 2
+    ints: 2
+  }
+  arg {
+    name: "values"
+    floats: 1.0
+    floats: 1.0
+    floats: 1.0
+    floats: 1.0
+  }
+}
+op {
+  output: "conv_b"
+  type: "GivenTensorFill"
+  arg {
+    name: "shape"
+    ints: 1
+  }
+  arg {
+    name: "values"
+    floats: 2.0
+  }
+}
+
diff --git a/tests/models/caffe2Models/conv_nhwc_predict_net.pbtxt b/tests/models/caffe2Models/conv_nhwc_predict_net.pbtxt
@@ -0,0 +1,29 @@
+name: "conv_test"
+op {
+  input: "inputs"
+  input: "conv_w"
+  input: "conv_b"
+  output: "conv_out"
+  type: "Conv"
+  arg {
+    name: "order"
+    s: "NHWC"
+  }
+  arg {
+    name: "kernel"
+    i: 2
+  }
+  arg {
+    name: "stride"
+    i: 1
+  }
+  arg {
+    name: "group"
+    i: 1
+  }
+  arg {
+    name: "pad"
+    i: 1
+  }
+}
+external_output: "conv_out"
diff --git a/tests/models/caffe2Models/maxpool_nhwc_predict_net.pbtxt b/tests/models/caffe2Models/maxpool_nhwc_predict_net.pbtxt
@@ -0,0 +1,23 @@
+name: "maxpool_test"
+op {
+  input: "inputs"
+  output: "output"
+  type: "MaxPool"
+  arg {
+    name: "order"
+    s: "NHWC"
+  }
+  arg {
+    name: "kernel"
+    i: 2
+  }
+  arg {
+    name: "stride"
+    i: 1
+  }
+  arg {
+    name: "pad"
+    i: 1
+  }
+}
+external_output: "output"
diff --git a/tests/models/caffe2Models/maxpool_predict_net.pbtxt b/tests/models/caffe2Models/maxpool_predict_net.pbtxt
@@ -0,0 +1,19 @@
+name: "maxpool_test"
+op {
+  input: "inputs"
+  output: "output"
+  type: "MaxPool"
+  arg {
+    name: "kernel"
+    i: 2
+  }
+  arg {
+    name: "stride"
+    i: 1
+  }
+  arg {
+    name: "pad"
+    i: 1
+  }
+}
+external_output: "output"