1 | /* Copyright 2017 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 | #ifndef TENSORFLOW_LITE_TOCO_TOOLING_UTIL_H_ |
16 | #define TENSORFLOW_LITE_TOCO_TOOLING_UTIL_H_ |
17 | |
18 | #include <algorithm> |
19 | #include <cmath> |
20 | #include <functional> |
21 | #include <iostream> |
22 | #include <limits> |
23 | #include <memory> |
24 | #include <string> |
25 | #include <vector> |
26 | |
27 | #include "absl/strings/string_view.h" |
28 | #include "tensorflow/core/lib/core/errors.h" |
29 | #include "tensorflow/core/lib/core/status.h" |
30 | #include "tensorflow/core/platform/logging.h" |
31 | #include "tensorflow/lite/kernels/internal/types.h" |
32 | #include "tensorflow/lite/toco/model.h" |
33 | #include "tensorflow/lite/toco/model_flags.pb.h" |
34 | #include "tensorflow/lite/toco/runtime/types.h" |
35 | #include "tensorflow/lite/toco/toco_flags.pb.h" |
36 | #include "tensorflow/lite/toco/types.pb.h" |
37 | |
38 | // TODO(aselle): Replace with using a container specific hash override instead. |
39 | namespace std { |
40 | template <> |
41 | struct hash<toco::OperatorType> { |
42 | size_t operator()(const toco::OperatorType& op) const { |
43 | return std::hash<size_t>()(static_cast<size_t>(op)); |
44 | } |
45 | }; |
46 | } // namespace std |
47 | |
48 | namespace toco { |
49 | |
50 | constexpr int kLogLevelModelChanged = 1; |
51 | constexpr int kLogLevelModelUnchanged = 2; |
52 | |
53 | absl::string_view FindLongestCommonPrefix(absl::string_view a, |
54 | absl::string_view b); |
55 | std::string LogName(const Operator& op); |
56 | |
57 | std::string ArrayDataTypeName(ArrayDataType data_type); |
58 | |
59 | // Returns true if the given array is specified as a model input array. |
60 | bool IsInputArray(const Model& model, const std::string& array_name); |
61 | // Returns true if the given array is specified as a model output array. |
62 | bool IsOutputArray(const Model& model, const std::string& array_name); |
63 | |
64 | bool IsArrayConsumed(const Model& model, const std::string& name); |
65 | int CountTrueOutputs(const Model& model, const Operator& op); |
66 | |
67 | int CountOpsWithInput(const Model& model, const std::string& array_name); |
68 | bool DeleteArrayIfUnused(const std::string& array_name, Model* model); |
69 | |
70 | // Deletes the op and any of its input and output arrays if they are unused |
71 | // after the op has been deleted. |
72 | void DeleteOpAndArrays(Model* model, const Operator* op); |
73 | |
74 | std::vector<std::unique_ptr<Operator>>::const_iterator FindOpWithOutput( |
75 | const Model& model, const std::string& array_name); |
76 | Operator* GetOpWithOutput(const Model& model, const std::string& array_name); |
77 | |
78 | std::vector<std::unique_ptr<Operator>>::iterator FindOpWithOutput( |
79 | Model& model, const std::string& array_name); |
80 | |
81 | std::vector<std::unique_ptr<Operator>>::const_iterator FindOpWithInput( |
82 | const Model& model, const std::string& array_name); |
83 | |
84 | std::vector<std::unique_ptr<Operator>>::iterator FindOpWithInput( |
85 | Model& model, const std::string& array_name); |
86 | |
87 | Operator* GetOpWithInput(const Model& model, const std::string& array_name); |
88 | Operator* GetFirstOpWithInput(const Model& model, |
89 | const std::string& array_name); |
90 | |
91 | // Replaces all uses of the |old_array_name| with the |new_array_name|. |
92 | void ReplaceArrayUsage(Model* model, const std::string& old_array_name, |
93 | const std::string& new_array_name); |
94 | |
95 | std::vector<std::unique_ptr<Operator>>::const_iterator FindOp( |
96 | const Model& model, const Operator* op); |
97 | std::vector<std::unique_ptr<Operator>>::iterator FindOp(Model& model, |
98 | const Operator* op); |
99 | |
100 | const char* OperatorTypeName(OperatorType type); |
101 | std::string HelpfulOperatorTypeName(const Operator& op); |
102 | |
103 | // Whether the operator can be fused with an activation function. Note that this |
104 | // will return false by default for new operators; fusing support is opt-in. |
105 | bool OperatorSupportsFusedActivation(OperatorType type); |
106 | |
107 | void DumpGraphvizVideoFrame(const Model& model); |
108 | void LogDump(int log_level, const std::string& message, const Model& model); |
109 | void LogSummary(int log_level, const std::string& message, const Model& model); |
110 | |
111 | // TODO(b/36075966): Clean up when dims superseded by array shape. |
112 | void ExtendShape(Shape* shape, int new_shape_size); |
113 | |
114 | // TODO(b/36075966): Clean up when dims superseded by array shape. |
115 | void UnextendShape(Shape* shape, int new_shape_size); |
116 | |
117 | // Checks that all dimensions of 'shape' are at least 1. Note that scalars, |
118 | // lacking dimensions, satisfy this condition and are considered non-empty. |
119 | bool IsNonEmpty(const Shape& shape); |
120 | |
121 | // Given two shapes with potentially different dimensionality and dimension |
122 | // arrays d0 and d1. Without loss of generality, assume that shape0 may have |
123 | // higher dimensionality (length(d0) >= length(d1)). Then shape0 and shape1 |
124 | // "agree up to broadcasting" if: |
125 | // - When walking the d0 and d1 from back to front with indices i0, i1, |
126 | // d0[i0] == d1[i1] or d0[i0] == 1 or d1[i1] == 1, for each dimension until |
127 | // i1 == 0 (inclusive). |
128 | bool ShapesAgreeUpToBroadcasting(const Shape& shape0, const Shape& shape1); |
129 | |
130 | // A stricter constraint than ShapesAgreeUpToBroadcasting(). |
131 | // |
132 | // Given two shapes with potentially different dimensionality and dimension |
133 | // arrays d0 and d1. Without loss of generality, assume that shape0 may have |
134 | // higher dimensionality (length(d0) >= length(d1)). Then shape0 and shape1 |
135 | // "agree up to extending" if: |
136 | // - When walking the d0 and d1 from back to front with indices i0, i1, |
137 | // d0[i0] == d1[i1] for each dimension until i1 == 0 (inclusive). |
138 | // - For the remaining indices [0..i0), d0[i0] == 1. |
139 | bool ShapesAgreeUpToExtending(const Shape& shape0, const Shape& shape1); |
140 | |
141 | inline ::tflite::RuntimeShape ToRuntimeShape(const Shape& shape) { |
142 | return ::tflite::RuntimeShape(shape.dimensions_count(), shape.dims().data()); |
143 | } |
144 | |
145 | bool IsArrayFullyConnectedWeights(const Model& model, const std::string& name); |
146 | |
147 | // If there is a wildcard dimension (-1), this may return a negative value. |
148 | int RequiredBufferSizeForShape(const Shape& shape); |
149 | |
150 | bool IsConstantParameterArray(const Model& model, const std::string& name); |
151 | |
152 | // Compares two constant parameter arrays for exact equality. |
153 | bool CompareConstantArrays(const Array& lhs_array, const Array& rhs_array); |
154 | |
155 | void CheckNoMissingArray(const Model& model); |
156 | void CheckInvariants(const Model& model); |
157 | |
158 | void CheckModelCounts(const Model& model); |
159 | |
160 | void FixOperatorOrdering(Model* model); |
161 | void FixNoMissingArray(Model* model); |
162 | void FixNoOrphanedArray(Model* model); |
163 | |
164 | // Fixes input/output arrays that may have issues during export or inference. |
165 | void FixEdgeArrays(Model* model); |
166 | |
167 | // Finds and deduplicates large constant arrays in the model. |
168 | // After constant propagation runs it's possible to end up with several of the |
169 | // same large array (whether they be zeros or otherwise). |
170 | // |
171 | // |min_size| is used to adjust the minimum size in bytes of an array before |
172 | // it's considered for deduping. As deduping can make the graphs more difficult |
173 | // to read this helps prevent small arrays from spidering out. |
174 | void DedupeConstantArrays(Model* model, size_t min_size); |
175 | |
176 | // Copies the contents of an array into another. |
177 | // Expects that the shape and data type match. |
178 | template <ArrayDataType A> |
179 | void CopyArrayBuffer(const Array& source_array, Array* target_array) { |
180 | int source_buffer_size = RequiredBufferSizeForShape(source_array.shape()); |
181 | int target_buffer_size = RequiredBufferSizeForShape(target_array->shape()); |
182 | CHECK_EQ(source_buffer_size, target_buffer_size) |
183 | << "Buffer sizes must match in element count" ; |
184 | CHECK(source_array.data_type == target_array->data_type) |
185 | << "Data types must match" ; |
186 | if (source_array.buffer) { |
187 | const auto& source_buffer = source_array.GetBuffer<A>(); |
188 | auto& target_buffer = target_array->GetMutableBuffer<A>(); |
189 | target_buffer.data = source_buffer.data; |
190 | } |
191 | } |
192 | |
193 | // Inserts a no-op reshape operator between the source array and the target |
194 | // array. This effectively just copies the data. |
195 | void InsertCopyOperator(Model* model, const std::string& source_array_name, |
196 | const std::string& target_array_name); |
197 | |
198 | // Clones an array with all data and parameters. |
199 | void CloneArray(Model* model, const std::string& source_array_name, |
200 | const std::string& target_array_name); |
201 | |
202 | void ResolveModelFlags(const ModelFlags& model_flags, Model* model); |
203 | |
204 | template <typename T> |
205 | T ConvertOperator(Operator* o, OperatorType type) { |
206 | if (o != nullptr && o->type == type) { |
207 | return static_cast<T>(o); |
208 | } |
209 | |
210 | return nullptr; |
211 | } |
212 | |
213 | void CheckIsReadyForQuantization(const Model& model); |
214 | |
215 | bool ReshapeIsEquivalentToTranspose(const Model& model, |
216 | const TensorFlowReshapeOperator* op, |
217 | bool ); |
218 | |
219 | inline int Offset(const Shape& shape, const std::vector<int>& indices) { |
220 | DCHECK_EQ(shape.dimensions_count(), indices.size()); |
221 | const int dims_count = shape.dimensions_count(); |
222 | int offset = 0; |
223 | for (int i = 0; i < dims_count; i++) { |
224 | const int index = indices[i]; |
225 | DCHECK(index >= 0 && index < shape.dims(i)); |
226 | offset *= shape.dims(i); |
227 | offset += index; |
228 | } |
229 | return offset; |
230 | } |
231 | |
232 | inline std::vector<int> ReverseOffset(const Shape& shape, int index) { |
233 | DCHECK_GE(index, 0); |
234 | DCHECK_LT(index, RequiredBufferSizeForShape(shape)); |
235 | const int dims_count = shape.dimensions_count(); |
236 | std::vector<int> indices(dims_count); |
237 | int residual = index; |
238 | for (int i = dims_count - 1; i >= 0; i--) { |
239 | indices[i] = residual % shape.dims(i); |
240 | residual /= shape.dims(i); |
241 | } |
242 | return indices; |
243 | } |
244 | |
245 | int ElementSize(ArrayDataType data_type); |
246 | |
247 | void DropMinMax(Model* model, const std::string& array_name); |
248 | |
249 | bool IsAllocatableTransientArray(const Model& model, |
250 | const std::string& array_name); |
251 | |
252 | void CreateOrCheckRnnStateArray(const std::string& name, int size, |
253 | int state_num_dims, Model* model); |
254 | |
255 | std::string AvailableArrayName(const Model& model, const std::string& name); |
256 | |
257 | // Formats a shape as a string: [ dims(0), dims(1), ..., dims(num_dims-1) ]. |
258 | std::string ShapeToString(const Shape& shape); |
259 | |
260 | void PrintArrayShape(Model* model, const std::string& name); |
261 | |
262 | void MakeArrayDims(int num_dims, int batch, int height, int width, int depth, |
263 | std::vector<int>* out_dims); |
264 | |
265 | // Defines a constant int32 array with the provided values formatted for use |
266 | // as op parameters. |
267 | std::string CreateInt32Array(Model* model, const std::string& param_name, |
268 | const std::vector<int>& value); |
269 | |
270 | bool EstimateArithmeticOpsCount(const Model& model, const Operator& op, |
271 | int64_t* result); |
272 | bool EstimateArithmeticOpsCount(const Model& model, int64_t* result); |
273 | std::string FormattedNumber(int64_t x); |
274 | |
275 | int AxesCount(AxesOrder axes_order); |
276 | |
277 | // Returns the permutation of the dimensions based on the input axes order and |
278 | // output axes order. |
279 | void GetShuffleShape(AxesOrder input_axes_order, AxesOrder output_axes_order, |
280 | std::vector<int>* shuffle); |
281 | |
282 | // Extend shuffle is designed to match ExtendShape, which pads the shape with |
283 | // unit dimensions at the beginning. |
284 | void ExtendShuffle(const std::vector<int>& input_shuffle, int newdim, |
285 | std::vector<int>* extended_shuffle); |
286 | |
287 | void ShuffleDims(const Shape& input_shape, AxesOrder input_axes_order, |
288 | AxesOrder output_axes_order, Shape* output_shape); |
289 | void ShuffleArray(const Shape& input_shape, AxesOrder input_axes_order, |
290 | AxesOrder output_axes_order, const Shape& output_shape, |
291 | const float* input_data, float* output_data); |
292 | void ShuffleArray(const Shape& input_shape, AxesOrder input_axes_order, |
293 | AxesOrder output_axes_order, const Shape& output_shape, |
294 | const uint8* input_data, uint8* output_data); |
295 | |
296 | // Returns true if it may be OK for any graph transformation to ever discard |
297 | // that array. The idea is that we can't ever discard arrays that are either |
298 | // an input or an output of the whole graph, or that appear in RNN back-edges, |
299 | // as that would undercut explicit flags that the user might pass. |
300 | bool IsDiscardableArray(const Model& model, const std::string& array_name); |
301 | |
302 | void CheckFinalDataTypesSatisfied(const Model& model); |
303 | |
304 | ArrayDataType ConvertIODataTypeToArrayDataType(IODataType type); |
305 | |
306 | // The process of building models varies according to the import format. |
307 | // |
308 | // (a) In some cases, such as model-proto format, the model should be fully |
309 | // specified. In these cases, no extra action should be taken by this function. |
310 | // (b) In other cases, such as TF graphdef format, the desired types of RNN |
311 | // arrays are not specified directly in the model, neither can they be inferred. |
312 | // However, we can set the types of RNN destination arrays to float. This breaks |
313 | // any cycles such as when resolution of the type of an RNN source array depends |
314 | // on the type of its destination array. |
315 | // |
316 | // This function is applied after the main import, after resolution of flags and |
317 | // after application of ArraysExtraInfo. It only defaults destination RNN arrays |
318 | // to float. If the model is subsequently quantized, it is assumed that the |
319 | // model contains sufficient information for that to be completed. If it is |
320 | // already quantized, then case (a) should hold. |
321 | void FinishBuildingRNNStates(Model* model); |
322 | |
323 | void (Model* model, bool quantize_output); |
324 | |
325 | // Calculates the number of elements in tensor given a shape. Shape elements |
326 | // are assumed to be of type T, while the result total is of type U. If U |
327 | // doesn't have enough range to represent the sum of elements, an error is |
328 | // returned. |
329 | template <typename T, typename U> |
330 | tensorflow::Status NumElements(const std::vector<T>& shape, U* num_elements) { |
331 | static_assert( |
332 | std::numeric_limits<T>::max() <= std::numeric_limits<uint64_t>::max(), |
333 | "vector type exceed capabilities of NumElements" ); |
334 | |
335 | *num_elements = 1; |
336 | for (const T& dim : shape) { |
337 | if (dim < 0) { |
338 | // TensorFlow's shapes sometimes include -1 to represent an "unknown" |
339 | // size but TOCO isn't able to create arrays of unknown sizes and will |
340 | // crash in RequiredBufferSizeForShape(). |
341 | return tensorflow::errors::InvalidArgument( |
342 | "Tensor shape should not include negative values" ); |
343 | } |
344 | if (*num_elements != 0 && |
345 | static_cast<uint64_t>(dim) > |
346 | std::numeric_limits<U>::max() / *num_elements) { |
347 | *num_elements = 0; |
348 | return tensorflow::errors::InvalidArgument("Tensor shape is too large" ); |
349 | } |
350 | *num_elements *= dim; |
351 | } |
352 | return ::tensorflow::OkStatus(); |
353 | } |
354 | |
355 | // A model file may have shuffled FC weights. |
356 | // When that happens, we want to de-shuffle them immediately on import, |
357 | // so that the rest of toco doesn't need to know about shuffled weights. |
358 | void UndoWeightsShuffling(Model* model); |
359 | |
360 | // Copies minmax, quantization_params, and narrow_range. |
361 | void CopyMinMaxAndQuantizationRelatedFields(const Array& src, Array* dst); |
362 | |
363 | // Delete Array if it's discardable and not referenced as input or output array |
364 | // by any other op than the specified op. |
365 | bool DeleteArrayIfUnusedOutsideOfOp(const std::string& array_name, |
366 | const Operator* op, Model* model); |
367 | |
368 | } // namespace toco |
369 | |
370 | #endif // TENSORFLOW_LITE_TOCO_TOOLING_UTIL_H_ |
371 | |