1 /**
2 * Copyright 2021 Huawei Technologies Co., Ltd
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 #include "src/delegate/tensorrt/op/pool_tensorrt.h"
18 #include "src/delegate/tensorrt/op/activation_tensorrt.h"
19 #include "src/delegate/tensorrt/tensorrt_utils.h"
20
21 namespace mindspore::lite {
IsSupport(const mindspore::schema::Primitive * primitive,const std::vector<mindspore::MSTensor> & in_tensors,const std::vector<mindspore::MSTensor> & out_tensors)22 int PoolTensorRT::IsSupport(const mindspore::schema::Primitive *primitive,
23 const std::vector<mindspore::MSTensor> &in_tensors,
24 const std::vector<mindspore::MSTensor> &out_tensors) {
25 if (!IsShapeKnown()) {
26 MS_LOG(ERROR) << "Unsupported input tensor unknown shape: " << op_name_;
27 return RET_ERROR;
28 }
29 if (in_tensors.size() != 1) {
30 MS_LOG(ERROR) << "Unsupported input tensor size, size is " << in_tensors.size();
31 return RET_ERROR;
32 }
33 if (out_tensors.size() != 1) {
34 MS_LOG(ERROR) << "Unsupported output tensor size, size is " << out_tensors.size();
35 return RET_ERROR;
36 }
37 if (in_tensors[0].format() != Format::NHWC && in_tensors[0].format() != Format::NCHW) {
38 MS_LOG(ERROR) << "Unsupported input tensor format of " << in_tensors[0].format();
39 return RET_ERROR;
40 }
41 return RET_OK;
42 }
43
AddInnerOp(nvinfer1::INetworkDefinition * network)44 int PoolTensorRT::AddInnerOp(nvinfer1::INetworkDefinition *network) {
45 if (tensorrt_in_tensors_.size() != 1) {
46 MS_LOG(ERROR) << "invalid input tensor size: " << tensorrt_in_tensors_.size();
47 return RET_ERROR;
48 }
49 int ret = ParseParams();
50 if (ret != RET_OK) {
51 MS_LOG(ERROR) << "ParseParams failed for : " << op_name_;
52 return RET_ERROR;
53 }
54 MS_LOG(DEBUG) << "before transpose "
55 << GetTensorFormat(tensorrt_in_tensors_[0].trt_tensor_, tensorrt_in_tensors_[0].format_);
56
57 nvinfer1::ITensor *pool_input = tensorrt_in_tensors_[0].trt_tensor_;
58 if (tensorrt_in_tensors_[0].trt_tensor_->getDimensions().nbDims == DIMENSION_4D &&
59 tensorrt_in_tensors_[0].format_ == Format::NHWC) {
60 // transpose: NHWC->NCHW
61 nvinfer1::IShuffleLayer *transpose_layer_in = NHWC2NCHW(network, *tensorrt_in_tensors_[0].trt_tensor_);
62 if (transpose_layer_in == nullptr) {
63 MS_LOG(ERROR) << "transpose: NHWC->NCHW failed";
64 return RET_ERROR;
65 }
66 transpose_layer_in->setName((op_name_ + "_transpose2NCHW").c_str());
67 pool_input = transpose_layer_in->getOutput(0);
68 }
69
70 // pooling layer
71 nvinfer1::Dims windowSize = lite::ConvertCudaDims(kernel_size_);
72 nvinfer1::IPoolingLayer *pooling_layer = network->addPoolingNd(*pool_input, pooling_type_, windowSize);
73 if (pooling_layer == nullptr) {
74 MS_LOG(ERROR) << "addPoolingNd failed for TensorRT.";
75 return RET_ERROR;
76 }
77 AddParams(pooling_layer);
78 pooling_layer->setName(op_name_.c_str());
79
80 // add activation
81 nvinfer1::ILayer *activation_layer = nullptr;
82 if (activation_type_ == schema::ActivationType::ActivationType_NO_ACTIVATION) {
83 activation_layer = pooling_layer;
84 } else {
85 activation_layer = ActivationTensorRT::AddActivation(network, activation_type_, 0, pooling_layer->getOutput(0));
86 if (activation_layer == nullptr) {
87 MS_LOG(ERROR) << "addActivation for pool failed";
88 return RET_ERROR;
89 }
90 activation_layer->setName((op_name_ + "_activation").c_str());
91 }
92 activation_layer->getOutput(0)->setName((op_name_ + "_output").c_str());
93 this->AddInnerOutTensors(ITensorHelper{activation_layer->getOutput(0), Format::NCHW});
94 MS_LOG(DEBUG) << "output " << GetTensorFormat(activation_layer->getOutput(0), Format::NCHW);
95 return RET_OK;
96 }
97
ParseParams()98 int PoolTensorRT::ParseParams() {
99 switch (type_) {
100 case (schema::PrimitiveType_AvgPoolFusion): {
101 const schema::AvgPoolFusion *pool_primitive = this->GetPrimitive()->value_as_AvgPoolFusion();
102 if (pool_primitive == nullptr) {
103 MS_LOG(ERROR) << "convert PoolFusion failed: " << op_name_;
104 return RET_ERROR;
105 }
106 pooling_type_ = nvinfer1::PoolingType::kAVERAGE;
107
108 auto kernel_size = pool_primitive->kernel_size();
109 if (kernel_size == nullptr) {
110 MS_LOG(ERROR) << "get kernel size failed: " << op_name_;
111 return RET_ERROR;
112 }
113 kernel_size_ = std::vector<int64_t>(kernel_size->begin(), kernel_size->end());
114
115 auto stride = pool_primitive->strides();
116 if (stride == nullptr) {
117 MS_LOG(ERROR) << "get stride failed: " << op_name_;
118 return RET_ERROR;
119 }
120 stride_ = std::vector<int64_t>(stride->begin(), stride->end());
121
122 auto padding = pool_primitive->pad();
123 if (padding == nullptr) {
124 MS_LOG(ERROR) << "get padding failed: " << op_name_;
125 return RET_ERROR;
126 }
127 padding_ = std::vector<int64_t>(padding->begin(), padding->end());
128
129 pad_mode_ = pool_primitive->pad_mode();
130 activation_type_ = pool_primitive->activation_type();
131 break;
132 }
133 case (schema::PrimitiveType_MaxPoolFusion): {
134 const schema::MaxPoolFusion *pool_primitive = this->GetPrimitive()->value_as_MaxPoolFusion();
135 if (pool_primitive == nullptr) {
136 MS_LOG(ERROR) << "convert PoolFusion failed: " << op_name_;
137 return RET_ERROR;
138 }
139 pooling_type_ = nvinfer1::PoolingType::kMAX;
140
141 auto kernel_size = pool_primitive->kernel_size();
142 if (kernel_size == nullptr) {
143 MS_LOG(ERROR) << "get kernel size failed: " << op_name_;
144 return RET_ERROR;
145 }
146 kernel_size_ = std::vector<int64_t>(kernel_size->begin(), kernel_size->end());
147
148 auto stride = pool_primitive->strides();
149 if (stride == nullptr) {
150 MS_LOG(ERROR) << "get stride failed: " << op_name_;
151 return RET_ERROR;
152 }
153 stride_ = std::vector<int64_t>(stride->begin(), stride->end());
154
155 auto padding = pool_primitive->pad();
156 if (padding == nullptr) {
157 MS_LOG(ERROR) << "get padding failed: " << op_name_;
158 return RET_ERROR;
159 }
160 padding_ = std::vector<int64_t>(padding->begin(), padding->end());
161
162 pad_mode_ = pool_primitive->pad_mode();
163 activation_type_ = pool_primitive->activation_type();
164 break;
165 }
166 default: {
167 MS_LOG(ERROR) << "unsupported primitive type of " << type_ << " for node: " << op_name_;
168 return RET_ERROR;
169 }
170 }
171 return RET_OK;
172 }
173
AddParams(nvinfer1::IPoolingLayer * pooling_layer)174 void PoolTensorRT::AddParams(nvinfer1::IPoolingLayer *pooling_layer) {
175 nvinfer1::Dims stride_dims = ConvertCudaDims(stride_);
176 pooling_layer->setStrideNd(stride_dims);
177 if (pad_mode_ == schema::PadMode::PadMode_SAME) {
178 pooling_layer->setPaddingMode(nvinfer1::PaddingMode::kSAME_UPPER);
179 } else {
180 nvinfer1::Dims dims{};
181 dims.nbDims = DIMENSION_2D;
182 dims.d[0] = padding_[0];
183 dims.d[1] = padding_[DIMENSION_2D];
184 pooling_layer->setPaddingNd(dims);
185 }
186 }
187 } // namespace mindspore::lite
188