Continue splitting SignatureDefFunction from ConcreteFunction.
PiperOrigin-RevId: 328682987 Change-Id: I4d32d48bbfb5f461c5e3edc50513057808c8c47f
This commit is contained in:
Brian Zhao
TensorFlower Gardener
parent
c5e13d84ca
commit
ecc6fe9066
@ -81,3 +81,26 @@ cc_library(
|
|||||||
"@com_google_absl//absl/types:span",
|
"@com_google_absl//absl/types:span",
|
||||||
],
|
],
|
||||||
)
|
)
|
||||||
|
|
||||||
|
cc_library(
|
||||||
|
name = "tf_signature_def_function",
|
||||||
|
srcs = [
|
||||||
|
"tf_signature_def_function.cc",
|
||||||
|
],
|
||||||
|
hdrs = [
|
||||||
|
"tf_signature_def_function.h",
|
||||||
|
],
|
||||||
|
deps = [
|
||||||
|
":tensorhandle_convertible",
|
||||||
|
"//tensorflow/c/eager:abstract_tensor_handle",
|
||||||
|
"//tensorflow/c/eager:immediate_execution_context",
|
||||||
|
"//tensorflow/c/eager:immediate_execution_operation",
|
||||||
|
"//tensorflow/c/eager:immediate_execution_tensor_handle",
|
||||||
|
"//tensorflow/c/experimental/saved_model/core:signature_def_function",
|
||||||
|
"//tensorflow/c/experimental/saved_model/core:signature_def_function_metadata",
|
||||||
|
"//tensorflow/core:lib",
|
||||||
|
"//tensorflow/core:protos_all_cc",
|
||||||
|
"//tensorflow/core/common_runtime/eager:context",
|
||||||
|
"@com_google_absl//absl/types:span",
|
||||||
|
],
|
||||||
|
)
|
||||||
|
|||||||
@ -0,0 +1,96 @@
|
|||||||
|
/* 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/experimental/saved_model/core/revived_types/tf_signature_def_function.h"
|
||||||
|
|
||||||
|
#include <memory>
|
||||||
|
#include <string>
|
||||||
|
|
||||||
|
#include "absl/types/span.h"
|
||||||
|
#include "tensorflow/c/eager/abstract_tensor_handle.h"
|
||||||
|
#include "tensorflow/c/eager/immediate_execution_operation.h"
|
||||||
|
#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
|
||||||
|
#include "tensorflow/c/experimental/saved_model/core/revived_types/tensorhandle_convertible.h"
|
||||||
|
#include "tensorflow/core/common_runtime/eager/context.h"
|
||||||
|
#include "tensorflow/core/framework/function.pb.h"
|
||||||
|
#include "tensorflow/core/platform/errors.h"
|
||||||
|
#include "tensorflow/core/platform/logging.h"
|
||||||
|
#include "tensorflow/core/platform/status.h"
|
||||||
|
#include "tensorflow/core/protobuf/saved_object_graph.pb.h"
|
||||||
|
#include "tensorflow/core/protobuf/struct.pb.h"
|
||||||
|
|
||||||
|
namespace tensorflow {
|
||||||
|
|
||||||
|
TFSignatureDefFunction::TFSignatureDefFunction(
|
||||||
|
const std::string& name,
|
||||||
|
std::vector<ImmediateExecutionTensorHandle*> captures,
|
||||||
|
SignatureDefFunctionMetadata metadata, ImmediateExecutionContext* ctx)
|
||||||
|
: name_(name),
|
||||||
|
captures_(std::move(captures)),
|
||||||
|
metadata_(std::move(metadata)),
|
||||||
|
ctx_(ctx) {}
|
||||||
|
|
||||||
|
TFSignatureDefFunction::~TFSignatureDefFunction() {
|
||||||
|
Status status = ctx_->RemoveFunction(name_);
|
||||||
|
if (!status.ok()) {
|
||||||
|
LOG(ERROR) << "Failed to remove functiondef " << name_ << ". "
|
||||||
|
<< status.error_message();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
Status TFSignatureDefFunction::Create(
|
||||||
|
const FunctionDef* function_def,
|
||||||
|
std::vector<ImmediateExecutionTensorHandle*> captures,
|
||||||
|
SignatureDefFunctionMetadata metadata, ImmediateExecutionContext* ctx,
|
||||||
|
std::unique_ptr<TFSignatureDefFunction>* out) {
|
||||||
|
TF_RETURN_IF_ERROR(ctx->AddFunctionDef(*function_def));
|
||||||
|
out->reset(new TFSignatureDefFunction(function_def->signature().name(),
|
||||||
|
std::move(captures),
|
||||||
|
std::move(metadata), ctx));
|
||||||
|
return Status();
|
||||||
|
}
|
||||||
|
|
||||||
|
const SignatureDefFunctionMetadata&
|
||||||
|
TFSignatureDefFunction::GetFunctionMetadata() const {
|
||||||
|
return metadata_;
|
||||||
|
}
|
||||||
|
|
||||||
|
Status TFSignatureDefFunction::MakeCallOp(
|
||||||
|
absl::Span<AbstractTensorHandle* const> inputs, ImmediateOpPtr* out) const {
|
||||||
|
out->reset(ctx_->CreateOperation());
|
||||||
|
// In eager mode, TF2 python executes functions by constructing an op with
|
||||||
|
// the name of the functiondef:
|
||||||
|
// https://github.com/tensorflow/tensorflow/blob/66668ec0ca432e2f38a575b814f45b6d299d01ed/tensorflow/python/eager/function.py#L545
|
||||||
|
// In graph mode, we create a PartitionedCallOp instead:
|
||||||
|
// https://github.com/tensorflow/tensorflow/blob/66668ec0ca432e2f38a575b814f45b6d299d01ed/tensorflow/python/eager/function.py#L573
|
||||||
|
|
||||||
|
// TODO(bmzhao): After discussing with Allen, we should execute this via a
|
||||||
|
// PartitionedCallOp for compatibility with "tooling that assumes functions in
|
||||||
|
// graphs are PartitionedCallOps".
|
||||||
|
TF_RETURN_IF_ERROR((*out)->Reset(name_.c_str(), nullptr));
|
||||||
|
|
||||||
|
// Adding the user-provided inputs to the function.
|
||||||
|
TF_RETURN_IF_ERROR((*out)->AddInputList(inputs));
|
||||||
|
|
||||||
|
absl::Span<AbstractTensorHandle* const> captures(
|
||||||
|
reinterpret_cast<AbstractTensorHandle* const*>(captures_.data()),
|
||||||
|
captures_.size());
|
||||||
|
|
||||||
|
// Adding the captures of the function.
|
||||||
|
TF_RETURN_IF_ERROR((*out)->AddInputList(captures));
|
||||||
|
return Status();
|
||||||
|
}
|
||||||
|
|
||||||
|
} // namespace tensorflow
|
||||||
@ -0,0 +1,90 @@
|
|||||||
|
/* 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.
|
||||||
|
==============================================================================*/
|
||||||
|
|
||||||
|
#ifndef TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_TF_SIGNATURE_DEF_FUNCTION_H_
|
||||||
|
#define TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_TF_SIGNATURE_DEF_FUNCTION_H_
|
||||||
|
|
||||||
|
#include <functional>
|
||||||
|
#include <memory>
|
||||||
|
#include <string>
|
||||||
|
#include <unordered_map>
|
||||||
|
#include <vector>
|
||||||
|
|
||||||
|
#include "tensorflow/c/eager/immediate_execution_context.h"
|
||||||
|
#include "tensorflow/c/eager/immediate_execution_operation.h"
|
||||||
|
#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
|
||||||
|
#include "tensorflow/c/experimental/saved_model/core/revived_types/tensorhandle_convertible.h"
|
||||||
|
#include "tensorflow/c/experimental/saved_model/core/signature_def_function.h"
|
||||||
|
#include "tensorflow/c/experimental/saved_model/core/signature_def_function_metadata.h"
|
||||||
|
#include "tensorflow/core/framework/function.pb.h"
|
||||||
|
#include "tensorflow/core/protobuf/saved_object_graph.pb.h"
|
||||||
|
|
||||||
|
namespace tensorflow {
|
||||||
|
|
||||||
|
// This is the TF eager runtime implementation of SignatureDefFunction (separate
|
||||||
|
// from the TFRT implementation). The user-facing API of SignatureDefFunctions
|
||||||
|
// and their semantic differences from ConcreteFunction are described here:
|
||||||
|
// https://github.com/tensorflow/tensorflow/blob/e2db60c9d9598ebae0b7741587ce6f5d473584d9/tensorflow/cc/saved_model/experimental/public/signature_def_function.h#L30-L59
|
||||||
|
// Additional implementation notes are available here:
|
||||||
|
// https://github.com/tensorflow/tensorflow/blob/e2db60c9d9598ebae0b7741587ce6f5d473584d9/tensorflow/c/experimental/saved_model/core/signature_def_function.h#L31-L48
|
||||||
|
class TFSignatureDefFunction : public SignatureDefFunction {
|
||||||
|
public:
|
||||||
|
// Factory function for creating a TFSignatureDefFunction.
|
||||||
|
//
|
||||||
|
// Params:
|
||||||
|
// function_def - The function_def associated with the created
|
||||||
|
// TFSignatureDefFunction. TFSignatureDefFunction will
|
||||||
|
// register this function_def with `ctx` on creation, and
|
||||||
|
// de-register it on destruction. function_def must be
|
||||||
|
// non-null, but otherwise has no lifetime requirements.
|
||||||
|
// captures - The captured TensorHandles associated with this
|
||||||
|
// TFConcreteFunction.
|
||||||
|
// metadata - FunctionMetadata associated with this TFSignatureDefFunction.
|
||||||
|
// ctx - A handle to the Tensorflow runtime. This MUST be non-null and
|
||||||
|
// outlive TFSignatureDefFunction.
|
||||||
|
// out - The output TFSignatureDefFunction.
|
||||||
|
static Status Create(const FunctionDef* function_def,
|
||||||
|
std::vector<ImmediateExecutionTensorHandle*> captures,
|
||||||
|
SignatureDefFunctionMetadata metadata,
|
||||||
|
ImmediateExecutionContext* ctx,
|
||||||
|
std::unique_ptr<TFSignatureDefFunction>* out);
|
||||||
|
|
||||||
|
// This method creates a "Call" Op used to execute the function.
|
||||||
|
Status MakeCallOp(absl::Span<AbstractTensorHandle* const> inputs,
|
||||||
|
ImmediateOpPtr* out) const override;
|
||||||
|
|
||||||
|
const SignatureDefFunctionMetadata& GetFunctionMetadata() const override;
|
||||||
|
|
||||||
|
~TFSignatureDefFunction() override;
|
||||||
|
|
||||||
|
private:
|
||||||
|
TFSignatureDefFunction(const std::string& name,
|
||||||
|
std::vector<ImmediateExecutionTensorHandle*> captures,
|
||||||
|
SignatureDefFunctionMetadata metadata,
|
||||||
|
ImmediateExecutionContext* ctx);
|
||||||
|
|
||||||
|
TFSignatureDefFunction(const TFSignatureDefFunction&) = delete;
|
||||||
|
TFSignatureDefFunction& operator=(const TFSignatureDefFunction&) = delete;
|
||||||
|
|
||||||
|
// Name of the FunctionDef corresponding to this TFSignatureDefFunction
|
||||||
|
std::string name_;
|
||||||
|
std::vector<ImmediateExecutionTensorHandle*> captures_;
|
||||||
|
SignatureDefFunctionMetadata metadata_;
|
||||||
|
ImmediateExecutionContext* ctx_;
|
||||||
|
};
|
||||||
|
|
||||||
|
} // namespace tensorflow
|
||||||
|
|
||||||
|
#endif // TENSORFLOW_C_EXPERIMENTAL_SAVED_MODEL_CORE_TF_SIGNATURE_DEF_FUNCTION_H_
|
||||||
Loading…
Reference in New Issue
Block a user