@@ -13,25 +13,57 @@ See the License for the specific language governing permissions and
13
13
limitations under the License. */
14
14
15
15
#include " paddle/fluid/inference/tensorrt/convert/op_converter.h"
16
+ #include " paddle/fluid/inference/tensorrt/plugin/avg_pool_op_plugin.h"
16
17
17
18
namespace paddle {
18
19
namespace inference {
19
20
namespace tensorrt {
20
21
22
+ void DealCeilMode (const nvinfer1::Dims &input_shape, std::vector<int > ksize,
23
+ std::vector<int > strides, std::vector<int > paddings,
24
+ nvinfer1::DimsHW *pre_pad, nvinfer1::DimsHW *post_pad,
25
+ int input_dims) {
26
+ int input_height = input_shape.d [input_dims - 2 ];
27
+ int input_width = input_shape.d [input_dims - 1 ];
28
+ int floor_h_output_size =
29
+ (input_height - ksize[0 ] + 2 * paddings[0 ]) / strides[0 ] + 1 ;
30
+ int ceil_h_output_size =
31
+ (input_height - ksize[0 ] + 2 * paddings[0 ] + strides[0 ] - 1 ) /
32
+ strides[0 ] +
33
+ 1 ;
34
+
35
+ int floor_w_output_size =
36
+ (input_width - ksize[1 ] + 2 * paddings[1 ]) / strides[1 ] + 1 ;
37
+ int ceil_w_output_size =
38
+ (input_width - ksize[1 ] + 2 * paddings[1 ] + strides[1 ] - 1 ) / strides[1 ] +
39
+ 1 ;
40
+ if (floor_h_output_size != ceil_h_output_size) {
41
+ post_pad->h () = strides[0 ] - 1 ;
42
+ }
43
+
44
+ if (floor_w_output_size != ceil_w_output_size) {
45
+ post_pad->w () = strides[1 ] - 1 ;
46
+ }
47
+ }
48
+
21
49
/*
22
50
* Pool2dOp, IPoolingLayer in TRT. This Layer doesn't has weights.
23
51
*/
24
52
class Pool2dOpConverter : public OpConverter {
25
53
public:
26
- void operator ()(const framework::proto::OpDesc& op,
27
- const framework::Scope& scope, bool test_mode) override {
28
- VLOG (3 )
54
+ void operator ()(const framework::proto::OpDesc & op,
55
+ const framework::Scope & scope, bool test_mode) override {
56
+ VLOG (40 )
29
57
<< " convert a fluid pool2d op to tensorrt pool2d layer without bias" ;
30
58
framework::OpDesc op_desc (op, nullptr );
31
59
// Declare inputs
32
60
PADDLE_ENFORCE_EQ (op_desc.Input (" X" ).size (), 1 );
33
61
PADDLE_ENFORCE_EQ (op_desc.Output (" Out" ).size (), 1 );
34
- auto * input1 = engine_->GetITensor (op_desc.Input (" X" )[0 ]);
62
+ auto *input1 = engine_->GetITensor (op_desc.Input (" X" )[0 ]);
63
+ nvinfer1::Dims input_shape = input1->getDimensions ();
64
+ int input_dims = input_shape.nbDims ;
65
+
66
+ PADDLE_ENFORCE_EQ (input_dims, 3UL );
35
67
36
68
bool global_pooling = boost::get<bool >(op_desc.GetAttr (" global_pooling" ));
37
69
std::string pool_type =
@@ -44,23 +76,6 @@ class Pool2dOpConverter : public OpConverter {
44
76
boost::get<std::vector<int >>(op_desc.GetAttr (" paddings" ));
45
77
bool ceil_mode = boost::get<bool >(op_desc.GetAttr (" ceil_mode" ));
46
78
47
- nvinfer1::Dims input_shape = input1->getDimensions ();
48
- int nbDims = input_shape.nbDims ;
49
- nvinfer1::DimsHW nv_ksize (ksize[0 ], ksize[1 ]);
50
- nvinfer1::DimsHW nv_strides (strides[0 ], strides[1 ]);
51
- nvinfer1::DimsHW nv_paddings (paddings[0 ], paddings[1 ]);
52
-
53
- if (global_pooling == true ) {
54
- nv_ksize.d [0 ] = input_shape.d [nbDims - 2 ];
55
- nv_ksize.d [1 ] = input_shape.d [nbDims - 1 ];
56
- nv_strides.h () = 1 ;
57
- nv_strides.w () = 1 ;
58
- nv_paddings.h () = 0 ;
59
- nv_paddings.w () = 0 ;
60
- }
61
-
62
- PADDLE_ENFORCE_EQ (input1->getDimensions ().nbDims , 3UL );
63
-
64
79
nvinfer1::PoolingType nv_pool_type = nvinfer1::PoolingType::kMAX ;
65
80
if (pool_type == " max" ) {
66
81
nv_pool_type = nvinfer1::PoolingType::kMAX ;
@@ -70,48 +85,71 @@ class Pool2dOpConverter : public OpConverter {
70
85
PADDLE_THROW (" TensorRT unsupported pooling type!" );
71
86
}
72
87
73
- if (ceil_mode) {
74
- nvinfer1::DimsHW pre_pad ( 0 , 0 );
75
- nvinfer1::DimsHW post_pad ( 0 , 0 );
76
- int input_height = input_shape. d [nbDims - 2 ];
77
- int input_width = input_shape. d [nbDims - 1 ] ;
78
- int floor_h_output_size =
79
- (input_height - ksize[ 0 ] + 2 * paddings[ 0 ]) / strides[ 0 ] + 1 ;
80
- int ceil_h_output_size =
81
- (input_height - ksize[ 0 ] + 2 * paddings[ 0 ] + strides[ 0 ] - 1 ) /
82
- strides[ 0 ] +
83
- 1 ;
84
-
85
- int floor_w_output_size =
86
- (input_width - ksize[ 1 ] + 2 * paddings[ 1 ]) / strides[ 1 ] + 1 ;
87
- int ceil_w_output_size =
88
- (input_width - ksize[ 1 ] + 2 * paddings[ 1 ] + strides[ 1 ] - 1 ) /
89
- strides[ 1 ] +
90
- 1 ;
91
- if (floor_h_output_size != ceil_h_output_size ) {
92
- post_pad. h () = strides[ 0 ] - 1 ;
88
+ nvinfer1::DimsHW nv_ksize (ksize[ 0 ], ksize[ 1 ]);
89
+ nvinfer1::DimsHW nv_strides (strides[ 0 ], strides[ 1 ] );
90
+ nvinfer1::DimsHW nv_paddings (paddings[ 0 ], paddings[ 1 ] );
91
+
92
+ nvinfer1::ILayer *layer = nullptr ;
93
+
94
+ if (global_pooling == true ) {
95
+ nv_ksize. d [ 0 ] = input_shape. d [input_dims - 2 ];
96
+ nv_ksize. d [ 1 ] = input_shape. d [input_dims - 1 ];
97
+ auto *layer = TRT_ENGINE_ADD_LAYER (
98
+ engine_, Pooling, * const_cast <nvinfer1::ITensor *>(input1),
99
+ nv_pool_type, nv_ksize);
100
+ PADDLE_ENFORCE_NOT_NULL (layer, " pool layer could not be created. " );
101
+ auto output_name = op_desc. Output ( " Out " )[ 0 ] ;
102
+ layer-> setName (( " pool2d (Output: " + output_name + " ) " ). c_str ());
103
+ layer-> getOutput ( 0 )-> setName (output_name. c_str ());
104
+ engine_-> SetITensor (output_name, layer-> getOutput ( 0 ));
105
+ if (test_mode ||
106
+ output_name == " patch_6_pool1.avg_pool.output.1.tmp_0702 " ) {
107
+ engine_-> DeclareOutput (output_name) ;
93
108
}
109
+ return ;
110
+ }
94
111
95
- if (floor_w_output_size != ceil_w_output_size) {
96
- post_pad.w () = strides[1 ] - 1 ;
112
+ if (pool_type == " max" ) {
113
+ nvinfer1::DimsHW pre_pad (paddings[0 ], paddings[1 ]);
114
+ nvinfer1::DimsHW post_pad (paddings[0 ], paddings[1 ]);
115
+ if (ceil_mode) {
116
+ // If ceil mode is true, we will pad the appropriate size to the input.
117
+ DealCeilMode (input_shape, ksize, strides, paddings, &pre_pad, &post_pad,
118
+ input_dims);
119
+ auto *pad_layer = TRT_ENGINE_ADD_LAYER (
120
+ engine_, Padding, *const_cast <nvinfer1::ITensor *>(input1), pre_pad,
121
+ post_pad);
122
+ PADDLE_ENFORCE_NOT_NULL (
123
+ pad_layer, " pad layer in poolOp converter could not be created." );
124
+ input1 = pad_layer->getOutput (0 );
125
+ }
126
+ auto *pool_layer = TRT_ENGINE_ADD_LAYER (
127
+ engine_, Pooling, *const_cast <nvinfer1::ITensor *>(input1),
128
+ nv_pool_type, nv_ksize);
129
+ PADDLE_ENFORCE_NOT_NULL (pool_layer, " pool layer could not be created." );
130
+ pool_layer->setStride (nv_strides);
131
+ pool_layer->setPadding (nv_paddings);
132
+ layer = pool_layer;
133
+ } else {
134
+ // Average pooling needs to exclude the padding pixels from the average
135
+ // mean.
136
+ // It is not supported well by TRT, we use a plugin here.
137
+ std::vector<int > input_shape_v;
138
+ for (int i = 0 ; i < input_dims; i++) {
139
+ input_shape_v.push_back (input_shape.d [i]);
97
140
}
98
- auto * layer = TRT_ENGINE_ADD_LAYER (
99
- engine_, Padding, * const_cast <nvinfer1::ITensor*>(input1), pre_pad,
100
- post_pad );
101
- input1 = layer-> getOutput ( 0 ) ;
141
+ AvgPoolPlugin *plugin =
142
+ new AvgPoolPlugin (ceil_mode, ksize, strides, paddings, input_shape_v);
143
+ auto *avg_pool_layer = engine_-> AddPlugin (&input1, 1 , plugin );
144
+ layer = avg_pool_layer ;
102
145
}
103
- auto * layer = TRT_ENGINE_ADD_LAYER (engine_, Pooling,
104
- *const_cast <nvinfer1::ITensor*>(input1),
105
- nv_pool_type, nv_ksize);
106
- PADDLE_ENFORCE_NOT_NULL (layer, " pool layer could not be created." );
107
- layer->setStride (nv_strides);
108
- layer->setPadding (nv_paddings);
109
146
110
147
auto output_name = op_desc.Output (" Out" )[0 ];
111
148
layer->setName ((" pool2d (Output: " + output_name + " )" ).c_str ());
112
149
layer->getOutput (0 )->setName (output_name.c_str ());
113
150
engine_->SetITensor (output_name, layer->getOutput (0 ));
114
- if (test_mode) {
151
+ if (test_mode ||
152
+ output_name == " patch_6_pool1.avg_pool.output.1.tmp_0702" ) {
115
153
engine_->DeclareOutput (output_name);
116
154
}
117
155
}
0 commit comments