Skip to content

Commit f85f2e8

Browse files
zhangjunzhangjun
authored
fix trt multiclass_nms3 (PaddlePaddle#45166)
* update * update * update
1 parent d9fac78 commit f85f2e8

File tree

6 files changed

+334
-44
lines changed

6 files changed

+334
-44
lines changed

‎paddle/fluid/inference/tensorrt/convert/multiclass_nms3_op.cc

Lines changed: 50 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -54,18 +54,34 @@ class MultiClassNMS3OpConverter : public OpConverter {
5454
PADDLE_GET_CONST(float, op_desc.GetAttr("nms_threshold"));
5555
int keep_top_k = PADDLE_GET_CONST(int, op_desc.GetAttr("keep_top_k"));
5656
bool normalized = PADDLE_GET_CONST(bool, op_desc.GetAttr("normalized"));
57-
int num_classes = scores_tensor->getDimensions().d[0];
57+
int class_index = engine_->with_dynamic_shape() ? 1 : 0;
58+
int num_classes = scores_tensor->getDimensions().d[class_index];
5859

5960
auto bboxes_dims = bboxes_tensor->getDimensions();
60-
nvinfer1::Dims3 bboxes_expand_dims(bboxes_dims.d[0], 1, bboxes_dims.d[1]);
61-
auto* bboxes_expand_layer =
62-
TRT_ENGINE_ADD_LAYER(engine_, Shuffle, *bboxes_tensor);
63-
bboxes_expand_layer->setReshapeDimensions(bboxes_expand_dims);
64-
65-
nvinfer1::Permutation permutation{1, 0};
66-
auto* scores_transpose_layer =
67-
TRT_ENGINE_ADD_LAYER(engine_, Shuffle, *scores_tensor);
68-
scores_transpose_layer->setFirstTranspose(permutation);
61+
nvinfer1::IShuffleLayer* bboxes_expand_layer = nullptr;
62+
nvinfer1::IShuffleLayer* scores_transpose_layer = nullptr;
63+
if (engine_->with_dynamic_shape()) {
64+
nvinfer1::Dims4 bboxes_expand_dims(
65+
bboxes_dims.d[0], bboxes_dims.d[1], 1, bboxes_dims.d[2]);
66+
bboxes_expand_layer =
67+
TRT_ENGINE_ADD_LAYER(engine_, Shuffle, *bboxes_tensor);
68+
bboxes_expand_layer->setReshapeDimensions(bboxes_expand_dims);
69+
70+
nvinfer1::Permutation permutation{0, 2, 1};
71+
scores_transpose_layer =
72+
TRT_ENGINE_ADD_LAYER(engine_, Shuffle, *scores_tensor);
73+
scores_transpose_layer->setFirstTranspose(permutation);
74+
} else {
75+
nvinfer1::Dims3 bboxes_expand_dims(bboxes_dims.d[0], 1, bboxes_dims.d[1]);
76+
bboxes_expand_layer =
77+
TRT_ENGINE_ADD_LAYER(engine_, Shuffle, *bboxes_tensor);
78+
bboxes_expand_layer->setReshapeDimensions(bboxes_expand_dims);
79+
80+
nvinfer1::Permutation permutation{1, 0};
81+
scores_transpose_layer =
82+
TRT_ENGINE_ADD_LAYER(engine_, Shuffle, *scores_tensor);
83+
scores_transpose_layer->setFirstTranspose(permutation);
84+
}
6985

7086
std::vector<nvinfer1::ITensor*> batch_nms_inputs;
7187
batch_nms_inputs.push_back(bboxes_expand_layer->getOutput(0));
@@ -101,35 +117,50 @@ class MultiClassNMS3OpConverter : public OpConverter {
101117
fields.size() * sizeof(nvinfer1::PluginField)));
102118
plugin_collections->nbFields = static_cast<int>(fields.size());
103119
plugin_collections->fields = fields.data();
104-
105-
auto creator = GetPluginRegistry()->getPluginCreator("BatchedNMS_TRT", "1");
120+
std::string nms_plugin_name = "BatchedNMS_TRT";
121+
if (engine_->with_dynamic_shape()) {
122+
nms_plugin_name = "BatchedNMSDynamic_TRT";
123+
}
124+
auto creator =
125+
GetPluginRegistry()->getPluginCreator(nms_plugin_name.c_str(), "1");
106126
auto batch_nms_plugin =
107-
creator->createPlugin("BatchNMSPlugin", plugin_collections);
127+
creator->createPlugin(nms_plugin_name.c_str(), plugin_collections);
108128
free(plugin_collections);
109129

110130
auto batch_nms_layer = engine_->network()->addPluginV2(
111131
batch_nms_inputs.data(), batch_nms_inputs.size(), *batch_nms_plugin);
132+
// static shape: [keep_topk, 4], [keep_topk], [keep_topk]
133+
// dynamic shape: [bs, keep_topk, 4], [bs, keep_topk], [bs, keep_topk]
112134
auto nmsed_boxes = batch_nms_layer->getOutput(1);
113135
auto nmsed_scores = batch_nms_layer->getOutput(2);
114136
auto nmsed_classes = batch_nms_layer->getOutput(3);
115137

116138
auto nmsed_scores_transpose_layer =
117139
TRT_ENGINE_ADD_LAYER(engine_, Shuffle, *nmsed_scores);
118-
nmsed_scores_transpose_layer->setReshapeDimensions(
119-
nvinfer1::Dims2(keep_top_k, 1));
120140
auto nmsed_classes_reshape_layer =
121141
TRT_ENGINE_ADD_LAYER(engine_, Shuffle, *nmsed_classes);
122-
nmsed_classes_reshape_layer->setReshapeDimensions(
123-
nvinfer1::Dims2(keep_top_k, 1));
124-
142+
if (engine_->with_dynamic_shape()) {
143+
nmsed_scores_transpose_layer->setReshapeDimensions(
144+
nvinfer1::Dims3(bboxes_dims.d[0], keep_top_k, 1));
145+
146+
nmsed_classes_reshape_layer->setReshapeDimensions(
147+
nvinfer1::Dims3(bboxes_dims.d[0], keep_top_k, 1));
148+
} else {
149+
nmsed_scores_transpose_layer->setReshapeDimensions(
150+
nvinfer1::Dims2(keep_top_k, 1));
151+
152+
nmsed_classes_reshape_layer->setReshapeDimensions(
153+
nvinfer1::Dims2(keep_top_k, 1));
154+
}
125155
std::vector<nvinfer1::ITensor*> concat_inputs;
126156
concat_inputs.push_back(nmsed_classes_reshape_layer->getOutput(0));
127157
concat_inputs.push_back(nmsed_scores_transpose_layer->getOutput(0));
128158
concat_inputs.push_back(nmsed_boxes);
129159

130160
auto nms_concat_layer = TRT_ENGINE_ADD_LAYER(
131161
engine_, Concatenation, concat_inputs.data(), concat_inputs.size());
132-
nms_concat_layer->setAxis(1);
162+
int axis_index = engine_->with_dynamic_shape() ? 1 : 0;
163+
nms_concat_layer->setAxis(axis_index + 1);
133164

134165
// add fake index as output to be consistent with the outputs of
135166
// multiclass_nms3

‎paddle/fluid/inference/tensorrt/convert/multiclass_nms_op.cc

Lines changed: 48 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -52,18 +52,34 @@ class MultiClassNMSOpConverter : public OpConverter {
5252
PADDLE_GET_CONST(float, op_desc.GetAttr("nms_threshold"));
5353
int keep_top_k = PADDLE_GET_CONST(int, op_desc.GetAttr("keep_top_k"));
5454
bool normalized = PADDLE_GET_CONST(bool, op_desc.GetAttr("normalized"));
55-
int num_classes = scores_tensor->getDimensions().d[0];
55+
int class_index = engine_->with_dynamic_shape() ? 1 : 0;
56+
int num_classes = scores_tensor->getDimensions().d[class_index];
5657

5758
auto bboxes_dims = bboxes_tensor->getDimensions();
58-
nvinfer1::Dims3 bboxes_expand_dims(bboxes_dims.d[0], 1, bboxes_dims.d[1]);
59-
auto* bboxes_expand_layer =
60-
TRT_ENGINE_ADD_LAYER(engine_, Shuffle, *bboxes_tensor);
61-
bboxes_expand_layer->setReshapeDimensions(bboxes_expand_dims);
62-
63-
nvinfer1::Permutation permutation{1, 0};
64-
auto* scores_transpose_layer =
65-
TRT_ENGINE_ADD_LAYER(engine_, Shuffle, *scores_tensor);
66-
scores_transpose_layer->setFirstTranspose(permutation);
59+
nvinfer1::IShuffleLayer* bboxes_expand_layer = nullptr;
60+
nvinfer1::IShuffleLayer* scores_transpose_layer = nullptr;
61+
if (engine_->with_dynamic_shape()) {
62+
nvinfer1::Dims4 bboxes_expand_dims(
63+
bboxes_dims.d[0], bboxes_dims.d[1], 1, bboxes_dims.d[2]);
64+
bboxes_expand_layer =
65+
TRT_ENGINE_ADD_LAYER(engine_, Shuffle, *bboxes_tensor);
66+
bboxes_expand_layer->setReshapeDimensions(bboxes_expand_dims);
67+
68+
nvinfer1::Permutation permutation{0, 2, 1};
69+
scores_transpose_layer =
70+
TRT_ENGINE_ADD_LAYER(engine_, Shuffle, *scores_tensor);
71+
scores_transpose_layer->setFirstTranspose(permutation);
72+
} else {
73+
nvinfer1::Dims3 bboxes_expand_dims(bboxes_dims.d[0], 1, bboxes_dims.d[1]);
74+
bboxes_expand_layer =
75+
TRT_ENGINE_ADD_LAYER(engine_, Shuffle, *bboxes_tensor);
76+
bboxes_expand_layer->setReshapeDimensions(bboxes_expand_dims);
77+
78+
nvinfer1::Permutation permutation{1, 0};
79+
scores_transpose_layer =
80+
TRT_ENGINE_ADD_LAYER(engine_, Shuffle, *scores_tensor);
81+
scores_transpose_layer->setFirstTranspose(permutation);
82+
}
6783

6884
std::vector<nvinfer1::ITensor*> batch_nms_inputs;
6985
batch_nms_inputs.push_back(bboxes_expand_layer->getOutput(0));
@@ -100,9 +116,14 @@ class MultiClassNMSOpConverter : public OpConverter {
100116
plugin_collections->nbFields = static_cast<int>(fields.size());
101117
plugin_collections->fields = fields.data();
102118

103-
auto creator = GetPluginRegistry()->getPluginCreator("BatchedNMS_TRT", "1");
119+
std::string nms_plugin_name = "BatchedNMS_TRT";
120+
if (engine_->with_dynamic_shape()) {
121+
nms_plugin_name = "BatchedNMSDynamic_TRT";
122+
}
123+
auto creator =
124+
GetPluginRegistry()->getPluginCreator(nms_plugin_name.c_str(), "1");
104125
auto batch_nms_plugin =
105-
creator->createPlugin("BatchNMSPlugin", plugin_collections);
126+
creator->createPlugin(nms_plugin_name.c_str(), plugin_collections);
106127
free(plugin_collections);
107128

108129
auto batch_nms_layer = engine_->network()->addPluginV2(
@@ -113,12 +134,21 @@ class MultiClassNMSOpConverter : public OpConverter {
113134

114135
auto nmsed_scores_transpose_layer =
115136
TRT_ENGINE_ADD_LAYER(engine_, Shuffle, *nmsed_scores);
116-
nmsed_scores_transpose_layer->setReshapeDimensions(
117-
nvinfer1::Dims2(keep_top_k, 1));
118137
auto nmsed_classes_reshape_layer =
119138
TRT_ENGINE_ADD_LAYER(engine_, Shuffle, *nmsed_classes);
120-
nmsed_classes_reshape_layer->setReshapeDimensions(
121-
nvinfer1::Dims2(keep_top_k, 1));
139+
if (engine_->with_dynamic_shape()) {
140+
nmsed_scores_transpose_layer->setReshapeDimensions(
141+
nvinfer1::Dims3(bboxes_dims.d[0], keep_top_k, 1));
142+
143+
nmsed_classes_reshape_layer->setReshapeDimensions(
144+
nvinfer1::Dims3(bboxes_dims.d[0], keep_top_k, 1));
145+
} else {
146+
nmsed_scores_transpose_layer->setReshapeDimensions(
147+
nvinfer1::Dims2(keep_top_k, 1));
148+
149+
nmsed_classes_reshape_layer->setReshapeDimensions(
150+
nvinfer1::Dims2(keep_top_k, 1));
151+
}
122152

123153
std::vector<nvinfer1::ITensor*> concat_inputs;
124154
concat_inputs.push_back(nmsed_classes_reshape_layer->getOutput(0));
@@ -127,7 +157,8 @@ class MultiClassNMSOpConverter : public OpConverter {
127157

128158
auto nms_concat_layer = TRT_ENGINE_ADD_LAYER(
129159
engine_, Concatenation, concat_inputs.data(), concat_inputs.size());
130-
nms_concat_layer->setAxis(1);
160+
int axis_index = engine_->with_dynamic_shape() ? 1 : 0;
161+
nms_concat_layer->setAxis(axis_index + 1);
131162

132163
RreplenishLayerAndOutput(
133164
nms_concat_layer, "multiclass_nms", {output_name}, test_mode);

‎paddle/fluid/inference/tensorrt/op_teller.cc

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,10 @@ namespace tensorrt {
3333
struct SimpleOpTypeSetTeller : public Teller {
3434
SimpleOpTypeSetTeller() {
3535
#if IS_TRT_VERSION_GE(7130)
36+
// use TensorRT plugin
3637
teller_set.insert("group_norm");
38+
teller_set.insert("multiclass_nms3");
39+
teller_set.insert("multiclass_nms");
3740
#endif
3841
#if IS_TRT_VERSION_GE(7000)
3942
teller_set.insert("tile");
@@ -278,7 +281,6 @@ struct SimpleOpTypeSetTeller : public Teller {
278281
"c_allreduce_prod",
279282
"roll",
280283
"cast",
281-
"multiclass_nms3",
282284
"transformer_input_convert",
283285
"recover_padding",
284286
"remove_padding",
@@ -847,7 +849,6 @@ bool OpTeller::Tell(const framework::ir::Node* node,
847849
}
848850

849851
if (op_type == "multiclass_nms" || op_type == "multiclass_nms3") {
850-
if (with_dynamic_shape) return false;
851852
auto* block = desc.Block();
852853
if (block == nullptr) {
853854
VLOG(3) << "The block desc is nullptr, we can't continue to analyze. "

‎paddle/fluid/inference/tests/infer_ut/test_ppyolo_mbv3.cc

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -73,7 +73,7 @@ TEST(tensorrt_tester_ppyolo_mbv3, multi_thread4_trt_fp32_bz2) {
7373
FLAGS_modeldir + "/model.pdiparams");
7474
config.EnableUseGpu(100, 0);
7575
config.EnableTensorRtEngine(
76-
1 << 20, 2, 3, paddle_infer::PrecisionType::kFloat32, false, false);
76+
1 << 25, 2, 3, paddle_infer::PrecisionType::kFloat32, false, false);
7777
LOG(INFO) << config.Summary();
7878
// get groudtruth by disbale ir
7979
paddle_infer::services::PredictorPool pred_pool_no_ir(config_no_ir, 1);

0 commit comments

Comments
 (0)