Implementation for equivalence of tf.moments

src/operator/nn/moments-inl.h

@@ -0,0 +1,254 @@
+/*
+ * 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) 2019 by Contributors
+ * \file moments-inl.h
+ * \brief Moments operator
+ * \author Hao Jin
+*/
+
+#ifndef MXNET_OPERATOR_NN_MOMENTS_INL_H_
+#define MXNET_OPERATOR_NN_MOMENTS_INL_H_
+
+#include <vector>
+#include "../tensor/broadcast_reduce_op.h"
+
+namespace mxnet {
+namespace op {
+
+struct MomentsParam : public dmlc::Parameter<MomentsParam> {
+  dmlc::optional<mxnet::TShape> axes;
+  bool keepdims;
+  DMLC_DECLARE_PARAMETER(MomentsParam) {
+    DMLC_DECLARE_FIELD(axes).set_default(dmlc::optional<mxnet::TShape>())
+      .describe("Array of ints. Axes along which to compute mean and variance.");
+    DMLC_DECLARE_FIELD(keepdims).set_default(false)
+      .describe("produce moments with the same dimensionality as the input.");
+  }
+};
+
+inline bool MomentsShape(const nnvm::NodeAttrs& attrs,
+                         mxnet::ShapeVector* in_attrs,
+                         mxnet::ShapeVector* out_attrs) {
+  const MomentsParam& param = nnvm::get<MomentsParam>(attrs.parsed);
+  CHECK_EQ(in_attrs->size(), 1U);
+  CHECK_EQ(out_attrs->size(), 2U);
+
+  mxnet::TShape out_shape =
+    ReduceAxesShapeImpl((*in_attrs)[0], param.axes, param.keepdims, false);
+  if (!param.axes.has_value() || param.axes.value().ndim() == 0) {
+    LOG(FATAL) << "Empty axes is not supported, if you would like to do global moments, "
+               << "please pass all axes to axes argument";
+  }
+  SHAPE_ASSIGN_CHECK(*out_attrs, 0, out_shape);
+  SHAPE_ASSIGN_CHECK(*out_attrs, 1, out_shape);
+  return true;
+}
+
+inline bool MomentsType(const nnvm::NodeAttrs& attrs,
+                        std::vector<int>* in_attrs,
+                        std::vector<int>* out_attrs) {
+  CHECK_EQ(in_attrs->size(), 1U);
+  CHECK_EQ(out_attrs->size(), 2U);
+
+  TYPE_ASSIGN_CHECK(*out_attrs, 0, in_attrs->at(0));
+  TYPE_ASSIGN_CHECK(*out_attrs, 1, in_attrs->at(0));
+  TYPE_ASSIGN_CHECK(*in_attrs, 0, out_attrs->at(0));
+  TYPE_ASSIGN_CHECK(*in_attrs, 0, out_attrs->at(1));
+  return out_attrs->at(0) != -1 && out_attrs->at(1) != -1;
+}
+
+struct VarBroadcastKernel {
+  template<typename DType>
+  MSHADOW_XINLINE static void Map(int i,
+                                  DType *out,
+                                  const DType *data,
+                                  const DType *mean,
+                                  mshadow::Shape<6> data_shape,
+                                  mshadow::Shape<6> mean_shape) {
+    size_t data_idx = i;
+    size_t mean_idx = i;
+    size_t data_stride = 1;
+    size_t mean_stride = 1;
+    for (int axis = 5; axis >= 0; --axis) {
+      size_t axis_idx = data_idx % data_shape[axis];
+      mean_idx -= axis_idx * data_stride;
+      if (mean_shape[axis] != 1) {
+        mean_idx += axis_idx * mean_stride;
+      }
+      data_idx /= data_shape[axis];
+      data_stride *= data_shape[axis];
+      mean_stride *= mean_shape[axis];
+    }
+    DType res = (data[i] - mean[mean_idx]);
+    out[i] = res * res;
+  }
+};
+
+template<typename xpu>
+inline void MomentsForwardImpl(const OpContext& ctx,
+                               const std::vector<TBlob>& inputs,
+                               const std::vector<OpReqType>& req,
+                               const std::vector<TBlob>& outputs,
+                               const dmlc::optional<mxnet::TShape>& axes,
+                               const bool keepdims) {
+  using namespace mshadow;
+  using namespace mshadow_op;
+  using namespace mxnet_op;
+
+  Stream<xpu> *s = ctx.get_stream<xpu>();
+
+  const TBlob& data = inputs[0];
+  const TBlob& mean = outputs[0];
+  const TBlob& var = outputs[1];
+
+  mxnet::TShape small;
+  if (keepdims) {
+    small = outputs[0].shape_;
+  } else {
+    small = ReduceAxesShapeImpl(inputs[0].shape_, axes, true, false);
+  }
+
+  ReduceAxesComputeImpl<xpu, mshadow_op::sum, true, true>(ctx, {data}, {req[0]}, {mean}, small);
+  MSHADOW_TYPE_SWITCH(data.type_flag_, DType, {
+    Shape<6> data_shape, mean_shape;
+    for (int i = 0; i < 6; ++i) {
+      data_shape[i] = (i < data.shape_.ndim()) ? data.shape_[i] : 1;
+      mean_shape[i] = (i < small.ndim()) ? small[i] : 1;
+    }
+    Tensor<xpu, 1, DType> temp_data =
+      ctx.requested[0].get_space_typed<xpu, 1, DType>(Shape1(data.shape_.Size()), s);;
+    Kernel<VarBroadcastKernel, xpu>::Launch(s, data.shape_.Size(), temp_data.dptr_,
+      data.dptr<DType>(), mean.dptr<DType>(), data_shape, mean_shape);
+    ReduceAxesComputeImpl<xpu, mshadow_op::sum, true, true>(
+      ctx, {TBlob(temp_data).reshape(data.shape_)}, {kWriteTo}, {var}, small);
+  });
+}
+
+template<typename xpu>
+inline void MomentsForward(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_op;
+  using namespace mxnet_op;
+
+  CHECK_EQ(inputs.size(), 1U);
+  CHECK_EQ(outputs.size(), 2U);
+
+  const MomentsParam& param = nnvm::get<MomentsParam>(attrs.parsed);
+
+  MomentsForwardImpl<xpu>(ctx, inputs, req, outputs, param.axes, param.keepdims);
+}
+
+template<int req>
+struct VarBackwardKernel {
+  template<typename DType>
+  MSHADOW_XINLINE static void Map(int i,
+                                  DType *igrad,
+                                  const DType *ograd,
+                                  const DType *data,
+                                  const DType *mean,
+                                  mshadow::Shape<6> data_shape,
+                                  mshadow::Shape<6> mean_shape,
+                                  const float N,
+                                  const float ddof = 0.0f) {
+    size_t data_idx = i;
+    size_t mean_idx = i;
+    size_t data_stride = 1;
+    size_t mean_stride = 1;
+    for (int axis = 5; axis >= 0; --axis) {
+      size_t axis_idx = data_idx % data_shape[axis];
+      mean_idx -= axis_idx * data_stride;
+      if (mean_shape[axis] != 1) {
+        mean_idx += axis_idx * mean_stride;
+      }
+      data_idx /= data_shape[axis];
+      data_stride *= data_shape[axis];
+      mean_stride *= mean_shape[axis];
+    }
+    KERNEL_ASSIGN(igrad[i], req, ograd[mean_idx] * (data[i] - mean[mean_idx]) * 2 / (N - ddof));
+  }
+};
+
+template<typename xpu>
+inline void MomentsBackwardImpl(const nnvm::NodeAttrs& attrs,
+                                const OpContext& ctx,
+                                const std::vector<TBlob>& inputs,
+                                const std::vector<OpReqType>& req,
+                                const std::vector<TBlob>& outputs,
+                                const dmlc::optional<mxnet::TShape>& axes) {
+  using namespace mshadow;
+  using namespace mshadow::expr;
+  using namespace mshadow_op;
+  using namespace mxnet_op;
+
+  Stream<xpu> *s = ctx.get_stream<xpu>();
+
+  const TBlob& mean_grad = inputs[0];
+  const TBlob& var_grad = inputs[1];
+  const TBlob& data = inputs[2];
+  const TBlob& mean = inputs[3];
+  const TBlob& var = inputs[4];
+  const TBlob& data_grad = outputs[0];
+
+  mxnet::TShape small = ReduceAxesShapeImpl(data.shape_, axes, true, false);
+  BroadcastComputeImpl<xpu>(attrs, ctx, {mean_grad}, req, outputs, small);
+  MSHADOW_TYPE_SWITCH(outputs[0].type_flag_, DType, {
+    Tensor<xpu, 1, DType> igrad = outputs[0].FlatTo1D<xpu, DType>(s);
+    igrad /= scalar<DType>(outputs[0].Size()/inputs[0].Size());
+  });
+
+  Shape<6> data_shape, var_shape;
+  float N = data_grad.Size() / var.Size();
+  for (int i = 0; i < 6; ++i) {
+    data_shape[i] = (i < data.shape_.ndim()) ? data.shape_[i] : 1;
+    var_shape[i] = (i < small.ndim()) ? small[i] : 1;
+  }
+  MSHADOW_TYPE_SWITCH(data_grad.type_flag_, DType, {
+    Kernel<VarBackwardKernel<kAddTo>, xpu>::Launch(
+      s, data_grad.shape_.Size(), data_grad.dptr<DType>(), var_grad.dptr<DType>(),
+      data.dptr<DType>(), mean.dptr<DType>(), data_shape, var_shape, N);
+  });
+}
+
+template<typename xpu>
+inline void MomentsBackward(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_op;
+  using namespace mxnet_op;
+
+  CHECK_EQ(inputs.size(), 5U);
+  CHECK_EQ(outputs.size(), 1U);
+
+  const MomentsParam& param = nnvm::get<MomentsParam>(attrs.parsed);
+
+  MomentsBackwardImpl<xpu>(attrs, ctx, inputs, req, outputs, param.axes);
+}
+
+}  // namespace op
+}  // namespace mxnet
+#endif  // MXNET_OPERATOR_NN_MOMENTS_INL_H_

