Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix FlattenContiguousRangeOpConverter out dim error #42087

Merged
merged 2 commits into from
Apr 24, 2022
Merged

fix FlattenContiguousRangeOpConverter out dim error #42087

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
19f9f81
fix FlattenContiguousRangeOpConverter out dim error
baoachun Apr 21, 2022
db80221
update code
baoachun Apr 21, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
150 changes: 92 additions & 58 deletions paddle/fluid/inference/tensorrt/convert/flatten_contiguous_range_op.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,14 +30,17 @@ class FlattenContiguousRangeOpConverter : public OpConverter {
public:
void operator()(const framework::proto::OpDesc& op,
const framework::Scope& scope, bool test_mode) override {
VLOG(3) << "convert a fluid flatten_contiguous_range op to tensorrt layer";
framework::OpDesc op_desc(op, nullptr);
// Declare inputs
auto* input = engine_->GetITensor(op_desc.Input("X")[0]);
int dims = input->getDimensions().nbDims;
const auto input_dim = input->getDimensions();
const int dims = input_dim.nbDims;
int start_axis = BOOST_GET_CONST(int, op_desc.GetAttr("start_axis"));
int stop_axis = BOOST_GET_CONST(int, op_desc.GetAttr("stop_axis"));

nvinfer1::IShuffleLayer* layer = nullptr;
nvinfer1::IShuffleLayer* layer =
TRT_ENGINE_ADD_LAYER(engine_, Shuffle, *input);
if (!engine_->with_dynamic_shape()) {
if (start_axis < 0) start_axis += dims + 1;
if (stop_axis < 0) stop_axis += dims + 1;
Expand All @@ -46,7 +49,7 @@ class FlattenContiguousRangeOpConverter : public OpConverter {
flatten_dim.nbDims = dims - (stop_axis - start_axis);
for (int i = 0, j = 0; i < dims; ++i) {
if (start_axis <= i + 1 && i + 1 <= stop_axis) {
int dim_i = input->getDimensions().d[i];
int dim_i = input_dim.d[i];
PADDLE_ENFORCE_GT(dim_i, 0, platform::errors::InvalidArgument(
"flatten_contiguous_range input dim "
"should be > 0, but got %d.",
Expand All @@ -56,72 +59,103 @@ class FlattenContiguousRangeOpConverter : public OpConverter {
flatten_dim.d[j++] = dim_prod;
}
} else {
flatten_dim.d[j++] = input->getDimensions().d[i];
flatten_dim.d[j++] = input_dim.d[i];
}
}
layer = TRT_ENGINE_ADD_LAYER(engine_, Shuffle, *input);
layer->setReshapeDimensions(flatten_dim);
} else {
if (start_axis < 0) start_axis += dims;
if (stop_axis < 0) stop_axis += dims;
auto* shape_layer = TRT_ENGINE_ADD_LAYER(engine_, Shape, *input);
auto* shape_layer_itensor = shape_layer->getOutput(0);

nvinfer1::Dims start_dim, size_dim, stride_dim;
start_dim.nbDims = 1;
size_dim.nbDims = 1;
stride_dim.nbDims = 1;
start_dim.d[0] = start_axis;
size_dim.d[0] = stop_axis - start_axis + 1;
stride_dim.d[0] = 1;
auto* slice_layer =
TRT_ENGINE_ADD_LAYER(engine_, Slice, *shape_layer_itensor, start_dim,
size_dim, stride_dim);
uint32_t reduce_dim = 1;
auto* reduce_prod_layer = TRT_ENGINE_ADD_LAYER(
engine_, Reduce, *(slice_layer->getOutput(0)),
nvinfer1::ReduceOperation::kPROD, reduce_dim, true);

nvinfer1::ITensor* input_shape = nullptr;
if (start_axis == 0 && stop_axis == dims - 1) {
input_shape = reduce_prod_layer->getOutput(0);
} else {
std::vector<nvinfer1::ITensor*> itensors;
if (start_axis > 0) {
nvinfer1::Dims left_start_dim, left_size_dim, left_stride_dim;
left_start_dim.nbDims = 1;
left_size_dim.nbDims = 1;
left_stride_dim.nbDims = 1;
left_start_dim.d[0] = 0;
left_size_dim.d[0] = start_axis;
left_stride_dim.d[0] = 1;
auto* slice_layer_left = TRT_ENGINE_ADD_LAYER(
engine_, Slice, *shape_layer_itensor, left_start_dim,
left_size_dim, left_stride_dim);
itensors.push_back(slice_layer_left->getOutput(0));
int dim_prod = 1;
int dim_negative = 0;
nvinfer1::Dims flatten_dim;
flatten_dim.nbDims = dims - (stop_axis - start_axis);
bool need_slice = false;
for (int i = 0, j = 0; i < dims; ++i) {
int dim_i = input_dim.d[i];
if (start_axis <= i && i <= stop_axis) {
if (dim_i < 0) {
need_slice = true;
break;
}
dim_prod *= dim_i;
if (i == stop_axis) {
flatten_dim.d[j++] = dim_prod;
}
} else {
if (dim_i < 0) dim_negative++;
if (dim_negative > 1) {
need_slice = true;
break;
}
flatten_dim.d[j++] = input_dim.d[i];
}
itensors.push_back(reduce_prod_layer->getOutput(0));
if (stop_axis < dims - 1) {
nvinfer1::Dims right_start_dim, right_size_dim, right_stride_dim;
right_start_dim.nbDims = 1;
right_size_dim.nbDims = 1;
right_stride_dim.nbDims = 1;
right_start_dim.d[0] = stop_axis + 1;
right_size_dim.d[0] = dims - stop_axis - 1;
right_stride_dim.d[0] = 1;
auto* slice_layer_right = TRT_ENGINE_ADD_LAYER(
engine_, Slice, *shape_layer_itensor, right_start_dim,
right_size_dim, right_stride_dim);
itensors.push_back(slice_layer_right->getOutput(0));
}

if (need_slice) {
VLOG(3) << "slice input dim when the input dimension has -1";
auto* shape_layer = TRT_ENGINE_ADD_LAYER(engine_, Shape, *input);
auto* shape_layer_itensor = shape_layer->getOutput(0);

nvinfer1::Dims start_dim, size_dim, stride_dim;
start_dim.nbDims = 1;
size_dim.nbDims = 1;
stride_dim.nbDims = 1;
start_dim.d[0] = start_axis;
size_dim.d[0] = stop_axis - start_axis + 1;
stride_dim.d[0] = 1;
auto* slice_layer =
TRT_ENGINE_ADD_LAYER(engine_, Slice, *shape_layer_itensor,
start_dim, size_dim, stride_dim);
uint32_t reduce_dim = 1;
auto* reduce_prod_layer = TRT_ENGINE_ADD_LAYER(
engine_, Reduce, *(slice_layer->getOutput(0)),
nvinfer1::ReduceOperation::kPROD, reduce_dim, true);

nvinfer1::ITensor* input_shape = nullptr;
if (start_axis == 0 && stop_axis == dims - 1) {
input_shape = reduce_prod_layer->getOutput(0);
} else {
std::vector<nvinfer1::ITensor*> itensors;
if (start_axis > 0) {
nvinfer1::Dims left_start_dim, left_size_dim, left_stride_dim;
left_start_dim.nbDims = 1;
left_size_dim.nbDims = 1;
left_stride_dim.nbDims = 1;
left_start_dim.d[0] = 0;
left_size_dim.d[0] = start_axis;
left_stride_dim.d[0] = 1;
auto* slice_layer_left = TRT_ENGINE_ADD_LAYER(
engine_, Slice, *shape_layer_itensor, left_start_dim,
left_size_dim, left_stride_dim);
itensors.push_back(slice_layer_left->getOutput(0));
}
itensors.push_back(reduce_prod_layer->getOutput(0));
if (stop_axis < dims - 1) {
nvinfer1::Dims right_start_dim, right_size_dim, right_stride_dim;
right_start_dim.nbDims = 1;
right_size_dim.nbDims = 1;
right_stride_dim.nbDims = 1;
right_start_dim.d[0] = stop_axis + 1;
right_size_dim.d[0] = dims - stop_axis - 1;
right_stride_dim.d[0] = 1;
auto* slice_layer_right = TRT_ENGINE_ADD_LAYER(
engine_, Slice, *shape_layer_itensor, right_start_dim,
right_size_dim, right_stride_dim);
itensors.push_back(slice_layer_right->getOutput(0));
}
auto* concat_layer = TRT_ENGINE_ADD_LAYER(
engine_, Concatenation, itensors.data(), itensors.size());
concat_layer->setAxis(0);
input_shape = concat_layer->getOutput(0);
}
auto* concat_layer = TRT_ENGINE_ADD_LAYER(
engine_, Concatenation, itensors.data(), itensors.size());
concat_layer->setAxis(0);
input_shape = concat_layer->getOutput(0);
layer->setInput(1, *input_shape);
} else {
layer->setReshapeDimensions(flatten_dim);
}
layer = TRT_ENGINE_ADD_LAYER(engine_, Shuffle, *input);
layer->setInput(1, *input_shape);
}

auto output_name = op_desc.Output("Out")[0];
RreplenishLayerAndOutput(layer, "flatten_contiguous_range", {output_name},
test_mode);
Expand Down