amp multicast support casting to narrowest type

cpplint

cpplint

Author: Jerryzcn

src/operator/contrib/bounding_box-inl.h

@@ -864,13 +864,17 @@ struct box_encode {
     out_masks[a_index + 2] = valid;
     out_masks[a_index + 3] = valid;
     out_targets[a_index + 0] = valid > static_cast<DType>(0.5) ?
-        ((ref_x - a_x) / a_width - static_cast<DType>(means[0])) / static_cast<DType>(stds[0]) : static_cast<DType>(0.0);
+        ((ref_x - a_x) / a_width - static_cast<DType>(means[0])) /
+        static_cast<DType>(stds[0]) : static_cast<DType>(0.0);
     out_targets[a_index + 1] = valid > static_cast<DType>(0.5) ?
-        ((ref_y - a_y) / a_height - static_cast<DType>(means[1])) / static_cast<DType>(stds[1]) : static_cast<DType>(0.0);
+        ((ref_y - a_y) / a_height - static_cast<DType>(means[1])) /
+        static_cast<DType>(stds[1]) : static_cast<DType>(0.0);
     out_targets[a_index + 2] = valid > static_cast<DType>(0.5) ?
-        (log(ref_width / a_width) - static_cast<DType>(means[2])) / static_cast<DType>(stds[2]) : static_cast<DType>(0.0);
+        (log(ref_width / a_width) - static_cast<DType>(means[2])) /
+        static_cast<DType>(stds[2]) : static_cast<DType>(0.0);
     out_targets[a_index + 3] = valid > static_cast<DType>(0.5) ?
-        (log(ref_height / a_height) - static_cast<DType>(means[3])) / static_cast<DType>(stds[3]) : static_cast<DType>(0.0);
+        (log(ref_height / a_height) - static_cast<DType>(means[3])) /
+        static_cast<DType>(stds[3]) : static_cast<DType>(0.0);
   }
 };
 

src/operator/contrib/bounding_box.cc

@@ -216,7 +216,8 @@ NNVM_REGISTER_OP(_contrib_box_encode)
 .set_attr<nnvm::FInferType>("FInferType", ElemwiseType<6, 2>)
 .set_attr<FCompute>("FCompute<cpu>", BoxEncodeForward<cpu>)
 .set_attr<nnvm::FGradient>("FGradient", MakeZeroGradNodes)
-.add_argument("samples", "NDArray-or-Symbol", "(B, N) value +1 (positive), -1 (negative), 0 (ignore)")
+.add_argument("samples", "NDArray-or-Symbol", "(B, N) value +1 (positive), -1 (negative), "
+              "0 (ignore)")
 .add_argument("matches", "NDArray-or-Symbol", "(B, N) value range [0, M)")
 .add_argument("anchors", "NDArray-or-Symbol", "(B, N, 4) encoded in corner")
 .add_argument("refs", "NDArray-or-Symbol", "(B, N, 4) encoded in corner")

src/operator/tensor/amp_cast.h

@@ -48,10 +48,13 @@ struct AMPCastParam : public dmlc::Parameter<AMPCastParam> {
 
 struct AMPMultiCastParam : public dmlc::Parameter<AMPMultiCastParam> {
   int num_outputs;
+  bool cast_narrow;
 
   DMLC_DECLARE_PARAMETER(AMPMultiCastParam) {
     DMLC_DECLARE_FIELD(num_outputs)
     .describe("Number of input/output pairs to be casted to the widest type.");
+    DMLC_DECLARE_FIELD(cast_narrow).set_default(false)
+    .describe("Whether to cast to the narrowest type");
   }
 };
 
@@ -80,10 +83,10 @@ inline bool AMPMultiCastType(const nnvm::NodeAttrs& attrs,
   CHECK_EQ(in_attrs->size(), param.num_outputs);
   CHECK_EQ(out_attrs->size(), param.num_outputs);
   bool ret = true;
-  int widest_type = kFloat16;
+  int widest_type = param.cast_narrow ? kFloat32 : kFloat16;
   for (int i = 0; i < param.num_outputs; ++i) {
     if ((*in_attrs)[i] == kFloat32 || (*out_attrs)[i] == kFloat32) {
-      widest_type = kFloat32;
+      widest_type = param.cast_narrow ? kFloat16 : kFloat32;
     }
   }
   for (int i = 0; i < param.num_outputs; ++i) {