Skip to content

Commit

Permalink
fix FlattenContiguousRangeOpConverter out dim error (#42087) (#42184)
Browse files Browse the repository at this point in the history 
* fix FlattenContiguousRangeOpConverter out dim error

* update code
  • Loading branch information
@baoachun
baoachun authored Apr 25, 2022
1 parent 8c3c6da commit 20e8bf1
Showing 1 changed file with 92 additions and 58 deletions.
150 changes: 92 additions & 58 deletions paddle/fluid/inference/tensorrt/convert/flatten_contiguous_range_op.cc
View file
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

0 comments on commit 20e8bf1

Please sign in to comment.