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 <numeric>
18 #include <functional>
19 #include "src/delegate/tensorrt/op/scale_tensorrt.h"
20 #include "src/delegate/tensorrt/op/activation_tensorrt.h"
21 #include "src/delegate/tensorrt/tensorrt_utils.h"
22
23 namespace mindspore::lite {
24 constexpr int SCALE_INDEX = 1;
25 constexpr int SHIFT_INDEX = 2;
26 constexpr int POWER_INDEX = 3;
27
IsSupport(const schema::Primitive * primitive,const std::vector<mindspore::MSTensor> & in_tensors,const std::vector<mindspore::MSTensor> & out_tensors)28 int ScaleTensorRT::IsSupport(const schema::Primitive *primitive, const std::vector<mindspore::MSTensor> &in_tensors,
29 const std::vector<mindspore::MSTensor> &out_tensors) {
30 if (!IsShapeKnown()) {
31 MS_LOG(ERROR) << "Unsupported input tensor unknown shape: " << op_name_;
32 return RET_ERROR;
33 }
34 if (in_tensors.size() != INPUT_SIZE2 && in_tensors.size() != INPUT_SIZE3 && in_tensors.size() != INPUT_SIZE4) {
35 MS_LOG(ERROR) << "Unsupported input tensor size, size is: " << in_tensors.size();
36 return RET_ERROR;
37 }
38 if (out_tensors.size() != 1) {
39 MS_LOG(ERROR) << "Unsupported output tensor size, size is: " << out_tensors.size();
40 return RET_ERROR;
41 }
42 return RET_OK;
43 }
44
AddInnerOp(nvinfer1::INetworkDefinition * network)45 int ScaleTensorRT::AddInnerOp(nvinfer1::INetworkDefinition *network) {
46 if (network == nullptr) {
47 MS_LOG(ERROR) << "network is invalid";
48 return RET_ERROR;
49 }
50 auto scale_op = op_primitive_->value_as_ScaleFusion();
51 if (scale_op == nullptr) {
52 MS_LOG(ERROR) << "convert failed";
53 return RET_ERROR;
54 }
55
56 schema::ActivationType activation_type = scale_op->activation_type();
57
58 // mode of scale
59 int64_t axis = scale_op->axis();
60 if (axis == -1) {
61 axis = static_cast<int64_t>(in_tensors_[0].Shape().size() - 1);
62 }
63 mode_ = GetScaleMode(axis);
64 out_format_ = tensorrt_in_tensors_[0].format_;
65 MS_LOG(DEBUG) << "before transpose " << GetTensorFormat(tensorrt_in_tensors_[0].trt_tensor_, out_format_);
66
67 nvinfer1::ITensor *scale_in_tensor = PreProcessInputTensor(network);
68 if (scale_in_tensor == nullptr) {
69 MS_LOG(ERROR) << "PreProcessInputTensor failed: " << op_name_;
70 return RET_ERROR;
71 }
72
73 MS_LOG(DEBUG) << "after transpose " << GetTensorFormat(scale_in_tensor, out_format_);
74
75 bool nd = false;
76 // (input * scale + shift) ^ power
77 nvinfer1::Weights power{nvinfer1::DataType::kFLOAT, nullptr, 0};
78 nvinfer1::Weights shift{nvinfer1::DataType::kFLOAT, nullptr, 0};
79 nvinfer1::Weights scale{nvinfer1::DataType::kFLOAT, nullptr, 0};
80 if (in_tensors_.size() > SCALE_INDEX) {
81 scale.values = in_tensors_[SCALE_INDEX].MutableData();
82 MS_ASSERT(scale.values);
83 scale.count = in_tensors_[SCALE_INDEX].ElementNum();
84 scale.type = ConvertDataType(in_tensors_[SCALE_INDEX].DataType());
85 shift.type = scale.type;
86 power.type = scale.type;
87 nd = in_tensors_[1].Shape().size() == 1 ? false : true;
88 }
89 if (in_tensors_.size() > SHIFT_INDEX) {
90 shift.values = in_tensors_[SHIFT_INDEX].MutableData();
91 MS_ASSERT(shift.values);
92 shift.count = in_tensors_[SHIFT_INDEX].ElementNum();
93 }
94 if (in_tensors_.size() > POWER_INDEX) {
95 power.values = in_tensors_[POWER_INDEX].MutableData();
96 MS_ASSERT(power.values);
97 power.count = in_tensors_[POWER_INDEX].ElementNum();
98 }
99 nvinfer1::IScaleLayer *cal_layer = nullptr;
100
101 if (nd) {
102 MS_LOG(WARNING) << "multi dims ScaleMode enter";
103 cal_layer = network->addScaleNd(*scale_in_tensor, mode_, shift, scale, power, axis);
104 } else {
105 cal_layer = network->addScale(*scale_in_tensor, mode_, shift, scale, power);
106 }
107
108 if (cal_layer == nullptr) {
109 MS_LOG(ERROR) << "addScaleNd failed for: " << op_name_;
110 return RET_ERROR;
111 }
112 cal_layer->setName(op_name_.c_str());
113
114 // add activation
115 nvinfer1::ITensor *activation_tensor = cal_layer->getOutput(0);
116 if (activation_type != schema::ActivationType::ActivationType_NO_ACTIVATION) {
117 auto activation_layer = ActivationTensorRT::AddActivation(network, activation_type, 0, cal_layer->getOutput(0));
118 if (activation_layer == nullptr) {
119 MS_LOG(ERROR) << "addActivation for scale failed";
120 return RET_ERROR;
121 }
122 activation_layer->setName((op_name_ + "_activation").c_str());
123 activation_tensor = activation_layer->getOutput(0);
124 }
125
126 // squeeze to origin dim
127 nvinfer1::ITensor *op_out_tensor = activation_tensor;
128 if (activation_tensor->getDimensions().nbDims > static_cast<int>(out_tensors_[0].Shape().size())) {
129 op_out_tensor = AddSqueezeOp(activation_tensor, network);
130 }
131 op_out_tensor->setName((op_name_ + "_output").c_str());
132 this->AddInnerOutTensors(ITensorHelper{op_out_tensor, out_format_});
133 MS_LOG(DEBUG) << "output " << GetTensorFormat(op_out_tensor, out_format_);
134 return RET_OK;
135 }
136
PreProcessInputTensor(nvinfer1::INetworkDefinition * network)137 nvinfer1::ITensor *ScaleTensorRT::PreProcessInputTensor(nvinfer1::INetworkDefinition *network) {
138 nvinfer1::ITensor *scale_in_tensor = tensorrt_in_tensors_[0].trt_tensor_;
139 if (in_tensors_[0].Shape().size() < INPUT_SIZE4) {
140 // unsqueeze input Itensor to 4 dims
141 scale_in_tensor = AddUnsqueezeOp(network);
142 if (scale_in_tensor == nullptr) {
143 MS_LOG(ERROR) << "AddUnsqueezeOp failed";
144 return nullptr;
145 }
146 } else if (tensorrt_in_tensors_[0].trt_tensor_->getDimensions().nbDims == DIMENSION_4D &&
147 mode_ == nvinfer1::ScaleMode::kCHANNEL && tensorrt_in_tensors_[0].format_ == Format::NHWC) {
148 // per channel input format should be nchw, otherwise should be same with scale nhwc
149 // transpose: NHWC->NCHW
150 nvinfer1::IShuffleLayer *transpose_layer_in = NHWC2NCHW(network, *tensorrt_in_tensors_[0].trt_tensor_);
151 if (transpose_layer_in == nullptr) {
152 MS_LOG(ERROR) << "op action convert failed";
153 return nullptr;
154 }
155 transpose_layer_in->setName((op_name_ + "_transpose2NCHW").c_str());
156 scale_in_tensor = transpose_layer_in->getOutput(0);
157 out_format_ = Format::NCHW;
158 } else if (tensorrt_in_tensors_[0].trt_tensor_->getDimensions().nbDims == DIMENSION_4D &&
159 tensorrt_in_tensors_[0].format_ == Format::NCHW && mode_ == nvinfer1::ScaleMode::kELEMENTWISE) {
160 // transpose: NCHW->NHWC
161 nvinfer1::IShuffleLayer *transpose_layer_in = NCHW2NHWC(network, *tensorrt_in_tensors_[0].trt_tensor_);
162 if (transpose_layer_in == nullptr) {
163 MS_LOG(ERROR) << "op action convert failed";
164 return nullptr;
165 }
166 transpose_layer_in->setName((op_name_ + "_transpose2NHWC").c_str());
167 scale_in_tensor = transpose_layer_in->getOutput(0);
168 out_format_ = Format::NHWC;
169 }
170 return scale_in_tensor;
171 }
172
GetScaleMode(int64_t axis)173 nvinfer1::ScaleMode ScaleTensorRT::GetScaleMode(int64_t axis) {
174 nvinfer1::ScaleMode mode = nvinfer1::ScaleMode::kUNIFORM;
175 auto input_data_shape = in_tensors_[0].Shape();
176 auto input_weight_shape = in_tensors_[1].Shape();
177 int total = std::accumulate(input_data_shape.begin(), input_data_shape.end(), 1, std::multiplies<int>());
178 MS_LOG(DEBUG) << "input tensor element cnt: " << total;
179 if (input_weight_shape.size() == 0 || (input_weight_shape.size() == 1 && input_weight_shape[0] == 1)) {
180 mode = nvinfer1::ScaleMode::kUNIFORM;
181 } else if ((axis < static_cast<int64_t>(input_data_shape.size()) && input_weight_shape.size() == 1 &&
182 input_data_shape[axis] == input_weight_shape[0]) ||
183 (input_data_shape.size() == DIMENSION_4D && axis == DIMENSION_3D)) {
184 mode = nvinfer1::ScaleMode::kCHANNEL;
185 } else if (input_weight_shape.size() == 1 && input_weight_shape[0] == total) {
186 mode = nvinfer1::ScaleMode::kELEMENTWISE;
187 } else {
188 MS_LOG(ERROR) << "ScaleMode create failed: " << op_name_;
189 return mode;
190 }
191 MS_LOG(DEBUG) << op_name_ << " ScaleMode(UNIFORM 0, CHANNEL 1, ELEMENTWISE 2): " << static_cast<int>(mode);
192 return mode;
193 }
194
AddUnsqueezeOp(nvinfer1::INetworkDefinition * network)195 nvinfer1::ITensor *ScaleTensorRT::AddUnsqueezeOp(nvinfer1::INetworkDefinition *network) {
196 nvinfer1::IShuffleLayer *unsqueeze_layer = network->addShuffle(*this->tensorrt_in_tensors_[0].trt_tensor_);
197 if (unsqueeze_layer == nullptr) {
198 MS_LOG(ERROR) << "addShuffle failed for: " << op_name_;
199 return nullptr;
200 }
201 unsqueeze_layer->setName((op_name_ + "_unsqueeze").c_str());
202 auto unsqueeze_shape = in_tensors_[0].Shape();
203 size_t unsqueeze_size = DIMENSION_4D - unsqueeze_shape.size();
204 for (size_t i = 0; i < unsqueeze_size; i++) {
205 unsqueeze_shape.push_back(1);
206 }
207 nvinfer1::Dims unsqueeze_dims = lite::ConvertCudaDims(unsqueeze_shape);
208 unsqueeze_layer->setReshapeDimensions(unsqueeze_dims);
209 return unsqueeze_layer->getOutput(0);
210 }
211
AddSqueezeOp(nvinfer1::ITensor * in_tensor,nvinfer1::INetworkDefinition * network)212 nvinfer1::ITensor *ScaleTensorRT::AddSqueezeOp(nvinfer1::ITensor *in_tensor, nvinfer1::INetworkDefinition *network) {
213 nvinfer1::IShuffleLayer *squeeze_layer = network->addShuffle(*in_tensor);
214 if (squeeze_layer == nullptr) {
215 MS_LOG(ERROR) << "addShuffle failed for: " << op_name_;
216 return nullptr;
217 }
218 squeeze_layer->setName((op_name_ + "_squeeze").c_str());
219 nvinfer1::Dims squeeze_dims = lite::ConvertCudaDims(out_tensors_[0].Shape());
220 MS_LOG(DEBUG) << "squeeze_dims cnt for scale: " << squeeze_dims.nbDims;
221 squeeze_layer->setReshapeDimensions(squeeze_dims);
222 return squeeze_layer->getOutput(0);
223 }
224 } // namespace mindspore::lite
225