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_MODEL_H_ 16 #define TENSORFLOW_LITE_TOCO_MODEL_H_ 17 18 #include <complex> 19 #include <functional> 20 #include <initializer_list> 21 #include <memory> 22 #include <string> 23 #include <unordered_map> 24 #include <unordered_set> 25 #include <vector> 26 27 #include "absl/types/optional.h" 28 #include "tensorflow/core/platform/logging.h" 29 #include "tensorflow/lite/toco/model_flags.pb.h" 30 #include "tensorflow/lite/toco/runtime/types.h" 31 #include "tensorflow/lite/toco/toco_port.h" 32 #include "tensorflow/lite/toco/toco_types.h" 33 34 namespace toco { 35 36 using tflite::QuantizationParams; 37 38 enum class OperatorType : uint8 { 39 kNone, 40 // General-purpose neural network operators. 41 kAdd, 42 kAddN, 43 kAveragePool, 44 kBatchMatMul, 45 kBatchNormalization, 46 kCeil, 47 kConv, 48 kConcatenation, 49 kCos, 50 kDepthwiseConv, 51 kDepthToSpace, 52 kSpaceToDepth, 53 kDequantize, 54 kDiv, 55 kExp, 56 kExpandDims, 57 kFill, 58 kFloorDiv, 59 kFloorMod, 60 kFullyConnected, 61 kL2Normalization, 62 kL2Pool, 63 kLstmCell, 64 kUnidirectionalSequenceLstm, 65 kLocalResponseNormalization, 66 kLog, 67 kLogistic, 68 kMaxPool, 69 kFakeQuant, 70 kMul, 71 kOneHot, 72 kRandomUniform, 73 kRange, 74 kRank, 75 kRelu, 76 kRelu1, 77 kRelu6, 78 kPRelu, 79 kHardSwish, 80 kSoftmax, 81 kLogSoftmax, 82 kSub, 83 kTanh, 84 kTransposeConv, 85 kCast, 86 kFloor, 87 kRound, 88 kGather, 89 kResizeBilinear, 90 kSin, 91 kSpaceToBatchND, 92 kPack, 93 kBatchToSpaceND, 94 kPad, 95 kPadV2, 96 kReduceProd, // Reduction product 97 kStridedSlice, 98 kSlice, 99 kSqueeze, 100 kMean, 101 kArgMax, 102 // The SVDF Op is a decomposition of a densely connected Op into 103 // low rank filters. For details: 104 // https://research.google.com/pubs/pub43813.html 105 kSvdf, 106 // Special operators used for importing TensorFlow nodes. 107 // The general intent is to have some graph transformation either 108 // drop them or rewrite them as general-purpose operators. 109 kAll, 110 kAssert, 111 kConcat, 112 kConcatV2, 113 kGreater, 114 kGreaterEqual, 115 kIdentity, 116 kLess, 117 kLessEqual, 118 kReduceMax, // Reduction Max 119 kMaximum, // Element-wise Maximum 120 kReduceMin, // Reduction Min 121 kMinimum, // Element-wise Minimum 122 kMatMul, 123 kMerge, 124 kNeg, 125 kReshape, 126 kRsqrt, 127 kShape, 128 kSplit, 129 kSplitV, 130 kSqrt, 131 kSquare, 132 kSquaredDifference, 133 kSum, 134 kSwitch, 135 kTile, 136 kTranspose, 137 kTopK_V2, 138 kDynamicPartition, 139 kDynamicStitch, 140 // An unsupported TF operation. It's only needed to be able to represent TF 141 // graph internally and is expected to be dropped by graph transformations. 142 kUnsupported, 143 // Finally, TensorFlow uses different conventions for axes ordering, 144 // see AxesOrder, and this cannot always be resolved at the time of importing 145 // nodes, as TensorFlow parameters may be constant-expression subgraphs 146 // instead of being given as plain constant arrays. So we need to insert 147 // special nodes in the graph to shuffle axes. 148 kReorderAxes, 149 kSegmentSum, 150 kSelect, 151 kSelectV2, 152 kSparseToDense, 153 kEqual, 154 kNotEqual, 155 kPow, 156 kArgMin, 157 kAny, 158 kLogicalAnd, 159 kLogicalNot, 160 kLogicalOr, 161 kCTCBeamSearchDecoder, 162 kUnpack, 163 kZerosLike, 164 kResizeNearestNeighbor, 165 kLeakyRelu, 166 kAbs, 167 kMirrorPad, 168 kUnique, 169 kUnidirectionalSequenceRnn, 170 kBidirectionalSequenceLstm, 171 kReverseV2, 172 kBidirectionalSequenceRnn, 173 kGatherNd, 174 kWhere, 175 kElu, 176 kReverseSequence, 177 kMatrixDiag, 178 kMatrixSetDiag, 179 kMatrixDiagV2, 180 kMatrixSetDiagV2, 181 kMatrixDiagV3, 182 kMatrixSetDiagV3, 183 // Debugging operators. 184 kNumericVerify 185 }; 186 187 // Helper to deal with TensorFlow arrays using a different ordering of 188 // dimensions 189 // ("axes") than our own. 190 // TODO(benoitjacob): Ultimately, we shouldn't have any "ordering" of axes, 191 // we should have associative arrays mapping symbolic axes identifiers (like 192 // "output_depth") to dimensions. We would then not need this anymore. 193 enum class AxesOrder { 194 kOneAxis, // one-dimensional array, one unique axis. 195 kCR, // column-major matrix storage order. Our standard. 196 kRC, // row-major matrix storage order. TensorFlow default. 197 kOHWI, // Our standard for conv weights 198 kHWIO, // TensorFlow conv weights 199 k1HWO, // Our standard for DepthwiseConv weights 200 kHWIM, // TensorFlow DepthwiseConv weights 201 kNHWC, // TensorFlow activations 202 kHWOI, // TensorFlow back-prop conv weights 203 }; 204 205 // The type of the scalars in an array. 206 // Note that the type does not by itself tell whether the values in the array 207 // are non-quantized (can be accessed directly) or quantized (must be 208 // interpreted in conjunction with QuantizationParams). 209 // 210 // In practice though: 211 // float values are never quantized 212 // uint8 values are always quantized 213 // int32 values are sometimes quantized (depending on whether 214 // QuantizationParams are present). 215 // complex values are never quantized 216 // other types are never quantized at the moment. 217 // 218 // kNone means that we don't know the data type yet, or that we don't care 219 // because we'll be dropping the array anyway (e.g. some exotic array types 220 // may be involved only in debug-only subgraphs that we may not be interested 221 // in actually supporting). 222 enum class ArrayDataType : uint8 { 223 kNone, // 0 224 kBool, 225 kFloat, 226 kInt8, 227 kUint8, 228 kInt16, // 5 229 kUint16, 230 kInt32, 231 kUint32, 232 kInt64, 233 kUint64, // 10 234 kString, 235 kComplex64, 236 kFloat16, 237 }; 238 239 // Compile-time logic to map ArrayDataType to the corresponding C++ scalar type 240 template <ArrayDataType A> 241 struct DataTypeImpl {}; 242 template <> 243 struct DataTypeImpl<ArrayDataType::kNone> { 244 typedef int Type; 245 }; 246 template <> 247 struct DataTypeImpl<ArrayDataType::kBool> { 248 typedef bool Type; 249 }; 250 template <> 251 struct DataTypeImpl<ArrayDataType::kFloat> { 252 typedef float Type; 253 }; 254 template <> 255 struct DataTypeImpl<ArrayDataType::kInt8> { 256 typedef int8 Type; 257 }; 258 template <> 259 struct DataTypeImpl<ArrayDataType::kUint8> { 260 typedef uint8 Type; 261 }; 262 template <> 263 struct DataTypeImpl<ArrayDataType::kInt16> { 264 typedef int16 Type; 265 }; 266 template <> 267 struct DataTypeImpl<ArrayDataType::kUint16> { 268 typedef uint16 Type; 269 }; 270 template <> 271 struct DataTypeImpl<ArrayDataType::kInt32> { 272 typedef int32 Type; 273 }; 274 template <> 275 struct DataTypeImpl<ArrayDataType::kUint32> { 276 typedef uint32 Type; 277 }; 278 template <> 279 struct DataTypeImpl<ArrayDataType::kInt64> { 280 typedef int64 Type; 281 }; 282 template <> 283 struct DataTypeImpl<ArrayDataType::kUint64> { 284 typedef uint64 Type; 285 }; 286 template <> 287 struct DataTypeImpl<ArrayDataType::kString> { 288 typedef string Type; 289 }; 290 template <> 291 struct DataTypeImpl<ArrayDataType::kComplex64> { 292 typedef std::complex<float> Type; 293 }; 294 295 template <ArrayDataType A> 296 using DataType = typename DataTypeImpl<A>::Type; 297 298 // Base class for type-specific buffer types. 299 struct GenericBuffer { 300 // Non-default-constructible: only ArrayDataType-specific subclass 301 // objects may be constructed. 302 GenericBuffer() = delete; 303 // Non-copyable-or-movable: we should only store pointers-to-Buffer 304 // in containers, not Operators themselves, so there should be no 305 // copy or move. 306 GenericBuffer(const GenericBuffer&) = delete; 307 GenericBuffer(const GenericBuffer&&) = delete; 308 309 // We need a virtual destructor so we can store pointers-to-Buffer 310 // in containers and have the containers call the right subclass destructor. 311 virtual ~GenericBuffer() {} 312 313 virtual int Length() const = 0; 314 315 const ArrayDataType type; 316 317 protected: 318 // Constructor used by subclasses for specific ArrayDataType's. 319 explicit GenericBuffer(ArrayDataType t) : type(t) {} 320 }; 321 322 // Type-specific buffer, containing type-specific storage. 323 template <ArrayDataType A> 324 struct Buffer : GenericBuffer { 325 Buffer() : GenericBuffer(A) {} 326 327 int Length() const override { return data.size(); } 328 329 std::vector<DataType<A>> data; 330 }; 331 332 class Shape { 333 public: 334 // For Shape, we stick to half-way encapsulation for now: 335 // we hide the raw dims_ member, but expose it raw by accessors 336 // because from some brainstorming, it's not at all easy to 337 // anticipate which flavor of more hermetic encapsulation would 338 // actually buy us future-proof-ness without being needlessly 339 // cumbersome. 340 Shape() {} 341 Shape(std::initializer_list<int> dim_list) : dims_(dim_list) {} 342 343 void ReplaceDims(std::initializer_list<int> dim_list) { 344 dims_ = std::vector<int>(dim_list); 345 } 346 347 const std::vector<int>& dims() const { return dims_; } 348 std::vector<int>* mutable_dims() { return &dims_; } 349 const int dimensions_count() const { return dims_.size(); } 350 351 // We still have that one convenience accessor to avoid 352 // the awkward double bracket issue: shape.dims()[i]. 353 int dims(int i) const { 354 // Always check for out-of-bounds accesses, even in optimized builds where 355 // standard assertions are disabled. Out-of-bounds access here is a common 356 // occurrence. 357 CHECK_GE(i, 0); 358 CHECK_GT(dims_.size(), i); 359 return dims_[i]; 360 } 361 362 bool operator==(const Shape& comp) const { 363 return (this->dims_ == comp.dims()); 364 } 365 366 bool operator!=(const Shape& comp) const { return !((*this) == comp); } 367 368 private: 369 std::vector<int> dims_; 370 }; 371 372 // Base class for all operator classes. 373 struct Operator { 374 // Non-default-constructible: only OperatorType-specific subclass 375 // objects may be constructed. 376 Operator() = delete; 377 // Non-copyable-or-movable: we should only store pointers-to-Operator 378 // in containers, not Operators themselves, so there should be no 379 // copy or move. 380 Operator(const Operator&) = delete; 381 Operator(const Operator&&) = delete; 382 383 // We need a virtual destructor so we can store pointers-to-Operator 384 // in containers and have the containers call the right subclass destructor. 385 virtual ~Operator() {} 386 387 // The specific type of operator. Corresponds 1:1 to subclasses. 388 const OperatorType type; 389 390 // The activation function that may be fused into this operator, 391 // or None if no activation function is fused. 392 FusedActivationFunctionType fused_activation_function; 393 394 // Input arrays: either activation arrays or constant array parameters. 395 // We refer to them by their name, not by their address; the mapping of 396 // names to addresses is given by the Model, which owns both Operator's and 397 // Array's. Thus, an Operator on its own doesn't contain much information, 398 // it is meant to be used in conjunction with the Model that owns it. 399 std::vector<string> inputs; 400 401 // Output activation arrays. Same comments as for inputs apply here too. 402 std::vector<string> outputs; 403 404 // If true, the operator has more outputs than are listed in the 'outputs' 405 // member. These need to be resolved by some graph transformation. 406 // This flag is only here to indicate that an operator should not be 407 // discarded as unused, even if from its 'outputs' member alone it 408 // looks unused. 409 bool unresolved_outputs = false; 410 411 // A serialized tensorflow::NodeDef string. 412 // The field is filled only when importing from TensorFlow. 413 // It's guaranteed to be filled for `TensorFlowUnsupportedOperator`. 414 // It's not guaranteed to be filled for other ops. Ops created by graph 415 // transformations won't have TensorFlow NodeDef. 416 string tensorflow_node_def; 417 418 protected: 419 // Constructor used by subclasses for specific OperatorType's. 420 explicit Operator(OperatorType t) 421 : type(t), 422 fused_activation_function(FusedActivationFunctionType::kNone) {} 423 }; 424 425 // Padding types for Conv-like operators. This is how padding is typically 426 // specified in model files. But for inference, we will need to resolve this 427 // to a FixedPadding, see below. 428 enum class PaddingType { kNone, kSame, kValid }; 429 430 // Padding as resolved for a specific layer shape, as needed for inference. 431 // For a given layer shape, a given padding type will resolve to a choice of 432 // a number of padding rows and columns, which we call the padding height and 433 // width respectively. 434 struct FixedPadding { 435 int width = 0; 436 int height = 0; 437 }; 438 439 // "Universal" padding struct containing both a generic PaddingType (as 440 // represented in a model file), and a FixedPadding (as needed for inference). 441 // The latter is resolved during the PropagateFixedSizes pass. 442 struct Padding { 443 FixedPadding& GetOrCreateFixedPadding() { 444 if (!fixed) { 445 FixedPadding* ptr = new FixedPadding; 446 fixed = std::unique_ptr<FixedPadding>(ptr); 447 } 448 return *fixed; 449 } 450 451 Padding() : type(PaddingType::kNone) {} 452 PaddingType type; 453 std::unique_ptr<FixedPadding> fixed; 454 }; 455 456 // "Convolutional" layer, as represented in model files. 457 // 458 // Inputs: 459 // inputs[0]: required: the input activations array 460 // inputs[1]: required: the Conv weights 461 // inputs[2]: optional: the bias vector, specifying the biases for each output 462 // channel. 463 // 464 // Outputs: 465 // outputs[0]: required: the output activations array 466 // outputs[1]: optional: the intermediate array of im2col-replicated input 467 // activations. Present when targeting implementations 468 // of Conv layers as Im2col+GEMM. 469 // 470 // TensorFlow equivalent: Conv2D 471 struct ConvOperator : Operator { 472 ConvOperator() : Operator(OperatorType::kConv) {} 473 Padding padding; 474 int stride_width = 0; 475 int stride_height = 0; 476 // A dilation_rate of 0 is invalid and this field is an optional attribute. 477 // Thus initializing it to 1 to allow default conv behavior when the 478 // attribute is not present. 479 int dilation_width_factor = 1; 480 int dilation_height_factor = 1; 481 }; 482 483 // CTCBeamSearchDecoder operator: 484 // 485 // Inputs: 486 // inputs[0]: required: the logits. 487 // inputs[1]: required: sequence length. 488 // inputs[2]: optional: beam width. 489 // inputs[3]: optional: top paths. 490 // inputs[4]: optional: merge repeated. 491 // 492 // Outputs: 493 // outputs[0]: deocoded. 494 // outputs[1]: log probability. 495 // 496 // TensorFlow equivalent: CTCBeamSearchDecoder 497 struct CTCBeamSearchDecoderOperator : Operator { 498 CTCBeamSearchDecoderOperator() 499 : Operator(OperatorType::kCTCBeamSearchDecoder) {} 500 int beam_width; 501 int top_paths; 502 bool merge_repeated = true; 503 }; 504 505 // Depthwise-separable convolution operator. 506 // 507 // Inputs: 508 // inputs[0]: required: the input activations array 509 // inputs[1]: required: the DepthwiseConv weights 510 // inputs[2]: optional: the bias vector, specifying the biases for each output 511 // channel. 512 // 513 // TensorFlow equivalent: DepthwiseConv2dNative 514 struct DepthwiseConvOperator : Operator { 515 DepthwiseConvOperator() : Operator(OperatorType::kDepthwiseConv) {} 516 Padding padding; 517 int stride_height = 0; 518 int stride_width = 0; 519 int depth_multiplier = 0; 520 // A dilation_rate of 0 is invalid and this field is an optional attribute. 521 // Thus initializing it to 1 to allow default conv behavior when the 522 // attribute is not present. 523 int dilation_width_factor = 1; 524 int dilation_height_factor = 1; 525 }; 526 527 // Depth-to-space transform operator. 528 // 529 // Inputs: 530 // inputs[0]: required: the input activations array 531 // 532 // TensorFlow equivalent: DepthToSpace 533 struct DepthToSpaceOperator : Operator { 534 DepthToSpaceOperator() : Operator(OperatorType::kDepthToSpace) {} 535 int block_size = 0; 536 }; 537 538 // Space-to-depth transform operator. 539 // 540 // Inputs: 541 // inputs[0]: required: the input activations array 542 // 543 // TensorFlow equivalent: SpaceToDepth 544 struct SpaceToDepthOperator : Operator { 545 SpaceToDepthOperator() : Operator(OperatorType::kSpaceToDepth) {} 546 int block_size = 0; 547 }; 548 549 // Fully-connected operator. 550 // 551 // Inputs: 552 // inputs[0]: required: the input activations array 553 // inputs[1]: required: the FullyConnected weights 554 // inputs[2]: optional: the bias vector, specifying the biases for each output 555 // channel. 556 // 557 // TensorFlow equivalent: a pair consisting of a Reshape node reshaping the 558 // input activations as a matrix, followed by a MatMul node. 559 struct FullyConnectedOperator : Operator { 560 FullyConnectedOperator() : Operator(OperatorType::kFullyConnected) {} 561 FullyConnectedWeightsFormat weights_format = 562 FullyConnectedWeightsFormat::kDefault; 563 564 // `keep_num_dims` is supported in the FullyConnected kernel version 5, but 565 // it's never supported by Toco. 566 bool keep_num_dims = false; 567 }; 568 569 // Dequantization operator, converting a quantized array of integers with 570 // quantization parameters specifying how these integers correspond to real 571 // numbers 572 // (see QuantizationParams) to an output activations array of floating-point 573 // values. 574 // 575 // In floating-point image models, there is typically a Dequantization operator 576 // at the very beginning, converting the input image RGB data, consisting of 577 // uint8 integer values, to floating-point input activations. That is where 578 // image model parameters such as "mean_value" and "std_value" are typically 579 // handled. 580 // 581 // This is the only operator type that converts from quantized to 582 // floating-point, 583 // and there is at the moment no operator type at all to convert from 584 // floating-point 585 // to quantized. Every other operator does either float->float or 586 // quantized->quantized. 587 // 588 // Inputs: 589 // inputs[0]: required: the input quantized activations array 590 // 591 // TensorFlow equivalent: Dequantize 592 struct DequantizeOperator : Operator { 593 DequantizeOperator() : Operator(OperatorType::kDequantize) {} 594 }; 595 596 // Numeric verification operator, converting a quantized array of integers with 597 // quantization parameters specifying how these integers correspond to real 598 // numbers 599 // (see QuantizationParams) and verify them with an array of floating-point 600 // values. 601 602 // Inputs: 603 // inputs[0]: required: the input quantized activations array 604 // inputs[1]: required: the input reference activations array 605 // 606 // TensorFlow equivalent: Dequantize 607 struct NumericVerifyOperator : Operator { 608 NumericVerifyOperator() : Operator(OperatorType::kNumericVerify) {} 609 }; 610 611 // Batch-normalization operator. 612 // 613 // We only support batch-normalization using pre-learned moments, so this is 614 // just 615 // computing (input - mean) * multiplier + offset. As such, this can be 616 // expressed as a combination of Add and Mul nodes, and indeed this is how 617 // we break it down during tooling for the purpose of fusing it into 618 // other operators. 619 // 620 // Inputs: 621 // inputs[0]: required: the input activations array 622 // inputs[1]: required: the learned mean array 623 // inputs[2]: required: the learned multiplier array 624 // inputs[3]: required: the learned offset array 625 // 626 // TensorFlow equivalent: a combination of Add and Mul nodes 627 struct BatchNormalizationOperator : Operator { 628 BatchNormalizationOperator() 629 : Operator(OperatorType::kBatchNormalization), 630 global_normalization(false) {} 631 bool global_normalization; 632 }; 633 634 // L2-normalization operator. 635 // 636 // Inputs: 637 // inputs[0]: required: the input activations array 638 // 639 // TensorFlow equivalent: none. In TensorFlow, L2 normalization is implemented 640 // by a sub-graph of operators implementing L2-normalization 641 // from lower-level arithmetic nodes; during tooling, we identify such 642 // sub-graphs 643 // and replace them by L2NormalizationOperator's. See IdentifyL2Normalization. 644 struct L2NormalizationOperator : Operator { 645 L2NormalizationOperator() : Operator(OperatorType::kL2Normalization) {} 646 }; 647 648 // LSTM Cell operator. 649 // 650 // Inputs: 651 // inputs[0]: required: the input data array 652 // inputs[1]: required: the previous output activations array 653 // inputs[2]: required: the learned weights array 654 // inputs[3]: required: the learned biases array 655 // inputs[4]: required: the previous output state 656 // outputs[0]: required: the output activations array 657 // outputs[1]: required: the new state array 658 // 659 // TensorFlow equivalent: none. In TensorFlow, an LSTM is implemented 660 // with a sub-graph of lower-level arithmetic nodes; during tooling, we identify 661 // such sub-graphs and replace them with LstmCells. See IdentifyLstmCell(). 662 struct LstmCellOperator : Operator { 663 enum Inputs { 664 DATA_INPUT = 0, 665 PREV_ACTIV_INPUT = 1, 666 WEIGHTS_INPUT = 2, 667 BIASES_INPUT = 3, 668 PREV_STATE_INPUT = 4, 669 NUM_INPUTS = 5 670 }; 671 enum Outputs { 672 ACTIV_OUTPUT = 0, 673 STATE_OUTPUT = 1, 674 CONCAT_TEMP = 2, 675 ACTIV_TEMP = 3, 676 NUM_OUTPUTS = 4 677 }; 678 enum KernelType { 679 KERNEL_BASIC = 0, 680 KERNEL_FULL = 1, 681 }; 682 683 LstmCellOperator() 684 : Operator(OperatorType::kLstmCell), kernel_type(KERNEL_BASIC) {} 685 686 KernelType kernel_type; 687 }; 688 689 struct UnidirectionalSequenceLstmOperator : Operator { 690 UnidirectionalSequenceLstmOperator() 691 : Operator(OperatorType::kUnidirectionalSequenceLstm) {} 692 }; 693 694 struct BidirectionalSequenceLstmOperator : Operator { 695 BidirectionalSequenceLstmOperator() 696 : Operator(OperatorType::kBidirectionalSequenceLstm) {} 697 bool merge_outputs; 698 }; 699 700 struct BidirectionalSequenceRnnOperator : Operator { 701 BidirectionalSequenceRnnOperator() 702 : Operator(OperatorType::kBidirectionalSequenceRnn) {} 703 bool merge_outputs; 704 }; 705 706 // Element-wise multiplication operator. 707 // 708 // Inputs: 709 // inputs[0]: required: the left-hand side array 710 // inputs[1]: required: the right-hand side array 711 // 712 // TensorFlow equivalent: Mul 713 struct MulOperator : Operator { 714 MulOperator() : Operator(OperatorType::kMul) {} 715 }; 716 717 // Element-wise Abs operator: 718 // x -> abs(x) 719 // 720 // Inputs: 721 // inputs[0]: required: the input array 722 // 723 // TensorFlow equivalent: abs 724 struct AbsOperator : Operator { 725 AbsOperator() : Operator(OperatorType::kAbs) {} 726 }; 727 728 // Element-wise HardSwish operator: 729 // x -> x * relu6(x+3)/6 730 // 731 // Inputs: 732 // inputs[0]: required: the input array 733 // 734 // TensorFlow equivalent: hard_swish 735 struct HardSwishOperator : Operator { 736 HardSwishOperator() : Operator(OperatorType::kHardSwish) {} 737 }; 738 739 // Elu 740 // f(x) -> exp(x) - 1 for x < 0, x for x >= 0. 741 // 742 // Inputs: 743 // inputs[0]: required: the input array 744 // 745 // TensorFlow equivalent: Elu 746 struct EluOperator : Operator { 747 EluOperator() : Operator(OperatorType::kElu) {} 748 }; 749 750 // Element-wise Relu operator: 751 // x -> max(0, x) 752 // 753 // Inputs: 754 // inputs[0]: required: the input array 755 // 756 // TensorFlow equivalent: Relu 757 struct ReluOperator : Operator { 758 ReluOperator() : Operator(OperatorType::kRelu) {} 759 }; 760 761 // Element-wise Relu1 operator: 762 // x -> min(max(x, -1), 1) 763 // 764 // Inputs: 765 // inputs[0]: required: the input array 766 // 767 // TensorFlow equivalent: none. We can construct the operator with Minimum 768 // and Maximum operations 769 struct Relu1Operator : Operator { 770 Relu1Operator() : Operator(OperatorType::kRelu1) {} 771 }; 772 773 // Element-wise Relu6 operator: 774 // x -> max(0, min(6, x)) 775 // 776 // Inputs: 777 // inputs[0]: required: the input array 778 // 779 // TensorFlow equivalent: Relu6 780 struct Relu6Operator : Operator { 781 Relu6Operator() : Operator(OperatorType::kRelu6) {} 782 }; 783 784 // PRelu 785 // f(x) = alpha * x for x < 0, f(x) = x for x >= 0. 786 // 787 // Inputs: 788 // inputs[0]: required: the input array 789 // inputs[1]: required: the alpha array 790 // 791 // Equivalent to keras.layers.PReLU. 792 struct PReluOperator : Operator { 793 PReluOperator() : Operator(OperatorType::kPRelu) {} 794 }; 795 796 // LeakyRelu 797 // x -> max(x, alpha * x) 798 // 799 // Inputs: 800 // inputs[0]: required: the input array 801 // 802 // TensorFlow equivalent: LeakyRelu 803 struct LeakyReluOperator : Operator { 804 LeakyReluOperator() : Operator(OperatorType::kLeakyRelu) {} 805 806 float alpha = 0.2f; // 0.2 matches the default value for the TF op attribute. 807 }; 808 809 // Element-wise Logistic operator: 810 // x -> Logistic(x) = 1 / (1 + exp(-x)) 811 // 812 // Inputs: 813 // inputs[0]: required: the input array 814 // 815 // TensorFlow equivalent: Sigmoid 816 struct LogisticOperator : Operator { 817 LogisticOperator() : Operator(OperatorType::kLogistic) {} 818 }; 819 820 // Element-wise natural log operator: 821 // x -> ln(x) 822 // 823 // Inputs: 824 // inputs[0]: required: the input array 825 // 826 // TensorFlow equivalent: Log 827 struct LogOperator : Operator { 828 LogOperator() : Operator(OperatorType::kLog) {} 829 }; 830 831 // Element-wise Tanh operator: 832 // x -> Tanh(x) = (exp(x) - exp(-x)) / (exp(x) + exp(-x)) 833 // 834 // Inputs: 835 // inputs[0]: required: the input array 836 // 837 // TensorFlow equivalent: Tanh 838 struct TanhOperator : Operator { 839 TanhOperator() : Operator(OperatorType::kTanh) {} 840 }; 841 842 // Element-wise Sin operator: 843 // x -> Sin(x) = sin(x) 844 // 845 // Inputs: 846 // inputs[0]: required: the input array 847 // 848 // TensorFlow equivalent: Sin 849 struct SinOperator : Operator { 850 SinOperator() : Operator(OperatorType::kSin) {} 851 }; 852 853 // Element-wise addition operator. 854 // 855 // Inputs: 856 // inputs[0]: required: the left-hand side array 857 // inputs[1]: required: the right-hand side array 858 // 859 // TensorFlow equivalent: Add 860 struct AddOperator : Operator { 861 AddOperator() : Operator(OperatorType::kAdd) {} 862 }; 863 864 // Element-wise addition operator for N inputs. 865 // 866 // Inputs: 867 // inputs[i]: The i-th array to add together to form the output. 868 // 869 // TensorFlow equivalent: AddN 870 struct AddNOperator : Operator { 871 AddNOperator() : Operator(OperatorType::kAddN) {} 872 }; 873 874 // Concatenation operator: concatenates its inputs 875 // along the axis. 876 // 877 // Inputs: this operator accepts any number >= 1 of inputs. 878 // inputs[i]: the i-th array to concatenate. 879 // 880 // TensorFlow equivalent: Concat. 881 struct ConcatenationOperator : Operator { 882 ConcatenationOperator() : Operator(OperatorType::kConcatenation) {} 883 int axis = 0; 884 }; 885 886 // Reordering dimensions. Used only during tooling to transform graphs from 887 // the TensorFlow format. 888 // 889 // Inputs: 890 // inputs[0]: required: the input array 891 // 892 // TensorFlow equivalent: none. This is only useful to convert between formats. 893 struct ReorderAxesOperator : Operator { 894 ReorderAxesOperator() : Operator(OperatorType::kReorderAxes) {} 895 AxesOrder input_axes_order; 896 AxesOrder output_axes_order; 897 }; 898 899 // Average-pooling operator. 900 // 901 // Inputs: 902 // inputs[0]: required: the input array 903 // 904 // TensorFlow equivalent: AveragePool 905 struct AveragePoolOperator : Operator { 906 AveragePoolOperator() : Operator(OperatorType::kAveragePool) {} 907 Padding padding; 908 int stride_height = 0; 909 int stride_width = 0; 910 int kheight = 0; 911 int kwidth = 0; 912 }; 913 914 // Local response normalization operator. 915 // 916 // Inputs: 917 // inputs[0]: required: the input array 918 // 919 // TensorFlow equivalent: LRN 920 struct LocalResponseNormalizationOperator : Operator { 921 LocalResponseNormalizationOperator() 922 : Operator(OperatorType::kLocalResponseNormalization) {} 923 924 int range = 0; 925 float bias = 0.f; 926 float alpha = 0.f; 927 float beta = 0.f; 928 }; 929 930 // Max-pooling operator. 931 // 932 // Inputs: 933 // inputs[0]: required: the input array 934 // 935 // TensorFlow equivalent: MaxPool 936 struct MaxPoolOperator : Operator { 937 MaxPoolOperator() : Operator(OperatorType::kMaxPool) {} 938 Padding padding; 939 int stride_height = 0; 940 int stride_width = 0; 941 int kheight = 0; 942 int kwidth = 0; 943 }; 944 945 // L2-pooling operator. 946 // 947 // Inputs: 948 // inputs[0]: required: the input array 949 // 950 // TensorFlow equivalent: none. Can be shimmed by squaring+avgpool+sqrt. 951 struct L2PoolOperator : Operator { 952 L2PoolOperator() : Operator(OperatorType::kL2Pool) {} 953 Padding padding; 954 int stride_height = 0; 955 int stride_width = 0; 956 int kheight = 0; 957 int kwidth = 0; 958 }; 959 960 // The expected [min, max] range of values in a given array. 961 // Used for quantization only. 962 // This information typically comes from special nodes found in quantized 963 // models, see FakeQuantOperator, and is used during quantization to resolve 964 // actual quantization parameters (see QuantizationParams). 965 struct MinMax { 966 double min = 0.; 967 double max = 0.; 968 }; 969 970 inline bool operator==(const MinMax& m1, const MinMax& m2) { 971 return m1.min == m2.min && m1.max == m2.max; 972 } 973 974 inline bool operator!=(const MinMax& m1, const MinMax& m2) { 975 return m1.min != m2.min || m1.max != m2.max; 976 } 977 978 // Fake-quantization operator. This does two things: 979 // - Annotate its input and output arrays with MinMax information, 980 // - Arithmetic-wise, this operator rounds incoming activation values 981 // to the nearest representable value on the scale of 256 982 // values from the min to the max value dictated by its MinMax info. 983 // 984 // Inputs: 985 // inputs[0]: required: the input array 986 // inputs[1]: optional: the 'min' value, if it has not yet been resolved 987 // to a constant. 988 // inputs[2]: optional: the 'max' value, if it has not yet been resolved 989 // to a constant. 990 // 991 // TensorFlow equivalent: FakeQuantWithMinMaxVars, FakeQuantWithMinMaxArgs. 992 struct FakeQuantOperator : Operator { 993 FakeQuantOperator() : Operator(OperatorType::kFakeQuant) {} 994 std::unique_ptr<MinMax> minmax; 995 int num_bits = 8; 996 bool narrow_range = false; 997 }; 998 999 // Element-wise division operator. 1000 // 1001 // Inputs: 1002 // inputs[0]: required: the left-hand side array 1003 // inputs[1]: required: the right-hand side array 1004 // 1005 // TensorFlow equivalent: Div 1006 struct DivOperator : Operator { 1007 DivOperator() : Operator(OperatorType::kDiv) {} 1008 }; 1009 1010 // Element-wise identity (x->x) operator. 1011 // 1012 // Inputs: 1013 // inputs[0]: required: the input array 1014 // 1015 // TensorFlow equivalent: Identity 1016 struct TensorFlowIdentityOperator : Operator { 1017 TensorFlowIdentityOperator() : Operator(OperatorType::kIdentity) {} 1018 }; 1019 1020 // Batch matrix multiplication operator. This comes from a tf.matmul where one 1021 // of the operands has rank 3 or more. 1022 // 1023 // Inputs: 1024 // inputs[0]: required: the left-hand side matrix 1025 // inputs[1]: required: the right-hand side matrix 1026 // 1027 // TensorFlow equivalent: MatMul 1028 struct BatchMatMulOperator : Operator { 1029 BatchMatMulOperator() : Operator(OperatorType::kBatchMatMul) {} 1030 bool adj_x = false; 1031 bool adj_y = false; 1032 }; 1033 1034 // General matrix multiplication operator. We don't want to support general 1035 // matrix multiplication at inference time, so we resolve it during tooling 1036 // to more specific operator types, namely, FullyConnected. 1037 // 1038 // Inputs: 1039 // inputs[0]: required: the left-hand side matrix 1040 // inputs[1]: required: the right-hand side matrix 1041 // 1042 // TensorFlow equivalent: MatMul 1043 struct TensorFlowMatMulOperator : Operator { 1044 TensorFlowMatMulOperator() : Operator(OperatorType::kMatMul) {} 1045 bool transpose_a = false; 1046 bool transpose_b = false; 1047 }; 1048 1049 // Padding operator. Pads a tensor with zeros. 1050 // 1051 // Inputs: 1052 // inputs[0]: required: the input array 1053 // inputs[1]: required: the padding array 1054 // 1055 // This operation pads a `input` with zeros according to the `paddings` you 1056 // specify. `paddings` is an integer tensor with shape `[Dn, 2]`, where n is the 1057 // rank of `input`. For each dimension D of `input`, `paddings[D, 0]` indicates 1058 // how many zeros to add before the contents of `input` in that dimension, and 1059 // `paddings[D, 1]` indicates how many zeros to add after the contents of 1060 // `input` in that dimension. 1061 // 1062 // TensorFlow equivalent: Pad 1063 struct PadOperator : Operator { 1064 PadOperator() : Operator(OperatorType::kPad) {} 1065 1066 std::vector<int> left_padding; 1067 std::vector<int> right_padding; 1068 }; 1069 1070 // PaddingV2 operator. Pads a tensor with the given constant value. 1071 // 1072 // Inputs: 1073 // inputs[0]: required: the input array 1074 // inputs[1]: required: the padding array 1075 // inputs[2]: required: the scalar constant_values 1076 // 1077 // This operation pads input according to the paddings and constant_values you 1078 // specify. paddings is an integer tensor with shape [Dn, 2], where n is the 1079 // rank of input. For each dimension D of input, paddings[D, 0] indicates how 1080 // many padding values to add before the contents of input in that dimension, 1081 // and paddings[D, 1] indicates how many padding values to add after the 1082 // contents of input in that dimension. constant_values is a scalar tensor of 1083 // the same type as input that indicates the value to use for padding input. 1084 // 1085 // TensorFlow equivalent: PadV2 1086 struct PadV2Operator : Operator { 1087 PadV2Operator() : Operator(OperatorType::kPadV2) {} 1088 1089 std::vector<int> left_padding; 1090 std::vector<int> right_padding; 1091 }; 1092 1093 // Strided slice operator. 1094 // 1095 // Inputs: 1096 // inputs[0]: required: the input array 1097 // inputs[1]: required: the begin array 1098 // inputs[2]: required: the end array 1099 // inputs[3]: optional: the strides array 1100 // 1101 // TensorFlow equivalent: StridedSlice 1102 struct StridedSliceOperator : Operator { 1103 StridedSliceOperator() : Operator(OperatorType::kStridedSlice) {} 1104 1105 std::vector<int> start_indices; 1106 std::vector<int> stop_indices; 1107 std::vector<int> strides; 1108 1109 int begin_mask; 1110 int ellipsis_mask; 1111 int end_mask; 1112 int new_axis_mask; 1113 int shrink_axis_mask; 1114 1115 StridedSliceOperator(const StridedSliceOperator& other) 1116 : Operator(OperatorType::kStridedSlice) { 1117 inputs = other.inputs; 1118 outputs = other.outputs; 1119 1120 start_indices = other.start_indices; 1121 stop_indices = other.stop_indices; 1122 strides = other.strides; 1123 1124 begin_mask = other.begin_mask; 1125 ellipsis_mask = other.ellipsis_mask; 1126 end_mask = other.end_mask; 1127 new_axis_mask = other.new_axis_mask; 1128 shrink_axis_mask = other.shrink_axis_mask; 1129 } 1130 1131 void PadIndices(int dim_count) { 1132 // Add indices and mask bits to fully include extra dimensions 1133 CHECK_GE(dim_count, start_indices.size()); 1134 CHECK_EQ(start_indices.size(), stop_indices.size()); 1135 CHECK_EQ(stop_indices.size(), strides.size()); 1136 1137 for (int i = start_indices.size(); i < dim_count; i++) { 1138 start_indices.push_back(0); 1139 stop_indices.push_back(0); 1140 strides.push_back(1); 1141 begin_mask |= 1 << i; 1142 end_mask |= 1 << i; 1143 } 1144 } 1145 1146 void ReverseIndices() { 1147 CHECK_EQ(start_indices.size(), stop_indices.size()); 1148 CHECK_EQ(stop_indices.size(), strides.size()); 1149 1150 std::reverse(start_indices.begin(), start_indices.end()); 1151 std::reverse(stop_indices.begin(), stop_indices.end()); 1152 std::reverse(strides.begin(), strides.end()); 1153 1154 begin_mask = toco::port::ReverseBits32(static_cast<uint32>(begin_mask)) >> 1155 (32 - start_indices.size()); 1156 ellipsis_mask = 1157 toco::port::ReverseBits32(static_cast<uint32>(ellipsis_mask)) >> 1158 (32 - start_indices.size()); 1159 end_mask = toco::port::ReverseBits32(static_cast<uint32>(end_mask)) >> 1160 (32 - start_indices.size()); 1161 new_axis_mask = 1162 toco::port::ReverseBits32(static_cast<uint32>(new_axis_mask)) >> 1163 (32 - start_indices.size()); 1164 shrink_axis_mask = 1165 toco::port::ReverseBits32(static_cast<uint32>(shrink_axis_mask)) >> 1166 (32 - start_indices.size()); 1167 } 1168 }; 1169 1170 // Reshaping operator, reshaping its input array to a two-dimensional shape 1171 // (a "matrix"). This is used in the TensorFlow format, in conjunction with 1172 // MatMul nodes, to implement fully-connected layers. 1173 // 1174 // Inputs: 1175 // inputs[0]: required: the input array 1176 // inputs[1]: optional: the output tensor shape 1177 // 1178 // TensorFlow equivalent: Reshape --- except that we only support a special case 1179 // here, where the output shape is a matrix (2D) shape. 1180 struct TensorFlowReshapeOperator : Operator { 1181 TensorFlowReshapeOperator() : Operator(OperatorType::kReshape) {} 1182 std::vector<int> shape; 1183 }; 1184 1185 // Removes dimensions of size 1 from the shape of a tensor. 1186 // https://www.tensorflow.org/api_docs/python/tf/squeeze 1187 // 1188 // Inputs: 1189 // inputs[0]: required: the input array 1190 // 1191 // TensorFlow equivalent: Squeeze 1192 struct SqueezeOperator : Operator { 1193 SqueezeOperator() : Operator(OperatorType::kSqueeze) {} 1194 1195 std::vector<int> squeeze_dims; 1196 }; 1197 1198 // Inputs: 1199 // inputs[0]: required: the output shape 1200 // inputs[1]: required: the weights 1201 // inputs[2]: required: the input activations array 1202 // NOTE: The input activations is NOT the first input. 1203 // 1204 // 1205 // Outputs: 1206 // outputs[0]: required: the output activations array 1207 // 1208 // TensorFlow equivalent: Conv2DBackpropInput 1209 struct TransposeConvOperator : Operator { 1210 enum Inputs { 1211 OUTPUT_SHAPE = 0, 1212 WEIGHTS = 1, 1213 DATA_INPUT = 2, 1214 }; 1215 1216 TransposeConvOperator() : Operator(OperatorType::kTransposeConv) {} 1217 Padding padding; 1218 int stride_width = 0; 1219 int stride_height = 0; 1220 // Dilation is possible with transpose convolution, but Tensorflow does not 1221 // currently support it, so we omit it. 1222 }; 1223 1224 // Given a tensor input, this operation calculates element-wise exponential 1225 // (y = e^x). 1226 // 1227 // Inputs: 1228 // inputs[0]: required: input tensor 1229 // 1230 // TensorFlow equivalent: Exp 1231 struct ExpOperator : Operator { 1232 ExpOperator() : Operator(OperatorType::kExp) {} 1233 }; 1234 1235 // Given a tensor input, this operation calculates element-wise exponential 1236 // (y = cos(x)). 1237 // 1238 // Inputs: 1239 // inputs[0]: required: input tensor 1240 // 1241 // TensorFlow equivalent: Cos 1242 struct CosOperator : Operator { 1243 CosOperator() : Operator(OperatorType::kCos) {} 1244 }; 1245 1246 // Given a tensor input, this operation inserts a dimension of 1 at the 1247 // dimension index axis of input's shape. The dimension index axis starts at 1248 // zero; if you specify a negative number for axis it is counted backward from 1249 // the end. 1250 // 1251 // Inputs: 1252 // inputs[0]: required: input tensor 1253 // inputs[1]: required: 0-D (scalar). Specifies the dimension index at which 1254 // to expand the shape of input 1255 // 1256 // TensorFlow equivalent: ExpandDims 1257 struct ExpandDimsOperator : Operator { 1258 ExpandDimsOperator() : Operator(OperatorType::kExpandDims) {} 1259 }; 1260 1261 // Ceates a tensor of shape dims and fills it with the given scalar value. 1262 // Output type will be the same as the given scalar value. 1263 // 1264 // Inputs: 1265 // inputs[0]: required: 1-D (int32) - the shape of the output tensor 1266 // inputs[1]: required: 0-D (scalar) - value to fill the tensor with 1267 // 1268 // TensorFlow equivalent: Fill 1269 struct FillOperator : Operator { 1270 FillOperator() : Operator(OperatorType::kFill) {} 1271 }; 1272 1273 // Element-wise floor division operator. 1274 // 1275 // Inputs: 1276 // inputs[0]: required: the left-hand side array 1277 // inputs[1]: required: the right-hand side array 1278 // 1279 // TensorFlow equivalent: FloorDiv 1280 struct FloorDivOperator : Operator { 1281 FloorDivOperator() : Operator(OperatorType::kFloorDiv) {} 1282 }; 1283 1284 // Element-wise floor mod operator. 1285 // 1286 // Inputs: 1287 // inputs[0]: required: the left-hand side array 1288 // inputs[1]: required: the right-hand side array 1289 // 1290 // TensorFlow equivalent: FloorMod 1291 struct FloorModOperator : Operator { 1292 FloorModOperator() : Operator(OperatorType::kFloorMod) {} 1293 }; 1294 1295 struct RandomUniformOperator : Operator { 1296 RandomUniformOperator() : Operator(OperatorType::kRandomUniform) {} 1297 ArrayDataType dtype = ArrayDataType::kNone; 1298 int64 seed; 1299 int64 seed2; 1300 }; 1301 1302 // Creates a sequence of numbers that begins at start and extends by increments 1303 // of delta up to but not including limit. 1304 // 1305 // The dtype of the resulting tensor is inferred from the inputs unless it is 1306 // provided explicitly. 1307 // 1308 // Inputs: 1309 // inputs[0]: required: the start 1310 // inputs[1]: required: the limit 1311 // inputs[2]: required: the delta 1312 // 1313 // TensorFlow equivalent: Range 1314 struct RangeOperator : Operator { 1315 RangeOperator() : Operator(OperatorType::kRange) {} 1316 ArrayDataType dtype = ArrayDataType::kNone; 1317 }; 1318 1319 // Rank operator. Extracts the rank of the tensor. 1320 // 1321 // Inputs: 1322 // inputs[0]: required: the input array 1323 // 1324 // This operation outputs a 0-D int32 Tensor representing the rank of input. 1325 // 1326 // TensorFlow equivalent: Rank. 1327 struct TensorFlowRankOperator : Operator { 1328 TensorFlowRankOperator() : Operator(OperatorType::kRank) {} 1329 ArrayDataType output_data_type = ArrayDataType::kInt32; 1330 }; 1331 1332 // Element-wise negation (-x) operator. 1333 // 1334 // Inputs: 1335 // inputs[0]: required: the input array 1336 // 1337 // TensorFlow equivalent: Neg 1338 struct NegOperator : Operator { 1339 NegOperator() : Operator(OperatorType::kNeg) {} 1340 }; 1341 1342 // Element-wise select operator choosing elements from inputs[1] or input[2] 1343 // 1344 // Inputs: 1345 // inputs[0]: required: boolean mask per index 1346 // inputs[1]: required: tensor of values if true 1347 // inputs[2]: required: tensor of values if false 1348 // 1349 // TensorFlow equivalent: Select 1350 struct SelectOperator : Operator { 1351 SelectOperator() : Operator(OperatorType::kSelect) {} 1352 }; 1353 1354 // Element-wise reciprocal-square-root (x^-0.5) operator. 1355 // 1356 // Inputs: 1357 // inputs[0]: required: the input array 1358 // 1359 // TensorFlow equivalent: Rsqrt 1360 struct TensorFlowRsqrtOperator : Operator { 1361 TensorFlowRsqrtOperator() : Operator(OperatorType::kRsqrt) {} 1362 }; 1363 1364 // Stacks a list of rank-R tensors into one rank-(R+1) tensor. 1365 // 1366 // Packs the list of tensors in values into a tensor with rank one higher than 1367 // each tensor in values, by packing them along the axis dimension. Given a list 1368 // of length N of tensors of shape (A, B, C);. 1369 // 1370 // Inputs: this operator accepts any number >= 1 of inputs. 1371 // inputs[i]: the i-th array to merge. 1372 // 1373 // TensorFlow equivalent: Pack 1374 struct PackOperator : Operator { 1375 PackOperator() : Operator(OperatorType::kPack) {} 1376 int values_count; 1377 int axis = 0; 1378 ArrayDataType dtype = ArrayDataType::kNone; 1379 }; 1380 1381 // Shape operator. Extracts the shape of the tensor. 1382 // 1383 // Inputs: 1384 // inputs[0]: required: the input array 1385 // 1386 // This operation outputs a 1-D integer tensor representing the shape of 1387 // the input. 1388 // 1389 // TensorFlow equivalent: Shape. 1390 struct TensorFlowShapeOperator : Operator { 1391 TensorFlowShapeOperator() : Operator(OperatorType::kShape) {} 1392 ArrayDataType output_data_type = ArrayDataType::kInt32; 1393 }; 1394 1395 // Element-wise square-root (x^0.5) operator. 1396 // 1397 // Inputs: 1398 // inputs[0]: required: the input array 1399 // 1400 // TensorFlow equivalent: Sqrt 1401 struct TensorFlowSqrtOperator : Operator { 1402 TensorFlowSqrtOperator() : Operator(OperatorType::kSqrt) {} 1403 }; 1404 1405 // Element-wise square (x*x) operator. 1406 // 1407 // Inputs: 1408 // inputs[0]: required: the input array 1409 // 1410 // TensorFlow equivalent: Square 1411 struct TensorFlowSquareOperator : Operator { 1412 TensorFlowSquareOperator() : Operator(OperatorType::kSquare) {} 1413 }; 1414 1415 // Element-wise squared difference ((x-y)*(x-y)) operator. 1416 // 1417 // Inputs: 1418 // inputs[0]: required: the left-hand side array 1419 // inputs[1]: required: the right-hand side array 1420 // 1421 // TensorFlow equivalent: SquaredDifference 1422 struct SquaredDifferenceOperator : Operator { 1423 SquaredDifferenceOperator() : Operator(OperatorType::kSquaredDifference) {} 1424 }; 1425 1426 // Transposes a tensor. 1427 // 1428 // By default, this operation performs a regular matrix transpose on 2-D input 1429 // tensors. 1430 // 1431 // Inputs: 1432 // inputs[0]: required: the input array 1433 // 1434 // TensorFlow equivalent: Transpose 1435 struct TransposeOperator : Operator { 1436 TransposeOperator() : Operator(OperatorType::kTranspose) {} 1437 std::vector<int> perm; 1438 }; 1439 1440 // Element-wise subtraction operator. 1441 // 1442 // Inputs: 1443 // inputs[0]: required: the left-hand side array 1444 // inputs[1]: required: the right-hand side array 1445 // 1446 // TensorFlow equivalent: Sub 1447 struct SubOperator : Operator { 1448 SubOperator() : Operator(OperatorType::kSub) {} 1449 }; 1450 1451 // Sum reduction: computes the sum of all of entries across the axes. 1452 // 1453 // Inputs: 1454 // inputs[0]: required: the input array 1455 // 1456 // TensorFlow equivalent: Sum 1457 struct TensorFlowSumOperator : Operator { 1458 TensorFlowSumOperator() : Operator(OperatorType::kSum) {} 1459 std::vector<int> axis; 1460 bool keep_dims = false; 1461 }; 1462 1463 // Prod reduction: computes the product of all of entries across the axes. 1464 // 1465 // Inputs: 1466 // inputs[0]: required: the input array 1467 // 1468 // TensorFlow equivalent: Prod 1469 struct TensorFlowProdOperator : Operator { 1470 TensorFlowProdOperator() : Operator(OperatorType::kReduceProd) {} 1471 std::vector<int> axis; 1472 bool keep_dims = false; 1473 }; 1474 1475 // TensorFlow Tile equivalent. Refer to TensorFlow documentation for details. 1476 // 1477 // Inputs: 1478 // inputs[0]: required: the input array 1479 // inputs[1]: required: int array with length of rank(input[0]) 1480 struct TensorFlowTileOperator : Operator { 1481 TensorFlowTileOperator() : Operator(OperatorType::kTile) {} 1482 }; 1483 1484 // TensorFlow Slice equivalent. Refer to TensorFlow documentation for details. 1485 struct SliceOperator : Operator { 1486 SliceOperator() : Operator(OperatorType::kSlice) {} 1487 1488 std::vector<int> begin; 1489 std::vector<int> size; 1490 }; 1491 1492 // TensorFlow Split equivalent. Refer to TensorFlow documentation for details. 1493 // Not fully supported, just a placeholder to handle TensorFlow graphs and 1494 // support graph transformations to other operator types by matching sub-graphs. 1495 struct TensorFlowSplitOperator : Operator { 1496 TensorFlowSplitOperator() : Operator(OperatorType::kSplit) {} 1497 int num_split = 0; 1498 }; 1499 1500 // TensorFlow SplitV equivalent. Refer to TensorFlow documentation for details. 1501 struct TensorFlowSplitVOperator : Operator { 1502 TensorFlowSplitVOperator() : Operator(OperatorType::kSplitV) {} 1503 int num_split = 0; 1504 }; 1505 1506 // TensorFlow Concat equivalent. Refer to TensorFlow documentation for details. 1507 // Not fully supported, just a placeholder to handle TensorFlow graphs and 1508 // support graph transformations to other operator types by matching sub-graphs. 1509 // Concretely, once the concat dim becomes known, if it is the depth 1510 // dimension then we can change this op into a DepthConcatenation op. 1511 // Otherwise, we hope for some other graph transformation to drop this node. 1512 struct TensorFlowConcatOperator : Operator { 1513 TensorFlowConcatOperator() : Operator(OperatorType::kConcat) {} 1514 }; 1515 1516 // TensorFlow ConcatV2 equivalent. Refer to TensorFlow documentation for 1517 // details. 1518 // Not fully supported, just a placeholder to handle TensorFlow graphs and 1519 // support graph transformations to other operator types by matching sub-graphs. 1520 // Concretely, once the concat dim becomes known, if it is the depth 1521 // dimension then we can change this op into a DepthConcatenation op. 1522 // Otherwise, we hope for some other graph transformation to drop this node. 1523 struct TensorFlowConcatV2Operator : Operator { 1524 TensorFlowConcatV2Operator() : Operator(OperatorType::kConcatV2) {} 1525 }; 1526 1527 // TensorFlow Merge equivalent. Refer to TensorFlow documentation for details. 1528 // 1529 // Inputs: this operator accepts any number >= 1 of inputs. 1530 // inputs[i]: the i-th array to merge. 1531 // 1532 // It is expected that graph transformations will drop all but exactly one 1533 // of the inputs, at which point the Merge node will be equivalent to an 1534 // Identity node forwarding the remaining input. 1535 // 1536 // Note: We do not currently support runtime control flow: we only support 1537 // control flow that can be resolved at tooling time (independently of input 1538 // activations). 1539 struct TensorFlowMergeOperator : Operator { 1540 TensorFlowMergeOperator() : Operator(OperatorType::kMerge) {} 1541 }; 1542 1543 // TensorFlow Switch equivalent. Refer to TensorFlow documentation for details. 1544 // 1545 // Inputs: 1546 // inputs[0]: required: the input array 1547 // inputs[1]: required: the boolean predicate, given as an array of size 1 1548 // and of type kBool, will determine which output gets selected. 1549 // 1550 // Outputs: a TensorFlow Switch node always has exactly two outputs. Depending 1551 // on the boolean value that the input predicate resolves to (see note below), 1552 // one or the other of the outputs will be 'selected': the input array will be 1553 // forwarded to the 'selected output' as if by a Identity node, while the other 1554 // output will be discarded, and any graph edge connecting that discarded output 1555 // will be dropped. The rule for selecting outputs is as follows: 1556 // outputs[0] will be selected if the input predicate resolves to 'true'. 1557 // outputs[1] will be selected if the input predicate resolves to 'false'. 1558 // 1559 // Note: We do not currently support runtime control flow: we only support 1560 // control flow that can be resolved at tooling time (independently of input 1561 // activations). 1562 struct TensorFlowSwitchOperator : Operator { 1563 TensorFlowSwitchOperator() : Operator(OperatorType::kSwitch) {} 1564 }; 1565 1566 // TensorFlow All equivalent. Refer to TensorFlow documentation for details. 1567 // Not fully supported, just a placeholder to handle TensorFlow graphs and 1568 // support graph transformations to other operator types by matching sub-graphs. 1569 // Typically, this is only used as an input to an Assert node, so can be 1570 // removed as an unused node as we drop Assert nodes. 1571 struct TensorFlowAllOperator : Operator { 1572 TensorFlowAllOperator() : Operator(OperatorType::kAll) {} 1573 }; 1574 1575 // TensorFlow Assert equivalent. Refer to TensorFlow documentation for details. 1576 // Not fully supported, just a placeholder to handle TensorFlow graphs and 1577 // support graph transformations to other operator types by matching sub-graphs. 1578 // Typically, we just drop Assert nodes. 1579 struct TensorFlowAssertOperator : Operator { 1580 TensorFlowAssertOperator() : Operator(OperatorType::kAssert) {} 1581 }; 1582 1583 // TensorFlow Less equivalent. Refer to TensorFlow documentation for details. 1584 // Not fully supported, just a placeholder to handle TensorFlow graphs and 1585 // support graph transformations to other operator types by matching sub-graphs. 1586 // Typically, this is only used as an input to an Assert node, so can be 1587 // removed as an unused node as we drop Assert nodes. 1588 struct TensorFlowLessOperator : Operator { 1589 TensorFlowLessOperator() : Operator(OperatorType::kLess) {} 1590 }; 1591 1592 // TensorFlow LessEqual equivalent. Refer to TensorFlow documentation for 1593 // details. 1594 // Not fully supported, just a placeholder to handle TensorFlow graphs and 1595 // support graph transformations to other operator types by matching sub-graphs. 1596 // Typically, this is only used as an input to an Assert node, so can be 1597 // removed as an unused node as we drop Assert nodes. 1598 struct TensorFlowLessEqualOperator : Operator { 1599 TensorFlowLessEqualOperator() : Operator(OperatorType::kLessEqual) {} 1600 }; 1601 1602 // TensorFlow Less equivalent. Refer to TensorFlow documentation for details. 1603 // Not fully supported, just a placeholder to handle TensorFlow graphs and 1604 // support graph transformations to other operator types by matching sub-graphs. 1605 // Typically, this is only used as an input to an Assert node, so can be 1606 // removed as an unused node as we drop Assert nodes. 1607 struct TensorFlowGreaterOperator : Operator { 1608 TensorFlowGreaterOperator() : Operator(OperatorType::kGreater) {} 1609 }; 1610 1611 // TensorFlow GreaterEqual equivalent. Refer to TensorFlow documentation for 1612 // details. 1613 // Not fully supported, just a placeholder to handle TensorFlow graphs and 1614 // support graph transformations to other operator types by matching sub-graphs. 1615 // Typically, this is only used as an input to an Assert node, so can be 1616 // removed as an unused node as we drop Assert nodes. 1617 struct TensorFlowGreaterEqualOperator : Operator { 1618 TensorFlowGreaterEqualOperator() : Operator(OperatorType::kGreaterEqual) {} 1619 }; 1620 1621 // TensorFlow Equal equivalent. Refer to TensorFlow documentation for 1622 // details. 1623 // Not fully supported, just a placeholder to handle TensorFlow graphs and 1624 // support graph transformations to other operator types by matching sub-graphs. 1625 // Typically, this is only used as an input to an Assert node, so can be 1626 // removed as an unused node as we drop Assert nodes. 1627 struct TensorFlowEqualOperator : Operator { 1628 TensorFlowEqualOperator() : Operator(OperatorType::kEqual) {} 1629 }; 1630 1631 // TensorFlow Not Equal equivalent. Refer to TensorFlow documentation for 1632 // details. 1633 struct TensorFlowNotEqualOperator : Operator { 1634 TensorFlowNotEqualOperator() : Operator(OperatorType::kNotEqual) {} 1635 }; 1636 1637 // Max reduction: computes the max of all of entries across the axes. 1638 // 1639 // Inputs: 1640 // inputs[0]: required: the input array 1641 // 1642 // TensorFlow equivalent: Max 1643 struct TensorFlowMaxOperator : Operator { 1644 TensorFlowMaxOperator() : Operator(OperatorType::kReduceMax) {} 1645 std::vector<int> axis; 1646 bool keep_dims = false; 1647 }; 1648 1649 // Min reduction: computes the min of all of entries across the axes. 1650 // 1651 // Inputs: 1652 // inputs[0]: required: the input array 1653 // 1654 // TensorFlow equivalent: Min 1655 struct TensorFlowMinOperator : Operator { 1656 TensorFlowMinOperator() : Operator(OperatorType::kReduceMin) {} 1657 std::vector<int> axis; 1658 bool keep_dims = false; 1659 }; 1660 1661 // Element-wise maximum operator. Currently it only supports scalar as 1662 // the second operand. 1663 // 1664 // Inputs: 1665 // inputs[0]: required: the left-hand side array 1666 // inputs[1]: required: the right-hand side array 1667 // 1668 // TensorFlow equivalent: Maximum 1669 struct TensorFlowMaximumOperator : Operator { 1670 TensorFlowMaximumOperator() : Operator(OperatorType::kMaximum) {} 1671 }; 1672 1673 // Element-wise minimum operator. Currently it only supports scalar as 1674 // the second operand. 1675 // 1676 // Inputs: 1677 // inputs[0]: required: the left-hand side array 1678 // inputs[1]: required: the right-hand side array 1679 // 1680 // TensorFlow equivalent: Minimum 1681 struct TensorFlowMinimumOperator : Operator { 1682 TensorFlowMinimumOperator() : Operator(OperatorType::kMinimum) {} 1683 }; 1684 1685 // General TF operation, unsupported by tf.mini. Expected to be dropped by 1686 // graph transformations. 1687 struct TensorFlowUnsupportedOperator : Operator { 1688 TensorFlowUnsupportedOperator() : Operator(OperatorType::kUnsupported) {} 1689 1690 // The original TF operation type. Used for diagnostic purposes. 1691 string tensorflow_op; 1692 // A boolean indicating if the unsupported op should be treated as quantized. 1693 bool quantized = false; 1694 // A boolean indicating if the unsupported op output should allow float values 1695 // in quantized mode. 1696 bool support_output_type_float_in_quantized_op = false; 1697 // Output data types 1698 std::vector<ArrayDataType> output_data_types; 1699 // Output shapes. 1700 std::vector<Shape> output_shapes; 1701 }; 1702 1703 // Softmax activation function. 1704 // 1705 // Inputs: 1706 // inputs[0]: required: the input array 1707 // 1708 // TensorFlow equivalent: Softmax 1709 struct SoftmaxOperator : Operator { 1710 SoftmaxOperator() : Operator(OperatorType::kSoftmax) {} 1711 float beta = 0.f; 1712 }; 1713 1714 // LogSoftmax activation function. 1715 // 1716 // Inputs: 1717 // inputs[0]: required: the logits input array 1718 // 1719 // TensorFlow equivalent: LogSoftmax 1720 struct LogSoftmaxOperator : Operator { 1721 LogSoftmaxOperator() : Operator(OperatorType::kLogSoftmax) {} 1722 1723 // LogSoftmax can in principal have very large negative output, depending on 1724 // the input size. However, input x_i that is less than x_max-10 is 1725 // accumulated as exp(x_i-x_max), which is truncated to zero. 1726 // 1727 // Since we effectively disregard smallish inputs in the normalizing factor, 1728 // we also drop them in the output (set to minimum output), and in doing so 1729 // make better use of the quantization range / resolution. 1730 static constexpr float kOutputRangeMin = -16.0; 1731 }; 1732 1733 // Cast operator. 1734 // 1735 // Inputs: 1736 // inputs[0]: required: the input array 1737 // 1738 // TensorFlow equivalent: Cast 1739 struct CastOperator : Operator { 1740 CastOperator() : Operator(OperatorType::kCast) {} 1741 ArrayDataType src_data_type = ArrayDataType::kNone; 1742 ArrayDataType dst_data_type = ArrayDataType::kNone; 1743 }; 1744 1745 // Floor operator. 1746 // 1747 // Inputs: 1748 // inputs[0]: required: the input array 1749 // 1750 // TensorFlow equivalent: Floor 1751 struct FloorOperator : Operator { 1752 FloorOperator() : Operator(OperatorType::kFloor) {} 1753 }; 1754 1755 // Ceil operator. 1756 // 1757 // Inputs: 1758 // inputs[0]: required: the input array 1759 // 1760 // TensorFlow equivalent: Ceil 1761 struct CeilOperator : Operator { 1762 CeilOperator() : Operator(OperatorType::kCeil) {} 1763 }; 1764 1765 // Round operator. 1766 // 1767 // Inputs: 1768 // inputs[0]: required: the input array 1769 // 1770 // TensorFlow equivalent: Round 1771 struct RoundOperator : Operator { 1772 RoundOperator() : Operator(OperatorType::kRound) {} 1773 }; 1774 1775 // Gather operator. It gathers slices from params according to indices. 1776 // Only 1-D indices are supported at the moment. 1777 // 1778 // Inputs: 1779 // inputs[0]: required: the params array 1780 // inputs[1]: required: the indices to gather 1781 // inputs[2]: optional: axis 1782 // 1783 // TensorFlow equivalent: Gather 1784 struct GatherOperator : Operator { 1785 GatherOperator() : Operator(OperatorType::kGather) {} 1786 // Axis is populated explicitly or implicitly from the axis input by 1787 // ResolveGatherAttributes. An empty axis indicates that the axis has not yet 1788 // be resolved. 1789 absl::optional<int> axis; 1790 1791 // This field is not used by the standard TF Lite export but it is still need 1792 // for legacy Gather implementations. 1793 int input_rank = 0; 1794 }; 1795 1796 // GatherNd operator. It gathers slices from params according to indices. 1797 // 1798 // Inputs: 1799 // inputs[0]: required: the params array 1800 // inputs[1]: required: the indices to gather 1801 // 1802 // TensorFlow equivalent: GatherNd 1803 struct GatherNdOperator : Operator { 1804 GatherNdOperator() : Operator(OperatorType::kGatherNd) {} 1805 }; 1806 1807 // ArgMax operator. It returns the index of the maximum value along axis. 1808 // 1809 // Inputs: 1810 // inputs[0]: required: the input tensor 1811 // inputs[1]: optional: 0-D (scalar) axis 1812 // 1813 // TensorFlow equivalent: ArgMax 1814 struct ArgMaxOperator : Operator { 1815 ArgMaxOperator() : Operator(OperatorType::kArgMax) {} 1816 ArrayDataType output_data_type = ArrayDataType::kInt64; 1817 }; 1818 1819 // ArgMin operator. It returns the index of the minimum value along axis. 1820 // 1821 // Inputs: 1822 // inputs[0]: required: the input tensor 1823 // inputs[1]: optional: 0-D (scalar) axis 1824 // 1825 // TensorFlow equivalent: ArgMin 1826 struct ArgMinOperator : Operator { 1827 ArgMinOperator() : Operator(OperatorType::kArgMin) {} 1828 ArrayDataType output_data_type = ArrayDataType::kInt64; 1829 }; 1830 1831 // ResizeBilinear operator. It resizes input images with bilinear interpolation. 1832 // It does not support align_corners at the moment. 1833 // 1834 // Inputs: 1835 // inputs[0]: required: the input array 1836 // inputs[1]: required: the new image size 1837 // 1838 // TensorFlow equivalent: ResizeBilinear 1839 struct ResizeBilinearOperator : Operator { 1840 ResizeBilinearOperator() : Operator(OperatorType::kResizeBilinear) {} 1841 1842 bool align_corners = false; 1843 bool half_pixel_centers = false; 1844 }; 1845 1846 // ResizeNearestNeighborOperator operator. It resizes input images with nearest 1847 // neighbor interpolation. It does not support align_corners at the moment. 1848 // 1849 // Inputs: 1850 // inputs[0]: required: the input array 1851 // inputs[1]: required: the new image size 1852 // 1853 // TensorFlow equivalent: ResizeNearestNeighbor 1854 struct ResizeNearestNeighborOperator : Operator { 1855 ResizeNearestNeighborOperator() 1856 : Operator(OperatorType::kResizeNearestNeighbor) {} 1857 1858 bool align_corners = false; 1859 }; 1860 1861 // SpaceToBatchND operator. It divides spatial dimensions into a grid of 1862 // blocks and interleaves these blocks with the batch dimension. Currently, 1863 // only 2-d blocks are supported. 1864 // 1865 // Inputs: 1866 // inputs[0]: required: the input array 1867 // inputs[1]: required: the block shape 1868 // inputs[2]: required: the paddings 1869 // 1870 // TensorFlow equivalent: SpaceToBatchND 1871 struct SpaceToBatchNDOperator : Operator { 1872 SpaceToBatchNDOperator() : Operator(OperatorType::kSpaceToBatchND) {} 1873 1874 std::vector<int> block_shape; 1875 std::vector<int> before_paddings; 1876 std::vector<int> after_paddings; 1877 }; 1878 1879 // BatchToSpaceND operator. Rearranges data from batch into blocks of 1880 // spatial data. Currently, only 2-d blocks are supported. 1881 // 1882 // Inputs: 1883 // inputs[0]: required: the input array 1884 // inputs[1]: required: the block shape 1885 // inputs[2]: required: the crops 1886 // 1887 // TensorFlow equivalent: BatchToSpaceND 1888 struct BatchToSpaceNDOperator : Operator { 1889 BatchToSpaceNDOperator() : Operator(OperatorType::kBatchToSpaceND) {} 1890 1891 std::vector<int> block_shape; 1892 std::vector<int> before_crops; 1893 std::vector<int> after_crops; 1894 }; 1895 1896 // Mean operator. 1897 // 1898 // Inputs: 1899 // inputs[0]: required: the input array 1900 // 1901 // TensorFlow equivalent: Mean 1902 struct MeanOperator : Operator { 1903 MeanOperator() : Operator(OperatorType::kMean) {} 1904 1905 std::vector<int> axis; 1906 bool keep_dims = false; 1907 }; 1908 1909 // Svdf operator: 1910 // 1911 // Inputs: 1912 // inputs[0]: required: the input array 1913 // inputs[1]: required: weights_feature 1914 // inputs[2]: required: weights_time 1915 // inputs[3]: optional: bias 1916 struct SvdfOperator : Operator { 1917 SvdfOperator() : Operator(OperatorType::kSvdf) {} 1918 int rank; 1919 }; 1920 1921 // TopKV2 operator. 1922 // 1923 // Inputs: 1924 // input tensor and top_k scalar. 1925 struct TopKV2Operator : Operator { 1926 TopKV2Operator() : Operator(OperatorType::kTopK_V2) {} 1927 }; 1928 1929 // DynamicPartition operator: 1930 // 1931 // Inputs: 1932 // inputs[0]: required: data. 1933 // inputs[1]: required: partitions. 1934 // 1935 // TensorFlow equivalent: DynamicPartition 1936 struct DynamicPartitionOperator : Operator { 1937 DynamicPartitionOperator() : Operator(OperatorType::kDynamicPartition) {} 1938 int num_partitions; 1939 }; 1940 1941 // DynamicStitch operator: 1942 // 1943 // Inputs: 1944 // inputs[0,N): required: indices. 1945 // inputs[N,2N): required: data. 1946 // 1947 // TensorFlow equivalent: DynamicStitch/ParallelDynamicStitch 1948 struct DynamicStitchOperator : Operator { 1949 DynamicStitchOperator() : Operator(OperatorType::kDynamicStitch) {} 1950 int num_partitions; 1951 }; 1952 1953 // SparseToDense operator: 1954 // 1955 // Inputs: 1956 // Inputs[0]: required: sparse_indices. 1957 // Inputs[1]: required: output_shape. 1958 // Inputs[2]: required: sparse_values. 1959 // 1960 // TensorFlow equivalent: SparseToDense. 1961 struct SparseToDenseOperator : Operator { 1962 SparseToDenseOperator() : Operator(OperatorType::kSparseToDense) {} 1963 bool validate_indices; 1964 }; 1965 1966 // Pow operator: 1967 // 1968 // Inputs: 1969 // Inputs[0]: required: A tensor. 1970 // Inputs[1]: required: A tensor. 1971 // 1972 // TensorFlow equivalent: Pow. 1973 struct PowOperator : Operator { 1974 PowOperator() : Operator(OperatorType::kPow) {} 1975 }; 1976 1977 // Any operator: 1978 // 1979 // Inputs: 1980 // Inputs[0]: required: A boolean input tensor. 1981 // Inputs[1]: required: reduction_indices. 1982 // 1983 // TensorFlow equivalent: tf.reduce_any. 1984 struct TensorFlowAnyOperator : Operator { 1985 TensorFlowAnyOperator() : Operator(OperatorType::kAny) {} 1986 std::vector<int> axis; 1987 bool keep_dims = false; 1988 }; 1989 1990 // LogicalAnd operator: 1991 // 1992 // Inputs: 1993 // Inputs[0]: required: A boolean tensor. 1994 // Inputs[1]: required: A boolean tensor. 1995 // 1996 // TensorFlow equivalent: tf.logical_and. 1997 struct LogicalAndOperator : Operator { 1998 LogicalAndOperator() : Operator(OperatorType::kLogicalAnd) {} 1999 }; 2000 2001 // LogicalNot operator: 2002 // 2003 // Inputs: 2004 // Inputs[0]: required: A boolean tensor. 2005 // 2006 // TensorFlow equivalent: tf.logical_not. 2007 struct LogicalNotOperator : Operator { 2008 LogicalNotOperator() : Operator(OperatorType::kLogicalNot) {} 2009 }; 2010 2011 // OneHot operator: 2012 // 2013 // Inputs: 2014 // Inputs[0]: required: indices. 2015 // Inputs[1]: required: depth. 2016 // Inputs[2]: required: on_value. 2017 // Inputs[3]: required: off_value. 2018 // 2019 // TensorFlow equivalent: OneHot. 2020 struct OneHotOperator : Operator { 2021 enum Inputs { 2022 INDICES_INPUT = 0, 2023 DEPTH_INPUT = 1, 2024 ON_VALUE_INPUT = 2, 2025 OFF_VALUE_INPUT = 3, 2026 }; 2027 2028 OneHotOperator() : Operator(OperatorType::kOneHot) {} 2029 int axis = -1; 2030 }; 2031 2032 // LogicalOr operator: 2033 // 2034 // Inputs: 2035 // Inputs[0]: required: A Bool tensor. 2036 // Inputs[1]: required: A Bool tensor. 2037 // 2038 // TensorFlow equivalent: LogicalOr. 2039 struct LogicalOrOperator : Operator { 2040 LogicalOrOperator() : Operator(OperatorType::kLogicalOr) {} 2041 }; 2042 2043 // Unpack operator: 2044 // 2045 // Inputs: 2046 // Inputs[0]: required: A boolean input tensor. 2047 // Inputs[1]: required: reduction_indices. 2048 // 2049 // TensorFlow equivalent: tf.unstack. 2050 struct UnpackOperator : Operator { 2051 UnpackOperator() : Operator(OperatorType::kUnpack) {} 2052 int num; 2053 int axis; 2054 ArrayDataType dtype = ArrayDataType::kNone; 2055 }; 2056 2057 // ZerosLike operator: 2058 // 2059 // Inputs: 2060 // inputs[0]: required: the input array 2061 // 2062 // TensorFlow equivalent: tf.zeros_like 2063 struct TensorFlowZerosLikeOperator : Operator { 2064 TensorFlowZerosLikeOperator() : Operator(OperatorType::kZerosLike) {} 2065 }; 2066 2067 // ReverseV2 operator: 2068 // 2069 // Inputs: 2070 // Inputs[0]: required: the input array. 2071 // 2072 // TensorFlow equivalent: ReverseV2. 2073 struct ReverseV2Operator : Operator { 2074 ReverseV2Operator() : Operator(OperatorType::kReverseV2) {} 2075 }; 2076 2077 enum class MirrorPadMode { kNone, kSymmetric, kReflect }; 2078 2079 // MirrorPad Operator: 2080 // 2081 // Inputs: 2082 // Inputs[0]: required: input tensor to be padded. 2083 // Inputs[1]: required: 2 Column matrix specifying padding sizes. The number of 2084 // rows must be the same as the rank of the input. 2085 // Inputs[2]: required: REFLECT or SYMMETRIC. 2086 // 2087 // TensorFlow equivalent: MirrorPad. 2088 struct MirrorPadOperator : Operator { 2089 MirrorPadOperator() : Operator(OperatorType::kMirrorPad) {} 2090 // mode is either SYMMETRIC or REFLECT. 2091 MirrorPadMode mode; 2092 }; 2093 2094 // ReverseSequence operator: 2095 // 2096 // Inputs: 2097 // Inputs[0]: required: the input array. 2098 // Inputs[1]: required: the lengths of the elements to be reversed. 2099 // 2100 // TensorFlow equivalent: tf.reverse_sequence. 2101 struct ReverseSequenceOperator : Operator { 2102 ReverseSequenceOperator() : Operator(OperatorType::kReverseSequence) {} 2103 int seq_dim; 2104 int batch_dim = 0; 2105 }; 2106 2107 // Unique Operator: 2108 // 2109 // Inputs: 2110 // inputs[0]: required: the input array 2111 // 2112 // TensorFlow equivalent: Unique 2113 struct UniqueOperator : Operator { 2114 UniqueOperator() : Operator(OperatorType::kUnique) {} 2115 ArrayDataType idx_out_type = ArrayDataType::kInt32; 2116 }; 2117 2118 struct UnidirectionalSequenceRnnOperator : Operator { 2119 UnidirectionalSequenceRnnOperator() 2120 : Operator(OperatorType::kUnidirectionalSequenceRnn) {} 2121 bool time_major; 2122 FusedActivationFunctionType fused_activation_function; 2123 }; 2124 2125 // Where Operator: 2126 // Return the coordinates of the true values in condition tensor in row-major 2127 // order. 2128 // 2129 // Inputs: 2130 // inputs[0]: required: boolean condition tensor 2131 // 2132 // TensorFlow equivalent: Where 2133 struct WhereOperator : Operator { 2134 WhereOperator() : Operator(OperatorType::kWhere) {} 2135 }; 2136 2137 // Matrix Diag Operator: 2138 // Construct a batched diagonal tensor with given batched diagonal values. 2139 // Inputs: A tensor of values that will be on the diagonal of the returned 2140 // tensor. 2141 struct MatrixDiagOperator : Operator { 2142 MatrixDiagOperator() : Operator(OperatorType::kMatrixDiag) {} 2143 }; 2144 2145 // Matrix Diag Operator V2: 2146 // Construct a batched diagonal tensor with given batched diagonal values. 2147 // Not fully supported, contains 4 extra inputs compared to MatrixDiag. Behave 2148 // like MatrixDiag when default parameters are used. 2149 struct MatrixDiagV2Operator : Operator { 2150 MatrixDiagV2Operator() : Operator(OperatorType::kMatrixDiagV2) {} 2151 }; 2152 2153 // Matrix Diag Operator V3: 2154 // Construct a batched diagonal tensor with given batched diagonal values. 2155 // Not fully supported, contains 5 extra inputs compared to MatrixDiag. Behave 2156 // like MatrixDiag when default parameters are used. 2157 // V3 is only different from V2 because it has an extra attribute (align) which 2158 // controls the alignment of diagonals in the band matrix (compact) format. 2159 // The alignment in V2 contradicts with the default alignment in V3 so V2 is 2160 // skipped. (It has never been, and should never be, exposed in the public API.) 2161 struct MatrixDiagV3Operator : Operator { 2162 MatrixDiagV3Operator() : Operator(OperatorType::kMatrixDiagV3) {} 2163 }; 2164 2165 // Matrix Set Diag Operator: 2166 // Construct a batched diagonal tensor with given input and diagonal values. 2167 // Input is a rank (k+1) tensor of values. 2168 // diagonal is a rank (k) tensor of values that will be on the diagonal 2169 // of the returned output. Output is rank k+1. 2170 // tensor. 2171 struct MatrixSetDiagOperator : Operator { 2172 MatrixSetDiagOperator() : Operator(OperatorType::kMatrixSetDiag) {} 2173 }; 2174 2175 // Matrix Set Diag Operator V2: 2176 // Construct a batched diagonal tensor with given input and diagonal values. 2177 // Not fully supported, contains 1 extra inputs compared to MatrixSetDiag. 2178 // Behave like MatrixSetDiag when default parameters are used. 2179 struct MatrixSetDiagV2Operator : Operator { 2180 MatrixSetDiagV2Operator() : Operator(OperatorType::kMatrixSetDiagV2) {} 2181 }; 2182 2183 // Matrix Set Diag Operator V3: 2184 // Construct a batched diagonal tensor with given input and diagonal values. 2185 // Not fully supported, contains 2 extra inputs compared to MatrixSetDiag. 2186 // Behave like MatrixSetDiag when default parameters are used. 2187 // V3 is only different from V2 because it has an extra attribute (align) which 2188 // controls the alignment of diagonals in the band matrix (compact) format. 2189 // The alignment in V2 contradicts with the default alignment in V3 so V2 is 2190 // skipped. (It has never been, and should never be, exposed in the public API.) 2191 struct MatrixSetDiagV3Operator : Operator { 2192 MatrixSetDiagV3Operator() : Operator(OperatorType::kMatrixSetDiagV3) {} 2193 }; 2194 2195 struct SegmentSumOperator : Operator { 2196 SegmentSumOperator() : Operator(OperatorType::kSegmentSum) {} 2197 }; 2198 2199 // Alloc's are used for transient arrays only. An Alloc specifies which interval 2200 // of the "transient_data" workspace buffer passed to inference functions, is to 2201 // be used for the transient array at hand. The 'start' and 'end' values are 2202 // offsets from the start of the workspace buffer, expressed in bytes. 2203 struct Alloc { 2204 int64 start = 0; 2205 int64 end = 0; 2206 }; 2207 2208 inline bool operator<(const Alloc& a, const Alloc& b) { 2209 return a.start < b.start; 2210 } 2211 2212 // Array represents an array (either a constant parameter array or an 2213 // activations array) in a Model. 2214 struct Array { 2215 template <ArrayDataType A> 2216 const Buffer<A>& GetBuffer() const { 2217 DCHECK(buffer); 2218 DCHECK(buffer->type == A); 2219 return *static_cast<const Buffer<A>*>(buffer.get()); 2220 } 2221 template <ArrayDataType A> 2222 Buffer<A>& GetMutableBuffer() { 2223 if (!buffer) { 2224 Buffer<A>* ptr = new Buffer<A>; 2225 buffer = std::unique_ptr<GenericBuffer>(ptr); 2226 } 2227 DCHECK(buffer); 2228 DCHECK(buffer->type == A); 2229 return *static_cast<Buffer<A>*>(buffer.get()); 2230 } 2231 Alloc& GetOrCreateAlloc() { 2232 if (!alloc) { 2233 alloc = std::unique_ptr<Alloc>(new Alloc); 2234 } 2235 return *alloc; 2236 } 2237 MinMax& GetOrCreateMinMax() { 2238 if (!minmax) { 2239 minmax = std::unique_ptr<MinMax>(new MinMax); 2240 } 2241 return *minmax; 2242 } 2243 MinMax& GetMinMax() const { 2244 DCHECK(minmax); 2245 return *minmax; 2246 } 2247 QuantizationParams& GetOrCreateQuantizationParams() { 2248 if (!quantization_params) { 2249 quantization_params = 2250 std::unique_ptr<QuantizationParams>(new QuantizationParams); 2251 } 2252 return *quantization_params; 2253 } 2254 QuantizationParams& GetQuantizationParams() const { 2255 DCHECK(quantization_params); 2256 return *quantization_params; 2257 } 2258 2259 // The data type of the actual elements of this array, that is: 2260 // - If there is a buffer (see 'buffer' member), it must be of the same 2261 // type. 2262 // - If there is no buffer, meaning that this is a runtime (i.e. activations) 2263 // array, then this specifies the type of elements that there will be 2264 // at runtime. 2265 // 2266 // Note that this only specifies the storage type of elements; this does 2267 // not specify whether these are to be treated as 'real' or 'quantized' 2268 // values. 2269 // That is decided by whether the 'quantization_params' member is null. 2270 ArrayDataType data_type = ArrayDataType::kNone; 2271 // The final value that data_type should have at the end of graph 2272 // transformations 2273 ArrayDataType final_data_type = ArrayDataType::kNone; 2274 // The dimensions of this array --- this specifies both sizes and strides 2275 // (the storage layout). 2276 // 2277 // Issues with shape handling that remain include: 2278 // - No way to distinguish between 0-dimensional dims and missing dims. 2279 // - No way to describe dims that may be runtime-variable. 2280 // - Addressing of dims by integer index differs in different graph formats 2281 // (TensorFlow vs. other frameworks vs. what we have informally grown 2282 // within toco). 2283 // This is currently quite messy; see ReorderAxesOperator which is how we 2284 // bridge some of these discrepancies at the moment. This is overdue for 2285 // a redesign; I'm thinking that it would be nice to have more flexible 2286 // dims that allow mapping 1:1, cleanly, dims as they are in various 2287 // formats, 2288 // then explicitly convert between different conventions. 2289 2290 // Proto-style accessors 2291 bool has_shape() const { return array_shape != nullptr; } 2292 const Shape& shape() const { 2293 CHECK(has_shape()); 2294 return *array_shape; 2295 } 2296 Shape* mutable_shape() { 2297 if (!array_shape) { 2298 array_shape.reset(new Shape); 2299 } 2300 return array_shape.get(); 2301 } 2302 void copy_shape(const Shape& src_shape) { *mutable_shape() = src_shape; } 2303 void clear_shape() { array_shape = nullptr; } 2304 2305 // The constant buffer backing this array. This is non-null if and only if 2306 // this is a constant parameter array. Conversely, this is null for 2307 // activations arrays. 2308 // 2309 // Note that this buffer is pure storage. In the case of quantized values, 2310 // it only stores the quantized values, it does not know by itself about the 2311 // quantization parameters necessary to interprete these values, that is 2312 // in the separate 'quantization_params' field. In fact, this 'buffer' field 2313 // does no even know whether values are quantized. It only has a data_type, 2314 // which must equal the 'data_type' member here, and which only describes 2315 // the storage type of element, does not tell whether they are quantized i.e. 2316 // whether they are to be interpreted with quantization_params. 2317 std::unique_ptr<GenericBuffer> buffer; 2318 // Only for activation arrays (i.e. when 'buffer' is null). 2319 // Only for code generation. 2320 // 2321 // Describes the allocation of this array within the workspace buffer 2322 // allocated 2323 // for all transient arrays. 2324 std::unique_ptr<Alloc> alloc; 2325 // Describes the [min, max] range of values 2326 // to be assumed when determining quantization_params. 2327 // 2328 // Only used for quantization. In fact, only used for determining 2329 // quantization_params. 2330 // 2331 // Used for both constant arrays (those having a 'buffer') and non-constant 2332 // arrays (activations). Indeed, it is important to use the same min-max range 2333 // as was used during training, even if that min-max range is slightly wrong 2334 // w.r.t. actual buffer elements. Doing otherwise would defeat the point of 2335 // re-training for quantization. 2336 std::unique_ptr<MinMax> minmax; 2337 // Quantization parameters. The non-null-ness of this pointer is what 2338 // defines whether this array is quantized or not. 2339 // 2340 // If this is non-null, then these quantization parameters are to be used 2341 // to assign a meaning as real numbers to the elements of this array. 2342 std::unique_ptr<QuantizationParams> quantization_params; 2343 // narrow_range is a detail of how toco handles FakeQuant operators with 2344 // narrow_range, see 2345 // https://www.tensorflow.org/api_docs/python/tf/fake_quant_with_min_max_vars 2346 // 2347 // For more context about what that is useful for, see the big comment in 2348 // graph_transformations/ensure_uint8_weights_safe_for_fast_int8_kernels.cc 2349 // 2350 // The narrow_range flag applies only to quantized arrays, and changes 2351 // their quantization in the following way when it is set to 'true': 2352 // 1. The computation of {zero_point, scale} from {min, max} needs to be 2353 // amended so that the real min value will get quantized to 2354 // (min_quantized_value + 1) instead of just (min_quantized_value). 2355 // E.g. for uint8 quantization, the real min value should get quantized to 2356 // the uint8 value 1, not 0. 2357 // 2. Quantized values should get clamped to the interval 2358 // [min_quantized_value + 1, max_value]. Equivalently, the 2359 // min_quantized_value should get nudged to (min_quantized_value + 1). 2360 // The reason why 1. does not imply 2. is that real values may not belong to 2361 // the stated [min, max] interval. Concretely, weights recorded at the last 2362 // learning step may not fall in the [min, max] interval recorded over 2363 // previous learning steps, as the values evolve across learning steps. 2364 // 2365 // Rationale why this is directly a field on Array: 2366 // - This can't be just a field on FakeQuantOperator, because 2367 // FakeQuantOperators are gone (DropFakeQuant) before we get to using that 2368 // information (Quantize). We need a place to store that bit in the interim. 2369 // - This can't be in QuantizationParams because we need to record this 2370 // ahead of quantization, and QuantizationParams are only created during 2371 // quantization. 2372 // - This could be in MinMax, but that would be an abuse of what MinMax is 2373 // about, and would break existing code that assumes that a MinMax is just 2374 // a min and a max. Unlike MinMax which is agnostic as to the quantized 2375 // data type, narrow_range refers to values in the quantized data type. 2376 bool narrow_range = false; 2377 2378 private: 2379 std::unique_ptr<Shape> array_shape; 2380 }; 2381 2382 // Our Model struct, represents an entire model (our "top-level" struct). 2383 // Owns everything. 2384 class Model { 2385 public: 2386 using ArrayMap = std::unordered_map<string, std::unique_ptr<Array>>; 2387 2388 bool HasArray(const string& name) const { return arrays.count(name) > 0; } 2389 Array& GetArray(const string& name) const { 2390 DCHECK(HasArray(name)) << "Array not found: " << name; 2391 return *arrays.at(name); 2392 } 2393 Array& GetOrCreateArray(const string& name) { 2394 // Make sure name is not used by an optional array 2395 DCHECK(!optional_arrays.count(name)); 2396 if (!HasArray(name)) { 2397 Array* ptr = new Array; 2398 arrays[name] = std::unique_ptr<Array>(ptr); 2399 } 2400 Array& result = GetArray(name); 2401 return result; 2402 } 2403 void CreateOptionalArray(const string& name) { 2404 DCHECK(!arrays.count(name) && !optional_arrays.count(name)); 2405 optional_arrays.insert(name); 2406 } 2407 bool IsOptionalArray(const string& name) const { 2408 return optional_arrays.count(name); 2409 } 2410 2411 // Note that this invalidates all array iterators. 2412 void EraseArray(const string& name) { arrays.erase(name); } 2413 void EraseArrays(std::function<bool(const string&)> discardable) { 2414 for (auto it = arrays.begin(); it != arrays.end();) { 2415 if (discardable(it->first)) { 2416 it = arrays.erase(it); 2417 } else { 2418 ++it; 2419 } 2420 } 2421 } 2422 const ArrayMap& GetArrayMap() const { return arrays; } 2423 ArrayMap& GetMutableArrayMap() { return arrays; } 2424 2425 int64 ArithmeticOpsCount() const { return ops_count; } 2426 2427 void AddInvalidInputArray(string invalid_input_array) { 2428 invalid_input_arrays_.insert(invalid_input_array); 2429 } 2430 2431 const std::unordered_set<string>& GetInvalidInputArrays() const { 2432 return invalid_input_arrays_; 2433 } 2434 2435 // Optional arrays are used for optional tensors, 2436 // these tensors do not have data, but with reserved names as op inputs. 2437 std::set<string> optional_arrays; 2438 2439 // The list of operators. Notice how it's a list of unique_ptr's, implying 2440 // that the Model is what owns Operator's and keeps them alive. 2441 std::vector<std::unique_ptr<Operator>> operators; 2442 2443 // Generic flags, a place where we combine information passed to us via 2444 // command-line parameters (e.g. --input_width=N) with information that 2445 // we may or may not find in the input model file. 2446 ModelFlags flags; 2447 // For code-generation only: required size of the transient_data buffer 2448 std::size_t transient_data_size = 0; 2449 // For code-generation only: required alignment of the transient_data buffer 2450 std::size_t transient_data_alignment = 0; 2451 // Arithmetic operations performed in the model. 2452 int64 ops_count = 0; 2453 2454 private: 2455 // The associative array mapping names to Array's. 2456 // Notice how it's a container of unique_ptr's, implying 2457 // that the Model is what owns Array's and keeps them alive. 2458 // The Operator's refer to these Array's by their name strings, not by their 2459 // addresses. See Operator::inputs, Operator::outputs. 2460 std::unordered_map<string, std::unique_ptr<Array>> arrays; 2461 2462 // Invalid input arrays. 2463 std::unordered_set<string> invalid_input_arrays_; 2464 }; 2465 2466 // OperatorSignature contains the information required to making versioning 2467 // decisions. 2468 struct OperatorSignature { 2469 // The operator. 2470 const Operator* op; 2471 2472 // The model in which the operator resides. 2473 const Model* model; 2474 }; 2475 } // namespace toco 2476 2477 #endif // TENSORFLOW_LITE_TOCO_MODEL_H_ 2478