[mir_caffe2] Do not insert Transposes (#7139)

Remove insertion of Transpose operations, set the data format on operations instead.

Signed-off-by: Sergei Barannikov <s.barannikov@samsung.com>

diff --git a/compiler/mir-caffe2-importer/caffe2_op_creator.cpp b/compiler/mir-caffe2-importer/caffe2_op_creator.cpp
index 46b0125..fdf509a 100644 (file)
--- a/compiler/mir-caffe2-importer/caffe2_op_creator.cpp
+++ b/compiler/mir-caffe2-importer/caffe2_op_creator.cpp
@@ -214,20 +214,6 @@ static std::vector<std::int32_t> getWindowSize(const ::caffe2::OperatorDef &op,
   return {kernel_h, kernel_w};
 }
 
-mir::Operation::Output *Caffe2OpCreator::convertCaffeToMIR(mir::Operation::Output *arg)
-{
-  // NCHW -> NHWC
-  auto transpose = createOp<ops::TransposeOp>(arg, std::vector<std::size_t>{0, 2, 3, 1});
-  return transpose->getOutput(0);
-}
-
-mir::Operation::Output *Caffe2OpCreator::convertMIRToCaffe(mir::Operation::Output *arg)
-{
-  // NHWC -> NCHW
-  auto transpose = createOp<ops::TransposeOp>(arg, std::vector<std::size_t>{0, 3, 1, 2});
-  return transpose->getOutput(0);
-}
-
 //
 // Check functions
 //
@@ -325,15 +311,19 @@ std::vector<mir::Operation::Output *>
 Caffe2OpCreator::convertAdd(const std::vector<mir::Operation::Output *> &inputs,
                             const ::caffe2::OperatorDef &op)
 {
+  assert(inputs.size() == 2);
+  auto lhs = inputs[0];
+  auto rhs = inputs[1];
+
   if (getSingleArgument(op, "broadcast", 0) != 0)
   {
     // FIXME This only works when 'axis' == 1 and the second input is 1-D.
-    auto result = createOp<ops::AddOp>(convertCaffeToMIR(inputs[0]), inputs[1])->getOutput(0);
-
-    return {convertMIRToCaffe(result)};
+    rhs = createOp<ops::ReshapeOp>(rhs, Shape{1, rhs->getShape().dim(0), 1, 1})->getOutput(0);
+    auto result = createOp<ops::AddOp>(lhs, rhs)->getOutput(0);
+    return {result};
   }
 
-  auto result = createOp<ops::AddOp>(inputs[0], inputs[1])->getOutput(0);
+  auto result = createOp<ops::AddOp>(lhs, rhs)->getOutput(0);
   return {result};
 }
 
@@ -351,11 +341,9 @@ Caffe2OpCreator::convertAveragePool(const std::vector<mir::Operation::Output *>
   std::vector<std::int32_t> pad_before, pad_after;
   std::tie(pad_before, pad_after) = getPadding(op);
 
-  input = convertCaffeToMIR(input);
   auto result = createOp<ops::AvgPool2DOp>(input, window_size, strides, pad_before, pad_after,
-                                           false, mir::DataFormat::NHWC)
+                                           false, DataFormat::NCHW)
                     ->getOutput(0);
-  result = convertMIRToCaffe(result);
   return {result};
 }
 
@@ -386,8 +374,8 @@ Caffe2OpCreator::convertConv(const std::vector<mir::Operation::Output *> &inputs
     // TODO handle properly kernel with layer multiplier
     auto transposed_tensor = mir::transposeTensor<0, 1, 3, 2>(kernel_tensor);
     auto kernel = createOp<ops::ConstantOp>(transposed_tensor)->getOutput(0);
-    result = createOp<ops::DepthwiseConv2DOp>(convertCaffeToMIR(inputs[0]), kernel, stride_shape,
-                                              pad_before, pad_after)
+    result = createOp<ops::DepthwiseConv2DOp>(inputs[0], kernel, stride_shape, pad_before,
+                                              pad_after, DataFormat::NCHW)
                  ->getOutput(0);
   }
   else
@@ -397,8 +385,8 @@ Caffe2OpCreator::convertConv(const std::vector<mir::Operation::Output *> &inputs
       kernel_tensor = fixGroupedKernel(num_groups, kernel_tensor);
     kernel_tensor = transposeTensor<3, 0, 1, 2>(kernel_tensor);
     auto kernel = createOp<ops::ConstantOp>(kernel_tensor)->getOutput(0);
-    result = createOp<ops::Conv2DOp>(convertCaffeToMIR(inputs[0]), kernel, stride_shape, pad_before,
-                                     pad_after)
+    result = createOp<ops::Conv2DOp>(inputs[0], kernel, stride_shape, pad_before, pad_after,
+                                     DataFormat::NCHW)
                  ->getOutput(0);
   }
 
@@ -407,7 +395,7 @@ Caffe2OpCreator::convertConv(const std::vector<mir::Operation::Output *> &inputs
     result = createOp<ops::AddOp>(result, inputs[2])->getOutput(0);
   }
 
-  return {convertMIRToCaffe(result)};
+  return {result};
 }
 
 std::vector<mir::Operation::Output *>
@@ -470,11 +458,9 @@ Caffe2OpCreator::convertMaxPool(const std::vector<mir::Operation::Output *> &inp
   std::vector<std::int32_t> pad_before, pad_after;
   std::tie(pad_before, pad_after) = getPadding(op);
 
-  input = convertCaffeToMIR(input);
   auto result = createOp<ops::MaxPool2DOp>(input, window_size, strides, pad_before, pad_after,
-                                           mir::DataFormat::NHWC)
+                                           DataFormat::NCHW)
                     ->getOutput(0);
