You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@mxnet.apache.org by pa...@apache.org on 2020/06/02 12:23:31 UTC
[incubator-mxnet] branch v1.x updated: [1.x] Add BatchNormWithReLU
fp32/bf16 (#18160)
This is an automated email from the ASF dual-hosted git repository.
patriczhao pushed a commit to branch v1.x
in repository https://gitbox.apache.org/repos/asf/incubator-mxnet.git
The following commit(s) were added to refs/heads/v1.x by this push:
new 843f278 [1.x] Add BatchNormWithReLU fp32/bf16 (#18160)
843f278 is described below
commit 843f2786d9312fe1a072365c8560c870d8064b64
Author: Xinyu Chen <xi...@intel.com>
AuthorDate: Tue Jun 2 20:22:13 2020 +0800
[1.x] Add BatchNormWithReLU fp32/bf16 (#18160)
* add bnrelu bf16 into amp list
* [MKL-DNN] BatchNormRelu Fusion (#17679)
* support bnrelu fusion
* add param to gluon
* move to contrib
* fix gluon interface
* reuse bn param
* fix
* fix forward
* cache flags
* fix lint
* fix ut
* trigger
* trigger
* trigger
* inherit from BN base
* fix lint
Co-authored-by: no <no...@no.no>
* Remove chinese period which leads to utf-8 encoding problem (#18223)
* add bnrelu amp test
* trigger ci
Co-authored-by: no <no...@no.no>
Co-authored-by: damNull <19...@163.com>
---
python/mxnet/contrib/amp/lists/symbol_bf16.py | 4 +-
python/mxnet/gluon/contrib/nn/basic_layers.py | 3 +-
python/mxnet/gluon/nn/basic_layers.py | 142 +++++++++++-
src/operator/contrib/batch_norm_relu.cc | 298 +++++++++++++++++++++++++
src/operator/nn/batch_norm.cc | 8 +-
src/operator/nn/mkldnn/mkldnn_batch_norm-inl.h | 49 +++-
tests/python/mkl/test_bf16_operator.py | 12 +
tests/python/mkl/test_mkldnn.py | 60 +++++
8 files changed, 555 insertions(+), 21 deletions(-)
diff --git a/python/mxnet/contrib/amp/lists/symbol_bf16.py b/python/mxnet/contrib/amp/lists/symbol_bf16.py
index c7612fa..86edfe6 100644
--- a/python/mxnet/contrib/amp/lists/symbol_bf16.py
+++ b/python/mxnet/contrib/amp/lists/symbol_bf16.py
@@ -32,6 +32,7 @@ BF16_FP32_FUNCS = [
'abs',
'_add',
'BatchNorm',
+ 'BatchNormWithReLU',
'clip',
'Concat',
'concat',
@@ -47,7 +48,8 @@ BF16_FP32_FUNCS = [
# Functions that when running with Bfloat16, the params that still need float32.
BF16_USE_FP32_PARAMS = {
- 'BatchNorm': ["", "gamma", "beta", "moving_mean", "moving_var"]
+ 'BatchNormWithReLU': ["", "gamma", "beta", "moving_mean", "moving_var"],
+ 'BatchNorm': ["", "gamma", "beta", "moving_mean", "moving_var"],
}
# Functions that have to be cast to FP32 due to possible
diff --git a/python/mxnet/gluon/contrib/nn/basic_layers.py b/python/mxnet/gluon/contrib/nn/basic_layers.py
index 706e5e4..bc7c3ce 100644
--- a/python/mxnet/gluon/contrib/nn/basic_layers.py
+++ b/python/mxnet/gluon/contrib/nn/basic_layers.py
@@ -220,8 +220,9 @@ class SyncBatchNorm(BatchNorm):
center=True, scale=True, use_global_stats=False, beta_initializer='zeros',
gamma_initializer='ones', running_mean_initializer='zeros',
running_variance_initializer='ones', **kwargs):
+ fuse_relu = False
super(SyncBatchNorm, self).__init__(1, momentum, epsilon, center, scale, use_global_stats,
- beta_initializer, gamma_initializer,
+ fuse_relu, beta_initializer, gamma_initializer,
running_mean_initializer, running_variance_initializer,
in_channels, **kwargs)
num_devices = self._get_num_devices() if num_devices is None else num_devices
diff --git a/python/mxnet/gluon/nn/basic_layers.py b/python/mxnet/gluon/nn/basic_layers.py
index 61f9ea0..70b0a71 100644
--- a/python/mxnet/gluon/nn/basic_layers.py
+++ b/python/mxnet/gluon/nn/basic_layers.py
@@ -19,7 +19,7 @@
# pylint: disable= arguments-differ
"""Basic neural network layers."""
__all__ = ['Sequential', 'HybridSequential', 'Dense', 'Dropout', 'Embedding',
- 'BatchNorm', 'InstanceNorm', 'LayerNorm', 'GroupNorm',
+ 'BatchNorm', 'BatchNormReLU', 'InstanceNorm', 'LayerNorm', 'GroupNorm',
'Flatten', 'Lambda', 'HybridLambda']
import warnings
import numpy as np
@@ -279,8 +279,9 @@ class Dropout(HybridBlock):
**self.__dict__)
-class BatchNorm(HybridBlock):
- """Batch normalization layer (Ioffe and Szegedy, 2014).
+class _BatchNorm(HybridBlock):
+ """Abstract BatchNorm layer (private, used as implementation base).
+ Batch normalization layer (Ioffe and Szegedy, 2014).
Normalizes the input at each batch, i.e. applies a transformation
that maintains the mean activation close to 0 and the activation
standard deviation close to 1.
@@ -307,6 +308,8 @@ class BatchNorm(HybridBlock):
If True, use global moving statistics instead of local batch-norm. This will force
change batch-norm into a scale shift operator.
If False, use local batch-norm.
+ fuse_relu: bool, default False
+ If True, this operator is equal to `BN+ReLU`.
beta_initializer: str or `Initializer`, default 'zeros'
Initializer for the beta weight.
gamma_initializer: str or `Initializer`, default 'ones'
@@ -328,12 +331,14 @@ class BatchNorm(HybridBlock):
- **out**: output tensor with the same shape as `data`.
"""
def __init__(self, axis=1, momentum=0.9, epsilon=1e-5, center=True, scale=True,
- use_global_stats=False, beta_initializer='zeros', gamma_initializer='ones',
+ use_global_stats=False, fuse_relu=False,
+ beta_initializer='zeros', gamma_initializer='ones',
running_mean_initializer='zeros', running_variance_initializer='ones',
in_channels=0, **kwargs):
- super(BatchNorm, self).__init__(**kwargs)
+ super(_BatchNorm, self).__init__(**kwargs)
self._kwargs = {'axis': axis, 'eps': epsilon, 'momentum': momentum,
'fix_gamma': not scale, 'use_global_stats': use_global_stats}
+ self.fuse_relu = fuse_relu
if in_channels != 0:
self.in_channels = in_channels
@@ -359,10 +364,12 @@ class BatchNorm(HybridBlock):
def cast(self, dtype):
if np.dtype(dtype).name == 'float16':
dtype = 'float32'
- super(BatchNorm, self).cast(dtype)
+ super(_BatchNorm, self).cast(dtype)
def hybrid_forward(self, F, x, gamma, beta, running_mean, running_var):
batch_norm = F.npx.batch_norm if is_np_array() else F.BatchNorm
+ if (not is_np_array()) and self.fuse_relu:
+ batch_norm = F.contrib.BatchNormWithReLU
return batch_norm(x, gamma, beta, running_mean, running_var,
name='fwd', **self._kwargs)
@@ -375,6 +382,129 @@ class BatchNorm(HybridBlock):
content=', '.join(['='.join([k, v.__repr__()])
for k, v in self._kwargs.items()]))
+class BatchNorm(_BatchNorm):
+ """Batch normalization layer (Ioffe and Szegedy, 2014).
+ Normalizes the input at each batch, i.e. applies a transformation
+ that maintains the mean activation close to 0 and the activation
+ standard deviation close to 1.
+
+ Parameters
+ ----------
+ axis : int, default 1
+ The axis that should be normalized. This is typically the channels
+ (C) axis. For instance, after a `Conv2D` layer with `layout='NCHW'`,
+ set `axis=1` in `BatchNorm`. If `layout='NHWC'`, then set `axis=3`.
+ momentum: float, default 0.9
+ Momentum for the moving average.
+ epsilon: float, default 1e-5
+ Small float added to variance to avoid dividing by zero.
+ center: bool, default True
+ If True, add offset of `beta` to normalized tensor.
+ If False, `beta` is ignored.
+ scale: bool, default True
+ If True, multiply by `gamma`. If False, `gamma` is not used.
+ When the next layer is linear (also e.g. `nn.relu`),
+ this can be disabled since the scaling
+ will be done by the next layer.
+ use_global_stats: bool, default False
+ If True, use global moving statistics instead of local batch-norm. This will force
+ change batch-norm into a scale shift operator.
+ If False, use local batch-norm.
+ fuse_relu: False
+ beta_initializer: str or `Initializer`, default 'zeros'
+ Initializer for the beta weight.
+ gamma_initializer: str or `Initializer`, default 'ones'
+ Initializer for the gamma weight.
+ running_mean_initializer: str or `Initializer`, default 'zeros'
+ Initializer for the running mean.
+ running_variance_initializer: str or `Initializer`, default 'ones'
+ Initializer for the running variance.
+ in_channels : int, default 0
+ Number of channels (feature maps) in input data. If not specified,
+ initialization will be deferred to the first time `forward` is called
+ and `in_channels` will be inferred from the shape of input data.
+
+
+ Inputs:
+ - **data**: input tensor with arbitrary shape.
+
+ Outputs:
+ - **out**: output tensor with the same shape as `data`.
+ """
+ def __init__(self, axis=1, momentum=0.9, epsilon=1e-5, center=True, scale=True,
+ use_global_stats=False, fuse_relu=False,
+ beta_initializer='zeros', gamma_initializer='ones',
+ running_mean_initializer='zeros', running_variance_initializer='ones',
+ in_channels=0, **kwargs):
+ assert not fuse_relu, "Please use BatchNormReLU with Relu fusion"
+ super(BatchNorm, self).__init__(
+ axis=1, momentum=0.9, epsilon=1e-5, center=True, scale=True,
+ use_global_stats=False, fuse_relu=False,
+ beta_initializer='zeros', gamma_initializer='ones',
+ running_mean_initializer='zeros', running_variance_initializer='ones',
+ in_channels=0, **kwargs)
+
+class BatchNormReLU(_BatchNorm):
+ """Batch normalization layer (Ioffe and Szegedy, 2014).
+ Normalizes the input at each batch, i.e. applies a transformation
+ that maintains the mean activation close to 0 and the activation
+ standard deviation close to 1.
+
+ Parameters
+ ----------
+ axis : int, default 1
+ The axis that should be normalized. This is typically the channels
+ (C) axis. For instance, after a `Conv2D` layer with `layout='NCHW'`,
+ set `axis=1` in `BatchNorm`. If `layout='NHWC'`, then set `axis=3`.
+ momentum: float, default 0.9
+ Momentum for the moving average.
+ epsilon: float, default 1e-5
+ Small float added to variance to avoid dividing by zero.
+ center: bool, default True
+ If True, add offset of `beta` to normalized tensor.
+ If False, `beta` is ignored.
+ scale: bool, default True
+ If True, multiply by `gamma`. If False, `gamma` is not used.
+ When the next layer is linear (also e.g. `nn.relu`),
+ this can be disabled since the scaling
+ will be done by the next layer.
+ use_global_stats: bool, default False
+ If True, use global moving statistics instead of local batch-norm. This will force
+ change batch-norm into a scale shift operator.
+ If False, use local batch-norm.
+ fuse_relu: True
+ beta_initializer: str or `Initializer`, default 'zeros'
+ Initializer for the beta weight.
+ gamma_initializer: str or `Initializer`, default 'ones'
+ Initializer for the gamma weight.
+ running_mean_initializer: str or `Initializer`, default 'zeros'
+ Initializer for the running mean.
+ running_variance_initializer: str or `Initializer`, default 'ones'
+ Initializer for the running variance.
+ in_channels : int, default 0
+ Number of channels (feature maps) in input data. If not specified,
+ initialization will be deferred to the first time `forward` is called
+ and `in_channels` will be inferred from the shape of input data.
+
+
+ Inputs:
+ - **data**: input tensor with arbitrary shape.
+
+ Outputs:
+ - **out**: output tensor with the same shape as `data`.
+ """
+ def __init__(self, axis=1, momentum=0.9, epsilon=1e-5, center=True, scale=True,
+ use_global_stats=False, fuse_relu=True,
+ beta_initializer='zeros', gamma_initializer='ones',
+ running_mean_initializer='zeros', running_variance_initializer='ones',
+ in_channels=0, **kwargs):
+ assert fuse_relu, "Please use BatchNorm w/o Relu fusion"
+ super(BatchNormReLU, self).__init__(
+ axis=1, momentum=0.9, epsilon=1e-5, center=True, scale=True,
+ use_global_stats=False, fuse_relu=True,
+ beta_initializer='zeros', gamma_initializer='ones',
+ running_mean_initializer='zeros', running_variance_initializer='ones',
+ in_channels=0, **kwargs)
class Embedding(HybridBlock):
r"""Turns non-negative integers (indexes/tokens) into dense vectors
diff --git a/src/operator/contrib/batch_norm_relu.cc b/src/operator/contrib/batch_norm_relu.cc
new file mode 100644
index 0000000..14452cc
--- /dev/null
+++ b/src/operator/contrib/batch_norm_relu.cc
@@ -0,0 +1,298 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * Copyright (c) 2020 by Contributors
+ * \file batch_norm_relu.cc
+ * \brief
+ * \author Xinyu Chen
+*/
+
+#include "../nn/batch_norm-inl.h"
+#include <nnvm/op_attr_types.h>
+#include "../elemwise_op_common.h"
+#include "../operator_common.h"
+#if MXNET_USE_MKLDNN == 1
+#include "../nn/mkldnn/mkldnn_batch_norm-inl.h"
+#endif
+
+namespace mxnet {
+namespace op {
+
+namespace batchnormrelu {
+
+enum BatchNormWithReLUOpInputs {kData, kGamma, kBeta, kInMovingMean,
+ kInMovingVar}; // kGamma: weights, kBeta: biases
+enum BatchNormWithReLUOpOutputs {kOut, kMean, kVar, kWorkspace}; // req, out_data
+enum BatchNormWithReLUOpResource {kTempSpace};
+enum BatchNormWithReLUOpAuxiliary {kMovingMean, kMovingVar}; // aux_states
+
+/*! \brief Default channel axis if none specified in the params */
+constexpr int DEFAULT_AXIS = 1;
+} // namespace batchnormrelu
+
+static bool BatchNormWithReLUShape(const nnvm::NodeAttrs& attrs,
+ mxnet::ShapeVector *in_shape,
+ mxnet::ShapeVector *out_shape) {
+ const BatchNormParam& param = nnvm::get<BatchNormParam>(attrs.parsed);
+ using namespace mshadow;
+ CHECK_EQ(in_shape->size(), 5U) << "Input:[data, gamma, beta, MovingMean, MovingVar]";
+ CHECK_EQ(out_shape->size(), 4U);
+ const mxnet::TShape &dshape = in_shape->at(batchnormrelu::kData);
+
+ const size_t channelAxis = static_cast<size_t>(param.axis < 0
+ ? static_cast<int>(dshape.ndim()) + param.axis
+ : param.axis);
+ CHECK_LT(channelAxis, dshape.ndim()) << "Channel axis out of range: " << param.axis;
+
+ const int channelCount = dshape[channelAxis];
+
+ if (!mxnet::ndim_is_known(dshape)) {
+ return false;
+ }
+
+ in_shape->at(batchnormrelu::kGamma) = mxnet::TShape(Shape1(channelCount));
+ in_shape->at(batchnormrelu::kBeta) = mxnet::TShape(Shape1(channelCount));
+ in_shape->at(batchnormrelu::kInMovingMean) = mxnet::TShape(Shape1(channelCount)); // kMovingMean
+ in_shape->at(batchnormrelu::kInMovingVar) = mxnet::TShape(Shape1(channelCount)); // kMovingVar
+
+ out_shape->clear();
+ out_shape->push_back(dshape); // kOut
+ out_shape->push_back(Shape1(channelCount)); // kMean
+ out_shape->push_back(Shape1(channelCount)); // kVar
+ out_shape->push_back(dshape); // kWorkspace
+ return true;
+}
+
+static bool BatchNormWithReLUType(const nnvm::NodeAttrs& attrs,
+ std::vector<int> *in_type, std::vector<int> *out_type) {
+ using namespace mshadow;
+ CHECK_GE(in_type->size(), 1U);
+ const int dtype = (*in_type)[0];
+ CHECK_NE(dtype, -1) << "First input must have specified type";
+ // For float16 input type beta, gamma, mean, and average are stored in float32.
+ // For other input types, these parameters have the same type as input
+ // NOTE: This requirement is from cuDNN (v. 4 and 5)
+ int dtype_param;
+ MSHADOW_REAL_TYPE_SWITCH_EX(dtype, DTypeX, AccRealX, {
+ dtype_param = mshadow::DataType<AccRealX>::kFlag; });
+ std::vector<std::string> args{"data", "gamma", "beta", "mean", "var"};
+ CHECK_LE(in_type->size(), args.size());
+ for (size_t i = 1; i < in_type->size(); ++i) {
+ if ((*in_type)[i] == -1) {
+ (*in_type)[i] = dtype_param;
+ } else {
+ UNIFORM_TYPE_CHECK((*in_type)[i], dtype_param, args[i]);
+ }
+ }
+ const size_t n_out = 4;
+ out_type->clear();
+ out_type->push_back(dtype);
+ for (size_t i = 1; i < n_out; ++i) {
+ out_type->push_back(dtype_param);
+ }
+ return true;
+}
+
+#if MXNET_USE_MKLDNN == 1
+static inline bool SupportMKLDNNBNReLU(const NDArray &input, const BatchNormParam ¶m) {
+ mxnet::TShape shape = input.shape();
+ return SupportMKLDNN(input) && shape.ndim() == 4
+ && param.axis == mxnet::op::batchnormrelu::DEFAULT_AXIS;
+}
+
+void BatchNormWithReLUComputeExCPU(const nnvm::NodeAttrs &attrs,
+ const OpContext &ctx,
+ const std::vector<NDArray> &inputs,
+ const std::vector<OpReqType> &req,
+ const std::vector<NDArray> &outputs) {
+ CHECK_EQ(inputs.size(), 5U);
+ const BatchNormParam ¶m = nnvm::get<BatchNormParam>(attrs.parsed);
+ bool fuse_relu = true;
+ if (SupportMKLDNNBNReLU(inputs[0], param)) {
+ CHECK_GT(outputs.size(), 3U);
+ MKLDNN_OPCHECK_INIT(false, outputs.size(), inputs, outputs);
+ MKLDNN_REAL_TYPE_SWITCH(inputs[0].dtype(), DTYPE, {
+ MKLDNNBatchNormForward<DTYPE>(attrs, ctx, inputs, req, outputs, fuse_relu);
+ });
+ return;
+ }
+ LOG(FATAL) << "BatchNormWithReLU operator only supports MKL-DNN Backend.";
+}
+
+void BatchNormWithReLUGradComputeExCPU(const nnvm::NodeAttrs &attrs,
+ const OpContext &ctx,
+ const std::vector<NDArray> &inputs,
+ const std::vector<OpReqType> &req,
+ const std::vector<NDArray> &outputs) {
+ const BatchNormParam ¶m = nnvm::get<BatchNormParam>(attrs.parsed);
+ bool fuse_relu = true;
+ if (SupportMKLDNNBNReLU(inputs[0], param)) {
+ CHECK_EQ(inputs.size(), 9U);
+ MKLDNN_OPCHECK_INIT(true, outputs.size(), inputs, outputs);
+ MKLDNNBatchNormBackward<float>(attrs, ctx, inputs, req, outputs, fuse_relu);
+ return;
+ }
+ LOG(FATAL) << "BatchNormWithReLU operator only supports MKL-DNN Backend.";
+}
+#endif
+
+static inline bool BatchNormWithReLUStorageType(const nnvm::NodeAttrs &attrs,
+ const int dev_mask,
+ DispatchMode *dispatch_mode,
+ std::vector<int> *in_attrs,
+ std::vector<int> *out_attrs) {
+ const BatchNormParam ¶m = nnvm::get<BatchNormParam>(attrs.parsed);
+
+ bool dispatched = false;
+#if MXNET_USE_MKLDNN == 1
+ if (!dispatched) {
+ dispatched = MKLDNNStorageType(attrs, dev_mask, true, dispatch_mode,
+ in_attrs, out_attrs);
+ }
+ if (!MKLDNNEnvSet()) {
+ *dispatch_mode = DispatchMode::kFComputeFallback;
+ }
+#else
+ for (int& v : *in_attrs)
+ if (v == - 1) v = kDefaultStorage;
+ if (!dispatched && common::ContainsOnlyStorage(*in_attrs, kDefaultStorage)) {
+ dispatched = storage_type_assign(out_attrs, kDefaultStorage,
+ dispatch_mode, DispatchMode::kFCompute);
+ }
+ if (!dispatched) {
+ dispatched = dispatch_fallback(out_attrs, dispatch_mode);
+ }
+#endif
+ if (!common::ContainsOnlyStorage(*in_attrs, kDefaultStorage) && param.fix_gamma) {
+ LOG(FATAL) << "fix_gamma=True is not supported for sparse ndarrays. Tracked at #11647";
+ }
+ return dispatched;
+}
+
+std::vector<nnvm::NodeEntry> BatchNormWithReLUGrad(const nnvm::ObjectPtr& n,
+ const std::vector<nnvm::NodeEntry>& ograds) {
+ std::vector<nnvm::NodeEntry> out_data;
+ out_data.reserve(n->num_outputs());
+ for (size_t i = 0; i < n->num_outputs(); ++i)
+ out_data.emplace_back(n, i, 0);
+ std::vector<nnvm::NodeEntry> heads;
+ heads.reserve(9);
+ heads.emplace_back(ograds.at(0));
+ heads.emplace_back(out_data.at(batchnormrelu::kMean));
+ heads.emplace_back(out_data.at(batchnormrelu::kVar));
+ heads.emplace_back(n->inputs.at(batchnormrelu::kData));
+ heads.emplace_back(n->inputs.at(batchnormrelu::kGamma));
+ heads.emplace_back(n->inputs.at(batchnormrelu::kBeta));
+ heads.emplace_back(n->inputs.at(batchnormrelu::kInMovingMean));
+ heads.emplace_back(n->inputs.at(batchnormrelu::kInMovingVar));
+ heads.emplace_back(out_data.at(batchnormrelu::kWorkspace));
+
+ nnvm::ObjectPtr gnode = nnvm::Node::Create();
+ gnode->inputs = std::move(heads);
+ gnode->control_deps.emplace_back(n);
+ gnode->attrs = n->attrs;
+ gnode->attrs.op = nnvm::Op::Get("_backward_contrib_BatchNormWithReLU");
+ gnode->attrs.name = n->attrs.name + "_backward";
+ // The input of batchnorm
+ std::vector<nnvm::NodeEntry> in_grad;
+ in_grad.reserve(5);
+ for (size_t i = 0; i < 3; ++i)
+ in_grad.emplace_back(gnode, i, 0);
+ // attach no gradient node to forbid gradient on aux_state
+ nnvm::ObjectPtr ng = nnvm::Node::Create();
+ ng->attrs.op = Op::Get("_NoGradient");
+ ng->attrs.name = "NoGradient";
+ // the aux state of batchnorm
+ for (size_t i = 3; i < 5; ++i)
+ in_grad.emplace_back(ng);
+ return in_grad;
+}
+
+NNVM_REGISTER_OP(_contrib_BatchNormWithReLU)
+.describe(R"code(Batch normalization with ReLU fusion.
+
+An extented operator of Batch normalization which can fuse ReLU activation.
+
+)code" ADD_FILELINE)
+.set_num_inputs(5)
+.set_num_outputs(4)
+.set_attr_parser(ParamParser<BatchNormParam>)
+.set_attr<nnvm::FListInputNames>("FListInputNames",
+ [](const NodeAttrs& attrs) {
+ return std::vector<std::string>{"data", "gamma", "beta", "moving_mean", "moving_var"};
+})
+.set_attr<nnvm::FListOutputNames>("FListOutputNames",
+ [](const NodeAttrs& attrs) {
+ return std::vector<std::string>{"output", "mean", "var", "workspace"};
+})
+.set_attr<nnvm::FNumVisibleOutputs>("FNumVisibleOutputs",
+ [](const NodeAttrs& attrs) {
+ const BatchNormParam& param = nnvm::get<BatchNormParam>(attrs.parsed);
+ return param.output_mean_var ? 3 : 1;
+})
+.set_attr<nnvm::FMutateInputs>("FMutateInputs", [](const nnvm::NodeAttrs& attrs) {
+ return std::vector<uint32_t>{3, 4};
+})
+.set_attr<mxnet::FInferShape>("FInferShape", BatchNormWithReLUShape)
+.set_attr<nnvm::FInferType>("FInferType", BatchNormWithReLUType)
+.set_attr<FInferStorageType>("FInferStorageType", BatchNormWithReLUStorageType)
+#if MXNET_USE_MKLDNN == 1
+.set_attr<FComputeEx>("FComputeEx<cpu>", BatchNormWithReLUComputeExCPU)
+#endif
+.set_attr<nnvm::FGradient>("FGradient", BatchNormWithReLUGrad)
+#if MXNET_USE_MKLDNN == 1
+.set_attr<bool>("TIsMKLDNN", true)
+.set_attr<FResourceRequest>("FResourceRequest", [](const NodeAttrs& n) {
+ return std::vector<ResourceRequest>{ResourceRequest::kTempSpace};
+})
+#endif
+.add_argument("data", "NDArray-or-Symbol", "Input data to batch normalization")
+.add_argument("gamma", "NDArray-or-Symbol", "gamma array")
+.add_argument("beta", "NDArray-or-Symbol", "beta array")
+.add_argument("moving_mean", "NDArray-or-Symbol", "running mean of input")
+.add_argument("moving_var", "NDArray-or-Symbol", "running variance of input")
+.add_arguments(BatchNormParam::__FIELDS__())
+.set_attr<nnvm::FSetInputVarAttrOnCompose>(
+ "FSetInputVarAttrOnCompose",
+ [](const nnvm::NodeAttrs& attrs, nnvm::ObjectPtr var, const int index) {
+ if (var->attrs.dict.find("__init__") != var->attrs.dict.end()) return;
+ if (index == 3) {
+ var->attrs.dict["__init__"] = "[\"zero\", {}]";
+ } else if (index == 4) {
+ var->attrs.dict["__init__"] = "[\"one\", {}]";
+ }
+ });
+
+NNVM_REGISTER_OP(_backward_contrib_BatchNormWithReLU)
+.set_num_inputs(9)
+.set_num_outputs(3)
+.set_attr<nnvm::TIsBackward>("TIsBackward", true)
+.set_attr<FInferStorageType>("FInferStorageType", BatchNormWithReLUStorageType)
+#if MXNET_USE_MKLDNN == 1
+.set_attr<FResourceRequest>("FResourceRequest", [](const NodeAttrs& n) {
+ return std::vector<ResourceRequest>{ResourceRequest::kTempSpace};
+})
+.set_attr<bool>("TIsMKLDNN", true)
+.set_attr<FComputeEx>("FComputeEx<cpu>", BatchNormWithReLUGradComputeExCPU)
+#endif
+.set_attr_parser(ParamParser<BatchNormParam>);
+
+} // namespace op
+} // namespace mxnet
diff --git a/src/operator/nn/batch_norm.cc b/src/operator/nn/batch_norm.cc
index 97acced..ef38347 100644
--- a/src/operator/nn/batch_norm.cc
+++ b/src/operator/nn/batch_norm.cc
@@ -394,11 +394,12 @@ void BatchNormComputeExCPU(const nnvm::NodeAttrs &attrs,
const std::vector<NDArray> &outputs) {
CHECK_EQ(inputs.size(), 5U);
const BatchNormParam ¶m = nnvm::get<BatchNormParam>(attrs.parsed);
+ bool fuse_relu = false;
if (SupportMKLDNNBN(inputs[0], param)) {
MKLDNN_OPCHECK_INIT(false, outputs.size(), inputs, outputs);
MKLDNN_REAL_TYPE_SWITCH(inputs[0].dtype(), DTYPE, {
- MKLDNNRun(MKLDNNBatchNormForward<DTYPE>, attrs, ctx, inputs, req, outputs);
- });
+ MKLDNNBatchNormForward<DTYPE>(attrs, ctx, inputs, req, outputs, fuse_relu);
+ });
MKLDNN_OPCHECK_RUN(BatchNormCompute<cpu>, attrs, ctx, inputs, req, outputs);
return;
}
@@ -411,9 +412,10 @@ void BatchNormGradComputeExCPU(const nnvm::NodeAttrs &attrs,
const std::vector<OpReqType> &req,
const std::vector<NDArray> &outputs) {
const BatchNormParam ¶m = nnvm::get<BatchNormParam>(attrs.parsed);
+ bool fuse_relu = false;
if (SupportMKLDNNBN(inputs[0], param)) {
MKLDNN_OPCHECK_INIT(true, outputs.size(), inputs, outputs);
- MKLDNNRun(MKLDNNBatchNormBackward<float>, attrs, ctx, inputs, req, outputs);
+ MKLDNNBatchNormBackward<float>(attrs, ctx, inputs, req, outputs, fuse_relu);
MKLDNN_OPCHECK_RUN(BatchNormGradCompute<cpu>, attrs, ctx, inputs, req, outputs);
return;
}
diff --git a/src/operator/nn/mkldnn/mkldnn_batch_norm-inl.h b/src/operator/nn/mkldnn/mkldnn_batch_norm-inl.h
index d407d94..d740134 100644
--- a/src/operator/nn/mkldnn/mkldnn_batch_norm-inl.h
+++ b/src/operator/nn/mkldnn/mkldnn_batch_norm-inl.h
@@ -45,8 +45,9 @@ typedef mkldnn::batch_normalization_backward::primitive_desc t_bn_b_pdesc;
typedef mkldnn::batch_normalization_backward::desc t_bn_b_desc;
inline static mkldnn::normalization_flags _GetFlags(const std::vector<NDArray> &in_data,
- const std::vector<NDArray> &aux_states,
- const BatchNormParam ¶m, bool is_train_and_not_global_stats) {
+ const std::vector<NDArray> &aux_states,
+ bool is_train_and_not_global_stats,
+ bool fuse_relu) {
mkldnn::normalization_flags flags = static_cast<mkldnn::normalization_flags>(0U);
if (in_data.size() == 3U) {
flags |= mkldnn::normalization_flags::use_scale_shift;
@@ -57,6 +58,10 @@ inline static mkldnn::normalization_flags _GetFlags(const std::vector<NDArray> &
if (aux_states.size() == 2U && !is_train_and_not_global_stats) {
flags |= mkldnn::normalization_flags::use_global_stats;
}
+
+ if (fuse_relu) {
+ flags |= mkldnn::normalization_flags::fuse_norm_relu;
+ }
return flags;
}
@@ -128,6 +133,7 @@ static MKLDNNBNForward &GetBNForward(const BatchNormParam& param,
MKLDNNBNSignature key(param);
key.AddSign(ctx.is_train);
key.AddSign(*data_mem);
+ key.AddSign(static_cast<int>(flags));
auto it = fwds.find(key);
if (it == fwds.end()) {
@@ -149,15 +155,15 @@ static MKLDNNBNForward &GetBNForward(const BatchNormParam& param,
template <typename DType>
void MKLDNNBatchNormForward(const nnvm::NodeAttrs &attrs, const OpContext &ctx,
const std::vector<NDArray> &inputs, const std::vector<OpReqType> &req,
- const std::vector<NDArray> &outputs) {
+ const std::vector<NDArray> &outputs, bool fuse_relu) {
const BatchNormParam ¶m = nnvm::get<BatchNormParam>(attrs.parsed);
const std::vector<NDArray> in_data(inputs.begin(), inputs.begin() + batchnorm::kInMovingMean);
const std::vector<NDArray> aux_states(inputs.begin() + batchnorm::kInMovingMean, inputs.end());
TmpMemMgr::Get()->Init(ctx.requested[batchnorm::kTempSpace]);
mkldnn::normalization_flags flags = _GetFlags(in_data,
aux_states,
- param,
- ctx.is_train && !param.use_global_stats);
+ ctx.is_train && !param.use_global_stats,
+ fuse_relu);
const NDArray &data = in_data[batchnorm::kData];
auto &fwd = GetBNForward<DType>(param, ctx, data, flags);
const NDArray &out = outputs[batchnorm::kOut];
@@ -201,7 +207,17 @@ void MKLDNNBatchNormForward(const nnvm::NodeAttrs &attrs, const OpContext &ctx,
net_args[MKLDNN_ARG_SRC] = *data.GetMKLDNNData();
net_args[MKLDNN_ARG_SCALE_SHIFT] = weight_mem;
net_args[MKLDNN_ARG_DST] = *out_mem;
-
+ if (fuse_relu) {
+ const NDArray *workspace = nullptr;
+ workspace = &outputs[3];
+ auto engine = CpuEngine::Get()->get_engine();
+ if (workspace == nullptr) {
+ LOG(FATAL) << "MKLDNN BatchNorm: incorrect workspace input";
+ }
+ auto ws = std::make_shared<mkldnn::memory>(fwd.GetPd().workspace_desc(),
+ engine, workspace->GetMKLDNNData()->get_data_handle());
+ net_args[MKLDNN_ARG_WORKSPACE] = *ws;
+ }
if (!ctx.is_train || param.use_global_stats) {
float* omean = outputs[batchnorm::kMean].data().dptr<float>();
float* ovar = outputs[batchnorm::kVar].data().dptr<float>();
@@ -269,6 +285,7 @@ static MKLDNNBNBackward &GetBNBackward(
MKLDNNBNSignature key(param);
key.AddSign(in_data);
key.AddSign(diff_data);
+ key.AddSign(static_cast<int>(flags));
auto it = bwds.find(key);
if (it == bwds.end()) {
@@ -282,8 +299,12 @@ static MKLDNNBNBackward &GetBNBackward(
template <typename DType>
void MKLDNNBatchNormBackward(const nnvm::NodeAttrs &attrs, const OpContext &ctx,
const std::vector<NDArray> &inputs, const std::vector<OpReqType> &req,
- const std::vector<NDArray> &outputs) {
- CHECK_EQ(inputs.size(), 8U);
+ const std::vector<NDArray> &outputs, bool fuse_relu) {
+ if (fuse_relu) {
+ CHECK_EQ(inputs.size(), 9U);
+ } else {
+ CHECK_EQ(inputs.size(), 8U);
+ }
const BatchNormParam ¶m = nnvm::get<BatchNormParam>(attrs.parsed);
std::vector<NDArray> out_grad(1);
std::vector<NDArray> out_data(3);
@@ -301,8 +322,8 @@ void MKLDNNBatchNormBackward(const nnvm::NodeAttrs &attrs, const OpContext &ctx,
TmpMemMgr::Get()->Init(ctx.requested[batchnorm::kTempSpace]);
mkldnn::normalization_flags flags = _GetFlags(in_data,
aux_states,
- param,
- ctx.is_train && !param.use_global_stats);
+ ctx.is_train && !param.use_global_stats,
+ fuse_relu);
const NDArray &data = in_data[batchnorm::kData];
const NDArray &diff = out_grad[batchnorm::kOut];
@@ -352,6 +373,14 @@ void MKLDNNBatchNormBackward(const nnvm::NodeAttrs &attrs, const OpContext &ctx,
net_args[MKLDNN_ARG_DIFF_SCALE_SHIFT] = bwd.GetGradw();
net_args[MKLDNN_ARG_DIFF_DST] = *diff_mem;
+ if (fuse_relu) {
+ const NDArray *workspace = nullptr;
+ workspace = &inputs[8];
+ if (workspace != nullptr) {
+ net_args[MKLDNN_ARG_WORKSPACE] = *(workspace->GetMKLDNNData());
+ }
+ }
+
// training but no input mean and variance
if (ctx.is_train && !param.use_global_stats) {
DType* moving_mean_ptr = moving_mean.data().dptr<DType>();
diff --git a/tests/python/mkl/test_bf16_operator.py b/tests/python/mkl/test_bf16_operator.py
index b275c96..a301b46 100644
--- a/tests/python/mkl/test_bf16_operator.py
+++ b/tests/python/mkl/test_bf16_operator.py
@@ -130,6 +130,18 @@ def test_bf16_bn():
check_operator_accuracy(sym_fp32=bn_fp32, sym_bf16=bn_bf16, data_shape=(32, 16, 64, 64), bf16_use_fp32_params=True, etol=1e-2)
@with_seed()
+def test_bf16_bnrelu():
+ data_sym_fp32 = mx.sym.Variable(name='data')
+ data_sym_bf16 = mx.sym.Variable(name='data', dtype=bfloat16)
+
+ bnrelu_params = {"eps": 2e-05, "fix_gamma": False, "use_global_stats": True, "name": "bn"}
+ bnrelu_fp32 = mx.sym.contrib.BatchNormWithReLU(data_sym_fp32, **bnrelu_params)
+
+ bnrelu_bf16 = mx.sym.contrib.BatchNormWithReLU(data_sym_bf16, **bnrelu_params)
+ check_operator_accuracy(sym_fp32=bnrelu_fp32, sym_bf16=bnrelu_bf16, data_shape=(3, 32, 28, 28), bf16_use_fp32_params=True, etol=1e-2)
+ check_operator_accuracy(sym_fp32=bnrelu_fp32, sym_bf16=bnrelu_bf16, data_shape=(32, 16, 64, 64), bf16_use_fp32_params=True, etol=1e-2)
+
+@with_seed()
def test_bf16_conv():
data_sym_fp32 = mx.sym.Variable(name='data')
data_sym_bf16 = mx.sym.Variable(name='data', dtype=bfloat16)
diff --git a/tests/python/mkl/test_mkldnn.py b/tests/python/mkl/test_mkldnn.py
index b52bb03..731a761 100644
--- a/tests/python/mkl/test_mkldnn.py
+++ b/tests/python/mkl/test_mkldnn.py
@@ -295,6 +295,66 @@ def test_batchnorm():
for stype in stypes:
check_batchnorm_training(stype)
+@with_seed()
+def test_batchnorm_relu_fusion():
+ def check_batchnorm_relu_fusion(shape):
+ x = mx.sym.Variable('x')
+ in_data = mx.nd.random.normal(shape=shape)
+ grad_out = mx.nd.random.uniform(0, 1, shape)
+ bn = mx.sym.BatchNorm(data=x, fix_gamma=False)
+ relu = mx.sym.Activation(data=bn, act_type='relu', name='relu')
+ exe = relu.simple_bind(ctx=mx.cpu(), x=shape, grad_req='write')
+ exe.arg_arrays[0][:] = in_data
+ exe.forward(is_train=True)
+ exe.backward(grad_out)
+ no_fuse_outputs = exe.outputs
+ no_fuse_grads = exe.grad_arrays
+
+ bnrelu = mx.sym.contrib.BatchNormWithReLU(data=x, fix_gamma=False)
+ exe_fuse = bnrelu.simple_bind(ctx=mx.cpu(), x=shape, grad_req='write')
+ exe_fuse.arg_arrays[0][:] = in_data
+ exe_fuse.forward(is_train=True)
+ exe_fuse.backward(grad_out)
+ fuse_outputs = exe_fuse.outputs
+ fuse_grads = exe_fuse.grad_arrays
+
+ for i in range(len(no_fuse_outputs)):
+ assert_almost_equal(no_fuse_outputs[i], fuse_outputs[i])
+ for i in range(len(no_fuse_grads)):
+ assert_almost_equal(no_fuse_grads[i], fuse_grads[i])
+
+ def check_batchnorm_relu_fusion_gluon(shape):
+ class BNNet(gluon.HybridBlock):
+ def __init__(self, fuse_relu):
+ super(BNNet, self).__init__()
+ self.fuse_relu = fuse_relu
+ with self.name_scope():
+ if self.fuse_relu:
+ self.bn = gluon.nn.BatchNormReLU()
+ else:
+ self.bn = gluon.nn.BatchNorm()
+ self.relu = gluon.nn.Activation('relu')
+
+ def forward(self, x):
+ y = self.bn(x)
+ if not self.fuse_relu:
+ y = self.relu(y)
+ return y
+ fused_net = BNNet(fuse_relu=True)
+ unfused_net = BNNet(fuse_relu=False)
+ fused_net.collect_params().initialize()
+ unfused_net.collect_params().initialize()
+ in_data = mx.nd.random.normal(shape=shape)
+ no_fuse_outputs = unfused_net.forward(in_data)
+ fuse_outputs = fused_net.forward(in_data)
+
+ for i in range(len(no_fuse_outputs)):
+ assert_almost_equal(no_fuse_outputs[i], fuse_outputs[i])
+
+ check_batchnorm_relu_fusion((1, 3, 224, 224))
+ check_batchnorm_relu_fusion((8, 3, 224, 224))
+ check_batchnorm_relu_fusion_gluon((1, 3, 224, 224))
+ check_batchnorm_relu_fusion_gluon((8, 3, 224, 224))
@with_seed()
def test_softmax():