wehub/tensorflow--tensorflow
1
0
Fork 0
Code Issues Pull Requests Actions 3 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
cffconvert / validate (push) Has been skipped
Details
License Check / license-check (push) Failing after 2s
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:14:16 +08:00
commit 8a852e4b4e
36502 changed files with 9277225 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tensorflow/compiler/tf2tensorrt/kernels/get_calibration_data_op.cc
+69
View File
@@ -0,0 +1,69 @@
/* Copyright 2018 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 <memory>
#include <vector>
#include "tensorflow/compiler/tf2tensorrt/utils/trt_lru_cache.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/resource_mgr.h"
#include "tensorflow/core/lib/core/refcount.h"
#if GOOGLE_CUDA && GOOGLE_TENSORRT
namespace tensorflow {
namespace tensorrt {
class GetCalibrationDataOp : public OpKernel {
public:
explicit GetCalibrationDataOp(OpKernelConstruction* context)
: OpKernel(context) {}
~GetCalibrationDataOp() override {}
void Compute(OpKernelContext* context) override {
// TODO(laigd): it will allocate the tensor on the device and copy the
// serialized string to that tensor, and later sess.run() will copy it back
// to host. We need to optimize this.
const string& resource_name = context->input(0).scalar<tstring>()();
// Get the resource.
TRTEngineCacheResource* resource = nullptr;
OP_REQUIRES_OK(context, context->resource_manager()->Lookup(
std::string(kTfTrtContainerName), resource_name,
&resource));
core::ScopedUnref sc(resource);
// Serialize the resource as output.
string serialized_resource = resource->calib_ctx_->TerminateCalibration();
OP_REQUIRES(context, !serialized_resource.empty(),
errors::Unknown("Calibration table is empty."));
Tensor* output = nullptr;
OP_REQUIRES_OK(context,
context->allocate_output(0, TensorShape({}), &output));
output->scalar<tstring>()() = serialized_resource;
}
};
REGISTER_KERNEL_BUILDER(Name("GetCalibrationDataOp").Device(DEVICE_GPU),
GetCalibrationDataOp);
} // namespace tensorrt
} // namespace tensorflow
#endif // GOOGLE_CUDA && GOOGLE_TENSORRT
tensorflow/compiler/tf2tensorrt/kernels/trt_engine_op.cc
+1448
View File
File diff suppressed because it is too large Load Diff
tensorflow/compiler/tf2tensorrt/kernels/trt_engine_op_test.cc
+377
View File
@@ -0,0 +1,377 @@
/* Copyright 2019 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 <memory>
#include <numeric>
#include <utility>
#include <vector>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "absl/container/inlined_vector.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "absl/types/span.h"
#include "tensorflow/cc/framework/scope.h"
#include "tensorflow/cc/ops/function_ops.h"
#include "tensorflow/cc/ops/math_ops.h"
#include "tensorflow/compiler/tf2tensorrt/convert/convert_graph.h"
#include "tensorflow/compiler/tf2tensorrt/utils/trt_lru_cache.h"
#include "xla/tsl/framework/fixedpoint/FixedPoint.h"
#include "tensorflow/core/common_runtime/device.h"
#include "tensorflow/core/common_runtime/device_factory.h"
#include "tensorflow/core/common_runtime/process_function_library_runtime.h"
#include "tensorflow/core/framework/attr_value.pb.h"
#include "tensorflow/core/framework/fake_input.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/node_def_builder.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/resource_mgr.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/kernels/ops_testutil.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/refcount.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/public/version.h"
#if GOOGLE_CUDA && GOOGLE_TENSORRT
namespace tensorflow {
namespace tensorrt {
using ::absl::StrCat;
using ::testing::ElementsAre;
struct TestParam {
bool static_engine;
};
class TRTEngineOpTestBase : public OpsTestBase {
public:
void AddSimpleTrtOp(DataType dtype, int max_cached_engines_count = 1,
PartialTensorShape shape = PartialTensorShape({-1, -1}),
bool use_implicit_batch = true,
bool allow_build_at_runtime = true,
bool static_engine = false) {
// Create the GPU device.
std::unique_ptr<Device> device(
DeviceFactory::NewDevice("GPU", {}, "/job:worker/replica:0/task:0"));
// Create simple TF graph.
Scope s = Scope::NewRootScope();
auto feed = ops::_Arg(s.WithOpName("TensorRTInputPH_0"), dtype, 0);
auto add = ops::Add(s.WithOpName("add"), feed, feed);
ops::_Retval give_me_a_name(s.WithOpName("TensorRTOutputPH_0"), add, 0);
// Serialize the graph. TRTEngineOp will convert it using dynamic mode.
GraphDef graph_def;
TF_ASSERT_OK(s.ToGraphDef(&graph_def));
Graph* graph = s.graph();
TF_ASSERT_OK(convert::RegisterGraphToFunctionLibrary(graph_def, graph,
std::string(kOpName)));
TF_ASSERT_OK(flib_def_->AddLibrary(graph->flib_def()));
string segment_string;
if (static_engine) {
convert::TRTOptimizationPass::ConversionParams params;
convert::EngineInfo info;
info.segment_graph_def.CopyFrom(graph_def);
info.precision_mode = TrtPrecisionMode::FP32;
info.max_workspace_size_bytes = 1 << 20;
info.engine_name = "TRTEngineOP_000_000";
params.use_implicit_batch = use_implicit_batch;
params.trt_logger_name = "DefaultLogger";
TrtShapeOptimizationProfile profile;
// We set the input mask to true (no resource inputs)
std::vector<bool> input_mask = {true};
profile.SetInputMask(input_mask);
// We set profile 0 to be incompatible with the input used in the test.
// This way we ensure that profile selection is tested.
TensorShape my_shape;
TF_CHECK_OK(
TensorShapeUtils::MakeShape(std::vector<int32>{4, 2}, &my_shape));
profile.AddShape({my_shape, {}});
TF_CHECK_OK(
TensorShapeUtils::MakeShape(std::vector<int32>{1, 2}, &my_shape));
profile.AddShape({my_shape, {}});
profile.InitProfiles({shape}, ProfileStrategy::kOptimal);
std::vector<PartialTensorShape> shape_vec{shape, {}};
TF_CHECK_OK(convert::CreateStaticEngine(
params, info, 1, shape_vec, &profile, &segment_string, nullptr));
}
// Create the op.
// In implicit batch mode, the input shapes that we specify here are not
// used for engine creation, we use the concrete shapes during inference
// time for creating the engine.
// In explicit batch mode, the input shapes attribute is used to define
// the network for the TensorRT engine.
OpsTestBase::SetDevice(DEVICE_GPU, std::move(device));
NameAttrList function;
function.set_name(StrCat(std::string(kOpName), "_native_segment"));
// We disable allow_soft_placement when executing the native segment of the
// TRTEngineOp for the following reasons:
// OpsTestBase only allow one device in the device manager.
// We need to define the GPU device to test TRTEngineOp.
// When allow_soft_placement is true, the TensorFlow runtime produces an
// error if a CPU device is not defined
// (see ProcessFunctionLibraryRuntime::InstantiateMultiDevice).
TF_ASSERT_OK(NodeDefBuilder(std::string(kOpName), "TRTEngineOp")
.Input(FakeInput(1, dtype))
.Attr("input_shapes", {shape})
.Attr("output_shapes", {shape})
.Attr("static_engine", static_engine)
.Attr("segment_func", function)
.Attr("serialized_segment", segment_string)
.Attr("calibration_data", "")
.Attr("max_cached_engines_count", max_cached_engines_count)
.Attr("workspace_size_bytes", 1 << 20)
.Attr("precision_mode", "FP32")
.Attr("use_calibration", false)
.Attr("profile_strategy", "optimal")
.Attr("_use_implicit_batch", use_implicit_batch)
.Attr("_allow_build_at_runtime", allow_build_at_runtime)
.Attr("_allow_soft_placement", false)
.Attr("OutT", {dtype})
.Finalize(OpsTestBase::node_def()));
TF_ASSERT_OK(InitOpWithFunctionLibrary());
}
static const absl::string_view kOpName;
template <typename T>
void AddSimpleInput(const TensorShape& shape) {
std::vector<T> input(shape.num_elements());
std::iota(input.begin(), input.end(), T(0));
OpsTestBase::AddInputFromArray<T>(shape, input);
}
void ResetInputs() {
inputs_.clear();
for (auto& temp : tensors_) {
delete temp;
}
tensors_.clear();
}
private:
Status InitOpWithFunctionLibrary() {
OpKernel* kernel = nullptr;
auto flr = pflr_->GetFLR(device_->name());
std::shared_ptr<const NodeProperties> props;
Status status = NodeProperties::CreateFromNodeDef(
node_def_, flr->GetFunctionLibraryDefinition(), &props);
if (status.ok()) {
status.Update(CreateOpKernel(device_type_, device_, allocator(), flr,
props, TF_GRAPH_DEF_VERSION, &kernel));
}
kernel_ = std::unique_ptr<OpKernel>(kernel);
if (kernel_ != nullptr) input_types_ = kernel_->input_types();
return status;
}
};
class TRTEngineOpTestWithParam
: public TRTEngineOpTestBase,
public ::testing::WithParamInterface<TestParam> {
public:
TRTEngineOpTestWithParam() : param_(GetParam()) {}
protected:
TestParam param_;
};
const absl::string_view TRTEngineOpTestBase::kOpName = "myop";
constexpr std::array<TestParam, 2> TestParameters{TestParam{false},
TestParam{true}};
INSTANTIATE_TEST_CASE_P(TRTEngineOpTestInstantiation, TRTEngineOpTestWithParam,
::testing::ValuesIn(TestParameters));
TEST_F(TRTEngineOpTestBase, DynamicEngines) {
// Test dynamic engine creation during inference time
TRTEngineOpTestBase::AddSimpleTrtOp(DT_FLOAT, /*max_cached_engines_count=*/4);
// Execute the op with batch size > 1.
TRTEngineOpTestBase::AddSimpleInput<float>(TensorShape({2, 2}));
TF_ASSERT_OK(OpsTestBase::RunOpKernel());
// Get the engine cache.
TRTEngineCacheResource* cache_resource = nullptr;
TF_ASSERT_OK(device_->resource_manager()->Lookup(
std::string(kTfTrtContainerName), std::string(kOpName), &cache_resource));
core::ScopedUnref sc(cache_resource);
// It should contain only one engine.
auto cache = &cache_resource->cache_;
EXPECT_EQ(1, cache->size());
EXPECT_EQ(1, cache->count({TensorShape({2, 2})}));
// Execute the op with batch size 1. It should reuse existing engine to
// execute.
ResetInputs();
TRTEngineOpTestBase::AddSimpleInput<float>(TensorShape({1, 2}));
TF_ASSERT_OK(OpsTestBase::RunOpKernel());
EXPECT_EQ(1, cache->size());
EXPECT_EQ(1, cache->count({TensorShape({2, 2})}));
// Execute the op with a larger batch size.
ResetInputs();
TRTEngineOpTestBase::AddSimpleInput<float>(TensorShape({3, 2}));
TF_ASSERT_OK(OpsTestBase::RunOpKernel());
EXPECT_EQ(2, cache->size());
EXPECT_EQ(1, cache->count({TensorShape({2, 2})}));
EXPECT_EQ(1, cache->count({TensorShape({3, 2})}));
// Execute the op with an input that has different non-batch dimension.
ResetInputs();
TRTEngineOpTestBase::AddSimpleInput<float>(TensorShape({10, 10}));
TF_ASSERT_OK(OpsTestBase::RunOpKernel());
// Execute it again with an input that has the same non-batch dimension but
// smallest batch size. It should find the correct engine to use.
ResetInputs();
TRTEngineOpTestBase::AddSimpleInput<float>(TensorShape({1, 10}));
TF_ASSERT_OK(OpsTestBase::RunOpKernel());
EXPECT_EQ(3, cache->size()); // Should only create 3 engines in total.
EXPECT_EQ(1, cache->count({TensorShape({2, 2})}));
EXPECT_EQ(1, cache->count({TensorShape({3, 2})}));
EXPECT_EQ(1, cache->count({TensorShape({10, 10})}));
}
TEST_F(TRTEngineOpTestBase, AllowBuildAtRuntime) {
TRTEngineOpTestBase::AddSimpleTrtOp(DT_FLOAT, /*max_cached_engines_count=*/1,
PartialTensorShape({-1, -1}),
/*use_implicit_batch=*/true,
/*allow_build_at_runtime=*/false);
// Execute the op
TensorShape input_shape({2, 2});
TRTEngineOpTestBase::AddSimpleInput<float>(input_shape);
TF_ASSERT_OK(OpsTestBase::RunOpKernel());
// Get the engine cache.
TRTEngineCacheResource* cache_resource = nullptr;
TF_ASSERT_OK(device_->resource_manager()->Lookup(
std::string(kTfTrtContainerName), std::string(kOpName), &cache_resource));
core::ScopedUnref sc(cache_resource);
// It should contain a placeholder with an empty cuda_engine (to mark that
// engine creation was not successful for the given input shape).
auto cache = &cache_resource->cache_;
EXPECT_EQ(1, cache->size());
ASSERT_EQ(1, cache->count({input_shape}));
EngineContext* ectx = cache->at({input_shape}).get();
EXPECT_EQ(ectx->GetCudaEngine(), nullptr);
}
TEST_P(TRTEngineOpTestWithParam, ExplicitBatch) {
// Test inference in explicit batch mode with static input shapes. Static
// shapes in this context means that the TensorRT knows all the input shapes
// during engine creation time.
TRTEngineOpTestBase::AddSimpleTrtOp(DT_FLOAT, /*max_cached_engines_count=*/1,
/*shape=*/PartialTensorShape({1, 2}),
/*use_implicit_batch=*/false,
/*allow_build_at_runtime=*/true,
/*static_engine=*/param_.static_engine);
TensorShape input_shape({1, 2});
TRTEngineOpTestBase::AddSimpleInput<float>(input_shape);
TF_ASSERT_OK(OpsTestBase::RunOpKernel());
// Get the engine cache.
TRTEngineCacheResource* cache_resource = nullptr;
TF_ASSERT_OK(device_->resource_manager()->Lookup(
std::string(kTfTrtContainerName), std::string(kOpName), &cache_resource));
core::ScopedUnref sc(cache_resource);
auto cache = &cache_resource->cache_;
EXPECT_EQ(1, cache->size());
ASSERT_EQ(1, cache->count({input_shape}));
EngineContext* ectx = cache->at({input_shape}).get();
EXPECT_NE(ectx->GetCudaEngine(), nullptr);
}
TEST_P(TRTEngineOpTestWithParam, DynamicShapes) {
// Test inference in explicit batch mode with dynamic input shapes. Dynamic
// shapes in this context means that some input shapes for TensorRT are
// unknown during engine creation time. When we create the network, the
// unknow shapes are repsesented as -1. Before we run inference, these shapes
// have to be specified by calling setBindingDimensions.
TRTEngineOpTestBase::AddSimpleTrtOp(DT_FLOAT, /*max_cached_engines_count=*/1,
/*shape=*/PartialTensorShape({-1, -1}),
/*use_implicit_batch=*/false,
/*allow_build_at_runtime=*/true,
param_.static_engine);
TensorShape input_shape({1, 2});
TRTEngineOpTestBase::AddSimpleInput<float>(input_shape);
TF_ASSERT_OK(OpsTestBase::RunOpKernel());
// Get the engine cache.
TRTEngineCacheResource* cache_resource = nullptr;
TF_ASSERT_OK(device_->resource_manager()->Lookup(
std::string(kTfTrtContainerName), std::string(kOpName), &cache_resource));
core::ScopedUnref sc(cache_resource);
auto cache = &cache_resource->cache_;
EXPECT_EQ(1, cache->size());
ASSERT_EQ(1, cache->count({input_shape}));
EngineContext* ectx = cache->at({input_shape}).get();
EXPECT_NE(ectx->GetCudaEngine(), nullptr);
// Execute the op with an incompatible shape.
ResetInputs();
TRTEngineOpTestBase::AddSimpleInput<float>(TensorShape({1, 37}));
// Test that the op runs. This should fall back to native segment.
TF_ASSERT_OK(OpsTestBase::RunOpKernel());
// We should still have a single engine that is not compatible with the input.
EXPECT_EQ(1, cache->size());
EXPECT_EQ(0, cache->count({TensorShape({1, 37})}));
}
template <typename T>
class TRTEngineOpTest : public TRTEngineOpTestBase {};
using TypeList = ::testing::Types<float, Eigen::half>;
TYPED_TEST_SUITE(TRTEngineOpTest, TypeList);
TYPED_TEST(TRTEngineOpTest, Basic) {
TRTEngineOpTestBase::AddSimpleTrtOp(DataTypeToEnum<TypeParam>::v());
// Execute the op.
OpsTestBase::AddInputFromArray<TypeParam>(TensorShape({1, 2}),
{TypeParam(0.0f), TypeParam(1.0f)});
TF_ASSERT_OK(OpsTestBase::RunOpKernel());
// Verify the result.
Tensor* output = OpsTestBase::GetOutput(0);
EXPECT_THAT(
absl::Span<const TypeParam>(output->template flat<TypeParam>().data(),
output->NumElements()),
ElementsAre(TypeParam(0.0f), TypeParam(2.0f)));
}
} // namespace tensorrt
} // namespace tensorflow
#endif // GOOGLE_CUDA && GOOGLE_TENSORRT
tensorflow/compiler/tf2tensorrt/kernels/trt_engine_resource_ops.cc
+310
View File
@@ -0,0 +1,310 @@
/* Copyright 2018 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 <algorithm>
#include <memory>
#include <vector>
#include "absl/memory/memory.h"
#include "absl/strings/string_view.h"
#include "tensorflow/compiler/tf2tensorrt/utils/trt_allocator.h"
#include "tensorflow/compiler/tf2tensorrt/utils/trt_engine_instance.pb.h" // NOLINT
#include "tensorflow/compiler/tf2tensorrt/utils/trt_logger.h"
#include "tensorflow/compiler/tf2tensorrt/utils/trt_lru_cache.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/resource_mgr.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/lib/core/refcount.h"
#include "tensorflow/core/lib/io/record_reader.h"
#include "tensorflow/core/lib/io/record_writer.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/mutex.h"
#include "tensorflow/core/platform/thread_annotations.h"
#include "tensorflow/core/profiler/lib/traceme.h"
#if GOOGLE_CUDA && GOOGLE_TENSORRT
#include "third_party/tensorrt/NvInfer.h"
namespace tensorflow {
namespace tensorrt {
using ::nvinfer1::IRuntime;
class CreateTRTResourceHandle : public OpKernel {
public:
explicit CreateTRTResourceHandle(OpKernelConstruction* ctx) : OpKernel(ctx) {
OP_REQUIRES_OK(ctx, ctx->GetAttr("resource_name", &resource_name_));
}
void Compute(OpKernelContext* ctx) override {
tensorflow::profiler::TraceMe activity(
"CreateTRTResourceHandle::Compute",
tensorflow::profiler::TraceMeLevel::kInfo);
{
mutex_lock l(mutex_);
if (!initialized_) {
AllocatorAttributes attr;
attr.set_on_host(true);
OP_REQUIRES_OK(ctx, ctx->allocate_temp(DT_RESOURCE, TensorShape({}),
&handle_, attr));
VLOG(1) << "Creating TRT engine cache resource handle for op "
<< resource_name_ << " on device " << ctx->device()->name();
handle_.scalar<ResourceHandle>()() =
MakeResourceHandle<TRTEngineCacheResource>(
ctx, std::string(kTfTrtContainerName), resource_name_);
initialized_ = true;
}
}
ctx->set_output(0, handle_);
}
private:
string resource_name_;
Tensor handle_;
mutex mutex_;
bool initialized_ TF_GUARDED_BY(mutex_) = false;
CreateTRTResourceHandle(const CreateTRTResourceHandle&) = delete;
void operator=(const CreateTRTResourceHandle&) = delete;
};
REGISTER_KERNEL_BUILDER(Name("CreateTRTResourceHandle")
.Device(DEVICE_GPU)
.HostMemory("resource_handle"),
CreateTRTResourceHandle);
class InitializeTRTResource : public OpKernel {
public:
explicit InitializeTRTResource(OpKernelConstruction* ctx) : OpKernel(ctx) {
OP_REQUIRES_OK(
ctx, ctx->GetAttr("max_cached_engines_count", &max_cached_engines_));
}
void Compute(OpKernelContext* ctx) override {
tensorflow::profiler::TraceMe activity(
"InitializeTRTResource::Compute",
tensorflow::profiler::TraceMeLevel::kInfo);
ResourceHandle handle = HandleFromInput(ctx, 0);
core::RefCountPtr<TRTEngineCacheResource> resource;
OP_REQUIRES_OK(
ctx, LookupOrCreateResource<TRTEngineCacheResource>(
ctx, handle, &resource,
[this, ctx](TRTEngineCacheResource** resource) -> Status {
*resource = new TRTEngineCacheResource(
ctx, this->max_cached_engines_);
return OkStatus();
}));
auto allocator = resource->allocator_.get();
OP_REQUIRES(ctx, allocator != nullptr,
errors::Internal("Not able to initialize TRT engine cache when "
"GPU allocator is empty."));
OP_REQUIRES(ctx, resource->cache_.size() == 0,
errors::Internal("Expect engine cache to be empty, but got ",
resource->cache_.size(), " entries."));
// Get the file name.
const string& filename = ctx->input(1).scalar<tstring>()();
OP_REQUIRES(ctx, !filename.empty(),
errors::InvalidArgument("filename cannot be empty."));
// Parse the serialized engines and add them to the cache.
std::unique_ptr<RandomAccessFile> file;
OP_REQUIRES_OK(ctx, ctx->env()->NewRandomAccessFile(filename, &file));
auto reader = std::make_unique<io::RecordReader>(file.get());
uint64 offset = 0;
int num_loaded_engine = 0;
do {
tstring record;
Status status = reader->ReadRecord(&offset, &record);
if (absl::IsOutOfRange(status)) break;
TRTEngineInstance engine_instance;
engine_instance.ParseFromString(record);
std::vector<TensorShape> engine_input_shapes;
const auto& input_shapes = engine_instance.input_shapes();
engine_input_shapes.reserve(input_shapes.size());
for (const TensorShapeProto& shape : input_shapes) {
engine_input_shapes.emplace_back(shape);
}
TrtUniquePtrType<IRuntime> infer(
nvinfer1::createInferRuntime(TRTEngineCacheResource::GetLogger()));
infer->setGpuAllocator(allocator);
TrtUniquePtrType<nvinfer1::ICudaEngine> engine(
infer->deserializeCudaEngine(
engine_instance.serialized_engine().c_str(),
engine_instance.serialized_engine().size(), nullptr));
auto raw_engine = engine.get();
std::vector<ExecutionContext> ctx_vec;
if (num_loaded_engine == 0) {
// Restore profiles if there are any. Currently only 1 engine is allowed
// in dynamic mode therefore we call this only for the 0th engine.
// it is a no-op in implicit batch mode.
OP_REQUIRES_OK(ctx, resource->profiles_.RestoreProfiles(
raw_engine, engine_input_shapes.size()));
OP_REQUIRES_OK(ctx, resource->profiles_.CreateExecutionContexts(
raw_engine, &ctx_vec));
} else {
// Multiple engines are only available in static mode. For each engine
// we have only a single execution context.
ctx_vec.push_back(ExecutionContext::Create(raw_engine));
}
resource->cache_.emplace(engine_input_shapes,
std::make_unique<EngineContext>(
std::move(engine), std::move(ctx_vec)));
++num_loaded_engine;
} while (1);
VLOG(1) << "Loaded " << num_loaded_engine << " TRT engines for op "
<< handle.name() << " on device " << ctx->device()->name()
<< " from file " << filename;
}
private:
// Maximum number of cached engines
int max_cached_engines_;
InitializeTRTResource(const InitializeTRTResource&) = delete;
void operator=(const InitializeTRTResource&) = delete;
};
REGISTER_KERNEL_BUILDER(Name("InitializeTRTResource")
.Device(DEVICE_GPU)
.HostMemory("resource_handle"),
InitializeTRTResource);
class SerializeTRTResource : public OpKernel {
public:
explicit SerializeTRTResource(OpKernelConstruction* ctx) : OpKernel(ctx) {
OP_REQUIRES_OK(ctx, ctx->GetAttr("delete_resource", &delete_resource_));
OP_REQUIRES_OK(ctx, ctx->GetAttr("save_gpu_specific_engines",
&save_gpu_specific_engines_));
}
void Compute(OpKernelContext* ctx) override {
tensorflow::profiler::TraceMe activity(
"SerializeTRTResource::Compute",
tensorflow::profiler::TraceMeLevel::kInfo);
const string& resource_name = ctx->input(0).scalar<tstring>()();
const string& filename = ctx->input(1).scalar<tstring>()();
OP_REQUIRES(ctx, !filename.empty(),
errors::InvalidArgument("filename cannot be empty."));
// Lookup engine cache resource.
TRTEngineCacheResource* resource = nullptr;
OP_REQUIRES(
ctx,
ctx->resource_manager()
->Lookup(std::string(kTfTrtContainerName), resource_name, &resource)
.ok(),
errors::NotFound("TRTEngineCacheResource not yet created"));
core::ScopedUnref unref_me(resource);
// Terminate the calibration if any.
if (resource->calib_ctx_) resource->calib_ctx_->TerminateCalibration();
// Serialize the engines and write them to file.
std::unique_ptr<WritableFile> file;
OP_REQUIRES_OK(ctx, ctx->env()->NewWritableFile(filename, &file));
auto writer = std::make_unique<io::RecordWriter>(file.get());
int num_serialized_engines = 0;
if (save_gpu_specific_engines_) {
// If user requests TRT engines export, recursively create
// requisite directories.
const char* export_trt_engines_env =
getenv("TF_TRT_EXPORT_TRT_ENGINES_PATH");
if (export_trt_engines_env) {
VLOG(1) << "Exporting TRT engines to directory: "
<< export_trt_engines_env;
OP_REQUIRES_OK(
ctx, ctx->env()->RecursivelyCreateDir(export_trt_engines_env));
}
for (const auto& pair : resource->cache_) {
// Ignore engines that failed to build.
const std::unique_ptr<EngineContext>& engine = pair.second;
if (!engine || !engine->GetCudaEngine()) continue;
TRTEngineInstance engine_instance;
// Add input shapes.
const std::vector<TensorShape>& engine_input_shapes = pair.first;
for (const TensorShape& shape : engine_input_shapes) {
shape.AsProto(engine_instance.add_input_shapes());
}
// Add the serialized engine.
TrtUniquePtrType<nvinfer1::IHostMemory> engine_data(
engine->GetCudaEngine()->serialize());
engine_instance.set_serialized_engine(engine_data->data(),
engine_data->size());
if (export_trt_engines_env) {
const std::string engine_filename =
std::string(export_trt_engines_env) + "/" + resource_name;
std::unique_ptr<WritableFile> engine_file;
OP_REQUIRES_OK(
ctx, ctx->env()->NewWritableFile(engine_filename, &engine_file));
OP_REQUIRES_OK(ctx, engine_file->Append(StringPiece(
static_cast<char*>(engine_data->data()),
engine_data->size())));
const std::string dims_filename =
std::string(export_trt_engines_env) + "/dims-" + resource_name;
std::unique_ptr<WritableFile> dims_file;
OP_REQUIRES_OK(
ctx, ctx->env()->NewWritableFile(dims_filename, &dims_file));
for (const TensorShape& shape : engine_input_shapes) {
OP_REQUIRES_OK(ctx,
dims_file->Append(StringPiece(shape.DebugString())));
}
}
OP_REQUIRES_OK(
ctx, writer->WriteRecord(engine_instance.SerializeAsString()));
++num_serialized_engines;
}
} else {
VLOG(1) << "TRT Engines are not serialized for op: " << resource_name;
}
VLOG(1) << "Serialized " << num_serialized_engines << " TRT engines for op "
<< resource_name << " on device " << ctx->device()->name()
<< " to file " << filename;
if (delete_resource_) {
VLOG(1) << "Destroying TRT engine cache resource for op " << resource_name
<< " on device " << ctx->device()->name();
OP_REQUIRES_OK(ctx,
ctx->resource_manager()->Delete<TRTEngineCacheResource>(
std::string(kTfTrtContainerName), resource_name));
}
}
private:
bool delete_resource_ = false;
bool save_gpu_specific_engines_ = true;
SerializeTRTResource(const SerializeTRTResource&) = delete;
void operator=(const SerializeTRTResource&) = delete;
};
REGISTER_KERNEL_BUILDER(Name("SerializeTRTResource").Device(DEVICE_GPU),
SerializeTRTResource);
} // namespace tensorrt
} // namespace tensorflow
#endif // GOOGLE_CUDA && GOOGLE_TENSORRT
tensorflow/compiler/tf2tensorrt/kernels/trt_engine_resource_ops_test.cc
+412
View File
@@ -0,0 +1,412 @@
/* Copyright 2019 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 <memory>
#include <string>
#include <utility>
#include <vector>
#include "absl/container/inlined_vector.h"
#include "absl/memory/memory.h"
#include "absl/strings/str_join.h"
#include "tensorflow/compiler/tf2tensorrt/common/datavec.h"
#include "tensorflow/compiler/tf2tensorrt/common/utils.h"
#include "tensorflow/compiler/tf2tensorrt/convert/utils.h"
#include "tensorflow/compiler/tf2tensorrt/utils/trt_engine_instance.pb.h" // NOLINT
#include "tensorflow/compiler/tf2tensorrt/utils/trt_logger.h"
#include "tensorflow/compiler/tf2tensorrt/utils/trt_lru_cache.h"
#include "tensorflow/core/common_runtime/device.h"
#include "tensorflow/core/common_runtime/device_factory.h"
#include "tensorflow/core/framework/fake_input.h"
#include "tensorflow/core/framework/node_def_builder.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/resource_mgr.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/framework/tensor_shape.pb.h"
#include "tensorflow/core/framework/tensor_types.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/kernels/ops_testutil.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/lib/io/record_reader.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/file_system.h"
#include "tensorflow/core/platform/path.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/tstring.h"
#include "tensorflow/core/platform/types.h"
#if GOOGLE_CUDA && GOOGLE_TENSORRT
namespace tensorflow {
namespace tensorrt {
struct TestParam {
nvinfer1::Dims dims;
bool dynamic_shape;
int n_inputs;
};
class TRTEngineResourceOpsTest
: public OpsTestBase,
public ::testing::WithParamInterface<TestParam> {
public:
TRTEngineResourceOpsTest() : param_(GetParam()) {}
protected:
void Reset() {
for (auto& temp : tensors_) {
delete temp;
}
for (auto& temp : managed_outputs_) {
delete temp;
}
tensors_.clear();
managed_outputs_.clear();
inputs_.clear();
}
ITensorProxyPtr NetworkWith1Input(nvinfer1::INetworkDefinition* network,
ITensorProxyPtr input) {
// Add a unary layer.
nvinfer1::IUnaryLayer* layer =
network->addUnary(*input->trt_tensor(), nvinfer1::UnaryOperation::kEXP);
EXPECT_NE(nullptr, layer);
return layer->getOutput(0);
}
// Constructs a network with two inputs, where the second input is a shape
// tensor. We take a slice of the first input with the size of the slice
// specified by the second input, assuming the first input is a 2D tensor.
// We then add the slice to itself to produce the output of the network.
ITensorProxyPtr NetworkWith2Inputs(nvinfer1::INetworkDefinition* network,
ITensorProxyPtr input) {
nvinfer1::Dims dims2{1, {2}};
ITensorProxyPtr input2 =
network->addInput(absl::StrCat(IONamePrefixes::kInputPHName, 1).c_str(),
nvinfer1::DataType::kINT32, dims2);
EXPECT_NE(nullptr, input2->trt_tensor());
nvinfer1::Dims start{2, {0, 0}};
nvinfer1::Dims stride{2, {1, 1}};
auto slice_layer =
network->addSlice(*input->trt_tensor(), start, stride, stride);
EXPECT_NE(nullptr, slice_layer);
slice_layer->setInput(2, *input2->trt_tensor());
ITensorProxyPtr sliced_input = slice_layer->getOutput(0);
EXPECT_NE(nullptr, sliced_input->trt_tensor());
auto layer = network->addElementWise(*sliced_input->trt_tensor(),
*sliced_input->trt_tensor(),
nvinfer1::ElementWiseOperation::kSUM);
EXPECT_NE(nullptr, layer);
return layer->getOutput(0);
}
TrtUniquePtrType<nvinfer1::ICudaEngine> CreateTRTEngine() {
TrtUniquePtrType<nvinfer1::IBuilder> builder(
nvinfer1::createInferBuilder(logger_));
TrtUniquePtrType<nvinfer1::INetworkDefinition> network;
#if IS_TRT_VERSION_GE(8, 0, 0, 0)
network =
TrtUniquePtrType<nvinfer1::INetworkDefinition>(builder->createNetworkV2(
1U << static_cast<int>(
nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH)));
#else
network =
TrtUniquePtrType<nvinfer1::INetworkDefinition>(builder->createNetworkV2(
1U << static_cast<int>(
nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH)));
#endif
// Add the input.
nvinfer1::Dims dims = this->param_.dims;
if (this->param_.dynamic_shape) {
std::fill(dims.d, dims.d + dims.nbDims, -1);
}
const std::string in_name = StrCat(IONamePrefixes::kInputPHName, 0);
ITensorProxyPtr input =
network->addInput(in_name.c_str(), nvinfer1::DataType::kFLOAT, dims);
EXPECT_NE(nullptr, input->trt_tensor());
// Mark the output.
ITensorProxyPtr output =
this->param_.n_inputs == 1
? this->NetworkWith1Input(network.get(), input)
: this->NetworkWith2Inputs(network.get(), input);
output->setName("output");
network->markOutput(*output->trt_tensor());
// Build the engine
TrtUniquePtrType<nvinfer1::IBuilderConfig> builder_config(
builder->createBuilderConfig());
builder_config->setMaxWorkspaceSize(1 << 10);
builder->setMaxBatchSize(1);
if (this->param_.dynamic_shape) {
TrtShapeOptimizationProfile profile;
profile.SetShapeTensorMask(network.get());
const int n_input = param_.n_inputs;
// Set the input mask to true (no resource input)
std::vector<bool> input_mask(n_input, true);
profile.SetInputMask(input_mask);
// The for loop defines three optimization profiles for the network.
for (int i = 1; i <= 3; i++) {
std::vector<TensorShape> shape_vec(n_input);
// Define a shape with all dimensions set to 3*i.
std::vector<int> dimvec(this->param_.dims.nbDims, 3 * i);
TensorShape shape;
TF_CHECK_OK(
TensorShapeUtils::MakeShape(dimvec.data(), dimvec.size(), &shape));
const ITensorProxyPtr input = network->getInput(0);
const char* name = input->getName();
VLOG(2) << "Defining profile for input " << name;
shape_vec[0] = shape;
if (this->param_.n_inputs == 2) {
// The shape of the shape tensor.
TF_CHECK_OK(TensorShapeUtils::MakeShape(
std::vector<int32>{param_.dims.nbDims}, &shape));
shape_vec[1] = shape;
// Values of the shape tensor
Tensor shape_tensor(DT_INT32, shape);
// Define shape values {1, i}, where 1 is the value of the first dim,
// and i is the value of the second dimension.
std::vector<int32> vals{1, i};
std::copy_n(vals.data(), vals.size(),
shape_tensor.flat<int32_t>().data());
DataVec shape_values{{"one", {}}, {"two", shape_tensor}};
TF_CHECK_OK(profile.CollectShapeValues(shape_values));
} else {
TF_CHECK_OK(profile.CollectShapeValues({{"one", {}}}));
}
profile.AddShape(shape_vec);
}
std::vector<PartialTensorShape> input_partial_shapes;
TF_CHECK_OK(GetNetworkInputShapes(network.get(), &input_partial_shapes));
profile.InitProfiles(input_partial_shapes, ProfileStrategy::kOptimal);
// Configure and build engine
TF_CHECK_OK(profile.ConfigureBuilder(builder.get(), builder_config.get(),
network.get()));
}
VLOG(2) << "ConfigureBuilder Finished";
TrtUniquePtrType<nvinfer1::ICudaEngine> engine(
builder->buildEngineWithConfig(*network, *builder_config));
VLOG(2) << "Engine constructed";
EXPECT_NE(nullptr, engine);
return engine;
}
Logger& logger_ = *Logger::GetLogger();
TestParam param_;
};
#if IS_TRT_VERSION_GE(7, 1, 3, 0)
constexpr std::array<TestParam, 3> TestParameters = {
TestParam{nvinfer1::Dims{1, {1}}, false, 1},
TestParam{nvinfer1::Dims{1, {1}}, true, 1},
TestParam{nvinfer1::Dims{2, {3, 3}}, true, 2}};
#else
constexpr std::array<TestParam, 2> TestParameters = {
TestParam{nvinfer1::Dims{1, {1}}, false, 1},
TestParam{nvinfer1::Dims{1, {1}}, true, 1}};
#endif
INSTANTIATE_TEST_CASE_P(EngineResourceOpsTestInstantiation,
TRTEngineResourceOpsTest,
::testing::ValuesIn(TestParameters));
TEST_P(TRTEngineResourceOpsTest, Basic) {
// Create the GPU device.
std::unique_ptr<Device> device(
DeviceFactory::NewDevice("GPU", {}, "/job:worker/replica:0/task:0"));
ResourceMgr* rm = device->resource_manager();
SetDevice(DEVICE_GPU, std::move(device));
// Create a resource handle.
const string container(kTfTrtContainerName);
const string resource_name = "myresource";
Reset();
TF_ASSERT_OK(NodeDefBuilder("op", "CreateTRTResourceHandle")
.Attr("resource_name", resource_name)
.Finalize(node_def()));
TF_ASSERT_OK(InitOp());
TF_ASSERT_OK(RunOpKernel());
ResourceHandle handle =
context_->mutable_output(0)->scalar<ResourceHandle>()();
// Check that a resource hasn't been created yet.
TRTEngineCacheResource* resource = nullptr;
EXPECT_TRUE(
errors::IsNotFound(rm->Lookup(container, resource_name, &resource)));
// Create a resource and use an empty file to initialize the resource.
Reset();
Env* env = Env::Default();
const string filename = io::JoinPath(testing::TmpDir(), "trt_engine_file");
{
std::unique_ptr<WritableFile> file;
TF_ASSERT_OK(env->NewWritableFile(filename, &file));
}
TF_ASSERT_OK(NodeDefBuilder("op", "InitializeTRTResource")
.Input(FakeInput(DT_RESOURCE))
.Input(FakeInput(DT_STRING))
.Attr("max_cached_engines_count", 1)
.Finalize(node_def()));
TF_ASSERT_OK(InitOp());
AddInputFromArray<ResourceHandle>(TensorShape({}), {handle});
AddInputFromArray<tstring>(TensorShape({}), {filename});
TF_ASSERT_OK(RunOpKernel());
// Check that the resource is registered with the resource manager and the
// cache of the resource is empty.
EXPECT_TRUE(rm->Lookup(container, resource_name, &resource).ok());
EXPECT_EQ(0, resource->cache_.size());
// Create an engine and add it to the cache of the resource.
TrtUniquePtrType<nvinfer1::ICudaEngine> engine = CreateTRTEngine();
ExecutionContext context = ExecutionContext::Create(engine.get());
std::vector<TensorShape> engine_input_shape(1);
TF_ASSERT_OK(DimsAdapter(param_.dims).TensorShape(&(engine_input_shape[0])));
if (param_.n_inputs > 1) {
engine_input_shape.push_back(TensorShape({1, 1}));
}
resource->cache_.emplace(
engine_input_shape,
std::make_unique<EngineContext>(std::move(engine), std::move(context)));
// Check that the resource has multiple references before it is unregistered
// from the resource manager.
EXPECT_FALSE(resource->RefCountIsOne());
// Serialize the engine to a file and unregistered the resource from the
// resource manager.
Reset();
TF_ASSERT_OK(NodeDefBuilder("op", "SerializeTRTResource")
.Attr("delete_resource", true)
.Input(FakeInput(DT_STRING))
.Input(FakeInput(DT_STRING))
.Finalize(node_def()));
TF_ASSERT_OK(InitOp());
AddInputFromArray<tstring>(TensorShape({}), {resource_name});
AddInputFromArray<tstring>(TensorShape({}), {filename});
TF_ASSERT_OK(RunOpKernel());
// Check that the resource now has only one reference. Detach the reference
// to the resource to destroy the resource.
EXPECT_TRUE(resource->RefCountIsOne());
resource->Unref();
// Check that unregistering the resource from the resource manager returns
// an error as the resource has already been unregistered.
Reset();
TF_ASSERT_OK(NodeDefBuilder("op", "DestroyResourceOp")
.Attr("ignore_lookup_error", false)
.Input(FakeInput(DT_RESOURCE))
.Finalize(node_def()));
TF_ASSERT_OK(InitOp());
AddInputFromArray<ResourceHandle>(TensorShape({}), {handle});
EXPECT_TRUE(errors::IsNotFound(RunOpKernel()));
// Verify the file for the serialized engine.
std::unique_ptr<RandomAccessFile> file;
TF_ASSERT_OK(env->NewRandomAccessFile(filename, &file));
auto reader = std::make_unique<io::RecordReader>(file.get());
uint64 offset = 0;
tstring record;
TF_ASSERT_OK(reader->ReadRecord(&offset, &record));
TRTEngineInstance engine_instance;
engine_instance.ParseFromString(record);
EXPECT_EQ(param_.n_inputs, engine_instance.input_shapes_size());
EXPECT_EQ(param_.dims.nbDims, engine_instance.input_shapes(0).dim_size());
for (int i = 0; i < param_.dims.nbDims; i++) {
EXPECT_EQ(param_.dims.d[i], engine_instance.input_shapes(0).dim(i).size());
}
EXPECT_TRUE(errors::IsOutOfRange(reader->ReadRecord(&offset, &record)));
// Recreate the resource and use the file with the serialized engine to
// initialize the resource.
Reset();
TF_ASSERT_OK(NodeDefBuilder("op", "InitializeTRTResource")
.Input(FakeInput(DT_RESOURCE))
.Input(FakeInput(DT_STRING))
.Attr("max_cached_engines_count", 1)
.Finalize(node_def()));
TF_ASSERT_OK(InitOp());
AddInputFromArray<ResourceHandle>(TensorShape({}), {handle});
AddInputFromArray<tstring>(TensorShape({}), {filename});
TF_ASSERT_OK(RunOpKernel());
// Check that the resource is registered with the resource manager again and
// the cache of the resource is not empty.
EXPECT_TRUE(rm->Lookup(container, resource_name, &resource).ok());
EXPECT_EQ(1, resource->cache_.size());
if (this->param_.dynamic_shape) {
EXPECT_EQ(3, resource->profiles_.GetNumProfiles());
EXPECT_EQ(3, resource->cache_.begin()->second->GetNumContexts());
if (this->param_.n_inputs == 1) {
// Check if profiles are restored correctly.
std::vector<TensorShape> shapes(1);
// We create a shape vector that matches only profile 1.
TF_CHECK_OK(
TensorShapeUtils::MakeShape(std::vector<int32>{6}, &shapes[0]));
EXPECT_EQ(1, resource->profiles_.GetProfileNumber(shapes));
} else {
// Check if shape values are restored correctly.
std::vector<TensorShape> shapes(2);
// We create a shape vector that matches only profile 2.
TF_CHECK_OK(
TensorShapeUtils::MakeShape(std::vector<int32>{9, 9}, &shapes[0]));
TF_CHECK_OK(
TensorShapeUtils::MakeShape(std::vector<int32>{2}, &shapes[1]));
Tensor shape_tensor(DT_INT32, shapes[1]);
std::vector<int32> vals{1, 3};
std::copy_n(vals.data(), vals.size(),
shape_tensor.flat<int32_t>().data());
// DataVec names are not in used CollectShapeValues, only the order
// matters.
DataVec shape_values{{"one", {}}, {"two", shape_tensor}};
TF_CHECK_OK(resource->profiles_.CollectShapeValues(shape_values));
EXPECT_EQ(2, resource->profiles_.GetProfileNumber(shapes));
}
}
// Check that the resource has multiple references before it is unregistered
// from the resource manager.
EXPECT_FALSE(resource->RefCountIsOne());
// Unregister the resource from the resource manager two times, expect that
// the second time produces an error.
Reset();
TF_ASSERT_OK(NodeDefBuilder("op", "DestroyResourceOp")
.Attr("ignore_lookup_error", false)
.Input(FakeInput(DT_RESOURCE))
.Finalize(node_def()));
TF_ASSERT_OK(InitOp());
AddInputFromArray<ResourceHandle>(TensorShape({}), {handle});
TF_ASSERT_OK(RunOpKernel());
EXPECT_TRUE(errors::IsNotFound(RunOpKernel()));
// Check that the resource now has only one reference. Detach the reference
// to the resource to destroy resource.
EXPECT_TRUE(resource->RefCountIsOne());
resource->Unref();
}
} // namespace tensorrt
} // namespace tensorflow
#endif // GOOGLE_CUDA && GOOGLE_TENSORRT