diff --git a/tensorflow/compiler/tf2tensorrt/convert/convert_nodes.cc b/tensorflow/compiler/tf2tensorrt/convert/convert_nodes.cc index e791ff9ff60..4aca5efd6c9 100644 --- a/tensorflow/compiler/tf2tensorrt/convert/convert_nodes.cc +++ b/tensorflow/compiler/tf2tensorrt/convert/convert_nodes.cc @@ -404,6 +404,18 @@ Status GetTrtBroadcastShape(const TRT_TensorOrWeights& operand_l, // Compare broadcast feasibility if (check_feasibility) { for (int i = 0; i < broadcast_num_dims; ++i) { + if (!use_implicit_batch && (output_l[i] == -1 || output_r[i] == -1)) { + // If the condition is true then we are in explicit batch mode and (at + // least) one of the input dimensions are unknown. In other words we + // are in dynamic shape mode. During conversion time we only see -1 for + // the unknown shapes, therefore we cannot decide on the feasibility of + // broadcast over the unknown dimensions. Therefore we just continue for + // the next dimension. In dynamic shape mode TRT can only check the + // feasibility of the broadcast when the actual input dimensions are + // specified by SetTrtEngineInputs and the inference job is launched by + // TrtEnque. + continue; + } if ((output_l[i] != output_r[i]) && (output_l[i] != 1) && (output_r[i] != 1)) { return errors::InvalidArgument("Infeasible broadcast scheme (", @@ -3888,11 +3900,26 @@ Status ConvertBiasAdd(OpConverterParams* params) { nvinfer1::Dims input_shape = inputs.at(0).GetTrtDims(); nvinfer1::Dims bias_shape = inputs.at(1).GetTrtDims(); - // If the input is NCHW, then we need to unsqueeze the bias such that its last - // dimensions are 1s (and the first dimension is C). + // The bias input arg is a 1-D tensor with length C. If the input is NCHW, + // then we need to unsqueeze the bias such that its shape is [1, C, 1, 1]. if (data_format == "NCHW") { - bias_shape.nbDims = input_shape.nbDims; - std::fill(bias_shape.d + 1, bias_shape.d + bias_shape.nbDims, 1); + if (params->use_implicit_batch) { + // The batch dim is not included in implicit batch mode, so the shape of + // the bias tensor is [C, 1, 1]. + bias_shape.nbDims = input_shape.nbDims; + std::fill(bias_shape.d + 1, bias_shape.d + bias_shape.nbDims, 1); + } else { + // In explicit batch mode we create a tensor with shape [1, C, 1, 1]. + std::vector bias_shape_vec(bias_shape.d, + bias_shape.d + bias_shape.nbDims); + // Insert 1 before for batch dim + bias_shape_vec.insert(bias_shape_vec.begin(), 1); + // Trail with 1s to match input_shape size + bias_shape_vec.insert(bias_shape_vec.end(), + input_shape.nbDims - bias_shape_vec.size(), 1); + TF_RETURN_IF_ERROR( + TensorShapeArrayToTrtDims(bias_shape_vec, &bias_shape)); + } } else { // Next, broadcast the bias across the input. TF_RETURN_IF_ERROR(GetTrtBroadcastShape(inputs.at(0), inputs.at(1), diff --git a/tensorflow/compiler/tf2tensorrt/convert/convert_nodes_test.cc b/tensorflow/compiler/tf2tensorrt/convert/convert_nodes_test.cc index d4badd1cc03..c7b69818d3d 100644 --- a/tensorflow/compiler/tf2tensorrt/convert/convert_nodes_test.cc +++ b/tensorflow/compiler/tf2tensorrt/convert/convert_nodes_test.cc @@ -1862,6 +1862,13 @@ INSTANTIATE_TEST_CASE_P( ::testing::Values(DT_FLOAT), ::testing::Values(TrtPrecisionMode::FP32))); +// Base class for tests that need to be tested for both FP32 and FP16. +class OpConverterTest2 : public ParameterizedOpConverterTestBase {}; +INSTANTIATE_TEST_CASE_P( + OpConvTestInstantiation, OpConverterTest2, + ::testing::Combine(::testing::ValuesIn(ValidTrtModes), + ::testing::Values(DT_FLOAT, DT_HALF), + ::testing::Values(TrtPrecisionMode::FP32))); template void CopyTensorElements(const Tensor& tensor, protobuf::RepeatedField* out) { out->Clear(); @@ -2396,91 +2403,70 @@ TEST_F(OpConverterTest, ConvertBatchMatMul) { TestMatMulHelper(this, get_batch_matmul_nodedef, "BatchMatMul"); } -template -void TestConvertBiasAdd(OpConverterTest* test) { +TEST_P(OpConverterTest2, ConvertBiasAdd) { + // Note that kINT32 is not supported by IScaleLayer, so we don't test + // DT_INT32 type here. DT_FLOAT and DT_HALF are tested. // Get the NodeDef for BiasAdd. - auto get_biasadd_nodedef = [](const string& data_format) -> NodeDef { + auto get_biasadd_nodedef = [](const string& data_format, + DataType tf_dtype) -> NodeDef { Scope s = Scope::NewRootScope(); - auto input = ops::Placeholder(s.WithOpName("input"), dtype); - auto weights = ops::Placeholder(s.WithOpName("weights"), dtype); + auto input = ops::Placeholder(s.WithOpName("input"), tf_dtype); + auto weights = ops::Placeholder(s.WithOpName("weights"), tf_dtype); const auto biasadd_attrs = ops::BiasAdd::DataFormat(data_format); auto biasadd = ops::BiasAdd(s.WithOpName("my_biasadd"), input, weights, biasadd_attrs); return biasadd.operation.node()->def(); }; - typedef typename EnumToDataType::Type CType; for (const string& data_format : {"NHWC", "NCHW"}) { for (const int trt_input_rank : {1, 2, 3, 4}) { - test->Reset(); - NodeDef node_def = get_biasadd_nodedef(data_format); + Reset(); + NodeDef node_def = get_biasadd_nodedef(data_format, tf_dtype); // Add input, dims_array will be like {2, 1, ..., 1, 3} - std::vector dims_array(trt_input_rank, 1); + std::vector dims_array(trt_input_rank + 1, 1); if (trt_input_rank == 1) { - dims_array[0] = (data_format == "NHWC" ? 3 : 2); + dims_array[1] = (data_format == "NHWC" ? 3 : 2); } else { - dims_array[0] = 2; - dims_array[trt_input_rank - 1] = 3; + dims_array[1] = 2; + dims_array[trt_input_rank] = 3; } - test->AddTestTensor("input", dims_array, /*batch_size=*/1, - TfDataTypeToTrt(dtype)); - - // Add bias weights. - const int channel_size = (data_format == "NHWC" ? 3 : 2); - std::vector bias(channel_size); - for (int i = 0; i < channel_size; ++i) { - bias[i] = CType(i + 1); // bias will be {1, 2, 3, ...} - } - test->AddTestWeights("weights", {channel_size}, bias); - - // Run the conversion. - test->RunValidationAndConversion(node_def); - TRT_TensorOrWeights output; - TF_EXPECT_OK(test->GetTensorOrWeights("my_biasadd", &output)); - ASSERT_TRUE(output.is_tensor()); - ExpectTrtDimsEqualsArray(dims_array, output.tensor()->getDimensions()); - - // Build and run the engine. const int num_input = TrtTensorDimsNumElements(GetTestDims(dims_array)); ASSERT_EQ(trt_input_rank > 1 ? 6 : (data_format == "NHWC" ? 3 : 2), num_input); + std::vector input_data(num_input, 0); + + AddTestTensor("input", dims_array, input_data); + + const int channel_size = (data_format == "NHWC" ? 3 : 2); + std::vector bias(channel_size); + for (int i = 0; i < channel_size; ++i) { + bias[i] = i + 1; // bias will be {1, 2, 3, ...} + } + AddTestWeights("weights", {channel_size}, bias, tf_dtype); + + // Build and run the engine. + std::vector output_data; - const DataVec input_data{ - {"input", test->ConstructTensor(num_input, CType(0))}}; - DataVec output_data{ - {"my_biasadd", test->ConstructTensor(num_input)}}; - TF_EXPECT_OK(test->BuildAndRun(input_data, &output_data)); if (trt_input_rank == 1) { if (data_format == "NHWC") { - EXPECT_THAT(GetSpanForData(output_data[0]), - ElementsAre(CType(1), CType(2), CType(3))); + output_data = {1, 2, 3}; } else { - EXPECT_THAT(GetSpanForData(output_data[0]), - ElementsAre(CType(1), CType(2))); + output_data = {1, 2}; } } else { if (data_format == "NHWC") { - EXPECT_THAT(GetSpanForData(output_data[0]), - ElementsAre(CType(1), CType(2), CType(3), CType(1), - CType(2), CType(3))); + output_data = {1, 2, 3, 1, 2, 3}; } else { - EXPECT_THAT(GetSpanForData(output_data[0]), - ElementsAre(CType(1), CType(1), CType(1), CType(2), - CType(2), CType(2))); + output_data = {1, 1, 1, 2, 2, 2}; } } + TestOpConverter("my_biasadd", node_def, dims_array, Status::OK(), + Status::OK(), ElementsAreArray(output_data)); } } } -TEST_F(OpConverterTest, ConvertBiasAdd) { - // OK. Note that kINT32 is not supported by IScaleLayer, so we don't test - // DT_INT32 type here. - TestConvertBiasAdd(this); - TestConvertBiasAdd(this); -} - template NodeDef GetBinaryOpNodeDef(const string& input_name_l, const string& input_name_r, DataType dtype) {