Caffe2 - C++ API
A deep learning, cross platform ML framework
net_dag_utils.h
1 
17 #ifndef CAFFE2_CORE_NET_DAG_UTILS_H_
18 #define CAFFE2_CORE_NET_DAG_UTILS_H_
19 
20 #include <atomic>
21 #include <climits>
22 #include <cstddef>
23 #include <thread> // NOLINT
24 #include <typeinfo>
25 #include <unordered_map>
26 #include <unordered_set>
27 #include <vector>
28 
29 #include "caffe2/core/blob.h"
30 #include "caffe2/core/common.h"
31 #include "caffe2/core/logging.h"
32 #include "caffe2/core/net.h"
33 #include "caffe2/core/observer.h"
34 #include "caffe2/core/operator_schema.h"
35 #include "caffe2/core/registry.h"
36 #include "caffe2/core/tensor.h"
37 #include "caffe2/core/workspace.h"
38 #include "caffe2/proto/caffe2.pb.h"
39 #include "caffe2/utils/simple_queue.h"
40 
41 namespace caffe2 {
42 namespace dag_utils {
43 
44 struct OperatorNode {
45  unique_ptr<OperatorBase> operator_;
46  vector<int> children_;
47  vector<int> parents_;
48  std::atomic<int> runtime_parent_count_;
49  bool is_chain_start_ = false;
50 };
51 
52 struct OpGraphNode {
53  vector<int> children_;
54  vector<int> parents_;
55  int visited_inputs = 0;
56  int num_orig_parents;
57 };
58 
59 using ExecutionChains = std::unordered_map<int, std::vector<int>>;
60 
61 ExecutionChains computeChains(std::vector<OperatorNode>& orig_nodes);
62 
63 ExecutionChains singleChains(std::vector<OperatorNode>& nodes);
64 
65 std::vector<OperatorNode> prepareOperatorNodes(
66  const std::shared_ptr<const NetDef>& net_def,
67  Workspace* ws);
68 
69 std::vector<OpGraphNode> prepareChainGraphNodes(
70  const std::vector<dag_utils::OperatorNode>& operator_nodes,
71  const std::vector<std::vector<int>>& execution_chains);
72 
73 } // namespace dag_utils
74 } // namespace caffe2
75 
76 #endif // CAFFE2_CORE_NET_DAG_UTILS_H_
Workspace is a class that holds all the related objects created during runtime: (1) all blobs...
Definition: workspace.h:63
Copyright (c) 2016-present, Facebook, Inc.