[mir_caffe] Refactor importing of BatchNorm (#6467)

Structurize the code, fix comments.

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

diff --git a/compiler/mir-caffe-importer/caffe_op_creator.cpp b/compiler/mir-caffe-importer/caffe_op_creator.cpp
index 0fb0cca..7878ae5 100644 (file)
--- a/compiler/mir-caffe-importer/caffe_op_creator.cpp
+++ b/compiler/mir-caffe-importer/caffe_op_creator.cpp
@@ -634,35 +634,37 @@ std::vector<mir::Operation::Output *>
 CaffeOpCreator::convertBatchNorm(const caffe::LayerParameter &layer,
                                  const std::vector<mir::Operation::Output *> &inputs)
 {
-  auto &opts = layer.batch_norm_param();
-  float eps = opts.eps();
-  auto scale_weight = convertBlob(layer.blobs(2));
-  float scale_factor = *reinterpret_cast<float *>(scale_weight.at(mir::Index{0}));
-  // Code below is taken from cpu caffe implementation:
-  // https://github.com/BVLC/caffe/blob/master/src/caffe/layers/batch_norm_layer.cpp#L100
+  const auto &opts = layer.batch_norm_param();
+  auto input = inputs[0];
+  auto mean_tensor = convertBlob(layer.blobs(0));
+  auto var_tensor = convertBlob(layer.blobs(1));
+  auto scale_tensor = convertBlob(layer.blobs(2));
+  const float eps = opts.eps();
+
+  float scale_factor = *reinterpret_cast<float *>(scale_tensor.at(mir::Index{0}));
+
+  // See https://github.com/BVLC/caffe/blob/master/src/caffe/layers/batch_norm_layer.cpp#L100
+  // Y = (X - mean / scale_factor) / sqrt(var / scale_factor + epsilon) =
+  //   = (X + C1) * C2
   if (scale_factor != 0.0f)
     scale_factor = 1.0f / scale_factor;
 
-  // create bias argument from mean:
-  // multiply elements of mean by scaleFactor and get opposite numbers
-  // to subtract mean from input via add operation
-  auto mean_weights = convertBlob(layer.blobs(0));
-  Tensor<float> bias_data(mean_weights);
-  for (Index idx : ShapeRange(bias_data.getShape()))
-    bias_data.at(idx) *= -scale_factor;
-  auto mean = createOp<ops::ConstantOp>("", mean_weights)->getOutput(0);
-  auto result = createOp<ops::AddOp>(layer.name() + ".bias", convertCaffeToMIR(inputs[0]), mean)
-                    ->getOutput(0);
-
-  // create scale argument from variance:
-  // multiply elements of variance by scaleFactor and
-  // normalize biased input using scale operation
-  auto variance_weights = convertBlob(layer.blobs(1));
-  Tensor<float> scale_data(variance_weights);
-  for (Index idx : ShapeRange(scale_data.getShape()))
-    scale_data.at(idx) = 1.0f / std::sqrt(scale_data.at(idx) * scale_factor + eps);
-  auto variance = createOp<ops::ConstantOp>("", variance_weights)->getOutput(0);
-  result = createOp<ops::MulOp>(layer.name() + ".scale", result, variance)->getOutput(0);
+  // C1 = -mean / scale_factor
+  Tensor<float> mean_accessor(mean_tensor);
+  for (const auto &idx : ShapeRange(mean_accessor.getShape()))
+    mean_accessor.at(idx) *= -scale_factor;
+  auto c1 = createOp<ops::ConstantOp>("", mean_tensor)->getOutput(0);
+
+  // C2 = 1 / sqrt(var / scale_factor + epsilon)
+  Tensor<float> var_accessor(var_tensor);
+  for (const auto &idx : ShapeRange(var_accessor.getShape()))
+    var_accessor.at(idx) = 1.0f / std::sqrt(var_accessor.at(idx) * scale_factor + eps);
+  auto c2 = createOp<ops::ConstantOp>("", var_tensor)->getOutput(0);
+
+  // Y = (X + C1) * C2
+  input = convertCaffeToMIR(input);
+  auto result = createOp<ops::AddOp>("", input, c1)->getOutput(0);
+  result = createOp<ops::MulOp>("", result, c2)->getOutput(0);
   return {convertMIRToCaffe(result)};
 }