Merge pull request #39156 from tfeher:trt_biasadd_dynamic_shape

PiperOrigin-RevId: 313402866
Change-Id: Iad1d5701f727973c8f892e20be090a2570baf1a9
This commit is contained in:
TensorFlower Gardener 2020-05-27 09:27:11 -07:00
commit be4d17c088
2 changed files with 70 additions and 57 deletions

View File

@ -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 (",
@ -3895,11 +3907,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<int> 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),

View File

@ -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 <typename T>
void CopyTensorElements(const Tensor& tensor, protobuf::RepeatedField<T>* out) {
out->Clear();
@ -2396,91 +2403,70 @@ TEST_F(OpConverterTest, ConvertBatchMatMul) {
TestMatMulHelper(this, get_batch_matmul_nodedef, "BatchMatMul");
}
template <DataType dtype>
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<dtype>::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<int32> dims_array(trt_input_rank, 1);
std::vector<int32> 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<CType> bias(channel_size);
for (int i = 0; i < channel_size; ++i) {
bias[i] = CType(i + 1); // bias will be {1, 2, 3, ...}
}
test->AddTestWeights<CType>("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<float> input_data(num_input, 0);
AddTestTensor("input", dims_array, input_data);
const int channel_size = (data_format == "NHWC" ? 3 : 2);
std::vector<float> 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<float> output_data;
const DataVec input_data{
{"input", test->ConstructTensor<CType>(num_input, CType(0))}};
DataVec output_data{
{"my_biasadd", test->ConstructTensor<CType>(num_input)}};
TF_EXPECT_OK(test->BuildAndRun(input_data, &output_data));
if (trt_input_rank == 1) {
if (data_format == "NHWC") {
EXPECT_THAT(GetSpanForData<CType>(output_data[0]),
ElementsAre(CType(1), CType(2), CType(3)));
output_data = {1, 2, 3};
} else {
EXPECT_THAT(GetSpanForData<CType>(output_data[0]),
ElementsAre(CType(1), CType(2)));
output_data = {1, 2};
}
} else {
if (data_format == "NHWC") {
EXPECT_THAT(GetSpanForData<CType>(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<CType>(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<DT_FLOAT>(this);
TestConvertBiasAdd<DT_HALF>(this);
}
template <typename OpType>
NodeDef GetBinaryOpNodeDef(const string& input_name_l,
const string& input_name_r, DataType dtype) {