Fix operators lying about their number of inputs (#17049)

* Add a check for number of inputs * Fix num inputs for backward_Deconvolution * Fix number of inputs to backward ROIAlign * Fix number of inputs to backward_SoftmaxOutput * Fix more operators lying about their number of inputs * Fix input number of backward NMS * Fixes * Fix dropout, RNN and upsampling backward number of inputs * Fix LeakyRelu number of inputs * Actually fix LeakyRelu * Fix pooling and concat * Fix Concat (attempt 2) * Fix from review * Incorporate Dick's changes * Add guard to MakeNonlossGradNode * Fix * Fix backward of SoftmaxActivation * Fix backward of np_prod and norm
apache · Jan 15, 2020 · bd67723 · bd67723
1 parent be9e17e
commit bd67723
Show file tree

Hide file tree

Showing 21 changed files with 105 additions and 30 deletions.
diff --git a/src/operator/contrib/bilinear_resize-inl.h b/src/operator/contrib/bilinear_resize-inl.h
@@ -328,15 +328,6 @@ inline uint16_t BilinearSampleOpNumInputs(const NodeAttrs& attrs) {
   }
 }
 
-inline uint16_t BilinearSampleOpNumBackwardInputs(const NodeAttrs& attrs) {
-  auto& param = nnvm::get<BilinearSampleParam>(attrs.parsed);
-  if (param.mode == bilinear_resize::like) {
-    return 3;
-  } else {
-    return 1;
-  }
-}
-
 inline uint16_t BilinearSampleOpNumBackwardOutputs(const NodeAttrs& attrs) {
   auto& param = nnvm::get<BilinearSampleParam>(attrs.parsed);
   if (param.mode == bilinear_resize::like) {

diff --git a/src/operator/contrib/bilinear_resize.cc b/src/operator/contrib/bilinear_resize.cc
@@ -232,7 +232,7 @@ for more details.
 
 NNVM_REGISTER_OP(_backward_contrib_BilinearResize2D)
 .set_attr_parser(ParamParser<BilinearSampleParam>)
-.set_num_inputs(BilinearSampleOpNumBackwardInputs)
+.set_num_inputs(1)
 .set_num_outputs(BilinearSampleOpNumBackwardOutputs)
 .set_attr<nnvm::TIsBackward>("TIsBackward", true)
 .set_attr<FCompute>("FCompute<cpu>", BilinearSampleOpBackward<cpu>);

diff --git a/src/operator/contrib/bounding_box.cc b/src/operator/contrib/bounding_box.cc
@@ -110,7 +110,7 @@ Examples::
 .add_arguments(BoxNMSParam::__FIELDS__());
 
 NNVM_REGISTER_OP(_backward_contrib_box_nms)
-.set_num_inputs(3)
+.set_num_inputs(4)
 .set_num_outputs(1)
 .set_attr_parser(ParamParser<BoxNMSParam>)
 .set_attr<nnvm::TIsBackward>("TIsBackward", true)

diff --git a/src/operator/contrib/roi_align.cc b/src/operator/contrib/roi_align.cc
@@ -621,6 +621,7 @@ He, Kaiming, et al. "Mask R-CNN." ICCV, 2017
 
 
 NNVM_REGISTER_OP(_backward_ROIAlign)
+.set_num_inputs(2)
 .set_num_outputs(2)
 .set_attr<nnvm::TIsBackward>("TIsBackward", true)
 .set_attr_parser(ParamParser<ROIAlignParam>)

diff --git a/src/operator/custom/custom.cc b/src/operator/custom/custom.cc
@@ -586,7 +586,7 @@ Please check the tutorial here: https://mxnet.incubator.apache.org/api/faq/new_o
 NNVM_REGISTER_OP(_backward_Custom)
 .set_num_inputs([](const NodeAttrs& attrs){
     const CustomParam& params = nnvm::get<CustomParam>(attrs.parsed);
-    return params.bwd_idx.size();
+    return params.bwd_idx.size() + params.num_auxs;
   })
 .set_num_outputs([](const NodeAttrs& attrs){
     const CustomParam& params = nnvm::get<CustomParam>(attrs.parsed);

diff --git a/src/operator/image/image_random.cc b/src/operator/image/image_random.cc
@@ -185,7 +185,7 @@ NNVM_REGISTER_OP(_image_normalize)
 
 NNVM_REGISTER_OP(_backward_image_normalize)
 .set_attr_parser(ParamParser<NormalizeParam>)
-.set_num_inputs(1)
+.set_num_inputs(2)
 .set_num_outputs(1)
 .set_attr<nnvm::TIsBackward>("TIsBackward", true)
 .set_attr<FCompute>("FCompute<cpu>", NormalizeOpBackward<cpu>);

diff --git a/src/operator/leaky_relu.cc b/src/operator/leaky_relu.cc
@@ -206,6 +206,19 @@ The following modified ReLU Activation functions are supported:
     });
 
 NNVM_REGISTER_OP(_backward_LeakyReLU)
+.set_num_inputs([](const NodeAttrs& attrs) {
+  const LeakyReLUParam& param = nnvm::get<LeakyReLUParam>(attrs.parsed);
+  if (param.act_type == leakyrelu::kPReLU) {
+    // forward has 2 inputs and 1 output
+    return 2 + 2 * 1;
+  } else if (param.act_type == leakyrelu::kRReLU) {
+    // forward has 1 input and 2 outputs
+    return 1 + 2 * 2;
+  } else {
+    // forward has 1 input and 1 output
+    return 1 + 2 * 1;
+  }
+})
 .set_num_outputs([](const NodeAttrs& attrs) {
   const LeakyReLUParam& param = nnvm::get<LeakyReLUParam>(attrs.parsed);
   return param.act_type == leakyrelu::kPReLU ? 2 : 1;

diff --git a/src/operator/nn/batch_norm.cc b/src/operator/nn/batch_norm.cc
@@ -593,6 +593,7 @@ then set ``gamma`` to 1 and its gradient to 0.
   });
 
 NNVM_REGISTER_OP(_backward_BatchNorm)
+.set_num_inputs(8)
 .set_num_outputs(3)
 .set_attr<nnvm::TIsBackward>("TIsBackward", true)
 .set_attr<FInferStorageType>("FInferStorageType", BatchNormStorageType)

diff --git a/src/operator/nn/concat.cc b/src/operator/nn/concat.cc
@@ -396,6 +396,14 @@ CONCAT_FORWARD_ATTRS
 .add_arguments(ConcatParam::__FIELDS__());
 
 NNVM_REGISTER_OP(_backward_Concat)
+.set_num_inputs([](const NodeAttrs& attrs) {
+#if MXNET_USE_MKLDNN == 1
+  const ConcatParam& params = nnvm::get<ConcatParam>(attrs.parsed);
+  return 1 + params.num_args;
+#else
+  return 1;
+#endif
+})
 .set_num_outputs([](const NodeAttrs& attrs) {
   const ConcatParam& params = nnvm::get<ConcatParam>(attrs.parsed);
   return params.num_args;

diff --git a/src/operator/nn/convolution.cc b/src/operator/nn/convolution.cc
@@ -510,6 +510,10 @@ There are other options to tune the performance.
 .add_arguments(ConvolutionParam::__FIELDS__());
 
 NNVM_REGISTER_OP(_backward_Convolution)
+.set_num_inputs([](const NodeAttrs& attrs) {
+  const ConvolutionParam& params = nnvm::get<ConvolutionParam>(attrs.parsed);
+  return params.no_bias ? 3 : 4;
+})
 .set_num_outputs([](const NodeAttrs& attrs) {
   const ConvolutionParam& params = nnvm::get<ConvolutionParam>(attrs.parsed);
   return params.no_bias ? 2 : 3;

diff --git a/src/operator/nn/ctc_loss.cc b/src/operator/nn/ctc_loss.cc
@@ -130,7 +130,7 @@ information on the definition and the algorithm.
 
 NNVM_REGISTER_OP(_backward_ctc_loss)
 .set_attr_parser(ParamParser<CTCLossOpParam>)
-.set_num_inputs(1)
+.set_num_inputs(4)
 .set_num_outputs(CTCLossOpNumInputs)
 .set_attr<nnvm::TIsBackward>("TIsBackward", true)
 .set_attr<FCompute>("FCompute<cpu>", CTCLossOpBackward<cpu>);

diff --git a/src/operator/nn/deconvolution.cc b/src/operator/nn/deconvolution.cc
@@ -445,6 +445,10 @@ NNVM_REGISTER_OP(Deconvolution)
 .add_arguments(DeconvolutionParam::__FIELDS__());
 
 NNVM_REGISTER_OP(_backward_Deconvolution)
+.set_num_inputs([](const NodeAttrs& attrs) {
+  const DeconvolutionParam& params = nnvm::get<DeconvolutionParam>(attrs.parsed);
+  return params.no_bias ? 3 : 4;
+})
 .set_num_outputs([](const NodeAttrs& attrs) {
   const DeconvolutionParam& params = nnvm::get<DeconvolutionParam>(attrs.parsed);
   return params.no_bias ? 2 : 3;

diff --git a/src/operator/nn/lrn.cc b/src/operator/nn/lrn.cc
@@ -184,6 +184,7 @@ number of kernels in the layer.
 .add_arguments(LRNParam::__FIELDS__());
 
 NNVM_REGISTER_OP(_backward_LRN)
+.set_num_inputs(3)
 .set_num_outputs(1)
 .set_attr_parser(ParamParser<LRNParam>)
 #if MXNET_USE_MKLDNN == 1

diff --git a/src/operator/nn/pooling.cc b/src/operator/nn/pooling.cc
@@ -453,6 +453,11 @@ For each window ``X``, the mathematical expression for Lp pooling is:
 .add_arguments(PoolingParam::__FIELDS__());
 
 NNVM_REGISTER_OP(_backward_Pooling)
+.set_num_inputs([](const NodeAttrs& attrs) {
+  const PoolingParam &param = nnvm::get<PoolingParam>(attrs.parsed);
+  // 1 input to fwd op and 2 * outputs from fwd op (fwd outputs and gradient inputs)
+  return 1 + 2 * GetNumOutputs(param);
+})
 .set_num_outputs(1)
 .set_attr<nnvm::TIsBackward>("TIsBackward", true)
 .set_attr<nnvm::FInplaceOption>(

diff --git a/src/operator/nn/softmax_activation.cc b/src/operator/nn/softmax_activation.cc
@@ -67,6 +67,7 @@ Example::
 .add_arguments(SoftmaxActivationParam::__FIELDS__());
 
 NNVM_REGISTER_OP(_backward_SoftmaxActivation)
+.set_num_inputs(2)
 .set_num_outputs(1)
 .set_attr<nnvm::TIsBackward>("TIsBackward", true)
 .set_attr<nnvm::FInplaceOption>("FInplaceOption", [](const NodeAttrs& attrs){

diff --git a/src/operator/nn/upsampling.cc b/src/operator/nn/upsampling.cc
@@ -211,6 +211,14 @@ Example::
     });
 
 NNVM_REGISTER_OP(_backward_UpSampling)
+.set_num_inputs([](const NodeAttrs& attrs) {
+  const UpSamplingParam& param_ = nnvm::get<UpSamplingParam>(attrs.parsed);
+  if (param_.sample_type != up_enum::kNearest) {
+    return 3;
+  } else {
+    return 1;
+  }
+})
 .set_num_outputs([](const NodeAttrs& attrs) {
   const UpSamplingParam& params = nnvm::get<UpSamplingParam>(attrs.parsed);
   return params.sample_type == up_enum::kNearest ? params.num_args : 2;

diff --git a/src/operator/numpy/np_broadcast_reduce_op_value.cc b/src/operator/numpy/np_broadcast_reduce_op_value.cc
@@ -238,7 +238,7 @@ NNVM_REGISTER_OP(_np_prod)
 .set_attr<nnvm::FGradient>("FGradient", ReduceGrad{"_backward_np_prod"});
 
 NNVM_REGISTER_OP(_backward_np_prod)
-.set_num_inputs(1)
+.set_num_inputs(3)
 .set_num_outputs(1)
 .set_attr_parser(ParamParser<NumpyReduceAxesParam>)
 .set_attr<nnvm::TIsBackward>("TIsBackward", true)

diff --git a/src/operator/operator_common.h b/src/operator/operator_common.h
@@ -359,12 +359,36 @@ inline bool dispatch_fallback(StorageTypeVector* stypes, DispatchMode* dispatch)
   return true;
 }
 
+inline std::vector<nnvm::NodeEntry>CreateNodeEntries(
+  nnvm::NodePtr pNode,
+  const std::vector<nnvm::NodeEntry>* pOgrads = nullptr,
+  const std::vector<nnvm::NodeEntry>* pInputs = nullptr) {
+  if (pOgrads)
+    pNode->inputs.insert(pNode->inputs.end(), pOgrads->begin(), pOgrads->end());
+
+  if (pInputs)
+    pNode->inputs.insert(pNode->inputs.end(), pInputs->begin(), pInputs->end());
+
+  if (!pNode->is_variable()) {
+    CHECK_EQ(pNode->num_inputs(), pNode->inputs.size())
+      << "Number of inputs to operator " << pNode->op()->name << " (" << pNode->num_inputs()
+      << ") does not match the actual number of inputs provided to operator "
+      << pNode->attrs.name << " (" << pNode->inputs.size() << ").";
+  }
+
+  std::vector<nnvm::NodeEntry> ret;
+  for (uint32_t i = 0; i < pNode->num_outputs(); ++i)
+    ret.emplace_back(nnvm::NodeEntry{pNode, i, 0});
+
+  return ret;
+}
+
 // make a new node with operator op_name. Inputs are not filled.
 inline nnvm::NodePtr MakeNode(
     const char* op_name, const std::string& name,
-    std::vector<nnvm::NodeEntry> const * inputs,
-    std::unordered_map<std::string, std::string> const * dict,
-    nnvm::NodePtr const * fwd_node) {
+    std::vector<nnvm::NodeEntry> const * inputs = nullptr,
+    std::unordered_map<std::string, std::string> const * dict = nullptr,
+    nnvm::NodePtr const * fwd_node = nullptr) {
   auto p = nnvm::Node::Create();
   p->attrs.op = nnvm::Op::Get(op_name);
   p->attrs.name = name;
@@ -376,6 +400,12 @@ inline nnvm::NodePtr MakeNode(
   if (p->op()->attr_parser != nullptr) {
     p->op()->attr_parser(&(p->attrs));
   }
+  if (inputs != nullptr) {
+    CHECK_EQ(p->num_inputs(), p->inputs.size())
+      << "Number of inputs to operator " << op_name << " (" << p->num_inputs()
+      << ") does not match the actual number of inputs provided to operator "
+      << name << " (" << p->inputs.size() << ").";
+  }
   return p;
 }
 
@@ -395,11 +425,8 @@ inline std::vector<nnvm::NodeEntry> MakeGradNode(
     const std::unordered_map<std::string, std::string>& dict) {
   auto p = MakeNode(op_name, n->attrs.name + "_backward",
                     &inputs, &dict, &n);
-  std::vector<nnvm::NodeEntry> ret;
-  for (uint32_t i = 0; i < p->num_outputs(); ++i) {
-    ret.emplace_back(p, i, 0);
-  }
-  return ret;
+
+  return CreateNodeEntries(p);
 }
 
 // quick helper to make gradient nodes that simply pass back zero. could be used in output ops.
@@ -446,13 +473,8 @@ inline std::vector<nnvm::NodeEntry> MakeNonlossGradNode(
     return MakeZeroGradNodes(n, ograds);
   auto p = MakeNode(op_name, n->attrs.name + "_backward",
                     nullptr, &dict, &n);
-  p->inputs.insert(p->inputs.end(), ograds.begin(), ograds.end());
-  p->inputs.insert(p->inputs.end(), inputs.begin(), inputs.end());
-  std::vector<nnvm::NodeEntry> ret;
-  for (uint32_t i = 0; i < p->num_outputs(); ++i) {
-    ret.emplace_back(p, i, 0);
-  }
-  return ret;
+
+  return CreateNodeEntries(p, &ograds, &inputs);
 }
 
 /*! \brief Parse keyword arguments as PType arguments and save to parsed */

diff --git a/src/operator/rnn.cc b/src/operator/rnn.cc
@@ -406,6 +406,20 @@ The definition of GRU here is slightly different from paper but compatible with
 .add_arguments(RNNParam::__FIELDS__());
 
 NNVM_REGISTER_OP(_backward_RNN)
+.set_num_inputs([](const NodeAttrs& attrs) {
+    const RNNParam& params = nnvm::get<RNNParam>(attrs.parsed);
+    int ret = 5;
+    if (params.state_outputs) {
+      ret += 2;
+    }
+    if (params.mode == rnn_enum::kLstm) {
+      ++ret;
+      if (params.state_outputs) {
+      ret += 2;
+      }
+    }
+    return ret;
+})
 .set_num_outputs([](const NodeAttrs& attrs) {
   const RNNParam& params = nnvm::get<RNNParam>(attrs.parsed);
   return GetNumInputArguments(params);

diff --git a/src/operator/softmax_output.cc b/src/operator/softmax_output.cc
@@ -258,6 +258,7 @@ NNVM_REGISTER_OP(SoftmaxOutput)
 NNVM_REGISTER_OP(SoftmaxOutput).add_alias("Softmax");
 
 NNVM_REGISTER_OP(_backward_SoftmaxOutput)
+.set_num_inputs(2)
 .set_num_outputs(2)
 .set_attr<nnvm::TIsBackward>("TIsBackward", true)
 .set_attr<nnvm::FInplaceOption>("FInplaceOption", [](const NodeAttrs& attrs){

diff --git a/src/operator/tensor/broadcast_reduce_norm_value.cc b/src/operator/tensor/broadcast_reduce_norm_value.cc
@@ -105,6 +105,7 @@ Examples::
 .add_arguments(NormParam::__FIELDS__());
 
 NNVM_REGISTER_OP(_backward_norm)
+.set_num_inputs(3)
 .set_num_outputs(1)
 .set_attr_parser(ParamParser<NormParam>)
 .set_attr<nnvm::TIsBackward>("TIsBackward", true)