Add aten mkldnn batch_norm operator

bddppq · facebook-github-bot · commit fb53c189b3e7 · 2019-04-26T13:41:48.000-07:00
Summary: Pull Request resolved: #19206 Reviewed By: dzhulgakov Differential Revision: D14887205 fbshipit-source-id: ea00c9e3205c449d08ab29535309164f951aab95
diff --git a/aten/src/ATen/native/mkldnn/Normalization.cpp b/aten/src/ATen/native/mkldnn/Normalization.cpp
@@ -0,0 +1,77 @@
+#include <ATen/ATen.h>
+#include <ATen/Config.h>
+#include <ATen/NativeFunctions.h>
+#include <tuple>
+
+#if !AT_MKLDNN_ENABLED()
+
+namespace at {
+namespace native {
+
+std::tuple<Tensor, Tensor, Tensor> mkldnn_batch_norm(
+    const Tensor& self,
+    const Tensor& weight,
+    const Tensor& bias,
+    const Tensor& running_mean,
+    const Tensor& running_var,
+    bool train,
+    double momentum,
+    double eps) {
+  AT_ERROR("mkldnn_batch_norm: ATen not compiled with MKLDNN support");
+}
+
+} // namespace native
+} // namespace at
+
+#else // AT_MKLDNN_EBABLED
+
+#include <ATen/native/mkldnn/MKLDNNCommon.h>
+
+namespace at {
+namespace native {
+
+std::tuple<Tensor, Tensor, Tensor> mkldnn_batch_norm(
+    const Tensor& input,
+    const Tensor& weight,
+    const Tensor& bias,
+    const Tensor& running_mean,
+    const Tensor& running_var,
+    bool train,
+    double momentum,
+    double eps) {
+  ideep::tensor& x = itensor_from_mkldnn(input);
+  ideep::tensor& w = itensor_from_mkldnn(weight);
+  ideep::tensor& b = itensor_from_mkldnn(bias);
+  ideep::tensor& m = itensor_from_mkldnn(running_mean);
+  ideep::tensor& v = itensor_from_mkldnn(running_var);
+
+  ideep::tensor y;
+
+  if (train) {
+    // TODO: support training
+    AT_ERROR("mkldnn_batch_norm: mkldnn training is not supported in yet.");
+
+    // ideep::tensor saved_mean;
+    // ideep::tensor saved_var;
+    // ideep::batch_normalization_forward_training::compute<AllocForMKLDNN>(
+    //     x, w, b, y, saved_mean, saved_var, m, v, momentum, eps);
+    // return std::make_tuple(
+    //     new_with_itensor_mkldnn(std::move(y), input.options()),
+    //     new_with_itensor_mkldnn(std::move(saved_mean), input.options()),
+    //     new_with_itensor_mkldnn(std::move(saved_var), input.options()));
+  } else {
+    AT_ASSERTM(input.dim() == 4 || input.dim() == 5,
+               "mkldnn_batch_norm: currently mkldnn only support 2d and 3d batchnorm");
+    ideep::batch_normalization_forward_inference::compute<AllocForMKLDNN>(
+        x, m, v, w, b, y, eps);
+    return std::make_tuple(
+        new_with_itensor_mkldnn(std::move(y), input.options()),
+        new_with_itensor_mkldnn(ideep::tensor{}, input.options()),
+        new_with_itensor_mkldnn(ideep::tensor{}, input.options()));
+  }
+}
+
+} // namespace native
+} // namespace at
+
+#endif // AT_MKLDNN_EBABLED
diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml
@@ -1307,6 +1307,7 @@
   dispatch:
     CPU: batch_norm_cpu
     CUDA: batch_norm_cuda
+    MkldnnCPU: mkldnn_batch_norm
 
 - func: batch_norm_stats(Tensor input, float eps) -> (Tensor, Tensor)
   dispatch:
diff --git a/test/test_mkldnn.py b/test/test_mkldnn.py
@@ -148,6 +148,19 @@ def test_avg_pool2d(self):
                 avg_pool2d(x),
                 avg_pool2d(x.to_mkldnn()).to_dense())
 
+    def test_batch_norm2d(self):
+        N = torch.randint(3, 10, (1,)).item()
+        C = torch.randint(3, 100, (1,)).item()
+        x = torch.randn(N, C, 35, 45, dtype=torch.float32) * 10
+
+        # TODO: support training
+        for train in [False]:
+            bn = torch.nn.BatchNorm2d(C).float().train(train)
+            mkldnn_bn = mkldnn_utils.to_mkldnn(copy.deepcopy(bn))
+            self.assertEqual(
+                bn(x),
+                mkldnn_bn(x.to_mkldnn()).to_dense())
+
 
 if __name__ == '__main__':
     run_tests()
