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

Cleanup trt engine op definition.

Browse Source 
PiperOrigin-RevId: 238484967
This commit is contained in:
Guangda Lai 2019-03-14 11:35:35 -07:00 committed by TensorFlower Gardener
parent 852d3364e6
commit 7239d65880
1 changed files with 7 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

17
tensorflow/compiler/tf2tensorrt/ops/trt_engine_op.cc
View File

@ -24,12 +24,9 @@ limitations under the License.
namespace tensorflow {
namespace shape_inference {
extern Status TRTEngineOpShapeInference(InferenceContext* c);
}
// NOTE: please try NOT to add/modify/remove attributes or inputs/outputs to the
// list below, this will break backward compatibility!
// NOTE: when making changes please follow
// https://www.tensorflow.org/guide/extend/op#backwards_compatibility to not
// break backward compatibility.
//
// TODO(laigd): consider making this op stateful. The only problem is it uses TF
// function which has to be stateless, but we can use function library as the
@ -41,8 +38,6 @@ REGISTER_OP("TRTEngineOp")
.Attr("segment_funcdef_name: string")
.Attr("InT: list({int8,float16,float32,int32})")
.Attr("OutT: list({int8,float16,float32,int32})")
.Attr("static_engine: bool = true")
.Attr("fixed_input_size: bool = true")
.Attr("cached_engine_batches: list(int) >= 0 = []")
.Attr("max_cached_engines_count: int = 1")
.Attr("workspace_size_bytes: int")
@ -57,8 +52,10 @@ REGISTER_OP("TRTEngineOp")
// implementation, we do require all input tensor to carry the same batch
// size, but this could change in the future). Hence we disable shape
// inference function as a workaround.
// .SetShapeFn(shape_inference::TRTEngineOpShapeInference);
.SetShapeFn(shape_inference::UnknownShape);
.SetShapeFn(shape_inference::UnknownShape)
// Deprecated attributes.
.Attr("fixed_input_size: bool = true")
.Attr("static_engine: bool = true");
} // namespace tensorflow
#endif // GOOGLE_TENSORRT