-  result = convertMIRToCaffe(result);
   return {result};
 }
 
@@ -482,15 +468,19 @@ std::vector<mir::Operation::Output *>
 Caffe2OpCreator::convertMul(const std::vector<mir::Operation::Output *> &inputs,
                             const ::caffe2::OperatorDef &op)
 {
+  assert(inputs.size() == 2);
+  auto lhs = inputs[0];
+  auto rhs = inputs[1];
+
   if (getSingleArgument(op, "broadcast", 0) != 0)
   {
     // FIXME This only works when `axis` == 1 and the second input is 1-D.
-    auto result = createOp<ops::MulOp>(convertCaffeToMIR(inputs[0]), inputs[1])->getOutput(0);
-
-    return {convertMIRToCaffe(result)};
+    rhs = createOp<ops::ReshapeOp>(rhs, Shape{1, rhs->getShape().dim(0), 1, 1})->getOutput(0);
+    auto result = createOp<ops::MulOp>(lhs, rhs)->getOutput(0);
+    return {result};
   }
 
-  auto result = createOp<ops::MulOp>(inputs[0], inputs[1])->getOutput(0);
+  auto result = createOp<ops::MulOp>(lhs, rhs)->getOutput(0);
   return {result};
 }
 
@@ -505,17 +495,17 @@ std::vector<mir::Operation::Output *>
 Caffe2OpCreator::convertResizeNearest(const std::vector<mir::Operation::Output *> &inputs,
                                       const ::caffe2::OperatorDef &op)
 {
-  // assume NCHW and convert to MIR (NHWC)
   std::vector<float> scales(4);
   assert(inputs[0]->getShape().rank() == 4 && "only 4d tensors is supported");
-  scales[0] = 1;
-  // default to noop
-  scales[1] = getSingleArgument(op, "height_scale", 1.0f);
-  scales[2] = getSingleArgument(op, "width_scale", 1.0f);
-  scales[3] = 1;
-  auto resize = createOp<ops::ResizeOp>(convertCaffeToMIR(inputs[0]),
-                                        ops::ResizeOp::ResizeMethod::nearestNeighbor, scales);
-  return {convertMIRToCaffe(resize->getOutput(0))};
+  // Assuming NCHW format.
+  scales[0] = 1.0f;
+  scales[1] = 1.0f;
+  scales[2] = getSingleArgument(op, "height_scale", 1.0f);
+  scales[3] = getSingleArgument(op, "width_scale", 1.0f);
+  auto result =
+      createOp<ops::ResizeOp>(inputs[0], ops::ResizeOp::ResizeMethod::nearestNeighbor, scales)
+          ->getOutput(0);
+  return {result};
 }
 
 std::vector<mir::Operation::Output *>
@@ -569,20 +559,23 @@ Caffe2OpCreator::convertSpatialBN(const std::vector<mir::Operation::Output *> &i
     bias_data.at(idx) *= -1;
 
   auto mean = createOp<ops::ConstantOp>(mean_tensor)->getOutput(0);
-  auto result = createOp<ops::AddOp>(convertCaffeToMIR(inputs[0]), mean)->getOutput(0);
+  mean = createOp<ops::ReshapeOp>(mean, Shape{1, mean->getShape().dim(0), 1, 1})->getOutput(0);
+  auto result = createOp<ops::AddOp>(inputs[0], mean)->getOutput(0);
 
   // res2 = res1 * scale / (var + epsilon)
   Tensor<float> multiplier(scale_tensor);
   for (auto &idx : ShapeRange(scale_tensor.getShape()))
     multiplier.at(idx) /= std::sqrt(*reinterpret_cast<float *>(var_tensor.at(idx)) + eps);
   auto scale = createOp<ops::ConstantOp>(scale_tensor)->getOutput(0);
+  scale = createOp<ops::ReshapeOp>(scale, Shape{1, scale->getShape().dim(0), 1, 1})->getOutput(0);
   result = createOp<ops::MulOp>(result, scale)->getOutput(0);
 
   // overall_res = res2 + bias
   auto bias = createOp<ops::ConstantOp>(bias_tensor)->getOutput(0);
+  bias = createOp<ops::ReshapeOp>(bias, Shape{1, bias->getShape().dim(0), 1, 1})->getOutput(0);
   result = createOp<ops::AddOp>(result, bias)->getOutput(0);
 
-  return {convertMIRToCaffe(result)};
+  return {result};
 }
 
 std::vector<mir::Operation::Output *>

diff --git a/compiler/mir-caffe2-importer/caffe2_op_creator.h b/compiler/mir-caffe2-importer/caffe2_op_creator.h
index 57424a2..686c6d2 100644 (file)
--- a/compiler/mir-caffe2-importer/caffe2_op_creator.h
+++ b/compiler/mir-caffe2-importer/caffe2_op_creator.h
@@ -104,10 +104,6 @@ public:
 private:
   mir::Graph *_graph = nullptr;
 
-  mir::Operation::Output *convertCaffeToMIR(mir::Operation::Output *arg);
-
-  mir::Operation::Output *convertMIRToCaffe(mir::Operation::Output *arg);
-
   template <typename OpType, typename... Types> mir::Operation *createOp(Types &&... args);
 };