You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@mxnet.apache.org by GitBox <gi...@apache.org> on 2018/06/30 21:04:34 UTC
[GitHub] zhanghang1989 closed pull request #10303: [MXNET-246] operators for
Synchronized BatchNorm
zhanghang1989 closed pull request #10303: [MXNET-246] operators for Synchronized BatchNorm
URL: https://github.com/apache/incubator-mxnet/pull/10303
This is a PR merged from a forked repository.
As GitHub hides the original diff on merge, it is displayed below for
the sake of provenance:
As this is a foreign pull request (from a fork), the diff is supplied
below (as it won't show otherwise due to GitHub magic):
diff --git a/src/imperative/imperative_utils.h b/src/imperative/imperative_utils.h
index 044ab823f77..fcb2b5817a9 100644
--- a/src/imperative/imperative_utils.h
+++ b/src/imperative/imperative_utils.h
@@ -164,7 +164,12 @@ inline void SetShapeType(const Context& ctx,
NDArrayStorageType storage_type = static_cast<NDArrayStorageType>(out_storage_types[i]);
if (outputs[i]->is_none()) {
if (storage_type == kDefaultStorage) {
- *outputs[i] = NDArray(out_shapes[i], ctx, true, out_types[i]);
+ if (outputs.size() == inputs.size()) {
+ // TODO FIXME (Hang Zhang), find a properate way to handle multi-device ctx
+ *outputs[i] = NDArray(out_shapes[i], inputs[i]->ctx(), true, out_types[i]);
+ } else {
+ *outputs[i] = NDArray(out_shapes[i], ctx, true, out_types[i]);
+ }
} else {
*outputs[i] = NDArray(storage_type, out_shapes[i], ctx, true, out_types[i]);
}
diff --git a/src/operator/allreduce-inl.h b/src/operator/allreduce-inl.h
new file mode 100644
index 00000000000..30215c1212c
--- /dev/null
+++ b/src/operator/allreduce-inl.h
@@ -0,0 +1,143 @@
+/*
+ * 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) 2018 by Contributors
+ * \file allreduce-inl.h
+ * \brief all reduce operator
+ * \author Hang Zhang
+ */
+#ifndef MXNET_OPERATOR_ALL_REDUCE_INL_H_
+#define MXNET_OPERATOR_ALL_REDUCE_INL_H_
+
+#include <dmlc/logging.h>
+#include <dmlc/parameter.h>
+#include <mxnet/operator.h>
+#include <mxnet/ndarray.h>
+#include <map>
+#include <vector>
+#include <string>
+#include <utility>
+#include "../ndarray/ndarray_function.h"
+#include "./operator_common.h"
+#include "./mxnet_op.h"
+#include "./mshadow_op.h"
+#include "../kvstore/comm.h"
+
+namespace mxnet {
+namespace op {
+
+struct AllReduceOpParam : public dmlc::Parameter<AllReduceOpParam> {
+ int num_args;
+ DMLC_DECLARE_PARAMETER(AllReduceOpParam) {
+ DMLC_DECLARE_FIELD(num_args).set_lower_bound(1)
+ .describe("Number of inputs to be allreduced.");
+ }
+}; // struct AllReduceOpParam
+
+template<typename xpu>
+inline void AllReduceOpForwardEx(const nnvm::NodeAttrs& attrs,
+ const OpContext &ctx,
+ const std::vector<NDArray> &inputs,
+ const std::vector<OpReqType> &req,
+ const std::vector<NDArray> &outputs) {
+ using namespace mshadow;
+ using namespace mshadow::expr;
+ CHECK_EQ(inputs.size(), outputs.size());
+ CHECK_EQ(inputs.size(), req.size());
+ //int priority = 0;
+ // create buf
+ std::vector<NDArray> reduce(inputs.size());
+ NDArray out(outputs[0].shape(), outputs[0].ctx(), false, outputs[0].dtype());
+ // copy to buf
+ for (size_t i = 0; i < inputs.size(); ++i) {
+ //inputs[i].WaitToRead();
+ reduce[i] = NDArray(
+ outputs[0].shape(), outputs[0].ctx(), false, outputs[0].dtype());
+ //CopyFromTo(inputs[i], &(reduce[i]), priority);
+ TBlob tmp = reduce[i].data();
+ ndarray::Copy<xpu, xpu>(inputs[i].data(), &tmp,
+ inputs[i].ctx(), reduce[i].ctx(), ctx.run_ctx);
+ }
+ // all reduce
+ std::vector<TBlob> source_tblob(reduce.size());
+ for (size_t i = 0; i < reduce.size(); ++i) {
+ source_tblob[i] = reduce[i].data();
+ }
+ TBlob tmp = out.data();
+ ndarray::ElementwiseSum<xpu>(source_tblob, &tmp, ctx.run_ctx);
+ // copy to each
+ for (size_t i = 0; i < outputs.size(); ++i) {
+ TBlob tmp = outputs[i].data();
+ ndarray::Copy<xpu, xpu>(out.data(), &tmp,
+ out.ctx(), outputs[i].ctx(), ctx.run_ctx);
+ }
+}
+
+
+inline bool AllReduceShape(const nnvm::NodeAttrs& attrs,
+ std::vector<TShape> *in_attrs,
+ std::vector<TShape> *out_attrs) {
+ CHECK_EQ(in_attrs->size(), out_attrs->size());
+ for (int i = 0; i < static_cast<int>(in_attrs->size()); ++i) {
+ TShape& ishape = (*in_attrs)[i];
+ SHAPE_ASSIGN_CHECK(*out_attrs, i, ishape);
+ }
+ for (int i = 0; i < static_cast<int>(in_attrs->size()); ++i) {
+ TShape& ishape = (*out_attrs)[i];
+ SHAPE_ASSIGN_CHECK(*in_attrs, i, ishape);
+ }
+ return true;
+}
+
+inline bool AllReduceType(const nnvm::NodeAttrs& attrs,
+ std::vector<int>* in_attrs,
+ std::vector<int>* out_attrs) {
+ int dtype = (*in_attrs)[0];
+ CHECK_NE(dtype, -1) << "First input must have specified type";
+ // assign
+ for (int i = 0; i < static_cast<int>(in_attrs->size()); ++i) {
+ dtype = (*in_attrs)[i];
+ TYPE_ASSIGN_CHECK(*out_attrs, i, dtype);
+ }
+ for (int i = 0; i < static_cast<int>(in_attrs->size()); ++i) {
+ dtype = (*out_attrs)[i];
+ TYPE_ASSIGN_CHECK(*in_attrs, i, dtype);
+ }
+ return true;
+}
+
+inline bool AllReduceStorageType(const nnvm::NodeAttrs& attrs,
+ const int dev_mask,
+ DispatchMode* dispatch_mode,
+ std::vector<int>* in_attrs,
+ std::vector<int>* out_attrs) {
+ CHECK_EQ(in_attrs->size(), out_attrs->size());
+ *dispatch_mode = DispatchMode::kFComputeEx;
+ for (int& v : *in_attrs) {
+ if (v == - 1) v = kDefaultStorage;
+ }
+ for (size_t i = 0; i < out_attrs->size(); i++) {
+ (*out_attrs)[i] = kDefaultStorage;
+ }
+ return true;
+}
+
+} // namespace op
+} // namespace mxnet
+#endif // MXNET_OPERATOR_ALL_REDUCE_INL_H_
diff --git a/src/operator/allreduce.cc b/src/operator/allreduce.cc
new file mode 100644
index 00000000000..269794a73fd
--- /dev/null
+++ b/src/operator/allreduce.cc
@@ -0,0 +1,84 @@
+/*
+ * 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) 2018 by Contributors
+ * \file allreduce.cc
+ * \brief all reduce operator
+ * \author Hang Zhang
+ */
+
+#include "./allreduce-inl.h"
+
+namespace mxnet {
+namespace op {
+
+DMLC_REGISTER_PARAMETER(AllReduceOpParam);
+
+NNVM_REGISTER_OP(AllReduce)
+.describe(R"code(TODO docs
+)code" ADD_FILELINE)
+.set_attr_parser(ParamParser<AllReduceOpParam>)
+.set_num_inputs([](const nnvm::NodeAttrs& attrs) {
+ uint32_t ret = dmlc::get<AllReduceOpParam>(attrs.parsed).num_args;
+ return ret;
+ })
+.set_num_outputs([](const nnvm::NodeAttrs& attrs) {
+ uint32_t ret = dmlc::get<AllReduceOpParam>(attrs.parsed).num_args;
+ return ret;
+ })
+.set_attr<nnvm::FInferShape>("FInferShape", AllReduceShape)
+.set_attr<nnvm::FInferType>("FInferType", AllReduceType)
+.set_attr<FInferStorageType>("FInferStorageType", AllReduceStorageType)
+.set_attr<std::string>("key_var_num_args", "num_args")
+.set_attr<FComputeEx>("FComputeEx<cpu>", AllReduceOpForwardEx<cpu>)
+.set_attr<nnvm::FGradient>("FGradient", ElemwiseGradUseNone{"_backward_AllReduce"})
+.set_attr<nnvm::FInplaceOption>("FInplaceOption", [](const NodeAttrs& attrs){
+ uint32_t n = dmlc::get<AllReduceOpParam>(attrs.parsed).num_args;
+ std::vector<std::pair<int, int> > ret;
+ for (uint32_t i = 0; i < n; i++) {
+ ret.push_back(std::pair<int, int>(i, i));
+ }
+ return ret;
+})
+.add_argument("data", "NDArray-or-Symbol[]", "List of arrays to allreduce");
+
+NNVM_REGISTER_OP(_backward_AllReduce)
+.set_attr_parser(ParamParser<AllReduceOpParam>)
+.set_num_inputs([](const nnvm::NodeAttrs& attrs) {
+ uint32_t ret = dmlc::get<AllReduceOpParam>(attrs.parsed).num_args;
+ return ret;
+ })
+.set_num_outputs([](const nnvm::NodeAttrs& attrs) {
+ uint32_t ret = dmlc::get<AllReduceOpParam>(attrs.parsed).num_args;
+ return ret;
+ })
+.set_attr<nnvm::TIsBackward>("TIsBackward", true)
+.set_attr<FInferStorageType>("FInferStorageType", AllReduceStorageType)
+.set_attr<nnvm::FInplaceOption>("FInplaceOption", [](const NodeAttrs& attrs){
+ uint32_t n = dmlc::get<AllReduceOpParam>(attrs.parsed).num_args;
+ std::vector<std::pair<int, int> > ret;
+ for (uint32_t i = 0; i < n; i++) {
+ ret.push_back(std::pair<int, int>(i, i));
+ }
+ return ret;
+})
+.set_attr<FComputeEx>("FComputeEx<cpu>", AllReduceOpForwardEx<cpu>);
+
+} // namespace op
+} // namespace mxnet
diff --git a/src/operator/allreduce.cu b/src/operator/allreduce.cu
new file mode 100644
index 00000000000..dcb0cb7c271
--- /dev/null
+++ b/src/operator/allreduce.cu
@@ -0,0 +1,38 @@
+/*
+ * 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) 2018 by Contributors
+ * \file allreduce.cu
+ * \brief all reduce operator
+ * \author Hang Zhang
+ */
+
+#include "./allreduce-inl.h"
+
+namespace mxnet {
+namespace op {
+
+NNVM_REGISTER_OP(AllReduce)
+.set_attr<FComputeEx>("FComputeEx<gpu>", AllReduceOpForwardEx<gpu>);
+
+NNVM_REGISTER_OP(_backward_AllReduce)
+.set_attr<FComputeEx>("FComputeEx<gpu>", AllReduceOpForwardEx<gpu>);
+
+} // namespace op
+} // namespace mxnet
diff --git a/src/operator/decouple_batch_norm-inl.h b/src/operator/decouple_batch_norm-inl.h
new file mode 100644
index 00000000000..fb69bfa4432
--- /dev/null
+++ b/src/operator/decouple_batch_norm-inl.h
@@ -0,0 +1,244 @@
+/*
+ * 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) 2018 by Contributors
+ * \file sync_batch_norm-inl.h
+ * \brief
+ * \author Hang Zhang
+ * Adapted from BatchNormV1
+ */
+#ifndef MXNET_OPERATOR_SYNC_BATCH_NORM_INL_H_
+#define MXNET_OPERATOR_SYNC_BATCH_NORM_INL_H_
+
+#include <dmlc/logging.h>
+#include <dmlc/parameter.h>
+//#include <mxnet/operator.h>
+#include <map>
+#include <vector>
+#include <string>
+#include <utility>
+#include "./operator_common.h"
+#include "./mshadow_op.h"
+
+namespace mxnet {
+namespace op {
+
+using namespace mshadow;
+template<typename xpu>
+inline void BNForward(Tensor<xpu, 4> data,
+ Tensor<xpu, 1> gamma,
+ Tensor<xpu, 1> beta,
+ Tensor<xpu, 1> mean,
+ Tensor<xpu, 1> std,
+ Tensor<xpu, 4> &out,
+ const std::vector<OpReqType> &req) {
+ using namespace mshadow::expr;
+ Assign(out, req[0], broadcast<1>(gamma / std, data.shape_) * data +
+ broadcast<1>(beta - (gamma * mean) / std, data.shape_));
+}
+
+template<typename xpu>
+inline void BNBackward(Tensor<xpu, 4> grad,
+ Tensor<xpu, 4> data,
+ Tensor<xpu, 1> mean,
+ Tensor<xpu, 1> std,
+ Tensor<xpu, 1> gamma,
+ Tensor<xpu, 4> &grad_in,
+ Tensor<xpu, 1> &gradGamma,
+ Tensor<xpu, 1> &gradBeta,
+ Tensor<xpu, 1> &gradMean,
+ Tensor<xpu, 1> &gradStd,
+ const std::vector<OpReqType> &req) {
+ using namespace mshadow::expr;
+ // the grad may not be zero originally? check this
+ Assign(gradMean, req[3], -1.f *
+ sumall_except_dim<1>(grad* broadcast<1>(gamma / std, data.shape_)));
+ Assign(gradStd, req[4],
+ sumall_except_dim<1>((grad * broadcast<1>(gamma, data.shape_)) *
+ (data - broadcast<1>(mean, data.shape_)) *
+ -1.f *
+ F<mshadow_op::power>(broadcast<1>(std, data.shape_),
+ -2.f)));
+ Assign(gradGamma, req[1],
+ sumall_except_dim<1>(
+ grad * (data - broadcast<1>(mean, data.shape_)) /
+ broadcast<1>(std, data.shape_)));
+ Assign(gradBeta, req[2], sumall_except_dim<1>(grad));
+ Assign(grad_in, req[0],
+ (grad * broadcast<1>(gamma / std, data.shape_)))
+}
+
+template<typename xpu>
+inline void DecoupleBNForward(const nnvm::NodeAttrs& attrs,
+ const OpContext &ctx,
+ const std::vector<TBlob> &inputs,
+ const std::vector<OpReqType> &req,
+ const std::vector<TBlob> &outputs) {
+ using namespace mshadow;
+ using namespace mshadow::expr;
+ CHECK_EQ(inputs.size(), 5U);
+ CHECK_EQ(outputs.size(), 1U);
+ CHECK_EQ(req.size(), 1U);
+ if (!ctx.is_train) {
+ CHECK_EQ(req[0], kWriteTo);
+ }
+ Stream<xpu> *s = ctx.get_stream<xpu>();
+ Tensor<xpu, 4> data;
+ Tensor<xpu, 4> out;
+ if (inputs[0].ndim() == 2) {
+ Shape<4> dshape = Shape4(inputs[0].shape_[0],
+ inputs[0].shape_[1], 1, 1);
+ data = inputs[0].get_with_shape<xpu, 4, real_t>(dshape, s);
+ out = outputs[0].get_with_shape<xpu, 4, real_t>(dshape, s);
+ } else {
+ data = inputs[0].get<xpu, 4, real_t>(s);
+ out = outputs[0].get<xpu, 4, real_t>(s);
+ }
+ Tensor<xpu, 1> gamma = inputs[1].get<xpu, 1, real_t>(s);
+ Tensor<xpu, 1> beta = inputs[2].get<xpu, 1, real_t>(s);
+ Tensor<xpu, 1> mean = inputs[3].get<xpu, 1, real_t>(s);
+ Tensor<xpu, 1> std = inputs[4].get<xpu, 1, real_t>(s);
+
+ BNForward<xpu>(data, gamma, beta, mean,std, out, req);
+}
+
+template<typename xpu>
+inline void DecoupleBNBackward(const nnvm::NodeAttrs& attrs,
+ const OpContext &ctx,
+ const std::vector<TBlob> &inputs,
+ const std::vector<OpReqType> &req,
+ const std::vector<TBlob> &outputs) {
+ using namespace mshadow;
+ using namespace mshadow::expr;
+ CHECK_EQ(inputs.size(), 7U);
+ CHECK_EQ(outputs.size(), 5U);
+ Stream<xpu> *s = ctx.get_stream<xpu>();
+ std::vector<TBlob> out_grad(inputs.begin(), inputs.begin() + 1);
+ std::vector<TBlob> in_data(inputs.begin() + 1,
+ inputs.begin() + 6);
+ //std::vector<TBlob> out_data(inputs.begin() + 6, inputs.end());
+ std::vector<TBlob> in_grad(outputs.begin(), outputs.begin() + 5);
+ Tensor<xpu, 4> data, grad, grad_in;
+ if (in_data[0].ndim() == 2) {
+ Shape<4> dshape = Shape4(out_grad[0].shape_[0],
+ out_grad[0].shape_[1], 1, 1);
+ data = in_data[0].get_with_shape<xpu, 4, real_t>(dshape, s);
+ grad = out_grad[0].get_with_shape<xpu, 4, real_t>(dshape, s);
+ grad_in = in_grad[0].get_with_shape<xpu, 4, real_t>(dshape, s);
+ } else {
+ data = in_data[0].get<xpu, 4, real_t>(s);
+ grad = out_grad[0].get<xpu, 4, real_t>(s);
+ grad_in = in_grad[0].get<xpu, 4, real_t>(s);
+ }
+
+ Tensor<xpu, 1> gamma = in_data[1].get<xpu, 1, real_t>(s);
+ //Tensor<xpu, 1> beta = in_data[2].get<xpu, 1, real_t>(s);
+ Tensor<xpu, 1> mean = in_data[3].get<xpu, 1, real_t>(s);
+ Tensor<xpu, 1> std = in_data[4].get<xpu, 1, real_t>(s);
+ Tensor<xpu, 1> gradGamma = in_grad[1].get<xpu, 1, real_t>(s);
+ Tensor<xpu, 1> gradBeta = in_grad[2].get<xpu, 1, real_t>(s);
+ Tensor<xpu, 1> gradMean = in_grad[3].get<xpu, 1, real_t>(s);
+ Tensor<xpu, 1> gradStd = in_grad[4].get<xpu, 1, real_t>(s);
+ BNBackward<xpu>(grad, data, mean, std, gamma,
+ grad_in, gradGamma, gradBeta, gradMean, gradStd, req);
+}
+
+
+static bool DecoupleBNInferShape(const nnvm::NodeAttrs& attrs,
+ std::vector<TShape> *in_shape,
+ std::vector<TShape> *out_shape) {
+ using namespace mshadow;
+ CHECK_EQ(in_shape->size(), 5U) << "Input:[data, gamma, beta, mean, std]";
+ const TShape &dshape = in_shape->at(0);
+ if (dshape.ndim() == 0) return false;
+ in_shape->at(1) = TShape(Shape1(dshape[1]));
+ in_shape->at(2) = TShape(Shape1(dshape[1]));
+ in_shape->at(3) = TShape(Shape1(dshape[1]));
+ in_shape->at(4) = TShape(Shape1(dshape[1]));
+ out_shape->clear();
+ out_shape->push_back(dshape);
+
+ return true;
+}
+
+static bool DecoupleBNInferType(const nnvm::NodeAttrs& attrs,
+ std::vector<int> *in_type,
+ std::vector<int> *out_type) {
+ using namespace mshadow;
+ CHECK_EQ(in_type->size(), 5U);
+ 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 = 0;
+ MSHADOW_REAL_TYPE_SWITCH_EX(dtype, DTypeX, AccRealX, {
+ dtype_param = mshadow::DataType<AccRealX>::kFlag; });
+ std::vector<std::string> args{"data", "gamma", "beta", "mean", "var"};
+ for (index_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]);
+ }
+ }
+ out_type->clear();
+ out_type->push_back(dtype);
+ return true;
+}
+
+static inline bool DecoupleBNStorageType(const nnvm::NodeAttrs &attrs,
+ const int dev_mask,
+ DispatchMode *dispatch_mode,
+ std::vector<int> *in_attrs,
+ std::vector<int> *out_attrs) {
+ CHECK_EQ(in_attrs->size(), 5);
+ CHECK_EQ(out_attrs->size(), 1);
+ *dispatch_mode = DispatchMode::kFCompute;
+ for (int& v : *in_attrs) {
+ if (v == - 1) v = kDefaultStorage;
+ }
+ for (size_t i = 0; i < out_attrs->size(); i++) {
+ (*out_attrs)[i] = kDefaultStorage;
+ }
+ return true;
+}
+
+static inline bool backward_DecoupleBNStorageType(const nnvm::NodeAttrs &attrs,
+ const int dev_mask,
+ DispatchMode *dispatch_mode,
+ std::vector<int> *in_attrs,
+ std::vector<int> *out_attrs) {
+ CHECK_EQ(in_attrs->size(), 7);
+ CHECK_EQ(out_attrs->size(), 5);
+ *dispatch_mode = DispatchMode::kFCompute;
+ for (int& v : *in_attrs) {
+ if (v == - 1) v = kDefaultStorage;
+ }
+ for (size_t i = 0; i < out_attrs->size(); i++) {
+ (*out_attrs)[i] = kDefaultStorage;
+ }
+ return true;
+}
+
+} // namespace op
+} // namespace mxnet
+
+#endif // MXNET_OPERATOR_SYNC_BATCH_NORM_INL_H_
diff --git a/src/operator/decouple_batch_norm.cc b/src/operator/decouple_batch_norm.cc
new file mode 100644
index 00000000000..ac7e8453c36
--- /dev/null
+++ b/src/operator/decouple_batch_norm.cc
@@ -0,0 +1,59 @@
+/*
+ * 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) 2015 by Contributors
+ * \file sync_batch_norm.cc
+ * \brief cpu sync BN operator
+ * \author Hang Zhang
+ * Adapted from BatchNormV1
+*/
+#include "decouple_batch_norm-inl.h"
+#include "elemwise_op_common.h"
+
+namespace mxnet {
+namespace op {
+
+
+NNVM_REGISTER_OP(DecoupleBatchNorm)
+.describe(R"code(Synchronized Cross-GPU Batch normalization.
+)code" ADD_FILELINE)
+.set_num_inputs(5)
+.set_num_outputs(1)
+.set_attr<nnvm::FInferShape>("FInferShape", DecoupleBNInferShape)
+.set_attr<nnvm::FInferType>("FInferType", DecoupleBNInferType)
+.set_attr<FInferStorageType>("FInferStorageType", DecoupleBNStorageType)
+.set_attr<FCompute>("FCompute<cpu>", DecoupleBNForward<cpu>)
+.set_attr<nnvm::FGradient>("FGradient", ElemwiseGradUseInOut{"_backward_DecoupleBatchNorm"})
+.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("mean", "NDArray-or-Symbol", "mean array")
+.add_argument("std", "NDArray-or-Symbol", "std array")
+;
+
+NNVM_REGISTER_OP(_backward_DecoupleBatchNorm)
+.set_num_outputs(5)
+.set_attr<nnvm::TIsBackward>("TIsBackward", true)
+.set_attr<FInferStorageType>("FInferStorageType", backward_DecoupleBNStorageType)
+.set_attr<FCompute>("FCompute<cpu>", DecoupleBNBackward<cpu>);
+;
+
+} // namespace op
+} // namespace mxnet
diff --git a/src/operator/decouple_batch_norm.cu b/src/operator/decouple_batch_norm.cu
new file mode 100644
index 00000000000..dc56dff4f40
--- /dev/null
+++ b/src/operator/decouple_batch_norm.cu
@@ -0,0 +1,40 @@
+/*
+ * 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) 2018 by Contributors
+ * \file sync_batch_norm.cu
+ * \brief CUDA Synchronized Multi-GPU Batch Normalization code
+ * \author Hang Zhang
+ * Adapted from BatchNormV1
+ */
+#include "decouple_batch_norm-inl.h"
+
+namespace mxnet {
+namespace op {
+
+NNVM_REGISTER_OP(DecoupleBatchNorm)
+.set_attr<FCompute>("FCompute<gpu>", DecoupleBNForward<gpu>);
+
+NNVM_REGISTER_OP(_backward_DecoupleBatchNorm)
+.set_attr<FCompute>("FCompute<gpu>", DecoupleBNBackward<gpu>);
+
+
+} // namespace op
+} // namespace mxnet
diff --git a/src/operator/sum_square-inl.h b/src/operator/sum_square-inl.h
new file mode 100644
index 00000000000..8b381bdcb10
--- /dev/null
+++ b/src/operator/sum_square-inl.h
@@ -0,0 +1,180 @@
+/*
+ * 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) 2018 by Contributors
+ * \file sum_square-inl.h
+ * \brief
+ * \author Hang Zhang
+ */
+
+#ifndef MXNET_OPERATOR_SUM_SQUARE_INL_H_
+#define MXNET_OPERATOR_SUM_SQUARE_INL_H_
+
+#include <dmlc/logging.h>
+#include <dmlc/parameter.h>
+#include <mxnet/operator.h>
+#include <map>
+#include <vector>
+#include <string>
+#include <utility>
+#include "./operator_common.h"
+#include "./mshadow_op.h"
+
+namespace mxnet {
+namespace op {
+
+template<typename xpu>
+inline void SumSquareForward(const nnvm::NodeAttrs& attrs,
+ const OpContext &ctx,
+ const std::vector<TBlob> &inputs,
+ const std::vector<OpReqType> &req,
+ const std::vector<TBlob> &outputs) {
+ using namespace mshadow;
+ using namespace mshadow::expr;
+ CHECK_EQ(inputs.size(), 1U);
+ CHECK_EQ(outputs.size(), 2U);
+ Stream<xpu> *s = ctx.get_stream<xpu>();
+ Tensor<xpu, 4> data;
+ if (inputs[0].ndim() == 2) {
+ Shape<4> dshape = Shape4(inputs[0].shape_[0],
+ inputs[0].shape_[1], 1, 1);
+ data = inputs[0].get_with_shape<xpu, 4, real_t>(dshape, s);
+ } else {
+ data = inputs[0].get<xpu, 4, real_t>(s);
+ }
+ Tensor<xpu, 1> sumx = outputs[0].get<xpu, 1, real_t>(s);
+ Tensor<xpu, 1> sumx2 = outputs[1].get<xpu, 1, real_t>(s);
+ Assign(sumx, req[0], sumall_except_dim<1>(data));
+ Assign(sumx2, req[1], sumall_except_dim<1>(F<mshadow_op::power>(data, 2.f)));
+}
+
+
+template<typename xpu>
+inline void SumSquareBackward(const nnvm::NodeAttrs& attrs,
+ const OpContext &ctx,
+ const std::vector<TBlob> &inputs,
+ const std::vector<OpReqType> &req,
+ const std::vector<TBlob> &outputs) {
+ using namespace mshadow;
+ using namespace mshadow::expr;
+ CHECK_EQ(inputs.size(), 5U);
+ CHECK_EQ(outputs.size(), 1U);
+ std::vector<TBlob> out_grad(inputs.begin(), inputs.begin() + 2);
+ std::vector<TBlob> in_data(inputs.begin() + 2,
+ inputs.begin() + 3);
+ std::vector<TBlob> in_grad(outputs.begin(), outputs.end());
+ Stream<xpu> *s = ctx.get_stream<xpu>();
+ Tensor<xpu, 4> data;
+ Tensor<xpu, 4> grad_in;
+ Tensor<xpu, 1> gradx;
+ Tensor<xpu, 1> gradx2;
+ if (in_data[0].ndim() == 2) {
+ Shape<4> dshape = Shape4(out_grad[0].shape_[0],
+ out_grad[0].shape_[1], 1, 1);
+ data = in_data[0].get_with_shape<xpu, 4, real_t>(dshape, s);
+ grad_in = in_grad[0].get_with_shape<xpu, 4, real_t>(dshape, s);
+ } else {
+ data = in_data[0].get<xpu, 4, real_t>(s);
+ grad_in = in_grad[0].get<xpu, 4, real_t>(s);
+ }
+ gradx = out_grad[0].get<xpu, 1, real_t>(s);
+ gradx2 = out_grad[1].get<xpu, 1, real_t>(s);
+ Assign(grad_in, req[0], broadcast<1>(gradx, data.shape_) +
+ data * broadcast<1>(gradx2 * 2.f, data.shape_));
+}
+
+static bool SumSquareInferShape(const nnvm::NodeAttrs& attrs,
+ std::vector<TShape> *in_shape,
+ std::vector<TShape> *out_shape) {
+ using namespace mshadow;
+ CHECK_EQ(in_shape->size(), 1U) << "Input:[data]";
+ const TShape &dshape = in_shape->at(0);
+ if (dshape.ndim() == 0) return false;
+ out_shape->clear();
+ out_shape->push_back(TShape(Shape1(dshape[1])));
+ out_shape->push_back(TShape(Shape1(dshape[1])));
+ return true;
+}
+
+static bool SumSquareInferType(const nnvm::NodeAttrs& attrs,
+ std::vector<int> *in_type,
+ std::vector<int> *out_type) {
+ using namespace mshadow;
+ CHECK_EQ(in_type->size(), 1U);
+ 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 = 0;
+ MSHADOW_REAL_TYPE_SWITCH_EX(dtype, DTypeX, AccRealX, {
+ dtype_param = mshadow::DataType<AccRealX>::kFlag; });
+ std::vector<std::string> args{"data"};
+ for (index_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]);
+ }
+ }
+ out_type->clear();
+ out_type->push_back(dtype);
+ out_type->push_back(dtype);
+ return true;
+}
+
+static inline bool SumSquareStorageType(const nnvm::NodeAttrs &attrs,
+ const int dev_mask,
+ DispatchMode *dispatch_mode,
+ std::vector<int> *in_attrs,
+ std::vector<int> *out_attrs) {
+ CHECK_EQ(in_attrs->size(), 1);
+ CHECK_EQ(out_attrs->size(), 2);
+ *dispatch_mode = DispatchMode::kFCompute;
+ for (int& v : *in_attrs) {
+ if (v == - 1) v = kDefaultStorage;
+ }
+ for (size_t i = 0; i < out_attrs->size(); i++) {
+ (*out_attrs)[i] = kDefaultStorage;
+ }
+ return true;
+}
+
+static inline bool backward_SumSquareStorageType(const nnvm::NodeAttrs &attrs,
+ const int dev_mask,
+ DispatchMode *dispatch_mode,
+ std::vector<int> *in_attrs,
+ std::vector<int> *out_attrs) {
+ CHECK_EQ(in_attrs->size(), 5);
+ CHECK_EQ(out_attrs->size(), 1);
+ *dispatch_mode = DispatchMode::kFCompute;
+ for (int& v : *in_attrs) {
+ if (v == - 1) v = kDefaultStorage;
+ }
+ for (size_t i = 0; i < out_attrs->size(); i++) {
+ (*out_attrs)[i] = kDefaultStorage;
+ }
+ return true;
+}
+
+
+} // namespace op
+} // namespace mxnet
+
+#endif // MXNET_OPERATOR_SUM_SQUARE_INL_H_
diff --git a/src/operator/sum_square.cc b/src/operator/sum_square.cc
new file mode 100644
index 00000000000..a82dcf61f85
--- /dev/null
+++ b/src/operator/sum_square.cc
@@ -0,0 +1,52 @@
+/*
+ * 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) 2018 by Contributors
+ * \file sum_square.cc
+ * \brief
+ * \author Hang Zhang
+*/
+#include "sum_square-inl.h"
+#include "elemwise_op_common.h"
+
+namespace mxnet {
+namespace op {
+
+NNVM_REGISTER_OP(SumSquare)
+.describe(R"code(In-device Sum and Sum of Square.
+)code" ADD_FILELINE)
+.set_num_inputs(1)
+.set_num_outputs(2)
+.set_attr<nnvm::FInferShape>("FInferShape", SumSquareInferShape)
+.set_attr<nnvm::FInferType>("FInferType", SumSquareInferType)
+.set_attr<FInferStorageType>("FInferStorageType", SumSquareStorageType)
+.set_attr<FCompute>("FCompute<cpu>", SumSquareForward<cpu>)
+.set_attr<nnvm::FGradient>("FGradient", ElemwiseGradUseInOut{"_backward_SumSquare"})
+.add_argument("data", "NDArray-or-Symbol", "Input data to batch normalization")
+;
+
+NNVM_REGISTER_OP(_backward_SumSquare)
+.set_num_outputs(1)
+.set_attr<nnvm::TIsBackward>("TIsBackward", true)
+.set_attr<FInferStorageType>("FInferStorageType", backward_SumSquareStorageType)
+.set_attr<FCompute>("FCompute<cpu>", SumSquareBackward<cpu>);
+;
+
+} // namespace op
+} // namespace mxnet
diff --git a/src/operator/sum_square.cu b/src/operator/sum_square.cu
new file mode 100644
index 00000000000..0ccf56b817a
--- /dev/null
+++ b/src/operator/sum_square.cu
@@ -0,0 +1,39 @@
+/*
+ * 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) 2018 by Contributors
+ * \file sync_batch_norm.cu
+ * \brief CUDA Synchronized Multi-GPU Batch Normalization code
+ * \author Hang Zhang
+ * Adapted from BatchNormV1
+ */
+#include "sum_square-inl.h"
+
+namespace mxnet {
+namespace op {
+
+NNVM_REGISTER_OP(SumSquare)
+.set_attr<FCompute>("FCompute<gpu>", SumSquareForward<gpu>);
+
+NNVM_REGISTER_OP(_backward_SumSquare)
+.set_attr<FCompute>("FCompute<gpu>", SumSquareBackward<gpu>);
+
+} // namespace op
+} // namespace mxnet
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services