mxnet
pass_functions.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 
28 #ifndef NNVM_PASS_FUNCTIONS_H_
29 #define NNVM_PASS_FUNCTIONS_H_
30 
31 #include <string>
32 #include <memory>
33 #include <vector>
34 #include <utility>
35 #include "base.h"
36 #include "pass.h"
37 #include "graph_attr_types.h"
38 
39 namespace nnvm {
40 namespace pass {
41 
47 inline Graph LoadJSON(const std::string& json_str) {
48  Graph ret;
49  ret.attrs["json"] = std::make_shared<any>(json_str);
50  return ApplyPass(ret, "LoadJSON");
51 }
52 
58 inline std::string SaveJSON(Graph graph) {
59  Graph ret = ApplyPass(std::move(graph), "SaveJSON");
60  return ret.GetAttr<std::string>("json");
61 }
62 
63 
69 inline std::string PrintGraphIR(Graph graph) {
70  Graph ret = ApplyPass(std::move(graph), "PrintGraphIR");
71  return ret.GetAttr<std::string>("graphir");
72 }
73 
84 inline Graph OrderMutation(Graph src) {
85  return ApplyPass(std::move(src), "OrderMutation");
86 }
87 
97 inline Graph InferShape(Graph graph,
98  ShapeVector shape_inputs,
99  std::string shape_attr_key = "") {
100  if (shape_inputs.size() != 0) {
101  graph.attrs["shape_inputs"] = std::make_shared<any>(std::move(shape_inputs));
102  }
103  if (shape_attr_key.length() != 0) {
104  graph.attrs["shape_attr_key"] = std::make_shared<any>(std::move(shape_attr_key));
105  }
106  return ApplyPass(std::move(graph), "InferShape");
107 }
108 
118 inline Graph InferType(Graph graph,
119  DTypeVector dtype_inputs,
120  std::string dtype_attr_key = "") {
121  if (dtype_inputs.size() != 0) {
122  graph.attrs["dtype_inputs"] = std::make_shared<any>(std::move(dtype_inputs));
123  }
124  if (dtype_attr_key.length() != 0) {
125  graph.attrs["dtype_attr_key"] = std::make_shared<any>(std::move(dtype_attr_key));
126  }
127  return ApplyPass(std::move(graph), "InferType");
128 }
129 
144 inline Graph PlaceDevice(Graph graph,
145  std::string device_group_attr_key,
146  DeviceAssignMap device_assign_map,
147  std::string device_copy_op) {
148  graph.attrs["device_group_attr_key"] = std::make_shared<any>(std::move(device_group_attr_key));
149  graph.attrs["device_assign_map"] = std::make_shared<any>(std::move(device_assign_map));
150  graph.attrs["device_copy_op"] = std::make_shared<any>(std::move(device_copy_op));
151  return ApplyPass(std::move(graph), "PlaceDevice");
152 }
153 
170  Graph graph,
171  std::vector<NodeEntry> ys,
172  std::vector<NodeEntry> xs,
173  std::vector<NodeEntry> ys_out_grad,
174  std::function<NodeEntry(std::vector<NodeEntry>&& inputs)> aggregate_fun = nullptr,
175  std::function<int(const Node& node)> mirror_fun = nullptr,
176  std::function<NodeEntry(const NodeEntry& src, const NodeEntry &like)>
177  attr_hint_fun = nullptr,
178  std::vector<const Op*> zero_ops = std::vector<const Op*>(),
179  std::string copy_op_str = std::string()) {
180  graph.attrs["grad_ys"] = std::make_shared<any>(std::move(ys));
181 
182  graph.attrs["grad_xs"] = std::make_shared<any>(std::move(xs));
183  graph.attrs["grad_ys_out_grad"] = std::make_shared<any>(std::move(ys_out_grad));
184  if (aggregate_fun != nullptr) {
185  graph.attrs["grad_aggregate_fun"] = std::make_shared<any>(aggregate_fun);
186  }
187 
188  if (mirror_fun != nullptr) {
189  graph.attrs["grad_mirror_fun"] = std::make_shared<any>(mirror_fun);
190  }
191 
192  if (attr_hint_fun != nullptr) {
193  graph.attrs["attr_hint_fun"] = std::make_shared<any>(attr_hint_fun);
194  }
195 
196  if (zero_ops.size()) {
197  graph.attrs["zero_ops"] = std::make_shared<any>(std::move(zero_ops));
198  }
199 
200  if (copy_op_str != std::string()) {
201  graph.attrs["copy_op"] = std::make_shared<any>(std::move(copy_op_str));
202  }
203 
204  return ApplyPass(std::move(graph), "Gradient");
205 }
206 
207 } // namespace pass
208 } // namespace nnvm
209 #endif // NNVM_PASS_FUNCTIONS_H_
Definition: base.h:35
Graph InferType(Graph graph, DTypeVector dtype_inputs, std::string dtype_attr_key="")
Infer types in the graph given the information.
Definition: pass_functions.h:118
Graph Gradient(Graph graph, std::vector< NodeEntry > ys, std::vector< NodeEntry > xs, std::vector< NodeEntry > ys_out_grad, std::function< NodeEntry(std::vector< NodeEntry > &&inputs)> aggregate_fun=nullptr, std::function< int(const Node &node)> mirror_fun=nullptr, std::function< NodeEntry(const NodeEntry &src, const NodeEntry &like)> attr_hint_fun=nullptr, std::vector< const Op * > zero_ops=std::vector< const Op * >(), std::string copy_op_str=std::string())
Get the gradient graph whose outputs are gradients of xs wrt to ys.
Definition: pass_functions.h:169
Graph OrderMutation(Graph src)
Add control flow dependencies between nodes.
Definition: pass_functions.h:84
Graph ApplyPass(Graph src, const std::string &pass)
Apply one pass to the graph.
Definition: pass.h:62
std::unordered_map< std::string, int > DeviceAssignMap
The result holder of device of each operator in the graph.
Definition: graph_attr_types.h:111
std::unordered_map< std::string, std::shared_ptr< any > > attrs
attributes of a graph Note that attribute is shared pointer and can be shared across graphs...
Definition: graph.h:60
Data structures that can appear in graph attributes.
Node represents an operation in a computation graph.
Definition: node.h:155
Graph PlaceDevice(Graph graph, std::string device_group_attr_key, DeviceAssignMap device_assign_map, std::string device_copy_op)
Place the devices for each operator in the graph.
Definition: pass_functions.h:144
Symbolic computation graph. This is the intermediate representation for optimization pass...
Definition: graph.h:46
std::string PrintGraphIR(Graph graph)
Print graph ir.
Definition: pass_functions.h:69
std::vector< TShape > ShapeVector
The result holder of shape of each NodeEntry in the graph.
Definition: graph_attr_types.h:60
an entry that represents output data from a node
Definition: node.h:51
std::vector< int > DTypeVector
The result holder of type of each NodeEntry in the graph.
Definition: graph_attr_types.h:75
Pass that can be applied to a graph.
const T & GetAttr(const std::string &attr_name) const
Get the immutable attribute from attrs.
Definition: graph.h:247
Graph LoadJSON(const std::string &json_str)
Load a graph from JSON string, redirects to "LoadJSON" pass.
Definition: pass_functions.h:47
Graph InferShape(Graph graph, ShapeVector shape_inputs, std::string shape_attr_key="")
Infer shapes in the graph given the information.
Definition: pass_functions.h:97
Configuration of nnvm as well as basic data structure.
std::string SaveJSON(Graph graph)
Save a graph to json, redirects to "SaveJSON" pass.
Definition: pass_functions.h:58