Merge pull request #39156 from tfeher:trt_biasadd_dynamic_shape
PiperOrigin-RevId: 313402866 Change-Id: Iad1d5701f727973c8f892e20be090a2570baf1a9
This commit is contained in:
commit
be4d17c088
@ -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),
|
||||
|
@ -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) {
|
||||
|
Loading…
x
Reference in New Issue
Block a user