245 lines
9.4 KiB
C++
245 lines
9.4 KiB
C++
/* Copyright 2020 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 "tensorflow/c/eager/mnist_gradients_testutil.h"
|
|
|
|
#include <memory>
|
|
|
|
#include "absl/container/flat_hash_set.h"
|
|
#include "absl/types/span.h"
|
|
#include "tensorflow/c/eager/abstract_tensor_handle.h"
|
|
#include "tensorflow/c/eager/c_api_experimental.h"
|
|
#include "tensorflow/c/eager/c_api_unified_experimental.h"
|
|
#include "tensorflow/c/eager/c_api_unified_experimental_internal.h"
|
|
#include "tensorflow/c/eager/gradients.h"
|
|
#include "tensorflow/c/eager/gradients_internal.h"
|
|
#include "tensorflow/c/eager/gradients_util.h"
|
|
#include "tensorflow/c/experimental/gradients/tape/tape_context.h"
|
|
#include "tensorflow/c/experimental/ops/array_ops.h"
|
|
#include "tensorflow/c/experimental/ops/math_ops.h"
|
|
#include "tensorflow/c/experimental/ops/nn_ops.h"
|
|
#include "tensorflow/core/lib/llvm_rtti/llvm_rtti.h"
|
|
|
|
namespace tensorflow {
|
|
namespace gradients {
|
|
namespace internal {
|
|
|
|
using std::vector;
|
|
|
|
//===================== Test Models to run =========================
|
|
|
|
// Computes
|
|
// y = inputs[0] + inputs[1]
|
|
// return grad(y, {inputs[0], inputs[1]})
|
|
Status AddGradModel(AbstractContext* ctx,
|
|
absl::Span<AbstractTensorHandle* const> inputs,
|
|
absl::Span<AbstractTensorHandle*> outputs,
|
|
const GradientRegistry& registry) {
|
|
auto tape = new Tape(/*persistent=*/false);
|
|
tape->Watch(inputs[0]); // Watch x.
|
|
tape->Watch(inputs[1]); // Watch y.
|
|
std::vector<AbstractTensorHandle*> add_outputs(1);
|
|
AbstractContextPtr tape_ctx(new TapeContext(ctx, tape, registry));
|
|
TF_RETURN_IF_ERROR(
|
|
ops::Add(tape_ctx.get(), inputs, absl::MakeSpan(add_outputs), "Add"));
|
|
TF_RETURN_IF_ERROR(tape->ComputeGradient(ctx, /*targets=*/add_outputs,
|
|
/*sources=*/inputs,
|
|
/*output_gradients=*/{}, outputs));
|
|
for (auto add_output : add_outputs) {
|
|
add_output->Unref();
|
|
}
|
|
delete tape;
|
|
return Status::OK();
|
|
}
|
|
|
|
// Computes
|
|
// y = inputs[0] * inputs[1]
|
|
// return grad(y, {inputs[0], inputs[1]})
|
|
Status MatMulGradModel(AbstractContext* ctx,
|
|
absl::Span<AbstractTensorHandle* const> inputs,
|
|
absl::Span<AbstractTensorHandle*> outputs,
|
|
const GradientRegistry& registry) {
|
|
auto tape = new Tape(/*persistent=*/false);
|
|
tape->Watch(inputs[0]); // Watch x.
|
|
tape->Watch(inputs[1]); // Watch y.
|
|
vector<AbstractTensorHandle*> mm_outputs(1);
|
|
AbstractContextPtr tape_ctx(new TapeContext(ctx, tape, registry));
|
|
TF_RETURN_IF_ERROR(ops::MatMul(tape_ctx.get(), inputs,
|
|
absl::MakeSpan(mm_outputs), "matmul0",
|
|
/*transpose_a=*/false,
|
|
/*transpose_b=*/false)); // Compute x*y.
|
|
|
|
TF_RETURN_IF_ERROR(tape->ComputeGradient(ctx, /*targets=*/mm_outputs,
|
|
/*sources=*/inputs,
|
|
/*output_gradients=*/{}, outputs));
|
|
for (auto mm_output : mm_outputs) {
|
|
mm_output->Unref();
|
|
}
|
|
delete tape;
|
|
return Status::OK();
|
|
}
|
|
|
|
// Model to run 2-layer net
|
|
Status MNISTForwardModel(AbstractContext* ctx,
|
|
absl::Span<AbstractTensorHandle* const> inputs,
|
|
absl::Span<AbstractTensorHandle*> outputs,
|
|
const GradientRegistry& registry) {
|
|
/**
|
|
* We will trace a 2-layer fully connected network for an MNIST model:
|
|
*
|
|
* def mnist_forward(X, W1, W2, y_labels):
|
|
* mm_out_1 = tf.matmul(X,W1)
|
|
* hidden_layer = tf.nn.relu(mm_out_1)
|
|
* scores = tf.matmul(hidden_layer,W2)
|
|
* softmax =
|
|
* tf.nn.sparse_softmax_cross_entropy_with_logits(scores,
|
|
* y_labels)
|
|
* return scores, softmax
|
|
*
|
|
* Use this convention for inputs:
|
|
*
|
|
* inputs = [X, W1, W2, y_labels]
|
|
*
|
|
*/
|
|
AbstractTensorHandle* X = inputs[0];
|
|
AbstractTensorHandle* W1 = inputs[1];
|
|
AbstractTensorHandle* W2 = inputs[2];
|
|
AbstractTensorHandle* y_labels = inputs[3];
|
|
|
|
vector<AbstractTensorHandle*> temp_outputs(1);
|
|
|
|
TF_RETURN_IF_ERROR(ops::MatMul(ctx, {X, W1}, absl::MakeSpan(temp_outputs),
|
|
"matmul0",
|
|
/*transpose_a=*/false,
|
|
/*transpose_b=*/false)); // Compute X*W1
|
|
|
|
TF_RETURN_IF_ERROR(ops::Relu(ctx, {temp_outputs[0]},
|
|
absl::MakeSpan(temp_outputs),
|
|
"relu")); // Compute Relu(X*W1)
|
|
|
|
TF_RETURN_IF_ERROR(ops::MatMul(
|
|
ctx, {temp_outputs[0], W2}, absl::MakeSpan(temp_outputs), "matmul1",
|
|
/*transpose_a=*/false, /*transpose_b=*/false)); // Compute W2*Relu(X*W1)
|
|
|
|
AbstractTensorHandle* scores = temp_outputs[0];
|
|
|
|
temp_outputs.resize(2);
|
|
TF_RETURN_IF_ERROR(ops::SparseSoftmaxCrossEntropyWithLogits(
|
|
ctx, {scores, y_labels}, absl::MakeSpan(temp_outputs),
|
|
"softmax_loss")); // Compute Softmax(Scores,labels)
|
|
|
|
AbstractTensorHandle* loss_vals = temp_outputs[0];
|
|
|
|
outputs[0] = scores;
|
|
outputs[1] = loss_vals;
|
|
return Status::OK();
|
|
}
|
|
|
|
Status MatMulTransposeModel(AbstractContext* ctx,
|
|
absl::Span<AbstractTensorHandle* const> inputs,
|
|
absl::Span<AbstractTensorHandle*> outputs,
|
|
const GradientRegistry& registry) {
|
|
AbstractTensorHandle* X = inputs[0];
|
|
AbstractTensorHandle* W1 = inputs[1];
|
|
|
|
TF_RETURN_IF_ERROR(ops::MatMul(ctx, {X, W1}, outputs, "matmul0",
|
|
/*transpose_a=*/true,
|
|
/*transpose_b=*/false)); // Compute X*W1
|
|
return Status::OK();
|
|
}
|
|
|
|
Status MNISTGradModel(AbstractContext* ctx,
|
|
absl::Span<AbstractTensorHandle* const> inputs,
|
|
absl::Span<AbstractTensorHandle*> outputs,
|
|
const GradientRegistry& registry) {
|
|
AbstractTensorHandle* X = inputs[0];
|
|
AbstractTensorHandle* W1 = inputs[1];
|
|
AbstractTensorHandle* W2 = inputs[2];
|
|
AbstractTensorHandle* y_labels = inputs[3];
|
|
|
|
auto tape = new Tape(/*persistent=*/true);
|
|
tape->Watch(X); // Watch X.
|
|
tape->Watch(W1); // Watch W1.
|
|
tape->Watch(W2); // Watch W1.
|
|
vector<AbstractTensorHandle*> temp_outputs(1);
|
|
AbstractContextPtr tape_ctx(new TapeContext(ctx, tape, registry));
|
|
TF_RETURN_IF_ERROR(ops::MatMul(tape_ctx.get(), {X, W1},
|
|
absl::MakeSpan(temp_outputs), "matmul0",
|
|
/*transpose_a=*/false,
|
|
/*transpose_b=*/false)); // Compute X*W1
|
|
|
|
AbstractTensorHandle* mm = temp_outputs[0];
|
|
|
|
TF_RETURN_IF_ERROR(ops::Relu(tape_ctx.get(), {mm},
|
|
absl::MakeSpan(temp_outputs), // Relu(X*W1)
|
|
"relu0"));
|
|
|
|
AbstractTensorHandle* hidden = temp_outputs[0];
|
|
|
|
TF_RETURN_IF_ERROR(ops::MatMul(
|
|
tape_ctx.get(), {hidden, W2}, absl::MakeSpan(temp_outputs), "matmul1",
|
|
/*transpose_a=*/false, /*transpose_b=*/false)); // W2*Relu(X*W1)
|
|
|
|
AbstractTensorHandle* scores = temp_outputs[0];
|
|
|
|
temp_outputs.resize(2);
|
|
TF_RETURN_IF_ERROR(ops::SparseSoftmaxCrossEntropyWithLogits(
|
|
tape_ctx.get(), {scores, y_labels}, absl::MakeSpan(temp_outputs),
|
|
"softmaxloss")); // W2*Relu(X*W1)
|
|
|
|
AbstractTensorHandle* loss = temp_outputs[0];
|
|
|
|
TF_RETURN_IF_ERROR(tape->ComputeGradient(ctx, /*targets=*/{loss},
|
|
/*sources=*/{W1, W2},
|
|
/*output_gradients=*/{},
|
|
outputs.subspan(0, 2)));
|
|
|
|
// Only release 2nd temp output as first holds loss values.
|
|
temp_outputs[1]->Unref();
|
|
outputs[2] = loss;
|
|
delete tape;
|
|
return Status::OK();
|
|
}
|
|
|
|
Status ScalarMulModel(AbstractContext* ctx,
|
|
absl::Span<AbstractTensorHandle* const> inputs,
|
|
absl::Span<AbstractTensorHandle*> outputs,
|
|
const GradientRegistry& registry) {
|
|
return ops::Mul(ctx, inputs, outputs,
|
|
"scalarMul0"); // Compute eta*A
|
|
}
|
|
|
|
Status MatMulModel(AbstractContext* ctx,
|
|
absl::Span<AbstractTensorHandle* const> inputs,
|
|
absl::Span<AbstractTensorHandle*> outputs,
|
|
const GradientRegistry& registry) {
|
|
return ops::MatMul(ctx, inputs, outputs, "matmul0",
|
|
/*transpose_a=*/false,
|
|
/*transpose_b=*/false); // Compute X*W1
|
|
}
|
|
|
|
Status MulModel(AbstractContext* ctx,
|
|
absl::Span<AbstractTensorHandle* const> inputs,
|
|
absl::Span<AbstractTensorHandle*> outputs,
|
|
const GradientRegistry& registry) {
|
|
return ops::Mul(ctx, inputs, outputs,
|
|
"mul0"); // Compute x*y
|
|
}
|
|
|
|
// ============================= End Models ================================
|
|
|
|
} // namespace internal
|
|
} // namespace gradients
|
|
} // namespace tensorflow
|