wehub-resource-sync
611 lines
25 KiB
C++
611 lines
25 KiB
C++
/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
|
|
|
|
Licensed 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.
|
|
==============================================================================*/
|
|
|
|
#include <cstdint>
|
|
#include <functional>
|
|
#include <string>
|
|
#include <vector>
|
|
|
|
#include "absl/status/status.h"
|
|
#include "absl/strings/string_view.h"
|
|
#include "tensorflow/cc/framework/grad_op_registry.h"
|
|
#include "tensorflow/cc/framework/gradients.h"
|
|
#include "tensorflow/cc/ops/nn_ops.h"
|
|
#include "tensorflow/cc/ops/nn_ops_internal.h"
|
|
#include "tensorflow/cc/ops/standard_ops.h"
|
|
#include "tensorflow/core/framework/types.pb.h"
|
|
|
|
namespace tensorflow {
|
|
namespace ops {
|
|
namespace {
|
|
|
|
absl::Status SoftmaxGrad(const Scope& scope, const Operation& op,
|
|
const std::vector<Output>& grad_inputs,
|
|
std::vector<Output>* grad_outputs) {
|
|
// Softmax gradient function.
|
|
// p = softmax(x) maps from [batch, n] to [batch, m]
|
|
// dp/dx = [dp0/dx0 ... dp0/dxn-1 ]
|
|
// [ ... ... ]
|
|
// [dpm-1/dx0 ... dpm-1/dxn-1]
|
|
// dL/dx = dp/dx * dL/dy
|
|
//
|
|
// Using alternative formula:
|
|
// dL/dx = dL/dy * y - sum(dL/dy * y) * y
|
|
// = (dL/dy - sum(dL/dy * y)) * y
|
|
auto y = op.output(0);
|
|
auto dyy = Mul(scope, grad_inputs[0], y);
|
|
auto sum = Sum(scope, dyy, /*axis=*/-1, Sum::KeepDims(true));
|
|
auto sub = Sub(scope, grad_inputs[0], sum);
|
|
auto dx = Mul(scope, sub, y);
|
|
grad_outputs->push_back(dx);
|
|
return scope.status();
|
|
}
|
|
REGISTER_GRADIENT_OP("Softmax", SoftmaxGrad);
|
|
|
|
bool IsZero(const Scope& scope, const Output& grad) {
|
|
std::string op_type_name = grad.op().node()->type_string();
|
|
if (op_type_name == "ZerosLike" || op_type_name == "Zeros") {
|
|
return true;
|
|
}
|
|
// The Operation we were provided is not named something obvious so
|
|
// we need to actually look at its contents.
|
|
// The original python code did this by calling a utility function called
|
|
// tensor_util.constant_value.
|
|
// There is no C++ equivalent to tensor_util.constant_value so we do nothing
|
|
// for the moment.
|
|
return false;
|
|
}
|
|
|
|
// Multiply after broadcasting vec to match dimensions of mat.
|
|
// Args:
|
|
// vec: A 1-D tensor of dimension [D0]
|
|
// mat: A 2-D tensor of dimension [D0, D1]
|
|
//
|
|
// Returns:
|
|
// A tensor of dimension [D0, D1], the result for vec * mat.
|
|
Output BroadcastMul(const Scope& scope, const Output& vec, const Output& mat) {
|
|
auto reshaped = ExpandDims(scope, vec, -1);
|
|
return Multiply(scope, reshaped, mat);
|
|
}
|
|
|
|
absl::Status SoftmaxCrossEntropyWithLogitsGrad(
|
|
const Scope& scope, const Operation& op,
|
|
const std::vector<Output>& grad_inputs, std::vector<Output>* grad_outputs) {
|
|
// Softmax gradient with cross entropy logits function.
|
|
// We multiply the backprop for cost with the gradients - op.output[1].
|
|
// There is no gradient for labels.
|
|
|
|
// The outputs of the network are at input index 0.
|
|
auto logits = op.input(0);
|
|
// The "truth" labels are at index 1.
|
|
auto softmax_grad = op.output(1);
|
|
|
|
// The loss is the output at index 0, and backprop is the output at index 1.
|
|
auto grad_loss = grad_inputs[0];
|
|
auto grad_grad = grad_inputs[1];
|
|
|
|
auto grad = BroadcastMul(scope, grad_loss, softmax_grad);
|
|
if (!IsZero(scope, grad_grad)) {
|
|
std::vector<int> axis;
|
|
auto logits_softmax = Softmax(scope, logits);
|
|
|
|
auto grad_grad_expand = ExpandDims(scope, grad_grad, 1);
|
|
auto logits_softmax_expand = ExpandDims(scope, logits_softmax, 2);
|
|
auto matmul_result =
|
|
BatchMatMul(scope, grad_grad_expand, logits_softmax_expand);
|
|
axis.push_back(1);
|
|
auto squeeze_result = Squeeze(scope, matmul_result, Squeeze::Axis(axis));
|
|
auto subtraction_result = Subtract(scope, grad_grad, squeeze_result);
|
|
auto multiply_result = Multiply(scope, subtraction_result, logits_softmax);
|
|
grad = Add(scope, grad, multiply_result);
|
|
}
|
|
auto minus_log_softmax = Multiply(scope, LogSoftmax(scope, logits), -1.0f);
|
|
grad_outputs->push_back(grad);
|
|
grad_outputs->push_back(BroadcastMul(scope, grad_loss, minus_log_softmax));
|
|
return scope.status();
|
|
}
|
|
REGISTER_GRADIENT_OP("SoftmaxCrossEntropyWithLogits",
|
|
SoftmaxCrossEntropyWithLogitsGrad);
|
|
|
|
absl::Status LogSoftmaxGrad(const Scope& scope, const Operation& op,
|
|
const std::vector<Output>& grad_inputs,
|
|
std::vector<Output>* grad_outputs) {
|
|
auto softmax = Exp(scope, op.output(0));
|
|
auto sum = Sum(scope, grad_inputs[0], {1}, Sum::KeepDims(true));
|
|
auto mul = Mul(scope, sum, softmax);
|
|
auto dx = Sub(scope, grad_inputs[0], mul);
|
|
grad_outputs->push_back(dx);
|
|
return scope.status();
|
|
}
|
|
REGISTER_GRADIENT_OP("LogSoftmax", LogSoftmaxGrad);
|
|
|
|
absl::Status ReluGradHelper(const Scope& scope, const Operation& op,
|
|
const std::vector<Output>& grad_inputs,
|
|
std::vector<Output>* grad_outputs) {
|
|
auto dx = internal::ReluGrad(scope, grad_inputs[0], op.input(0));
|
|
grad_outputs->push_back(dx);
|
|
return scope.status();
|
|
}
|
|
REGISTER_GRADIENT_OP("Relu", ReluGradHelper);
|
|
|
|
absl::Status Relu6GradHelper(const Scope& scope, const Operation& op,
|
|
const std::vector<Output>& grad_inputs,
|
|
std::vector<Output>* grad_outputs) {
|
|
auto dx = internal::Relu6Grad(scope, grad_inputs[0], op.input(0));
|
|
grad_outputs->push_back(dx);
|
|
return scope.status();
|
|
}
|
|
REGISTER_GRADIENT_OP("Relu6", Relu6GradHelper);
|
|
|
|
absl::Status LeakyReluGradHelper(const Scope& scope, const Operation& op,
|
|
const std::vector<Output>& grad_inputs,
|
|
std::vector<Output>* grad_outputs) {
|
|
float alpha;
|
|
TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "alpha", &alpha));
|
|
internal::LeakyReluGrad::Attrs attrs;
|
|
auto dx = internal::LeakyReluGrad(scope, grad_inputs[0], op.input(0),
|
|
attrs.Alpha(alpha));
|
|
grad_outputs->push_back(dx);
|
|
return scope.status();
|
|
}
|
|
REGISTER_GRADIENT_OP("LeakyRelu", LeakyReluGradHelper);
|
|
|
|
absl::Status LeakyReluGradGradHelper(const Scope& scope, const Operation& op,
|
|
const std::vector<Output>& grad_inputs,
|
|
std::vector<Output>* grad_outputs) {
|
|
float alpha;
|
|
TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "alpha", &alpha));
|
|
internal::LeakyReluGrad::Attrs attrs;
|
|
auto dx = internal::LeakyReluGrad(scope, grad_inputs[0], op.input(1),
|
|
attrs.Alpha(alpha));
|
|
grad_outputs->push_back(dx);
|
|
grad_outputs->push_back(NoGradient());
|
|
return scope.status();
|
|
}
|
|
REGISTER_GRADIENT_OP("LeakyReluGrad", LeakyReluGradGradHelper);
|
|
|
|
absl::Status EluGradHelper(const Scope& scope, const Operation& op,
|
|
const std::vector<Output>& grad_inputs,
|
|
std::vector<Output>* grad_outputs) {
|
|
auto dx = internal::EluGrad(scope, grad_inputs[0], op.output(0));
|
|
grad_outputs->push_back(dx);
|
|
return scope.status();
|
|
}
|
|
REGISTER_GRADIENT_OP("Elu", EluGradHelper);
|
|
|
|
absl::Status SeluGradHelper(const Scope& scope, const Operation& op,
|
|
const std::vector<Output>& grad_inputs,
|
|
std::vector<Output>* grad_outputs) {
|
|
auto dx = internal::SeluGrad(scope, grad_inputs[0], op.output(0));
|
|
grad_outputs->push_back(dx);
|
|
return scope.status();
|
|
}
|
|
REGISTER_GRADIENT_OP("Selu", SeluGradHelper);
|
|
|
|
absl::Status L2LossGrad(const Scope& scope, const Operation& op,
|
|
const std::vector<Output>& grad_inputs,
|
|
std::vector<Output>* grad_outputs) {
|
|
grad_outputs->push_back(Mul(scope, op.input(0), grad_inputs[0]));
|
|
return scope.status();
|
|
}
|
|
REGISTER_GRADIENT_OP("L2Loss", L2LossGrad);
|
|
|
|
absl::Status BiasAddGradHelper(const Scope& scope, const Operation& op,
|
|
const std::vector<Output>& grad_inputs,
|
|
std::vector<Output>* grad_outputs) {
|
|
std::string data_format;
|
|
TF_RETURN_IF_ERROR(
|
|
GetNodeAttr(op.output(0).node()->attrs(), "data_format", &data_format));
|
|
auto dx_1 =
|
|
BiasAddGrad(scope, grad_inputs[0], BiasAddGrad::DataFormat(data_format));
|
|
grad_outputs->push_back(Identity(scope, grad_inputs[0]));
|
|
grad_outputs->push_back(dx_1);
|
|
return scope.status();
|
|
}
|
|
REGISTER_GRADIENT_OP("BiasAdd", BiasAddGradHelper);
|
|
|
|
absl::Status Conv2DGrad(const Scope& scope, const Operation& op,
|
|
const std::vector<Output>& grad_inputs,
|
|
std::vector<Output>* grad_outputs) {
|
|
std::string data_format;
|
|
std::string padding;
|
|
std::vector<int32_t> strides;
|
|
bool use_cudnn_on_gpu;
|
|
auto attrs = op.output(0).node()->attrs();
|
|
TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "data_format", &data_format));
|
|
TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "padding", &padding));
|
|
TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "strides", &strides));
|
|
TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "use_cudnn_on_gpu", &use_cudnn_on_gpu));
|
|
auto dx_1 = Conv2DBackpropInput(scope, Shape(scope, op.input(0)), op.input(1),
|
|
grad_inputs[0], strides, padding,
|
|
Conv2DBackpropInput::DataFormat(data_format)
|
|
.UseCudnnOnGpu(use_cudnn_on_gpu));
|
|
grad_outputs->push_back(dx_1);
|
|
auto dx_2 =
|
|
Conv2DBackpropFilter(scope, op.input(0), Shape(scope, op.input(1)),
|
|
grad_inputs[0], strides, padding,
|
|
Conv2DBackpropFilter::DataFormat(data_format)
|
|
.UseCudnnOnGpu(use_cudnn_on_gpu));
|
|
grad_outputs->push_back(dx_2);
|
|
return scope.status();
|
|
}
|
|
REGISTER_GRADIENT_OP("Conv2D", Conv2DGrad);
|
|
|
|
absl::Status MaxPoolGradHelper(const Scope& scope, const Operation& op,
|
|
const std::vector<Output>& grad_inputs,
|
|
std::vector<Output>* grad_outputs) {
|
|
std::string data_format;
|
|
std::string padding;
|
|
std::vector<int32_t> strides;
|
|
std::vector<int32_t> ksize;
|
|
auto attrs = op.output(0).node()->attrs();
|
|
TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "data_format", &data_format));
|
|
TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "ksize", &ksize));
|
|
TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "padding", &padding));
|
|
TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "strides", &strides));
|
|
auto dx = internal::MaxPoolGrad(
|
|
scope, op.input(0), op.output(0), grad_inputs[0], ksize, strides, padding,
|
|
internal::MaxPoolGrad::DataFormat(data_format));
|
|
grad_outputs->push_back(dx);
|
|
return scope.status();
|
|
}
|
|
REGISTER_GRADIENT_OP("MaxPool", MaxPoolGradHelper);
|
|
|
|
absl::Status MaxPoolGradV2Helper(const Scope& scope, const Operation& op,
|
|
const std::vector<Output>& grad_inputs,
|
|
std::vector<Output>* grad_outputs) {
|
|
std::string data_format;
|
|
std::string padding;
|
|
auto attrs = op.output(0).node()->attrs();
|
|
TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "data_format", &data_format));
|
|
TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "padding", &padding));
|
|
auto dx = MaxPoolGradV2(scope, op.input(0), op.output(0), grad_inputs[0],
|
|
op.input(1), op.input(2), padding,
|
|
MaxPoolGradV2::DataFormat(data_format));
|
|
grad_outputs->push_back(dx);
|
|
grad_outputs->push_back(NoGradient());
|
|
grad_outputs->push_back(NoGradient());
|
|
return scope.status();
|
|
}
|
|
REGISTER_GRADIENT_OP("MaxPoolV2", MaxPoolGradV2Helper);
|
|
|
|
absl::Status MaxPool3DGradHelper(const Scope& scope, const Operation& op,
|
|
const std::vector<Output>& grad_inputs,
|
|
std::vector<Output>* grad_outputs) {
|
|
std::vector<int32_t> ksize;
|
|
std::vector<int32_t> strides;
|
|
std::string padding;
|
|
std::string data_format;
|
|
auto attrs = op.output(0).node()->attrs();
|
|
TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "ksize", &ksize));
|
|
TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "strides", &strides));
|
|
TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "padding", &padding));
|
|
TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "data_format", &data_format));
|
|
MaxPool3DGrad::Attrs grad_attrs;
|
|
auto dx =
|
|
MaxPool3DGrad(scope, op.input(0), op.output(0), grad_inputs[0], ksize,
|
|
strides, padding, grad_attrs.DataFormat(data_format));
|
|
grad_outputs->push_back(dx);
|
|
return scope.status();
|
|
}
|
|
REGISTER_GRADIENT_OP("MaxPool3D", MaxPool3DGradHelper);
|
|
|
|
absl::Status AvgPoolGradHelper(const Scope& scope, const Operation& op,
|
|
const std::vector<Output>& grad_inputs,
|
|
std::vector<Output>* grad_outputs) {
|
|
std::vector<int32_t> ksize;
|
|
std::vector<int32_t> strides;
|
|
std::string padding;
|
|
std::string data_format;
|
|
auto attrs = op.output(0).node()->attrs();
|
|
TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "ksize", &ksize));
|
|
TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "strides", &strides));
|
|
TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "padding", &padding));
|
|
TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "data_format", &data_format));
|
|
internal::AvgPoolGrad::Attrs grad_attrs;
|
|
auto dx = internal::AvgPoolGrad(scope, Shape(scope, op.input(0)),
|
|
grad_inputs[0], ksize, strides, padding,
|
|
grad_attrs.DataFormat(data_format));
|
|
grad_outputs->push_back(dx);
|
|
return scope.status();
|
|
}
|
|
REGISTER_GRADIENT_OP("AvgPool", AvgPoolGradHelper);
|
|
|
|
absl::Status AvgPool3DGradHelper(const Scope& scope, const Operation& op,
|
|
const std::vector<Output>& grad_inputs,
|
|
std::vector<Output>* grad_outputs) {
|
|
std::vector<int32_t> ksize;
|
|
std::vector<int32_t> strides;
|
|
std::string padding;
|
|
std::string data_format;
|
|
auto attrs = op.output(0).node()->attrs();
|
|
TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "ksize", &ksize));
|
|
TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "strides", &strides));
|
|
TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "padding", &padding));
|
|
TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "data_format", &data_format));
|
|
AvgPool3DGrad::Attrs grad_attrs;
|
|
auto dx =
|
|
AvgPool3DGrad(scope, Shape(scope, op.input(0)), grad_inputs[0], ksize,
|
|
strides, padding, grad_attrs.DataFormat(data_format));
|
|
grad_outputs->push_back(dx);
|
|
return scope.status();
|
|
}
|
|
REGISTER_GRADIENT_OP("AvgPool3D", AvgPool3DGradHelper);
|
|
|
|
absl::Status LRNGradHelper(const Scope& scope, const Operation& op,
|
|
const std::vector<Output>& grad_inputs,
|
|
std::vector<Output>* grad_outputs) {
|
|
auto dx = internal::LRNGrad(scope, grad_inputs[0], op.input(0), op.output(0));
|
|
grad_outputs->push_back(dx);
|
|
return scope.status();
|
|
}
|
|
REGISTER_GRADIENT_OP("LRN", LRNGradHelper);
|
|
|
|
absl::Status SoftplusGradHelper(const Scope& scope, const Operation& op,
|
|
const std::vector<Output>& grad_inputs,
|
|
std::vector<Output>* grad_outputs) {
|
|
auto dx = internal::SoftplusGrad(scope, grad_inputs[0], op.input(0));
|
|
grad_outputs->push_back(dx);
|
|
return scope.status();
|
|
}
|
|
REGISTER_GRADIENT_OP("Softplus", SoftplusGradHelper);
|
|
|
|
absl::Status SoftsignGradHelper(const Scope& scope, const Operation& op,
|
|
const std::vector<Output>& grad_inputs,
|
|
std::vector<Output>* grad_outputs) {
|
|
auto dx = internal::SoftsignGrad(scope, grad_inputs[0], op.input(0));
|
|
grad_outputs->push_back(dx);
|
|
return scope.status();
|
|
}
|
|
REGISTER_GRADIENT_OP("Softsign", SoftsignGradHelper);
|
|
|
|
absl::Status FractionalAvgPoolGradHelper(const Scope& scope,
|
|
const Operation& op,
|
|
const std::vector<Output>& grad_inputs,
|
|
std::vector<Output>* grad_outputs) {
|
|
bool overlapping;
|
|
TF_RETURN_IF_ERROR(
|
|
GetNodeAttr(op.output(0).node()->attrs(), "overlapping", &overlapping));
|
|
auto dx = internal::FractionalAvgPoolGrad(
|
|
scope, Shape(scope, op.input(0), Shape::OutType(DT_INT64)),
|
|
grad_inputs[0], op.output(1), op.output(2),
|
|
internal::FractionalAvgPoolGrad::Overlapping(overlapping));
|
|
grad_outputs->push_back(dx);
|
|
return scope.status();
|
|
}
|
|
REGISTER_GRADIENT_OP("FractionalAvgPool", FractionalAvgPoolGradHelper);
|
|
|
|
absl::Status FractionalMaxPoolGradHelper(const Scope& scope,
|
|
const Operation& op,
|
|
const std::vector<Output>& grad_inputs,
|
|
std::vector<Output>* grad_outputs) {
|
|
bool overlapping;
|
|
TF_RETURN_IF_ERROR(
|
|
GetNodeAttr(op.output(0).node()->attrs(), "overlapping", &overlapping));
|
|
auto dx = internal::FractionalMaxPoolGrad(
|
|
scope, op.input(0), op.output(0), grad_inputs[0], op.output(1),
|
|
op.output(2), internal::FractionalMaxPoolGrad::Overlapping(overlapping));
|
|
grad_outputs->push_back(dx);
|
|
return scope.status();
|
|
}
|
|
REGISTER_GRADIENT_OP("FractionalMaxPool", FractionalMaxPoolGradHelper);
|
|
|
|
// Templated constructor for FusedBatchNormGrad[..]::Attrs.
|
|
template <typename T>
|
|
T FusedBatchNormGradAttrs(float epsilon, absl::string_view 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<absl::Status(
|
|
const Scope&, Output x, Output grad_y, Output scale,
|
|
const std::vector<Output>& reserve_spaces, float epsilon,
|
|
absl::string_view data_format, bool is_training,
|
|
std::vector<Output>* grad_outputs)>;
|
|
|
|
absl::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 absl::InvalidArgumentError(
|
|
"FusedBatchNorm requires at least 5 outputs");
|
|
}
|
|
if (grad_inputs.empty()) {
|
|
return absl::InvalidArgumentError(
|
|
"FusedBatchNorm grad requires 1 grad input");
|
|
}
|
|
if (op.num_inputs() < 3) {
|
|
return absl::InvalidArgumentError("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 absl::InvalidArgumentError(
|
|
"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});
|
|
}
|
|
|
|
absl::string_view 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();
|
|
}
|
|
}
|
|
|
|
absl::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,
|
|
absl::string_view 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);
|
|
|
|
absl::Status Conv2DBackpropInputGrad(const Scope& scope, const Operation& op,
|
|
const std::vector<Output>& grad_inputs,
|
|
std::vector<Output>* grad_outputs) {
|
|
if (op.num_inputs() != 3) {
|
|
return absl::InvalidArgumentError("Conv2DBackpropInput requires 3 inputs.");
|
|
}
|
|
if (grad_inputs.empty()) {
|
|
return absl::InvalidArgumentError(
|
|
"Conv2DBackpropInput grad requires 1 grad input");
|
|
}
|
|
|
|
std::vector<int> dilations, strides, explicit_paddings;
|
|
bool use_cudnn_on_gpu;
|
|
std::string data_format, padding;
|
|
TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "dilations", &dilations));
|
|
TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "strides", &strides));
|
|
TF_RETURN_IF_ERROR(
|
|
GetNodeAttr(op.node()->attrs(), "explicit_paddings", &explicit_paddings));
|
|
TF_RETURN_IF_ERROR(
|
|
GetNodeAttr(op.node()->attrs(), "use_cudnn_on_gpu", &use_cudnn_on_gpu));
|
|
TF_RETURN_IF_ERROR(
|
|
GetNodeAttr(op.node()->attrs(), "data_format", &data_format));
|
|
TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "padding", &padding));
|
|
|
|
grad_outputs->push_back(NoGradient());
|
|
|
|
Conv2DBackpropFilter::Attrs filter_attrs;
|
|
filter_attrs.use_cudnn_on_gpu_ = use_cudnn_on_gpu;
|
|
filter_attrs.explicit_paddings_ = explicit_paddings;
|
|
filter_attrs.data_format_ = data_format;
|
|
filter_attrs.dilations_ = dilations;
|
|
grad_outputs->push_back(
|
|
Conv2DBackpropFilter(scope, grad_inputs[0], Shape(scope, op.input(1)),
|
|
op.input(2), strides, padding, filter_attrs));
|
|
|
|
Conv2D::Attrs conv_attrs;
|
|
conv_attrs.use_cudnn_on_gpu_ = use_cudnn_on_gpu;
|
|
conv_attrs.explicit_paddings_ = explicit_paddings;
|
|
conv_attrs.data_format_ = data_format;
|
|
conv_attrs.dilations_ = dilations;
|
|
grad_outputs->push_back(
|
|
Conv2D(scope, grad_inputs[0], op.input(1), strides, padding, conv_attrs));
|
|
return scope.status();
|
|
}
|
|
REGISTER_GRADIENT_OP("Conv2DBackpropInput", Conv2DBackpropInputGrad);
|
|
|
|
absl::Status DepthwiseConv2dNativeGrad(const Scope& scope, const Operation& op,
|
|
const std::vector<Output>& grad_inputs,
|
|
std::vector<Output>* grad_outputs) {
|
|
if (op.num_inputs() != 2) {
|
|
return absl::InvalidArgumentError(
|
|
"DepthwiseConv2dNative requires 2 inputs.");
|
|
}
|
|
if (grad_inputs.empty()) {
|
|
return absl::InvalidArgumentError(
|
|
"DepthwiseConv2dNative grad requires 1 grad input");
|
|
}
|
|
|
|
std::vector<int> dilations, strides, explicit_paddings;
|
|
std::string data_format, padding;
|
|
TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "dilations", &dilations));
|
|
TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "strides", &strides));
|
|
TF_RETURN_IF_ERROR(
|
|
GetNodeAttr(op.node()->attrs(), "explicit_paddings", &explicit_paddings));
|
|
TF_RETURN_IF_ERROR(
|
|
GetNodeAttr(op.node()->attrs(), "data_format", &data_format));
|
|
TF_RETURN_IF_ERROR(GetNodeAttr(op.node()->attrs(), "padding", &padding));
|
|
|
|
DepthwiseConv2dNativeBackpropInput::Attrs input_attrs;
|
|
input_attrs.explicit_paddings_ = explicit_paddings;
|
|
input_attrs.data_format_ = data_format;
|
|
input_attrs.dilations_ = dilations;
|
|
grad_outputs->push_back(DepthwiseConv2dNativeBackpropInput(
|
|
scope, Shape(scope, op.input(0)), op.input(1), grad_inputs[0], strides,
|
|
padding, input_attrs));
|
|
|
|
DepthwiseConv2dNativeBackpropFilter::Attrs filter_attrs;
|
|
filter_attrs.explicit_paddings_ = explicit_paddings;
|
|
filter_attrs.data_format_ = data_format;
|
|
filter_attrs.dilations_ = dilations;
|
|
grad_outputs->push_back(DepthwiseConv2dNativeBackpropFilter(
|
|
scope, op.input(0), Shape(scope, op.input(1)), grad_inputs[0], strides,
|
|
padding, filter_attrs));
|
|
return scope.status();
|
|
}
|
|
REGISTER_GRADIENT_OP("DepthwiseConv2dNative", DepthwiseConv2dNativeGrad);
|
|
|
|
} // anonymous namespace
|
|
} // namespace ops
|
|
} // namespace tensorflow
|