src/operator/nn/moments.cc

@@ -0,0 +1,85 @@
+/*
+ * 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) 2019 by Contributors
+ * \file moments.cc
+ * \brief Moments operator
+ * \author Hao Jin
+*/
+
+#include "./moments-inl.h"
+
+namespace mxnet {
+namespace op {
+
+DMLC_REGISTER_PARAMETER(MomentsParam);
+
+NNVM_REGISTER_OP(moments)
+.describe(R"code(
+Calculate the mean and variance of `data`.
+
+The mean and variance are calculated by aggregating the contents of data across axes.
+If x is 1-D and axes = [0] this is just the mean and variance of a vector.
+
+Example:
+
+     x = [[1, 2, 3], [4, 5, 6]]
+     mean, var = moments(data=x, axes=[0])
+     mean = [2.5, 3.5, 4.5]
+     var = [2.25, 2.25, 2.25]
+     mean, var = moments(data=x, axes=[1])
+     mean = [2.0, 5.0]
+     var = [0.66666667, 0.66666667]
+     mean, var = moments(data=x, axis=[0, 1])
+     mean = [3.5]
+     var = [2.9166667]
+
+)code" ADD_FILELINE)
+.set_attr_parser(ParamParser<MomentsParam>)
+.set_num_inputs(1)
+.set_num_outputs(2)
+.set_attr<nnvm::FListInputNames>("FListInputNames",
+  [](const NodeAttrs& attrs) {
+    return std::vector<std::string>{"data"};
+  })
+.set_attr<mxnet::FInferShape>("FInferShape", MomentsShape)
+.set_attr<nnvm::FInferType>("FInferType", MomentsType)
+.set_attr<FCompute>("FCompute<cpu>", MomentsForward<cpu>)
+.set_attr<FResourceRequest>("FResourceRequest",
+  [](const NodeAttrs& attrs) {
+    return std::vector<ResourceRequest>{ResourceRequest::kTempSpace};
+  })
+.set_attr<nnvm::FGradient>("FGradient", ElemwiseGradUseInOut{"_backward_moments"})
+.set_attr<nnvm::FInplaceOption>("FInplaceOption",
+  [](const NodeAttrs& attrs) {
+    return std::vector<std::pair<int, int> >{{0, 0}};
+  })
+.add_argument("data", "NDArray-or-Symbol", "Input ndarray")
+.add_arguments(MomentsParam::__FIELDS__());
+
+NNVM_REGISTER_OP(_backward_moments)
+.set_attr_parser(ParamParser<MomentsParam>)
+.set_num_inputs(5)
+.set_num_outputs(1)
+.set_attr<nnvm::TIsBackward>("TIsBackward", true)
+.set_attr<FCompute>("FCompute<cpu>", MomentsBackward<cpu>);
+
+}  // namespace op
+}  // namespace mxnet

src/operator/nn/moments.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) 2019 by Contributors
+ * \file moments.cu
+ * \brief Moments operator
+ * \author Hao Jin
+*/
+
+#include "./moments-inl.h"
+
+namespace mxnet {
+namespace op {
+
+NNVM_REGISTER_OP(moments)
+.set_attr<FCompute>("FCompute<gpu>", MomentsForward<gpu>);
+
+NNVM_REGISTER_OP(_backward_moments)
+.set_attr<FCompute>("FCompute<gpu>", MomentsBackward<gpu>);
+
+}  // namespace op
+}  // namespace mxnet