mxnet
imperative.h
Go to the documentation of this file.
1 /*
2  * Licensed to the Apache Software Foundation (ASF) under one
3  * or more contributor license agreements. See the NOTICE file
4  * distributed with this work for additional information
5  * regarding copyright ownership. The ASF licenses this file
6  * to you under the Apache License, Version 2.0 (the
7  * "License"); you may not use this file except in compliance
8  * with the License. You may obtain a copy of the License at
9  *
10  * http://www.apache.org/licenses/LICENSE-2.0
11  *
12  * Unless required by applicable law or agreed to in writing,
13  * software distributed under the License is distributed on an
14  * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15  * KIND, either express or implied. See the License for the
16  * specific language governing permissions and limitations
17  * under the License.
18  */
19 
20 #ifndef MXNET_IMPERATIVE_H_
21 #define MXNET_IMPERATIVE_H_
22 
23 #include <mxnet/op_attr_types.h>
24 #include <mxnet/graph_attr_types.h>
25 #include <mxnet/c_api.h>
26 #include <nnvm/symbolic.h>
27 #include <nnvm/op.h>
28 #include <nnvm/graph.h>
29 #include <vector>
30 #include <atomic>
31 #include <unordered_map>
32 
33 #include "./ndarray.h"
34 
35 namespace mxnet {
37 class Imperative {
38  public:
40  class AGInfo {
41  public:
45  std::vector<NDArray> outputs;
46  std::vector<NDArray> out_grads;
48 
49  AGInfo() :
50  grad_req(kNullOp), fresh_out_grad(false) {}
51 
52  static void Clear(const nnvm::NodePtr& node) {
53  if (node == nullptr || node->info.empty()) return;
54  AGInfo& info = Get(node);
55  if (info.grad_req != kNullOp) return;
56  node->info.clear();
57  }
58 
59  static AGInfo& Get(const nnvm::NodePtr& node) {
60  return dmlc::get<AGInfo>(node->info);
61  }
62 
63  static AGInfo& Create(const nnvm::NodePtr& node) {
64  node->info.construct<AGInfo>();
65  return Get(node);
66  }
67 
68  static bool IsNone(const NDArray& arr) {
69  return arr.entry_.node == nullptr || arr.entry_.node->info.empty();
70  }
71 
72  static bool IsVariable(const nnvm::NodePtr& node) {
73  AGInfo& info = Get(node);
74  return info.grad_req != kNullOp && info.outputs.size() == 1
75  && info.out_grads.size() == 1;
76  }
77  };
78  class CachedOp {
79  public:
80  explicit CachedOp(const nnvm::Symbol& sym);
81  uint32_t num_inputs() {
82  return fwd_graph_.indexed_graph().input_nodes().size();
83  }
84  uint32_t num_outputs() {
85  return fwd_graph_.outputs.size();
86  }
87  uint32_t num_backward_inputs() {
88  return bwd_ograd_dep_.size() + bwd_in_dep_.size() + bwd_out_dep_.size();
89  }
90  std::vector<bool>& save_inputs() {
91  return save_inputs_;
92  }
93  std::vector<bool>& save_outputs() {
94  return save_outputs_;
95  }
96  const std::unordered_set<uint32_t>& mutable_input_nodes() {
97  return fwd_graph_.indexed_graph().mutable_input_nodes();
98  }
99  nnvm::Graph GetForwardGraph(const bool recording,
100  const std::vector<NDArray*>& inputs);
101  nnvm::Graph GetBackwardGraph(const OpStatePtr& state,
102  const std::vector<OpReqType>& reqs,
103  const std::vector<NDArray*>& inputs);
104  std::vector<nnvm::NodeEntry> Gradient(const nnvm::NodePtr& node,
105  const std::vector<nnvm::NodeEntry>& ograds);
106  OpStatePtr Forward(const std::vector<NDArray*>& inputs,
107  const std::vector<NDArray*>& outputs);
108  void Backward(const bool retain_graph,
109  const OpStatePtr& state,
110  const std::vector<NDArray*>& inputs,
111  const std::vector<OpReqType>& reqs,
112  const std::vector<NDArray*>& outputs);
113 
114  private:
115  struct CachedOpState {
116  std::vector<NDArray> buff;
117  std::vector<OpStatePtr> states;
118  };
119  std::mutex mutex_;
120  nnvm::Graph fwd_graph_;
121  nnvm::Graph grad_graph_;
122  nnvm::Graph full_graph_;
123  std::vector<bool> curr_grad_req_;
124  std::vector<uint32_t> bwd_in_dep_, bwd_out_dep_, bwd_ograd_dep_;
125  std::vector<uint32_t> bwd_input_eid_;
126  std::vector<bool> save_inputs_, save_outputs_;
127  };
129  bool is_training() const {
130  return is_train_;
131  }
133  bool set_is_training(bool is_train) {
134  bool old = is_train_;
135  is_train_ = is_train;
136  return old;
137  }
139  bool is_recording() const {
140  return is_recording_;
141  }
144  bool old = is_recording_;
145  is_recording_ = is_recording;
146  return old;
147  }
149  void RecordOp(nnvm::NodeAttrs&& attrs,
150  const std::vector<NDArray*>& inputs,
151  const std::vector<NDArray*>& outputs,
152  const OpStatePtr& state = OpStatePtr(),
153  std::vector<bool>* p_save_inputs = nullptr,
154  std::vector<bool>* p_save_outputs = nullptr);
156  OpStatePtr Invoke(const Context& default_ctx,
157  const nnvm::NodeAttrs& attrs,
158  const std::vector<NDArray*>& inputs,
159  const std::vector<NDArray*>& outputs);
162  const nnvm::NodeAttrs& attrs,
163  const std::vector<NDArray*>& inputs,
164  const std::vector<NDArray*>& outputs,
165  const std::vector<OpReqType>& req,
166  const DispatchMode dispatch_mode,
169  void MarkVariables(const std::vector<NDArray*>& variables,
170  const std::vector<mx_uint>& grad_reqs,
171  const std::vector<NDArray*>& gradients);
173  std::vector<NDArray*> Backward(const std::vector<NDArray*>& outputs,
174  const std::vector<NDArray*>& ograds,
175  const std::vector<NDArray*>& variables,
176  bool is_train, bool retain_graph,
177  bool create_graph);
179  static Imperative* Get();
180 
181  private:
182  friend class NDArray;
184  Imperative() {}
186  void GetBackwardDependency(
187  const nnvm::NodePtr& node,
188  uint32_t num_inputs, uint32_t num_outputs,
189  std::vector<bool> *p_save_inputs,
190  std::vector<bool> *p_save_outputs);
191  void RunGraph(
192  const bool retain_graph,
193  const nnvm::IndexedGraph& idx,
194  const std::vector<NDArray*> arrays,
195  size_t node_start, size_t node_end,
196  std::vector<OpReqType>&& array_reqs,
197  std::vector<uint32_t>&& ref_count,
198  std::vector<OpStatePtr> *p_states,
199  const DispatchModeVector& dispatch_modes);
201 #if DMLC_CXX11_THREAD_LOCAL
202  static thread_local bool is_train_;
203  static thread_local bool is_recording_;
204 #else
205  static MX_THREAD_LOCAL bool is_train_;
206  static MX_THREAD_LOCAL bool is_recording_;
207 #endif
208 
209  std::atomic<uint64_t> node_count_{0};
211  std::atomic<uint64_t> variable_count_{0};
212 };
213 
214 using CachedOpPtr = std::shared_ptr<Imperative::CachedOp>;
215 
216 } // namespace mxnet
217 #endif // MXNET_IMPERATIVE_H_
C API of mxnet.
bool is_recording() const
whether operator recording is on.
Definition: imperative.h:139
static bool IsNone(const NDArray &arr)
Definition: imperative.h:68
uint32_t num_outputs()
Definition: imperative.h:84
static AGInfo & Create(const nnvm::NodePtr &node)
Definition: imperative.h:63
std::vector< DispatchMode > DispatchModeVector
The result holder of dispatch mode of each Node in the graph.
Definition: graph_attr_types.h:60
bool is_training() const
whether operator recording is on.
Definition: imperative.h:129
no operation, do not write anything
Definition: op_attr_types.h:46
bool set_is_training(bool is_train)
turn on or turn off operator recording for autograd.
Definition: imperative.h:133
namespace of mxnet
Definition: base.h:126
std::vector< NDArray * > Backward(const std::vector< NDArray * > &outputs, const std::vector< NDArray * > &ograds, const std::vector< NDArray * > &variables, bool is_train, bool retain_graph, bool create_graph)
compute the gradient of outputs w.r.t variables.
static void Clear(const nnvm::NodePtr &node)
Definition: imperative.h:52
Additional operator attributes beside the ones provided by NNVM.
OpStatePtr Invoke(const Context &default_ctx, const nnvm::NodeAttrs &attrs, const std::vector< NDArray * > &inputs, const std::vector< NDArray * > &outputs)
AGInfo()
Definition: imperative.h:49
DispatchMode
the dispatch mode of the operator
Definition: op_attr_types.h:106
std::vector< NDArray > outputs
Definition: imperative.h:45
Definition: imperative.h:40
uint32_t num_backward_inputs()
Definition: imperative.h:87
bool set_is_recording(bool is_recording)
turn on or turn off operator recording for autograd.
Definition: imperative.h:143
Definition: imperative.h:78
bool fresh_out_grad
Definition: imperative.h:47
OpStatePtr state
Definition: imperative.h:44
std::vector< NDArray > out_grads
Definition: imperative.h:46
Data structures that can appear in graph attributes.
void RecordOp(nnvm::NodeAttrs &&attrs, const std::vector< NDArray * > &inputs, const std::vector< NDArray * > &outputs, const OpStatePtr &state=OpStatePtr(), std::vector< bool > *p_save_inputs=nullptr, std::vector< bool > *p_save_outputs=nullptr)
to record operator, return corresponding node.
std::shared_ptr< Imperative::CachedOp > CachedOpPtr
Definition: imperative.h:214
OpReqType grad_req
Definition: imperative.h:43
Context ctx
Definition: imperative.h:42
OpReqType
operation request type to Forward and Backward
Definition: op_attr_types.h:44
std::vector< bool > & save_inputs()
Definition: imperative.h:90
OpStatePtr InvokeOp(const Context &ctx, const nnvm::NodeAttrs &attrs, const std::vector< NDArray * > &inputs, const std::vector< NDArray * > &outputs, const std::vector< OpReqType > &req, const DispatchMode dispatch_mode, OpStatePtr state=OpStatePtr())
static bool IsVariable(const nnvm::NodePtr &node)
Definition: imperative.h:72
runtime functions for NDArray
Definition: imperative.h:37
static AGInfo & Get(const nnvm::NodePtr &node)
Definition: imperative.h:59
const std::unordered_set< uint32_t > & mutable_input_nodes()
Definition: imperative.h:96
void MarkVariables(const std::vector< NDArray * > &variables, const std::vector< mx_uint > &grad_reqs, const std::vector< NDArray * > &gradients)
mark variables for computing gradients.
uint32_t num_inputs()
Definition: imperative.h:81
Context information about the execution environment.
Definition: base.h:141
ndarray interface
Definition: ndarray.h:69
std::vector< bool > & save_outputs()
Definition: imperative.h:93
Operator state. This is a pointer type, its content is mutable even if OpStatePtr is const...
Definition: op_attr_types.h:122