Apollo 11.0
自动驾驶开放平台
dfmb_psroi_align_plugin.h
浏览该文件的文档.
1/******************************************************************************
2 * Copyright 2020 The Apollo Authors. All Rights Reserved.
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 * http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 *****************************************************************************/
16
17#pragma once
18
19#include <NvInferVersion.h>
20
21#include "cyber/common/log.h"
22#include "modules/perception/common/inference/tensorrt/rt_common.h"
23
24namespace apollo {
25namespace perception {
26namespace inference {
27
28// TODO(chenjiahao): complete member functions
29// Custom layer for DFMBPSROIAlign operation,
30// i.e. DeForMaBle Position Sensitive ROI Align.
31// input0 dims: [C, H, W], input1 dims: [num_rois, 5, 1, 1]
32// input2 dims: [N, C2, H2, W2]
33#ifdef NV_TENSORRT_MAJOR
34#if NV_TENSORRT_MAJOR != 8
35class DFMBPSROIAlignPlugin : public nvinfer1::IPlugin {
36 public:
37 DFMBPSROIAlignPlugin(
38 const DFMBPSROIAlignParameter &dfmb_psroi_align_parameter,
39 nvinfer1::Dims *in_dims, int nbInputs) {
40 heat_map_a_ = dfmb_psroi_align_parameter.heat_map_a();
41 output_channel_ = dfmb_psroi_align_parameter.output_dim();
42 group_height_ = dfmb_psroi_align_parameter.group_height();
43 group_width_ = dfmb_psroi_align_parameter.group_width();
44 pooled_height_ = dfmb_psroi_align_parameter.pooled_height();
45 pooled_width_ = dfmb_psroi_align_parameter.pooled_width();
46 pad_ratio_ = dfmb_psroi_align_parameter.pad_ratio();
47 sample_per_part_ = dfmb_psroi_align_parameter.sample_per_part();
48
49 trans_std_ = dfmb_psroi_align_parameter.trans_std();
50 part_height_ = dfmb_psroi_align_parameter.part_height();
51 part_width_ = dfmb_psroi_align_parameter.part_width();
52 heat_map_b_ = dfmb_psroi_align_parameter.heat_map_b();
53 no_trans_ = (nbInputs < 3);
54 num_classes_ = no_trans_ ? 1 : in_dims[2].d[1];
55
56 CHECK_GT(heat_map_a_, 0);
57 CHECK_GE(heat_map_b_, 0);
58 CHECK_GE(pad_ratio_, 0);
59 CHECK_GT(output_channel_, 0);
60 CHECK_GT(sample_per_part_, 0);
61 CHECK_GT(group_height_, 0);
62 CHECK_GT(group_width_, 0);
63 CHECK_GT(pooled_height_, 0);
64 CHECK_GT(pooled_width_, 0);
65 CHECK_GE(part_height_, 0);
66 CHECK_GE(part_width_, 0);
67
68 channels_ = in_dims[0].d[0];
69 height_ = in_dims[0].d[1];
70 width_ = in_dims[0].d[2];
71 output_dims_ = nvinfer1::Dims4(in_dims[1].d[0], output_channel_,
72 pooled_height_, pooled_width_);
73 output_size_ =
74 in_dims[1].d[0] * output_channel_ * pooled_height_ * pooled_width_;
75
76 CHECK_EQ(channels_, output_channel_ * group_height_ * group_width_);
77 CHECK_EQ(in_dims[1].d[1], 5);
78 if (!no_trans_) {
79 CHECK_EQ(in_dims[2].d[1] % 2, 0);
80 int num_classes = in_dims[2].d[1] / 2;
81 CHECK_EQ(output_channel_ % num_classes, 0);
82 CHECK_EQ(part_height_, in_dims[2].d[2]);
83 CHECK_EQ(part_width_, in_dims[2].d[3]);
84 }
85 }
86
87 virtual ~DFMBPSROIAlignPlugin() {}
88
89 virtual int initialize() { return 0; }
90 virtual void terminate() {}
91 int getNbOutputs() const override { return 1; }
92
93 nvinfer1::Dims getOutputDimensions(int index, const nvinfer1::Dims *inputs,
94 int nbInputDims) override {
95 // TODO(chenjiahao): complete input dims assertion
96 return output_dims_;
97 }
98
99 void configure(const nvinfer1::Dims *inputDims, int nbInputs,
100 const nvinfer1::Dims *outputDims, int nbOutputs,
101 int maxBatchSize) override {}
102
103 size_t getWorkspaceSize(int maxBatchSize) const override { return 0; }
104
105 virtual int enqueue(int batchSize, const void *const *inputs, void **outputs,
106 void *workspace, cudaStream_t stream);
107
108 size_t getSerializationSize() override { return 0; }
109
110 void serialize(void *buffer) override {
111 char *d = reinterpret_cast<char *>(buffer), *a = d;
112 size_t size = getSerializationSize();
113 CHECK_EQ(d, a + size);
114 }
115
116 private:
117 const int thread_size_ = 512;
118 float heat_map_a_;
119 float heat_map_b_;
120 float pad_ratio_;
121
122 int output_channel_;
123 bool no_trans_;
124 float trans_std_;
125 int sample_per_part_;
126 int group_height_;
127 int group_width_;
128 int pooled_height_;
129 int pooled_width_;
130 int part_height_;
131 int part_width_;
132 int num_classes_;
133
134 int channels_;
135 int height_;
136 int width_;
137 int output_size_;
138
139 nvinfer1::Dims output_dims_;
140};
141
142#else
143class DFMBPSROIAlignPlugin : public nvinfer1::IPluginV2Ext {
144 public:
145 DFMBPSROIAlignPlugin(
146 const DFMBPSROIAlignParameter &dfmb_psroi_align_parameter,
147 nvinfer1::Dims *in_dims, int nbInputs) {
148 heat_map_a_ = dfmb_psroi_align_parameter.heat_map_a();
149 output_channel_ = dfmb_psroi_align_parameter.output_dim();
150 group_height_ = dfmb_psroi_align_parameter.group_height();
151 group_width_ = dfmb_psroi_align_parameter.group_width();
152 pooled_height_ = dfmb_psroi_align_parameter.pooled_height();
153 pooled_width_ = dfmb_psroi_align_parameter.pooled_width();
154 pad_ratio_ = dfmb_psroi_align_parameter.pad_ratio();
155 sample_per_part_ = dfmb_psroi_align_parameter.sample_per_part();
156
157 trans_std_ = dfmb_psroi_align_parameter.trans_std();
158 part_height_ = dfmb_psroi_align_parameter.part_height();
159 part_width_ = dfmb_psroi_align_parameter.part_width();
160 heat_map_b_ = dfmb_psroi_align_parameter.heat_map_b();
161 no_trans_ = (nbInputs < 3);
162 num_classes_ = no_trans_ ? 1 : in_dims[2].d[1];
163
164 CHECK_GT(heat_map_a_, 0);
165 CHECK_GE(heat_map_b_, 0);
166 CHECK_GE(pad_ratio_, 0);
167 CHECK_GT(output_channel_, 0);
168 CHECK_GT(sample_per_part_, 0);
169 CHECK_GT(group_height_, 0);
170 CHECK_GT(group_width_, 0);
171 CHECK_GT(pooled_height_, 0);
172 CHECK_GT(pooled_width_, 0);
173 CHECK_GE(part_height_, 0);
174 CHECK_GE(part_width_, 0);
175
176 channels_ = in_dims[0].d[0];
177 height_ = in_dims[0].d[1];
178 width_ = in_dims[0].d[2];
179 output_dims_ = nvinfer1::Dims4(in_dims[1].d[0], output_channel_,
180 pooled_height_, pooled_width_);
181 output_size_ =
182 in_dims[1].d[0] * output_channel_ * pooled_height_ * pooled_width_;
183
184 CHECK_EQ(channels_, output_channel_ * group_height_ * group_width_);
185 CHECK_EQ(in_dims[1].d[1], 5);
186 if (!no_trans_) {
187 CHECK_EQ(in_dims[2].d[1] % 2, 0);
188 int num_classes = in_dims[2].d[1] / 2;
189 CHECK_EQ(output_channel_ % num_classes, 0);
190 CHECK_EQ(part_height_, in_dims[2].d[2]);
191 CHECK_EQ(part_width_, in_dims[2].d[3]);
192 }
193 }
194
195 DFMBPSROIAlignPlugin() {}
196 virtual ~DFMBPSROIAlignPlugin() {}
197
198 virtual int initialize() noexcept { return 0; }
199 virtual void terminate() noexcept {}
200 int32_t getNbOutputs() const noexcept override { return 1; }
201
202 nvinfer1::Dims getOutputDimensions(int32_t index,
203 const nvinfer1::Dims *inputs,
204 int32_t nbInputDims) noexcept override {
205 // TODO(chenjiahao): complete input dims assertion
206 return output_dims_;
207 }
208
209 void configureWithFormat(const nvinfer1::Dims *inputDims, int32_t nbInputs,
210 const nvinfer1::Dims *outputDims, int32_t nbOutputs,
211 nvinfer1::DataType type, nvinfer1::PluginFormat format,
212 int32_t maxBatchSize) noexcept override {}
213
214 size_t getWorkspaceSize(int32_t maxBatchSize)
215 const noexcept override { return 0; }
216
217 int32_t enqueue(int32_t batchSize, const void *const *inputs,
218 void *const *outputs, void *workspace,
219 cudaStream_t stream) noexcept override;
220
221 size_t getSerializationSize() const noexcept override { return 0; }
222
223 void serialize(void *buffer) const noexcept override {
224 char *d = reinterpret_cast<char *>(buffer), *a = d;
225 size_t size = getSerializationSize();
226 CHECK_EQ(d, a + size);
227 }
228
229 nvinfer1::AsciiChar const* getPluginType() const noexcept override {
230 return plugin_type;
231 }
232
233 nvinfer1::AsciiChar const* getPluginVersion() const noexcept override {
234 return plugin_version;
235 }
236
237 void setPluginNamespace(const nvinfer1::AsciiChar* libNamespace)
238 noexcept override {
239 plugin_namespace = const_cast<nvinfer1::AsciiChar*>(libNamespace);
240 }
241
242 nvinfer1::AsciiChar const* getPluginNamespace() const noexcept override {
243 return const_cast<nvinfer1::AsciiChar*>(plugin_namespace);
244 }
245
246 bool supportsFormat(nvinfer1::DataType type,
247 nvinfer1::PluginFormat format) const noexcept override {
248 return true;
249 }
250
251 void destroy() noexcept override {
252 delete this;
253 }
254
255 bool isOutputBroadcastAcrossBatch(int32_t outputIndex,
256 bool const *inputIsBroadcasted,
257 int32_t nbInputs) const noexcept override {
258 return false;
259 }
260
261 bool canBroadcastInputAcrossBatch(int32_t inputIndex)
262 const noexcept override {
263 return false;
264 }
265
266 nvinfer1::DataType getOutputDataType(int32_t index,
267 nvinfer1::DataType const *inputTypes,
268 int32_t nbInputs) const noexcept {
269 return nvinfer1::DataType::kFLOAT;
270 }
271
272 void configurePlugin(
273 nvinfer1::Dims const *inputDims, int32_t nbInputs,
274 nvinfer1::Dims const *outputDims, int32_t nbOutputs,
275 nvinfer1::DataType const *inputTypes, nvinfer1::DataType const *outputTypes,
276 bool const *inputIsBroadcast, bool const *outputIsBroadcast,
277 nvinfer1::PluginFormat floatFormat, int32_t maxBatchSize) noexcept {}
278
279 nvinfer1::IPluginV2Ext* clone() const noexcept override {
280 DFMBPSROIAlignPlugin* p = new DFMBPSROIAlignPlugin();
281 p->heat_map_a_ = heat_map_a_;
282 p->heat_map_b_ = heat_map_b_;
283 p->pad_ratio_ = pad_ratio_;
284
285 p->output_channel_ = output_channel_;
286 p->no_trans_ = no_trans_;
287 p->trans_std_ = trans_std_;
288 p->sample_per_part_ = sample_per_part_;
289 p->group_height_ = group_height_;
290 p->group_width_ = group_width_;
291 p->pooled_height_ = pooled_height_;
292 p->pooled_width_ = pooled_width_;
293 p->part_height_ = part_height_;
294 p->part_width_ = part_width_;
295 p->num_classes_ = num_classes_;
296
297 p->channels_ = channels_;
298 p->height_ = height_;
299 p->width_ = width_;
300 p->output_size_ = output_size_;
301
302 p->output_dims_ = nvinfer1::Dims4(
303 output_size_ / (output_channel_ * pooled_height_ * pooled_width_),
304 output_channel_, pooled_height_, pooled_width_);
305
306 p->plugin_namespace = plugin_namespace;
307 return p;
308 }
309
310 private:
311 const int thread_size_ = 512;
312 float heat_map_a_;
313 float heat_map_b_;
314 float pad_ratio_;
315
316 int output_channel_;
317 bool no_trans_;
318 float trans_std_;
319 int sample_per_part_;
320 int group_height_;
321 int group_width_;
322 int pooled_height_;
323 int pooled_width_;
324 int part_height_;
325 int part_width_;
326 int num_classes_;
327
328 int channels_;
329 int height_;
330 int width_;
331 int output_size_;
332
333 nvinfer1::Dims output_dims_;
334
335 nvinfer1::AsciiChar* plugin_namespace;
336 const nvinfer1::AsciiChar* plugin_type = "";
337 const nvinfer1::AsciiChar* plugin_version = "";
338};
339
340#endif
341#endif
342} // namespace inference
343} // namespace perception
344} // namespace apollo
apollo::perception::inference::DFMBPSROIAlignPlugin::DFMBPSROIAlignPlugin
DFMBPSROIAlignPlugin(const DFMBPSROIAlignParameter &dfmb_psroi_align_parameter, Dims *in_dims, int nbInputs)
Definition dfmb_psroi_align_plugin.h:34
apollo::perception::inference::DFMBPSROIAlignPlugin::getWorkspaceSize
size_t getWorkspaceSize(int maxBatchSize) const override
Definition dfmb_psroi_align_plugin.h:99
apollo::perception::inference::DFMBPSROIAlignPlugin::getOutputDimensions
Dims getOutputDimensions(int index, const Dims *inputs, int nbInputDims) override
Definition dfmb_psroi_align_plugin.h:90
apollo::perception::inference::DFMBPSROIAlignPlugin::configure
void configure(const Dims *inputDims, int nbInputs, const Dims *outputDims, int nbOutputs, int maxBatchSize) override
Definition dfmb_psroi_align_plugin.h:96
apollo::perception::inference::DFMBPSROIAlignPlugin::enqueue
int enqueue(int batchSize, const void *const *inputs, void **outputs, void *workspace, cudaStream_t stream) override
apollo
class register implement
Definition arena_queue.h:37
rt_common.h