experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Improve message when engine creation fails. Remove spammy warning

Browse Source
This commit is contained in:
Trevor Morris 2019-03-12 09:32:34 -07:00
parent fe020fccbb
commit 3764ac1d6c
1 changed files with 9 additions and 14 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

23
tensorflow/compiler/tf2tensorrt/kernels/trt_engine_op.cc
View File

@ -376,9 +376,9 @@ void TRTEngineOp::ComputeAsync(OpKernelContext* ctx,
}
EngineContext* engine_context = GetEngine(input_shapes, ctx);
if (!engine_context->cuda_engine) {
LOG(WARNING) << "Engine retrieval for input shapes: "
<< TensorShapeUtils::ShapeListString(input_shapes)
<< " failed. Running native segment for " << name();
VLOG(1) << "Engine retrieval for input shapes: "
<< TensorShapeUtils::ShapeListString(input_shapes)
<< " failed. Running native segment for " << name();
ExecuteNativeSegment(ctx, helper);
return;
}
@ -625,17 +625,12 @@ EngineContext* TRTEngineOp::GetEngine(
partial_shapes, &logger, allocator, calibrator_.get(), &engine,
use_calibration_, &convert_successfully);
if (!status.ok()) {
if (convert_successfully) {
// This means it fail to build the engine even when the network is built
// successfully, probably due to internal issues. In this case we don't
// retry in the future.
cache.emplace(engine_input_shapes, absl::make_unique<EngineContext>());
}
LOG(WARNING) << "Engine creation for batch size " << batch_size
<< " failed " << status;
// Store an empty engine here so we don't try to build the same engine
// again.
cache.emplace(engine_input_shapes, empty_context);
LOG(WARNING) << "Engine creation for " << name() << " failed. "
<< "The native segment will be used instead. "
<< "Reason: " << status;
// Store an empty engine in the cache for these input shapes so we don't
// try to build the same failing engine again.
cache.emplace(engine_input_shapes, absl::make_unique<EngineContext>());
return &empty_context;
}
VLOG(1) << "Conversion is done";