Apollo 11.0
自动驾驶开放平台
softmax_plugin.h
浏览该文件的文档.
1/******************************************************************************
2 * Copyright 2018 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 <string>
22
23#include "modules/perception/common/inference/tensorrt/rt_common.h"
24namespace apollo {
25namespace perception {
26namespace inference {
27
28#ifdef NV_TENSORRT_MAJOR
29#if NV_TENSORRT_MAJOR != 8
30class SoftmaxPlugin : public nvinfer1::IPlugin {
31 public:
32 SoftmaxPlugin(const SoftmaxParameter &param, nvinfer1::Dims in_dims) {
33 input_dims_.nbDims = in_dims.nbDims;
34 for (int i = 0; i < in_dims.nbDims; i++) {
35 input_dims_.d[i] = in_dims.d[i];
36 input_dims_.type[i] = in_dims.type[i];
37 }
38 axis_ = param.axis() - 1;
39 CHECK_GE(axis_, 0);
40 CHECK_LE(axis_ + 1, input_dims_.nbDims);
41
42 inner_num_ = 1;
43 for (int i = axis_ + 1; i < input_dims_.nbDims; i++) {
44 inner_num_ *= input_dims_.d[i];
45 }
46 outer_num_ = 1;
47 for (int i = 0; i < axis_; i++) {
48 outer_num_ *= input_dims_.d[i];
49 }
50 cudnnCreateTensorDescriptor(&input_desc_);
51 cudnnCreateTensorDescriptor(&output_desc_);
52 }
53
54 SoftmaxPlugin() {}
55
56 ~SoftmaxPlugin() {
57 cudnnDestroyTensorDescriptor(input_desc_);
58 cudnnDestroyTensorDescriptor(output_desc_);
59 }
60 virtual int initialize() {
61 cudnnCreate(&cudnn_); // initialize cudnn and cublas
62 cublasCreate(&cublas_);
63 return 0;
64 }
65 virtual void terminate() {
66 cublasDestroy(cublas_);
67 cudnnDestroy(cudnn_);
68 }
69 int getNbOutputs() const override { return 1; }
70
71 nvinfer1::Dims getOutputDimensions(int index, const nvinfer1::Dims *inputs,
72 int nbInputDims) override {
73 nvinfer1::Dims out_dims = inputs[0];
74 return out_dims;
75 }
76
77 void configure(const nvinfer1::Dims *inputDims, int nbInputs,
78 const nvinfer1::Dims *outputDims, int nbOutputs,
79 int maxBatchSize) override {
80 input_dims_ = inputDims[0];
81 }
82
83 size_t getWorkspaceSize(int maxBatchSize) const override { return 0; }
84
85 int enqueue(int batchSize, const void *const *inputs, void **outputs,
86 void *workspace, cudaStream_t stream) override;
87
88 size_t getSerializationSize() override { return 0; }
89
90 void serialize(void *buffer) override {
91 char *d = reinterpret_cast<char *>(buffer), *a = d;
92 size_t size = getSerializationSize();
93 CHECK_EQ(d, a + size);
94 }
95
96 private:
97 cudnnHandle_t cudnn_;
98 cublasHandle_t cublas_;
99 nvinfer1::Dims input_dims_;
100 int axis_;
101 int inner_num_;
102 int outer_num_;
103 cudnnTensorDescriptor_t input_desc_;
104 cudnnTensorDescriptor_t output_desc_;
105};
106
107#else
108class SoftmaxPlugin : public nvinfer1::IPluginV2 {
109 public:
110 SoftmaxPlugin(const char *layer_name, const void *serial_data,
111 size_t serial_length) {
112 const char *p = reinterpret_cast<const char *>(serial_data);
113 // CHECK_EQ(getPluginType(), Name());
114 // p += (Name().size() + strlen(name_));
115
116 input_dims_.nbDims = reinterpret_cast<const int *>(p)[0];
117 p += sizeof(int);
118 for (int i = 0; i < input_dims_.nbDims; ++i) {
119 input_dims_.d[i] = reinterpret_cast<const int *>(p)[0];
120 p += sizeof(int);
121 // input_dims_.type[i] = static_cast<nvinfer1::DimensionType>(
122 // reinterpret_cast<const int *>(p)[0]);
123 // p += sizeof(int);
124 }
125
126 inner_num_ = reinterpret_cast<const int *>(p)[0];
127 p += sizeof(int);
128
129 outer_num_ = reinterpret_cast<const int *>(p)[0];
130 p += sizeof(int);
131
132 axis_ = reinterpret_cast<const int *>(p)[0];
133 CHECK_GT(axis_, 0);
134 // CHECK_LE(axis_ + 1, input_dims_.nbDims);
135 CHECK_LE(axis_, input_dims_.nbDims);
136 p += sizeof(int);
137
138 CHECK_EQ(p, reinterpret_cast<const char *>(serial_data) + serial_length);
139 CHECK_EQ(getSerializationSize(), serial_length);
140
141 cudnnCreateTensorDescriptor(&input_desc_);
142 cudnnCreateTensorDescriptor(&output_desc_);
143 }
144
145 SoftmaxPlugin(const SoftmaxParameter &param, nvinfer1::Dims in_dims) {
146 input_dims_.nbDims = in_dims.nbDims;
147 for (int i = 0; i < in_dims.nbDims; i++) {
148 input_dims_.d[i] = in_dims.d[i];
149 // input_dims_.type[i] = in_dims.type[i];
150 }
151 axis_ = param.axis() - 1;
152 CHECK_GE(axis_, 0);
153 CHECK_LE(axis_ + 1, input_dims_.nbDims);
154
155 inner_num_ = 1;
156 for (int i = axis_ + 1; i < input_dims_.nbDims; i++) {
157 inner_num_ *= input_dims_.d[i];
158 }
159 outer_num_ = 1;
160 for (int i = 0; i < axis_; i++) {
161 outer_num_ *= input_dims_.d[i];
162 }
163 cudnnCreateTensorDescriptor(&input_desc_);
164 cudnnCreateTensorDescriptor(&output_desc_);
165 }
166
167 SoftmaxPlugin(int axis, nvinfer1::Dims in_dims) {
168 input_dims_.nbDims = in_dims.nbDims;
169 for (int i = 0; i < in_dims.nbDims; i++) {
170 input_dims_.d[i] = in_dims.d[i];
171 // input_dims_.type[i] = in_dims.type[i];
172 }
173 axis_ = axis - 1;
174 CHECK_GE(axis_, 0);
175 CHECK_LE(axis_ + 1, input_dims_.nbDims);
176
177 inner_num_ = 1;
178 for (int i = axis_ + 1; i < input_dims_.nbDims; i++) {
179 inner_num_ *= input_dims_.d[i];
180 }
181 outer_num_ = 1;
182 for (int i = 0; i < axis_; i++) {
183 outer_num_ *= input_dims_.d[i];
184 }
185 cudnnCreateTensorDescriptor(&input_desc_);
186 cudnnCreateTensorDescriptor(&output_desc_);
187 }
188 SoftmaxPlugin() {}
189 ~SoftmaxPlugin() {
190 cudnnDestroyTensorDescriptor(input_desc_);
191 cudnnDestroyTensorDescriptor(output_desc_);
192 }
193
194 virtual int initialize() noexcept {
195 cudnnCreate(&cudnn_); // initialize cudnn and cublas
196 cublasCreate(&cublas_);
197 return 0;
198 }
199 virtual void terminate() noexcept {
200 cublasDestroy(cublas_);
201 cudnnDestroy(cudnn_);
202 }
203 int getNbOutputs() const noexcept override { return 1; }
204
205 nvinfer1::Dims getOutputDimensions(int index, const nvinfer1::Dims *inputs,
206 int nbInputDims) noexcept override {
207 nvinfer1::Dims out_dims = inputs[0];
208 return out_dims;
209 }
210
211 void configureWithFormat(const nvinfer1::Dims *inputDims, int nbInputs,
212 const nvinfer1::Dims *outputDims, int nbOutputs,
213 nvinfer1::DataType type,
214 nvinfer1::PluginFormat format,
215 int maxBatchSize) noexcept override {
216 input_dims_ = inputDims[0];
217 }
218 size_t getWorkspaceSize(int maxBatchSize) const noexcept override {
219 return 0;
220 }
221
222 int enqueue(int batchSize, const void *const *inputs, void *const *outputs,
223 void *workspace, cudaStream_t stream) noexcept override;
224
225 size_t getSerializationSize() const noexcept override {
226 /* return sizeof(size_t) + Name().size() + \ */
227 return sizeof(int) + sizeof(int) * 1 * input_dims_.nbDims +
228 sizeof(inner_num_) + sizeof(outer_num_) + sizeof(axis_);
229 }
230
231 void serialize(void *buffer) const noexcept override {
232 char *p = reinterpret_cast<char *>(buffer);
233 // setPluginType(&p, Name());
234
235 reinterpret_cast<int *>(p)[0] = input_dims_.nbDims;
236 p += sizeof(int);
237 for (int i = 0; i < input_dims_.nbDims; ++i) {
238 reinterpret_cast<int *>(p)[0] = input_dims_.d[i];
239 p += sizeof(int);
240 // reinterpret_cast<int *>(p)[0] = static_cast<int>(input_dims_.type[i]);
241 // p += sizeof(int);
242 }
243
244 reinterpret_cast<int *>(p)[0] = inner_num_;
245 p += sizeof(int);
246
247 reinterpret_cast<int *>(p)[0] = outer_num_;
248 p += sizeof(int);
249
250 reinterpret_cast<int *>(p)[0] = axis_;
251 p += sizeof(int);
252 }
253
254 std::string Name() const { return "softmax_plugin"; }
255
256 const char *getPluginType() const noexcept override {
257 return "softmax_plugin";
258 };
259
260 const char *getPluginVersion() const noexcept override { return "1"; }
261
262 void setPluginNamespace(const char *s_name) noexcept override {
263 name_ = s_name;
264 }
265
266 const char *getPluginNamespace() const noexcept override { return "camera"; }
267 bool supportsFormat(nvinfer1::DataType type,
268 nvinfer1::PluginFormat format) const noexcept override {
269 return true;
270 }
271
272 void destroy() noexcept override { delete this; }
273
274 nvinfer1::IPluginV2 *clone() const noexcept override {
275 SoftmaxPlugin *p = new SoftmaxPlugin();
276 cudnnCreate(&(p->cudnn_)); // initialize cudnn and cublas
277 cublasCreate(&(p->cublas_));
278 p->axis_ = axis_;
279 p->inner_num_ = inner_num_;
280 p->outer_num_ = outer_num_;
281 (p->input_dims_).nbDims = input_dims_.nbDims;
282 for (int i = 0; i < input_dims_.nbDims; i++) {
283 (p->input_dims_).d[i] = input_dims_.d[i];
284 }
285 cudnnCreateTensorDescriptor(&(p->input_desc_));
286 cudnnCreateTensorDescriptor(&(p->output_desc_));
287 return p;
288 }
289
290 nvinfer1::DataType getOutputDataType(int32_t index,
291 nvinfer1::DataType const *inputTypes,
292 int32_t nbInputs) const noexcept {
293 return nvinfer1::DataType::kFLOAT;
294 }
295
296 void configurePlugin(nvinfer1::Dims const *inputDims, int32_t nbInputs,
297 nvinfer1::Dims const *outputDims, int32_t nbOutputs,
298 nvinfer1::DataType const *inputTypes,
299 nvinfer1::DataType const *outputTypes,
300 bool const *inputIsBroadcast,
301 bool const *outputIsBroadcast,
302 nvinfer1::PluginFormat floatFormat,
303 int32_t maxBatchSize) noexcept {}
304
305 private:
306 const char *name_ = "camera";
307 cudnnHandle_t cudnn_;
308 cublasHandle_t cublas_;
309 nvinfer1::Dims input_dims_;
310 int axis_;
311 int inner_num_;
312 int outer_num_;
313 cudnnTensorDescriptor_t input_desc_;
314 cudnnTensorDescriptor_t output_desc_;
315};
316
317class SoftmaxPluginCreator : public nvinfer1::IPluginCreator {
318 public:
319 SoftmaxPluginCreator() {}
320 ~SoftmaxPluginCreator() {
321 // terminae();
322 }
323 nvinfer1::IPluginV2 *createPlugin(
324 const char *name,
325 const nvinfer1::PluginFieldCollection *fc) noexcept override {
326 int axis;
327 nvinfer1::Dims in_dims;
328 const nvinfer1::PluginField *fields = fc->fields;
329 for (int i = 0; i < fc->nbFields; ++i) {
330 const char *attrName = fields[i].name;
331 if (!strcmp(attrName, "axis")) {
332 // ASSERT(fields[i].type == nvinfer1::PluginFieldType::kINT32);
333 axis = *(static_cast<const int32_t *>(fields[i].data));
334 } else if (!strcmp(attrName, "in_dims")) {
335 // ASSERT(fields[i].type ==nvinfer1::PluginFieldType::kDIMS);
336 in_dims = *(static_cast<const nvinfer1::Dims *>(fields[i].data));
337 }
338 }
339
340 SoftmaxPlugin *obj = new SoftmaxPlugin(axis, in_dims);
341 obj->setPluginNamespace(mNamespaces.c_str());
342 return obj;
343 }
344 nvinfer1::IPluginV2 *deserializePlugin(
345 const char *name, const void *serial_data,
346 size_t serial_length) noexcept override {
347 return new SoftmaxPlugin(name, serial_data, serial_length);
348 }
349 const char *getPluginName() const noexcept override {
350 return "softmax_plugincamera";
351 }
352 const char *getPluginVersion() const noexcept override { return "1"; }
353 const char *getPluginNamespace() const noexcept override { return "camera"; }
354 void setPluginNamespace(const char *s_name) noexcept override {
355 mNamespaces = s_name;
356 }
357 const nvinfer1::PluginFieldCollection *getFieldNames() noexcept override {
358 return nullptr;
359 }
360
361 private:
362 std::string mNamespaces = "camera";
363};
364REGISTER_TENSORRT_PLUGIN(SoftmaxPluginCreator);
365#endif
366#endif
367
368} // namespace inference
369} // namespace perception
370} // namespace apollo
apollo
class register implement
Definition arena_queue.h:37
rt_common.h