...
- Have an internalized graph library that can be maintained as part as MXNet. This makes contributions easier and helps evolve the graph without potential design conflicts for a library that is now part of other project with different goals and use cases than MXNet (TVM). In this way we can evolve the library and fit it for MXNet use cases.
- Use a class passed as template argument to the graph for node and edge attributes, where the graph attributes are stored in a maintainable, type safe and compile-time checked way.
- The nodes of the graph also have templated Node data and attributes class which decouple the implementation of generic graph algorithms, traversal etc, from their usage. This will allow for more data fields and structured data and classes to be added to the graph nodes and edges in a propper way.
- The graph is type safe,extensible, reusable, and is debugging friendly.
- Developers can easily extend the graph implementing logic to dump and process the graph in a generic way without having to deal with too many implementation details.
- Generic facilities can be added such as serialization, optimization passes, operation fusion passes etc. which will add tremendous value to MXNet as a flexible framework going forward in the next steps of deep learning developments.
- AGInfo is simplified and such a vital part can be part of the Node attributes class (NodeData in the example below).
- This change would help support higher order gradient calculations, since it needs to do passes on the graph that generate additional backward nodes recursively.
- Better framework for advanced graph processing. (See other proposals affecting graph processing in this wiki node).
- Eliminate the need for an additional indexed graph data structure, saving on memory, CPU cost and complexity. Eliminate the need for equivalent data structures such as Node, NodePtr, NodeEntry and coalesce into a single one representing either a variable or an op.
- Pave the way for powerful autodiferentiation and hybridization of tensor operations using MXNet for statically compiled languages and further improve autograd in dynamic languages such as Python.
- Pave the way for powerful graph optimization and operator fusion such as what's done in Tensorflow XLA.
Examples of design changes:
...
class MXNetNodeData {
vector<NDArray*> inputs;
vector<NDArray*> outputs;
Autograd autograd;
Operator* op;
string name;
[...]
bool is_operator();
bool is_variable();
bool has_autograd();
};
class EdgeData {};
Graph<MXNetNodeData, EdgeData, size_t, uint16_t> g;
g.DFSVisit([](const MXNetNodeData& x) { ... });
Diferentiation, and optimization would be implemented as modular and flexible passes to the graph, which will be encapsulated and maintainable.
For example, a piece of code that would be greatly improved by this would be Imperative::Backward. This code as of today is breaking encapsulation of the classes that it's using. Holding indices to internal data structures in the graph as well as using node ids, pointers to arrays and variables. It's extremely complex to follow and reason. It can be broken into different passes that augment the graph and add nodes or mutate it when doing shape inference, adding backward nodes etc.
...