1 /* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
2
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations under the License.
14 ==============================================================================*/
15 #include <algorithm>
16 #include <cstdint>
17 #include <string>
18 #include <vector>
19
20 #include "tensorflow/core/lib/core/status.h"
21 #include "tensorflow/core/platform/logging.h"
22 #include "tensorflow/lite/toco/graph_transformations/graph_transformations.h"
23 #include "tensorflow/lite/toco/model.h"
24 #include "tensorflow/lite/toco/tooling_util.h"
25
26 // Sometimes, user may choose to use broadcast mul to perform nearest neighbor
27 // upsampling, how it works?
28 //
29 // If you want to upsample [batch, height, width, channel] to
30 // [batch, height * scale1, width * scale2, channel],
31 // you can do a broadcast mul like following:
32 // [batch, height, 1, width, 1 channel] with
33 // ones [1, 1, scale1, 1, scale2, channel].
34 //
35 // However, there's a 6-d tensor broadcasting mul operation which is not
36 // supported by the current runtime, also, it's less efficient for the quantized
37 // case if we can just do some memcpy.
38 // So we can transform this broadcast mul pattern to pack:
39 //
40 // [batch, height, 1, width, 1, channel]
41 // || Reshape
42 // \/
43 // [batch, height, width, channel]
44 // || Pack at axis 3, with scale2
45 // \/
46 // [batch, height, width, scale2, channel]
47 // || Pack at axis 2, with scale1
48 // \/
49 // [batch, height, scale1, width, scale2, channel]
50 namespace toco {
51 namespace {
52
CopyShape(const std::vector<int> & new_shape,Array * array)53 void CopyShape(const std::vector<int>& new_shape, Array* array) {
54 for (int dim : new_shape) {
55 array->mutable_shape()->mutable_dims()->push_back(dim);
56 }
57 }
58
59 template <ArrayDataType DataType, typename T>
HasSameValues(const Array & input_array,T value)60 bool HasSameValues(const Array& input_array, T value) {
61 auto& buffer = input_array.GetBuffer<DataType>();
62 for (int i = 0; i < buffer.Length(); ++i) {
63 if (buffer.data.at(i) != value) {
64 return false;
65 }
66 }
67 return true;
68 }
69
FindOperator(Model * model,const Operator & op)70 std::vector<std::unique_ptr<Operator>>::iterator FindOperator(
71 Model* model, const Operator& op) {
72 return std::find_if(
73 model->operators.begin(), model->operators.end(),
74 [&op](const std::unique_ptr<Operator>& ptr) { return ptr.get() == &op; });
75 }
76
77 } // namespace
78
79 // It's possible the model uses mul-broadcast to implement nearest neighbor
80 // upsample which may involve 5-d, 6-d tensors. We can actually change this
81 // pattern to be pack-based which is easier for us to handle.
Run(Model * model,std::size_t op_index,bool * modified)82 ::tensorflow::Status IdentifyNearestUpsample::Run(Model* model,
83 std::size_t op_index,
84 bool* modified) {
85 *modified = false;
86 auto op_it = model->operators.begin() + op_index;
87 auto* op = op_it->get();
88 if (op->type != OperatorType::kMul) {
89 return ::tensorflow::OkStatus();
90 }
91
92 // We only support one operand being constant.
93 const std::string& lhs = op->inputs.at(0);
94 const std::string& rhs = op->inputs.at(1);
95 const std::string& output = op->outputs.at(0);
96
97 Operator* next_op = GetOpWithOutput(*model, output);
98 if (next_op == nullptr) {
99 return ::tensorflow::OkStatus();
100 }
101
102 if (IsConstantParameterArray(*model, lhs) ==
103 IsConstantParameterArray(*model, rhs)) {
104 return ::tensorflow::OkStatus();
105 }
106
107 Array& const_array = IsConstantParameterArray(*model, lhs)
108 ? model->GetArray(lhs)
109 : model->GetArray(rhs);
110 Array& nonconst_array = IsConstantParameterArray(*model, lhs)
111 ? model->GetArray(rhs)
112 : model->GetArray(lhs);
113 Array& output_array = model->GetArray(output);
114
115 // Wait for shape propogation finished.
116 if (!const_array.has_shape() || !nonconst_array.has_shape() ||
117 !output_array.has_shape()) {
118 return ::tensorflow::OkStatus();
119 }
120
121 // We need to make sure they have same dimension count & the const parameter
122 // only contain ones.
123 if (const_array.shape().dimensions_count() !=
124 nonconst_array.shape().dimensions_count()) {
125 return ::tensorflow::OkStatus();
126 }
127
128 if (const_array.data_type == ArrayDataType::kFloat) {
129 if (!HasSameValues<ArrayDataType::kFloat, float>(const_array, 1))
130 return ::tensorflow::OkStatus();
131 } else if (const_array.data_type == ArrayDataType::kInt32) {
132 if (!HasSameValues<ArrayDataType::kInt32, int>(const_array, 1))
133 return ::tensorflow::OkStatus();
134 } else if (const_array.data_type == ArrayDataType::kInt8) {
135 if (!HasSameValues<ArrayDataType::kInt8, int8_t>(const_array, 127))
136 return ::tensorflow::OkStatus();
137 } else if (const_array.data_type == ArrayDataType::kUint8) {
138 if (!HasSameValues<ArrayDataType::kUint8, uint8_t>(const_array, 255))
139 return ::tensorflow::OkStatus();
140 } else {
141 return ::tensorflow::OkStatus();
142 }
143
144 // We're recognizing the following patterns:
145 // non-const shape: [..., M, 1, ..., N, 1, ...]
146 // const shape: [..., 1, scale_m, ..., 1, scale_n, ...]
147 // and base on that, we're imaging the original shape, which is
148 // [..., M, ... N, ...], then we're a serious of packing starting with the
149 // lowest axis.
150
151 // Scale_factors will store [scale_m, scale_n, ...].
152 std::vector<int> scale_factors;
153 // Pack_axis with store [axis_for_M + 1, axis_for_N + 1,...]
154 std::vector<int> pack_axis;
155 std::vector<int> imagined_original_shape;
156
157 const auto& current_const_shape = const_array.shape();
158 const auto& current_nonconst_shape = nonconst_array.shape();
159 int imaged_original_shape_current_axis = 0;
160 int i = 0;
161 for (; i < current_const_shape.dimensions_count() - 1;
162 ++i, ++imaged_original_shape_current_axis) {
163 if (current_const_shape.dims(i) == 1 &&
164 current_nonconst_shape.dims(i + 1) == 1 &&
165 current_nonconst_shape.dims(i) != 1) {
166 pack_axis.push_back(imaged_original_shape_current_axis + 1);
167 scale_factors.push_back(current_const_shape.dims(i + 1));
168 imagined_original_shape.push_back(current_nonconst_shape.dims(i));
169 ++i;
170 } else {
171 imagined_original_shape.push_back(current_nonconst_shape.dims(i));
172 }
173 }
174 // Push the rest dim.
175 for (; i < current_nonconst_shape.dimensions_count(); ++i) {
176 imagined_original_shape.push_back(current_nonconst_shape.dims(i));
177 }
178
179 if (pack_axis.empty()) {
180 return ::tensorflow::OkStatus();
181 }
182
183 std::vector<Operator*> to_be_inserted_ops;
184
185 // First put the reshape op.
186 // This reshape op will reshape the input tensor to what we imagined as the
187 // original shape.
188 auto* reshape_op = new TensorFlowReshapeOperator;
189 to_be_inserted_ops.push_back(reshape_op);
190 const std::string& original_array_name =
191 IsConstantParameterArray(*model, lhs) ? rhs : lhs;
192 reshape_op->inputs.push_back(original_array_name);
193
194 // Create shape param.
195 const std::string shape_array_name = AvailableArrayName(*model, "new_shape");
196 reshape_op->inputs.push_back(shape_array_name);
197 Array& shape_array = model->GetOrCreateArray(shape_array_name);
198 const int dim_size = imagined_original_shape.size();
199 *(shape_array.mutable_shape()->mutable_dims()) = {dim_size};
200 shape_array.data_type = ArrayDataType::kInt32;
201 auto& shape_buffer = shape_array.GetMutableBuffer<ArrayDataType::kInt32>();
202 // This is what imagined as the original shape.
203 for (size_t i = 0; i < imagined_original_shape.size(); ++i) {
204 shape_buffer.data.push_back(imagined_original_shape.at(i));
205 }
206
207 const std::string& reshape_output_name =
208 AvailableArrayName(*model, "reshape_output");
209 Array& reshape_output_array = model->GetOrCreateArray(reshape_output_name);
210 reshape_output_array.data_type = output_array.data_type;
211 CopyShape(imagined_original_shape, &reshape_output_array);
212
213 // Copy the quantization/minmax params if applicable.
214 if (output_array.minmax) {
215 reshape_output_array.GetOrCreateMinMax().min = output_array.minmax->min;
216 reshape_output_array.GetOrCreateMinMax().max = output_array.minmax->max;
217 }
218 if (output_array.quantization_params) {
219 reshape_output_array.GetOrCreateQuantizationParams().scale =
220 output_array.quantization_params->scale;
221 reshape_output_array.GetOrCreateQuantizationParams().zero_point =
222 output_array.quantization_params->zero_point;
223 }
224 reshape_op->outputs.push_back(reshape_output_name);
225
226 // Place the pack op as described in the file comment.
227 std::string current_pack_input_name = reshape_output_name;
228 std::vector<int> current_shape = imagined_original_shape;
229 for (int i = pack_axis.size() - 1; i >= 0; --i) {
230 auto* pack_op = new PackOperator;
231 int scale = scale_factors.at(i);
232 for (int j = 0; j < scale; ++j) {
233 pack_op->inputs.push_back(current_pack_input_name);
234 }
235 // The axis is computed before, the values count is actually the scale.
236 int axis = pack_axis.at(i);
237 pack_op->axis = axis;
238 pack_op->values_count = scale;
239 const std::string& pack_output_array_name =
240 AvailableArrayName(*model, absl::StrCat("pack_at_", axis));
241
242 // We need to copy the quantization/minmax params if applicable.
243 Array& pack_output_array = model->GetOrCreateArray(pack_output_array_name);
244 if (output_array.minmax) {
245 pack_output_array.GetOrCreateMinMax().min = output_array.minmax->min;
246 pack_output_array.GetOrCreateMinMax().max = output_array.minmax->max;
247 }
248 if (output_array.quantization_params) {
249 pack_output_array.GetOrCreateQuantizationParams().scale =
250 output_array.quantization_params->scale;
251 pack_output_array.GetOrCreateQuantizationParams().zero_point =
252 output_array.quantization_params->zero_point;
253 }
254
255 pack_output_array.data_type = nonconst_array.data_type;
256 current_shape.insert(current_shape.begin() + axis, scale);
257 CopyShape(current_shape, &pack_output_array);
258 pack_op->outputs.push_back(pack_output_array_name);
259
260 // The output is actually the input to the next pack op.
261 current_pack_input_name = pack_output_array_name;
262 to_be_inserted_ops.push_back(pack_op);
263 }
264
265 // Rewire the final pack op.
266 to_be_inserted_ops.at(to_be_inserted_ops.size() - 1)->outputs.clear();
267 to_be_inserted_ops.at(to_be_inserted_ops.size() - 1)
268 ->outputs.push_back(op->outputs.at(0));
269
270 // Insert all added ops in a reverse order.
271 for (int i = to_be_inserted_ops.size() - 1; i >= 0; --i) {
272 op_it = model->operators.emplace(op_it, to_be_inserted_ops.at(i));
273 }
274
275 // Delete the mul op.
276 model->operators.erase(FindOperator(model, *op));
277
278 *modified = true;
279 return ::tensorflow::OkStatus();
280 }
281
282 } // namespace toco
283