From 1033515e9ede0996f0213da293d0a7cff6dc094c Mon Sep 17 00:00:00 2001 From: Tamas Bela Feher Date: Thu, 7 May 2020 18:17:52 +0200 Subject: [PATCH] Add TF-TRT converter for Shape op. The unit test includes test cases when the TensorRT network is called with no input tensor. The helper routines are adjusted to handle this case. --- .../tf2tensorrt/convert/convert_nodes.cc | 35 ++++++++++ .../tf2tensorrt/convert/convert_nodes_test.cc | 69 +++++++++++++++++-- .../utils/trt_shape_optimization_profiles.cc | 17 +++-- 3 files changed, 107 insertions(+), 14 deletions(-) diff --git a/tensorflow/compiler/tf2tensorrt/convert/convert_nodes.cc b/tensorflow/compiler/tf2tensorrt/convert/convert_nodes.cc index 369b339d01a..6674081011f 100644 --- a/tensorflow/compiler/tf2tensorrt/convert/convert_nodes.cc +++ b/tensorflow/compiler/tf2tensorrt/convert/convert_nodes.cc @@ -2410,6 +2410,40 @@ Status ConvertTranspose(OpConverterParams* params) { return Status::OK(); } +Status ConvertShape(OpConverterParams* params) { + const auto& inputs = params->inputs; + TF_RETURN_IF_ERROR( + CheckInputsWeights(*params, {{"input", TrtInputArg::kBoth}})); + if (params->use_implicit_batch) { + return errors::Unimplemented( + "Shape is only supported for explicit batch mode."); + } + if (HasStaticShape(inputs.at(0).GetTrtDims())) { + if (params->validation_only) return Status::OK(); + nvinfer1::Dims input_dims = inputs.at(0).GetTrtDims(); + nvinfer1::Dims output_dims{1, {input_dims.nbDims}}; + // Create a const node with the values of output_dims + TRT_ShapedWeights weight = params->weight_store->GetTempWeights( + nvinfer1::DataType::kINT32, output_dims); + int32* values_ptr = static_cast(weight.GetValues()); + std::copy(input_dims.d, input_dims.d + input_dims.nbDims, values_ptr); + auto output = params->converter->CreateConstantLayer(weight, output_dims); + params->outputs->push_back(TRT_TensorOrWeights(output)); + return Status::OK(); + } +#if IS_TRT_VERSION_GE(6, 0, 0, 0) + if (params->validation_only) return Status::OK(); + nvinfer1::IShapeLayer* shape_layer = + params->converter->network()->addShape(*inputs.at(0).tensor()); + TFTRT_RETURN_ERROR_IF_NULLPTR(shape_layer, params->node_def.name()); + params->outputs->push_back(TRT_TensorOrWeights(shape_layer->getOutput(0))); + return Status::OK(); +#else + return errors::Unavailable( + "Shape op conversion requires TensorRT 6 or above"); +#endif +} + Status ConvertReshape(OpConverterParams* params) { const auto& inputs = params->inputs; TF_RETURN_IF_ERROR( @@ -5958,6 +5992,7 @@ static void RegisterValidatableOpConverters( (*registration)[pool_op_type] = ConvertPool3D; } #endif + (*registration)["Shape"] = ConvertShape; (*registration)["Rsqrt"] = ConvertRsqrt; (*registration)["Slice"] = ConvertSlice; (*registration)["Softmax"] = ConvertSoftmax; diff --git a/tensorflow/compiler/tf2tensorrt/convert/convert_nodes_test.cc b/tensorflow/compiler/tf2tensorrt/convert/convert_nodes_test.cc index 52d05ff8225..9ca1c8c4c9f 100644 --- a/tensorflow/compiler/tf2tensorrt/convert/convert_nodes_test.cc +++ b/tensorflow/compiler/tf2tensorrt/convert/convert_nodes_test.cc @@ -1781,7 +1781,8 @@ class ParameterizedOpConverterTestBase void BuildAndRun(const string& name, const std::vector>& expected_output_dims, const Status& expected_runtime_status, - const std::vector>>& matcher) { + const std::vector>>& matcher, + const std::vector& out_tf_types = {}) { TensorShape shape; const int n_output = expected_output_dims.size(); ASSERT_EQ(n_output, matcher.size()); @@ -1790,12 +1791,14 @@ class ParameterizedOpConverterTestBase TF_EXPECT_OK( TensorShapeUtils::MakeShape(expected_output_dims[i], &shape)); string out_name = (n_output == 1) ? name : StrCat(name, ":", i); - InputOutputData data{out_name, - ConstructTensor(shape.num_elements(), 0, tf_type)}; + DataType out_tf_type = + out_tf_types.size() > i ? out_tf_types[i] : tf_type; + InputOutputData data{ + out_name, ConstructTensor(shape.num_elements(), 0, out_tf_type)}; output_data.push_back(data); } - ASSERT_FALSE(input_data_.empty()); - const int batch_size = input_data_[0].tensor.shape().dim_size(0); + const int batch_size = + input_data_.empty() ? 1 : input_data_[0].tensor.shape().dim_size(0); Status stat = OpConverterTest::BuildAndRun(input_data_, &output_data, batch_size); ASSERT_EQ(expected_runtime_status.ok(), stat.ok()) @@ -1820,13 +1823,15 @@ class ParameterizedOpConverterTestBase const std::vector& expected_output_dims, const Status& expected_conversion_status, const Status& expected_runtime_status, - const Matcher>& matcher) { + const Matcher>& matcher, + const std::vector& out_tf_types = {}) { RunValidationAndConversion(node_def, expected_conversion_status, name.c_str(), expected_output_dims); if (expected_conversion_status.ok()) { BuildAndRun(name, std::vector>({expected_output_dims}), expected_runtime_status, - std::vector>>({matcher})); + std::vector>>({matcher}), + out_tf_types); } } @@ -2169,6 +2174,56 @@ TEST_F(OpConverterTest, ConvertReshape) { } } +TEST_P(OpConverterTest1, ConvertShape) { + // Get the NodeDef for Shape op. + Scope s = Scope::NewRootScope(); + auto input = ops::Placeholder(s.WithOpName("input"), tf_type); + auto shape = ops::Shape(s.WithOpName("my_shape"), input); + const NodeDef& node_def = shape.operation.node()->def(); + + Status conversion_status = + (trt_mode == TrtTestMode::kImplicitBatch) + ? errors::Unimplemented( + "Shape is only supported for explicit batch mode.") + : Status::OK(); + std::vector test_params = { + TestParamBase{{1, 2, 3}, {}, {3}, {}, conversion_status}, + // Add input as weight (we use non empty param ({1}) to trigger this). + TestParamBase{{1, 2, 3}, {}, {3}, {1}, conversion_status}, + }; + + auto input_is_weight = [](const TestParamBase p) { return !p.param.empty(); }; + for (auto p : test_params) { + SCOPED_TRACE(p); + Reset(); + // Number of elements of the input tensor. We leave it 0 in case we do + // not need to add an input tensor. This happens in explicit batch mode: the + // shape is known at conversion time and therefore the shape is added to the + // network as a constant layer. (In this case the single node network that + // we use for the unit test will have no actual input tensor when converted + // to a TensorRT network.) + int n_elements = 0; + // In explicit batch mode the shape is known at conversion time and + // therefore the shape is added to the network as a constant layer. As + // a result, the single node network that we use for this unit test will + // have no actual input tensor when converted to a TensorRT network. + if (input_is_weight(p) || trt_mode != TrtTestMode::kExplicitBatch) { + // Calculate the number of elements for adding input data. + n_elements = std::accumulate(p.input_dims.begin(), p.input_dims.end(), 1, + std::multiplies()); + } + std::vector input_val(n_elements, 1); + if (!input_is_weight(p)) { + AddTestTensor("input", p.input_dims, input_val); + } else { + AddTestWeights("input", p.input_dims, input_val, tf_type); + } + TestOpConverter("my_shape", node_def, p.expected_output_dims, p.status, + p.runtime_status, ElementsAreArray(p.input_dims), + {DT_INT32}); + } +} + // Helper function for testing MatMul and BatchMatMul // get_matmul corresponds to the function used to generate the node. It should // accept (DataType, transpose_a, transpose_b) as parameters. diff --git a/tensorflow/compiler/tf2tensorrt/utils/trt_shape_optimization_profiles.cc b/tensorflow/compiler/tf2tensorrt/utils/trt_shape_optimization_profiles.cc index 70a0a9a7b65..2f31865751f 100644 --- a/tensorflow/compiler/tf2tensorrt/utils/trt_shape_optimization_profiles.cc +++ b/tensorflow/compiler/tf2tensorrt/utils/trt_shape_optimization_profiles.cc @@ -18,6 +18,7 @@ limitations under the License. #include #include +#include "absl/algorithm/container.h" #include "tensorflow/compiler/tf2tensorrt/convert/utils.h" #if GOOGLE_CUDA && GOOGLE_TENSORRT @@ -35,14 +36,16 @@ void TrtShapeOptimizationProfile::InitProfiles() { << "for each input (min=opt=max)."; } for (auto& shape_vec : input_shapes_) { - std::vector dimvec; - for (auto& shape : shape_vec) { - dimvec.push_back(TensorShapeToTrtDims(shape, false)); + if (!shape_vec.empty()) { + std::vector dimvec(shape_vec.size()); + absl::c_transform(shape_vec, dimvec.begin(), [](TensorShape shape) { + return TensorShapeToTrtDims(shape, false); + }); + // Set min=opt=max. + OptimizationProfileConfig profConfig{dimvec, dimvec, dimvec}; + profiles_.push_back(std::move(profConfig)); + VLOG(1) << "Created profile " << profiles_.back().DebugString(); } - // We set min=opt=max. - OptimizationProfileConfig profConfig{dimvec, dimvec, dimvec}; - profiles_.push_back(std::move(profConfig)); - VLOG(1) << "Created profile " << profiles_.back().DebugString(); } }