[GitHub] TaoLv commented on a change in pull request #13699: add mkldnn softmax_output

[GitBox <gi...@apache.org>]

TaoLv commented on a change in pull request #13699: add mkldnn softmax_output
URL: https://github.com/apache/incubator-mxnet/pull/13699#discussion_r255939205
 
 

 ##########
 File path: src/operator/softmax_output.cc
 ##########
 @@ -21,30 +21,139 @@
  * Copyright (c) 2015 by Contributors
  * \file softmax_output.cc
  * \brief
- * \author Bing Xu
+ * \author Bing Xu, Zhang Rong A
 */
 #include "./softmax_output-inl.h"
-
+#if MXNET_USE_MKLDNN == 1
+#include "./nn/mkldnn/mkldnn_ops-inl.h"
+#endif
 namespace mxnet {
 namespace op {
-template<>
-Operator *CreateOp<cpu>(SoftmaxOutputParam param, int dtype) {
-  Operator *op = nullptr;
-  MSHADOW_REAL_TYPE_SWITCH(dtype, DType, {
-    op = new SoftmaxOutputOp<cpu, DType>(param);
-  })
-  return op;
+
+DMLC_REGISTER_PARAMETER(SoftmaxOutputParam);
+struct SoftmaxOutputGrad {
+  const char *op_name;
+  std::vector<nnvm::NodeEntry> operator()(const nnvm::NodePtr& n,
+                                          const std::vector<nnvm::NodeEntry>& ograds) const {
+  std::vector<nnvm::NodeEntry> out_data(n->num_outputs());
+  for (uint32_t i = 0; i < out_data.size(); ++i) {
+    out_data[i] = nnvm::NodeEntry{n, i, 0};
+  }
+  std::vector<nnvm::NodeEntry> heads;
+  heads.push_back(out_data[softmaxout_enum::kOut]);
+  heads.push_back(n->inputs[softmaxout_enum::kLabel]);
+
+  nnvm::NodePtr 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_SoftmaxOutput");
+  gnode->attrs.name = n->attrs.name + "_backward";
+  std::vector<nnvm::NodeEntry> in_grad(2);
+  in_grad[0] = nnvm::NodeEntry{gnode, 0, 0};
+  in_grad[1] = nnvm::NodeEntry{gnode, 1, 0};
+  return in_grad;
+  }
+};
+
+static inline std::vector<std::string> ListArguments() {
+  return {"data", "label"};
+}
+
+
+static bool SoftmaxOutputType(const nnvm::NodeAttrs& attrs,
+                              std::vector<int> *in_type,
+                              std::vector<int> *out_type) {
+  CHECK_EQ(in_type->size(), 2U);
+  int dtype = (*in_type)[0];
+  CHECK_NE(dtype, -1) << "First input must have specified type";
+  for (size_t i = 0; i < in_type->size(); ++i) {
+    if ((*in_type)[i] == -1) {
+      (*in_type)[i] = dtype;
+    } else {
+      UNIFORM_TYPE_CHECK((*in_type)[i], dtype, ListArguments()[i]);
+    }
+  }
+  out_type->clear();
+  out_type->push_back(dtype);
+  return true;
 }
 
-// DO_BIND_DISPATCH comes from operator_common.h
-Operator *SoftmaxOutputProp::CreateOperatorEx(Context ctx, std::vector<TShape> *in_shape,
-                                     std::vector<int> *in_type) const {
-  DO_BIND_DISPATCH(CreateOp, param_, (*in_type)[0]);
+static bool SoftmaxOutputShape(const nnvm::NodeAttrs& attrs,
+                               std::vector<TShape> *in_shape,
+                               std::vector<TShape> *out_shape) {
+  using namespace mshadow;
+  const SoftmaxOutputParam& param = nnvm::get<SoftmaxOutputParam>(attrs.parsed);
+  CHECK_EQ(in_shape->size(), 2U) << "Input:[data, label]";
+  const TShape &dshape = in_shape->at(0);
+  if (dshape.ndim() == 0) return false;
+
+  // label.shape == data.shape: use probability as label
+  if (dshape != (*in_shape)[softmaxout_enum::kLabel]) {
+    if (param.multi_output) {
+      TShape lshape1 = Shape2(dshape[0], dshape.Size()/dshape[0]/dshape[1]);
+      TShape lshape2(dshape.ndim() - 1);
+      lshape2[0] = dshape[0];
+      for (index_t i = 2; i < dshape.ndim(); ++i)
+        lshape2[i-1] = dshape[i];
+      TShape lshape3 = dshape;
+      lshape3[1] = 1;
+      if (in_shape->at(softmaxout_enum::kLabel).ndim() == 0) {
+        in_shape->at(softmaxout_enum::kLabel) = lshape1;
+      } else if (in_shape->at(softmaxout_enum::kLabel) == lshape1) {
+      } else if (in_shape->at(softmaxout_enum::kLabel) == lshape2) {
+      } else if (in_shape->at(softmaxout_enum::kLabel) == lshape3) {
+      } else {
+        std::ostringstream os;
+        os << "Expecting " << lshape1 << " or " << lshape2
+           << ". But got " << in_shape->at(softmaxout_enum::kLabel);
+        throw InferShapeError(os.str(), softmaxout_enum::kLabel);
+      }
+    } else {
+      TShape label_shape(dshape.ndim() - 1);
+      for (index_t i = 0; i + 1 < dshape.ndim(); ++i)
+        label_shape[i] = dshape[i];
+      SHAPE_ASSIGN_CHECK(*in_shape, softmaxout_enum::kLabel, label_shape);
+    }
+  }
+
+  out_shape->clear();
+  out_shape->push_back(dshape);
+  return true;
 }
 
-DMLC_REGISTER_PARAMETER(SoftmaxOutputParam);
+#if MXNET_USE_MKLDNN == 1
+inline static bool SoftmaxOutputStorageType(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(), 2);
+  CHECK_EQ(out_attrs->size(), 1);
 
-MXNET_REGISTER_OP_PROPERTY(SoftmaxOutput, SoftmaxOutputProp)
+  return MKLDNNStorageType(attrs, dev_mask, true, dispatch_mode, in_attrs,
+                           out_attrs);
+}
+
+void SoftmaxOutputComputeExCPU(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(), 2U);
+  const SoftmaxOutputParam &param = nnvm::get<SoftmaxOutputParam>(attrs.parsed);
+  if (SupportMKLDNN(inputs[0]) && !ctx.is_train && SupportMKLDNNSoftmaxOutput(param)) {
+    MKLDNN_OPCHECK_INIT(false, outputs.size(), inputs, outputs);
+    MKLDNNSoftmaxOutputForward(attrs, ctx, inputs, req, outputs);
+    MKLDNN_OPCHECK_RUN(SoftmaxOutputCompute<cpu>, attrs, ctx, inputs, req, outputs);
+    return;
+  }
+  FallBackCompute(SoftmaxOutputCompute<cpu>, attrs, ctx, inputs, req, outputs);
+}
+
 
 Review comment:
   remove blank line.

----------------------------------------------------------------
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