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)};
}