• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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_GRAPH_INFO_H_
16 #define TENSORFLOW_LITE_GRAPH_INFO_H_
17 
18 #include <stddef.h>
19 
20 #include <vector>
21 
22 #include "tensorflow/lite/c/common.h"
23 
24 namespace tflite {
25 
26 // Basic information about an inference graph, where execution nodes
27 // are connected via tensors.
28 class GraphInfo {
29  public:
~GraphInfo()30   virtual ~GraphInfo() {}
31 
32   // Total number of tensors in the graph.
33   virtual size_t num_tensors() const = 0;
34 
35   // Returns a tensor given its index which is expected to be between 0 and
36   // num_tensors().
37   virtual TfLiteTensor* tensor(size_t index) = 0;
38 
39   // Number of nodes in the current execution plan.
40   virtual size_t num_execution_nodes() const = 0;
41 
42   // Total number of known nodes, which may include nodes that are no longer in
43   // the execution plan. This happens in case of applying multiple delegates.
44   // Should be >= num_execution_nodes()
45   virtual size_t num_total_nodes() const = 0;
46 
47   // Returns a node given its index in the execution plan, which is expected to
48   // be between 0 and num_execution_nodes().
49   virtual const TfLiteNode& node(size_t index) const = 0;
50 
51   // Returns an implementation-specific node index which may be different from
52   // execution-plan index.
53   // Expected to be between 0 and num_total_nodes().
54   virtual size_t node_index(size_t index) const = 0;
55 
56   // Returns the indices of the input tensors.
57   virtual const std::vector<int>& inputs() const = 0;
58 
59   // Returns the indices of the output tensors.
60   virtual const std::vector<int>& outputs() const = 0;
61 
62   // Returns the indices of the variable tensors.
63   virtual const std::vector<int>& variables() const = 0;
64 };
65 
66 // Represents a subset of nodes in a TensorFlow Lite graph.
67 struct NodeSubset {
68   enum Type {
69     kTfUnexplored = 0,  // temporarily used during creation
70     kTfPartition,
71     kTfNonPartition
72   };
73   Type type = kTfUnexplored;
74   // Nodes within the node sub set
75   std::vector<int> nodes;
76   // Tensors that stride output from another node sub set that this depends on,
77   // or global inputs to the TensorFlow Lite full graph.
78   std::vector<int> input_tensors;
79   // Outputs that are consumed by other node sub sets or are global output
80   // tensors. All output tensors of the nodes in the node sub set that do not
81   // appear in this list are intermediate results that can be potentially
82   // elided.
83   std::vector<int> output_tensors;
84 };
85 
86 // Partitions a list of node indices `nodes_to_partition` into node sub sets.
87 // Each node sub set is in dependency order (i.e. all members of the node sub
88 // sets). `node_subsets` is assumed to be empty.
89 TfLiteStatus PartitionGraphIntoIndependentNodeSubsets(
90     const GraphInfo* info, const TfLiteIntArray* nodes_to_partition,
91     std::vector<NodeSubset>* node_subsets);
92 
93 }  // namespace tflite
94 
95 #endif  // TENSORFLOW_LITE_GRAPH_INFO_H_
96