[Op] Rollback c++ gradient op for FusedBatchNormV3Grad.

Author: liutongxuan

tensorflow/cc/gradients/nn_grad.cc

@@ -394,121 +394,6 @@ Status FractionalMaxPoolGradHelper(const Scope& scope, const Operation& op,
 }
 REGISTER_GRADIENT_OP("FractionalMaxPool", FractionalMaxPoolGradHelper);
 
-// Templated constructor for FusedBatchNormGrad[..]::Attrs.
-template <typename T>
-T FusedBatchNormGradAttrs(float epsilon, std::string data_format,
-                          bool is_training) {
-  T result;
-  result.epsilon_ = epsilon;
-  result.data_format_ = data_format;
-  result.is_training_ = is_training;
-  return result;
-}
-
-using BatchNormGradFn =
-    std::function<Status(const Scope&, Output x, Output grad_y, Output scale,
-                         const std::vector<Output>& reserve_spaces,
-                         float epsilon, std::string data_format,
-                         bool is_training, std::vector<Output>* grad_outputs)>;
-
-Status BaseFusedBatchNormGrad(const Scope& scope, const Operation& op,
-                              const std::vector<Output>& grad_inputs,
-                              BatchNormGradFn grad_fn,
-                              std::vector<Output>* grad_outputs) {
-  if (op.num_outputs() < 5) {
-    return errors::InvalidArgument(
-        "FusedBatchNorm requires at least 5 outputs");
-  }
-  if (grad_inputs.empty()) {
-    return errors::InvalidArgument("FusedBatchNorm grad requires 1 grad input");
-  }
-  if (op.num_inputs() < 3) {
-    return errors::InvalidArgument("FusedBatchNorm has too few inputs");
-  }
-
-  Output x = op.input(0);
-  Output grad_y = grad_inputs[0];
-  Output scale = op.input(1);
-  float epsilon;
-  std::string data_format;
-  bool is_training;
-  TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "epsilon", &epsilon));
-  TF_RETURN_IF_ERROR(
-      GetNodeAttr(op.node()->attrs(), "data_format", &data_format));
-  TF_RETURN_IF_ERROR(
-      GetNodeAttr(op.node()->attrs(), "is_training", &is_training));
-
-  std::vector<Output> reserve_spaces;
-  reserve_spaces.push_back(op.output(3));
-  reserve_spaces.push_back(op.output(4));
-  if (op.num_outputs() > 5) {
-    reserve_spaces.push_back(op.output(5));
-  }
-
-  if (is_training) {
-    return grad_fn(scope, x, grad_y, scale, reserve_spaces, epsilon,
-                   data_format, is_training, grad_outputs);
-  } else {
-    if (op.num_inputs() < 5) {
-      return errors::InvalidArgument(
-          "FusedBatchNorm requires 5 inputs in eval mode");
-    }
-
-    reserve_spaces[0] = op.input(3);  // pop_mean
-    reserve_spaces[1] = op.input(4);  // pop_var
-    if (data_format == "NCHW") {
-      x = Transpose(scope, x, {0, 2, 3, 1});
-      grad_y = Transpose(scope, grad_y, {0, 2, 3, 1});
-    } else if (data_format == "NCDHW") {
-      x = Transpose(scope, x, {0, 2, 3, 4, 1});
-      grad_y = Transpose(scope, grad_y, {0, 2, 3, 4, 1});
-    }
-
-    std::string target_data_format;
-    if (data_format == "NCHW" || data_format == "NHWC") {
-      target_data_format = "NHWC";
-    } else {
-      target_data_format = "NDHWC";
-    }
-
-    TF_RETURN_IF_ERROR(grad_fn(scope, x, grad_y, scale, reserve_spaces, epsilon,
-                               target_data_format, is_training, grad_outputs));
-    if (data_format == "NCHW") {
-      (*grad_outputs)[0] = Transpose(scope, (*grad_outputs)[0], {0, 3, 1, 2});
-    } else if (data_format == "NCDHW") {
-      (*grad_outputs)[0] =
-          Transpose(scope, (*grad_outputs)[0], {0, 4, 1, 2, 3});
-    }
-    return scope.status();
-  }
-}
-
-Status FusedBatchNormV3Grad(const Scope& scope, const Operation& op,
-                            const std::vector<Output>& grad_inputs,
-                            std::vector<Output>* grad_outputs) {
-  return BaseFusedBatchNormGrad(
-      scope, op, grad_inputs,
-      [](const Scope& scope, Output x, Output grad_y, Output scale,
-         const std::vector<Output>& reserve_spaces, float epsilon,
-         std::string data_format, bool is_training,
-         std::vector<Output>* grad_outputs) {
-        FusedBatchNormGradV3 grad(
-            scope, grad_y, x, scale, reserve_spaces[0], reserve_spaces[1],
-            reserve_spaces[2],
-            FusedBatchNormGradAttrs<FusedBatchNormGradV3::Attrs>(
-                epsilon, data_format, is_training));
-        grad_outputs->push_back(grad.x_backprop);
-        grad_outputs->push_back(grad.scale_backprop);
-        grad_outputs->push_back(grad.offset_backprop);
-        grad_outputs->push_back(NoGradient());
-        grad_outputs->push_back(NoGradient());
-        return scope.status();
-      },
-      grad_outputs);
-}
-
-REGISTER_GRADIENT_OP("FusedBatchNormV3", FusedBatchNormV3Grad);
-
 Status Conv2DBackpropInputGrad(const Scope& scope, const Operation& op,
                                const std::vector<Output>& grad_inputs,
                                std::vector<Output>* grad_outputs) {

tensorflow/cc/gradients/nn_grad_test.cc

@@ -35,7 +35,6 @@ using ops::Conv2DBackpropInput;
 using ops::Elu;
 using ops::FractionalAvgPool;
 using ops::FractionalMaxPool;
-using ops::FusedBatchNormV3;
 using ops::L2Loss;
 using ops::LogSoftmax;
 using ops::LRN;
@@ -344,40 +343,6 @@ TEST_F(NNGradTest, FractionalMaxPoolGradHelper) {
   RunTest(x, x_init_value, y.output, y_shape);
 }
 
-class FusedBatchNormGradTest : public NNGradTest,
-                               public ::testing::WithParamInterface<
-                                   std::tuple<bool, bool, TensorShape>> {};
-
-TEST_P(FusedBatchNormGradTest, FusedBatchNormV3Grad) {
-  FusedBatchNormV3::Attrs attrs;
-  attrs.is_training_ = std::get<0>(GetParam());
-  bool channel_first = std::get<1>(GetParam());
-  TensorShape shape = std::get<2>(GetParam());
-  int channel_dim = (channel_first) ? 1 : shape.dims() - 1;
-  TensorShape scale_shape({shape.dim_size(channel_dim)});
-  auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape));
-  auto scale = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(scale_shape));
-  auto offset = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(scale_shape));
-  auto mean = ops::ZerosLike(scope_, scale);
-  auto var = ops::OnesLike(scope_, scale);
-
-  if (!channel_first) {
-    attrs.data_format_ = (shape.dims() == 5) ? "NDHWC" : "NHWC";
-  } else {
-    attrs.data_format_ = (shape.dims() == 5) ? "NCDHW" : "NCHW";
-  }
-
-  auto y = FusedBatchNormV3(scope_, x, scale, offset, mean, var, attrs);
-  RunTest({x, scale, offset}, {shape, scale_shape, scale_shape}, {y.y},
-          {shape});
-}
-
-INSTANTIATE_TEST_SUITE_P(
-    FusedBatchNormGrad, FusedBatchNormGradTest,
-    ::testing::Combine(::testing::Bool(), ::testing::Bool(),
-                       ::testing::Values(TensorShape({2, 3, 4, 5}),
-                                         TensorShape({2, 3, 2, 2, 2}))));
-
 TEST_F(NNGradTest, Conv2DBackpropInputGrad) {
   TensorShape shape({1, 2, 2, 1});
   TensorShape filter_shape({1, 1, 1, 1});