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 allow_extra_unary_dims); 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 UseArraysExtraInfo(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