From 9ce0e29dc3e0cb572f0315a9371f96881f37ac01 Mon Sep 17 00:00:00 2001 From: Jacek Czaja Date: Fri, 6 Dec 2019 03:01:01 +0100 Subject: [PATCH] [MKL-DNN] Batch norm mkl-dnn NHWC support (#21553) * - BAtch norm mkl-dnn NHWC test=develop - compilation fix test=develop - UT fix - cosmetics test=develop - Fix to Batch Norm MKL-DNN NHWC UT test=develop Conflicts: paddle/fluid/operators/batch_norm_op.h * - Lint fixes test=develop --- paddle/fluid/operators/batch_norm_op.cc | 37 ++++++++++++++++++- paddle/fluid/operators/batch_norm_op.h | 4 ++ .../mkldnn/test_batch_norm_mkldnn_op.py | 7 ++++ .../tests/unittests/test_batch_norm_op.py | 15 ++++++++ 4 files changed, 61 insertions(+), 2 deletions(-) diff --git a/paddle/fluid/operators/batch_norm_op.cc b/paddle/fluid/operators/batch_norm_op.cc index a2447cb968..0708e104f2 100644 --- a/paddle/fluid/operators/batch_norm_op.cc +++ b/paddle/fluid/operators/batch_norm_op.cc @@ -94,8 +94,9 @@ void BatchNormOp::InferShape(framework::InferShapeContext *ctx) const { x_dims, x_dims.size()); const int64_t C = - (data_layout == DataLayout::kNCHW ? x_dims[1] - : x_dims[x_dims.size() - 1]); + ((this->IsMKLDNNType() == true) || (data_layout == DataLayout::kNCHW) + ? x_dims[1] + : x_dims[x_dims.size() - 1]); auto scale_dim = ctx->GetInputDim("Scale"); auto bias_dim = ctx->GetInputDim("Bias"); @@ -169,6 +170,32 @@ framework::OpKernelType BatchNormOp::GetExpectedKernelType( library); } +framework::OpKernelType BatchNormOp::GetKernelTypeForVar( + const std::string &var_name, const Tensor &tensor, + const framework::OpKernelType &expected_kernel_type) const { +#ifdef PADDLE_WITH_MKLDNN + // Only input require reshaping, weights and + // bias are having shape in NCHW order + if ((var_name == "X") && + (expected_kernel_type.data_layout_ == framework::DataLayout::kMKLDNN) && + (tensor.layout() != framework::DataLayout::kMKLDNN)) { + auto attrs = Attrs(); + auto ar = paddle::framework::AttrReader(attrs); + const std::string data_layout = ar.Get("data_layout"); + auto dl = framework::StringToDataLayout(data_layout); + // Some models may have intentionally set "AnyLayout" for pool + // op. Treat this as NCHW (default data_format value) + if (dl != framework::DataLayout::kAnyLayout) { + return framework::OpKernelType( + expected_kernel_type.data_type_, tensor.place(), + framework::StringToDataLayout(data_layout)); + } + } +#endif + return framework::OpKernelType(expected_kernel_type.data_type_, + tensor.place(), tensor.layout()); +} + void BatchNormOpMaker::Make() { AddAttr("is_test", "(bool, default false) Set to true for inference only, false " @@ -465,6 +492,12 @@ framework::OpKernelType BatchNormGradOp::GetExpectedKernelType( #ifdef PADDLE_WITH_MKLDNN if (library == framework::LibraryType::kPlain && platform::CanMKLDNNBeUsed(ctx)) { + // TODO(jczaja): Add support for NHWC + const std::string data_layout = ctx.Attr("data_layout"); + PADDLE_ENFORCE_NE( + data_layout, "NHWC", + platform::errors::Unimplemented( + "Batch Norm MKLDNN grad does not support NHWC data format yet")); library = framework::LibraryType::kMKLDNN; layout = framework::DataLayout::kMKLDNN; } diff --git a/paddle/fluid/operators/batch_norm_op.h b/paddle/fluid/operators/batch_norm_op.h index cb02194b6d..6a8ff6ceab 100644 --- a/paddle/fluid/operators/batch_norm_op.h +++ b/paddle/fluid/operators/batch_norm_op.h @@ -134,6 +134,10 @@ class BatchNormOp : public framework::OperatorWithKernel { protected: framework::OpKernelType GetExpectedKernelType( const framework::ExecutionContext& ctx) const override; + + framework::OpKernelType GetKernelTypeForVar( + const std::string& var_name, const Tensor& tensor, + const framework::OpKernelType& expected_kernel_type) const override; }; class BatchNormGradOp : public framework::OperatorWithKernel { diff --git a/python/paddle/fluid/tests/unittests/mkldnn/test_batch_norm_mkldnn_op.py b/python/paddle/fluid/tests/unittests/mkldnn/test_batch_norm_mkldnn_op.py index 2c8daf3817..4439d79d3d 100644 --- a/python/paddle/fluid/tests/unittests/mkldnn/test_batch_norm_mkldnn_op.py +++ b/python/paddle/fluid/tests/unittests/mkldnn/test_batch_norm_mkldnn_op.py @@ -86,6 +86,13 @@ class TestMKLDNNBatchNormOpInference(TestBatchNormOpInference): self.check_with_place(place, data_format, self.dtype, [2, 3, 4, 5]) +class TestMKLDNNBatchNormOpInference_NHWC(TestMKLDNNBatchNormOpInference): + def test_check_output(self): + place = core.CPUPlace() + data_format = "NHWC" + self.check_with_place(place, data_format, self.dtype, [2, 4, 5, 3]) + + class TestMKLDNNBatchNormOpWithReluInference(TestBatchNormOpInference): def init_kernel_type(self): self.use_mkldnn = True diff --git a/python/paddle/fluid/tests/unittests/test_batch_norm_op.py b/python/paddle/fluid/tests/unittests/test_batch_norm_op.py index 7c445594c7..e25a09385c 100644 --- a/python/paddle/fluid/tests/unittests/test_batch_norm_op.py +++ b/python/paddle/fluid/tests/unittests/test_batch_norm_op.py @@ -262,6 +262,21 @@ class TestBatchNormOpInference(unittest.TestCase): batch_norm_op.run(scope, place) + # When op is called without Executor then + # MKL-DNN Tensor is returned. For NHWC data layout + # dims will be in NCHW order as it is MKL-DNN way + # of memory descripting. So we need to convert NCHW + # dims into NHWC. + if data_layout == "NHWC" and self.use_mkldnn == True: + # Create executor to have MKL-DNN cache + # cleared after NHWC unit test + place = core.CPUPlace() + exe = fluid.Executor(place) + dims = y_tensor.shape() + c = dims.pop(1) + dims.append(c) + y_tensor._set_dims(dims) + # check inference result self.__assert_close( y_tensor, -- GitLab