大白话讲nnvm

大家好，欢迎来到IT知识分享网。

之前工作经验中，在某大厂，开发过机器学习框架，在和业务同学的合作下，取得还可以的成绩，但是一直觉得缺少了什么，最近在刷ai-system相关的公开课，才明白计算图的重要性，以往觉得不能理解的东西，现在突然都l理解了，工程能力可能真的要开始成为必备能力了，平时工作闲时加上周末时间，研究了nnvm这块的代码，和大家分享，有纰漏的地方，欢迎大家指出，谢谢。

NNVM是啥

还记得一到两年前突然，陈天奇突然推了个tinyflow，号称2000行的TensorFlow吗？（nnvm就是那个时候出来的，通常会带的title是“深度学习编译器”，由于本身背景原因，我其实不太理解深度学习编译器是个什么概念，因此经常看到相关的文章，尤其是pr文章时，经常云里雾里，这里呼吁下做pr的媒体同学，能在介绍知识时，不要高大上，介绍实质）。抛开其他一切的东西不谈，针对tinyflow，我认为（个人意见）主要就包括两个部分：一个是所有计算逻辑的计算（Op，如cnn、lstm、fc等等）这部分肯定不会包括在2000行里，是引入的torch lua的实现代码；另外一个就是nnvm，它的工作就是构建一个graph，这个graph就是现在深度学习相关框架都不能不提的计算图；计算图，顾名思义：就是将一个模型用有向无环图来表示，

那么，既然，模型可以用有向无环图表示，那么他的计算过程是否能够也由图的操作表示呢？当然是可以的，NNVM完成的就是这部分的意图。

图如何表示计算

计算包括前向计算和后向传播,前向计算的，把graph按post order DFS遍历一次， 就可以进行一次前向计算，很简单，比较复杂的就是后向传播，这里就要提下auto diff（这里要向弄清楚，可以看下auto diff，auto diff反向构造一个graiden的graph，去反向计算gradient， 比bp更有优势）：

如上图，会根据model的graph，去构造gradient graph，这样反向的梯度计算同样可以表示为一个graph，可以通过post order DFS，就可以计算gradient了，具体实现是通过构造一个名为gradient的pass function来实现。

计算图的抽象(graph.h/cc, node.h/cc, op.h/cc)

NNVM在构造graph的代码，就在graph.h/cc里， 抽象出node, 而node代表一个operation，如matmul、add等等，op在nnvm表示所有和graph本身计算逻辑的一个数据结构，是计算图得以完成forward、gradient计算的的基础。op定义了基本的attr，其实我觉得这里理解为接口会更合适，如op_attr_types.h中声明的各种函数，相当于整个graph中某个node根据你的op空出来很多槽，当选择某个node时，会填入对应的实现逻辑：

如conv2d的op注册，会将conv2d的计算和梯度计算、以及各种接口函数的实现注册到conv2d中，conv2d和其他的op不一样，（比如infer shape， infer type，这里记起来纪念前刚开始接触mxnet的时候为mxnet写inception-resnetv2的model时，当时惊讶于有infer shape的接口，而tf是没有的，现在看这些算是联想起来了，infershape其实也是类似于auto diff的逻辑，每一个op会注册一个infer shape的逻辑，当对graph作一个infer shape的操作时，其实是对整个graph 作一个poster dfs，然后每一个node 去infer shape）：

这里FGradient是通过新建一个conv2d_grad的op，conv2_grad的注册，注意TIsBackward为true：

稍微简单的一个op matmul的注册代码如下：

NNVM_REGISTER_OP(matmul) .describe(R"doc(Matrix multiplication of two arrays. ``dot``'s behavior depends on the input array dimensions: - 1-D arrays: inner product of vectors - 2-D arrays: matrix multiplication - N-D arrays: a sum product over the last axis of the first input and the first axis of the second input For example, given 3-D ``x`` with shape `(n,m,k)` and ``y`` with shape `(k,r,s)`, the result array will have shape `(n,m,r,s)`. It is computed by:: dot(x,y) = sum(x[i,j,:]*y[:,a,b]) )doc" NNVM_ADD_FILELINE) .set_support_level(1) .set_num_inputs(2) .set_num_outputs(1) .set_attr_parser(ParamParser 
  
    ) .set_attr 
   
     ("FGetAttrDict", ParamGetAttrDict 
    
      ) .add_arguments(MatMulParam::__FIELDS__()) .add_argument("lhs", "NDArray-or-Symbol", "The first input") .add_argument("rhs", "NDArray-or-Symbol", "The second input") .set_attr 
     
       ("FInferShape", DotShape) .set_attr 
      
        ("FInferType", ElemwiseType<2, 1>) .set_attr 
       
         ("FCorrectLayout", DotCorrectLayout) .set_attr 
        
          ( "FTVMCompute", [](const NodeAttrs& attrs, const Array 
         
           & inputs, const Array 
          
            & out_info) { const MatMulParam& param = nnvm::get 
           
             (attrs.parsed); return Array 
            
              { topi::matmul(inputs[0], inputs[1], param.transpose_a, param.transpose_b) }; }) .set_attr 
             
               ( "FGradient", [](const NodePtr& n, const std::vector 
              
                & ograds) { // z = x dot y // xshape (n,m,k), yshape (k,r,s) const MatMulParam& param = nnvm::get 
               
                 (n->attrs.parsed); bool Ta = param.transpose_a; bool Tb = param.transpose_b; // Ta = false, Tb = false // grad_x = grad_z dot y.T // grad_y = x.T dot grad_z if (!Ta && !Tb) { return std::vector 
                
                  { MakeNode("matmul", n->attrs.name + "_grad_0", {ograds[0], n->inputs[1]}, {{"transpose_a", "false"}, {"transpose_b", "true"}}), MakeNode("matmul", n->attrs.name + "_grad_1", {n->inputs[0], ograds[0]}, {{"transpose_a", "true"}, {"transpose_b", "false"}}) }; } else if (Ta && !Tb) { // Ta = true, Tb = false // grad_x = y dot grad_z.T // grad_y = x dot grad_z return std::vector 
                 
                   { MakeNode("matmul", n->attrs.name + "_grad_0", {n->inputs[1], ograds[0]}, {{"transpose_a", "false"}, {"transpose_b", "true"}}), MakeNode("matmul", n->attrs.name + "_grad_1", {n->inputs[0], ograds[0]}, {{"transpose_a", "false"}, {"transpose_b", "false"}}) }; } else if (!Ta && Tb) { // Ta = false, Tb = true // grad_x = grad_z dot y // grad_y = grad_z.T dot x return std::vector 
                  
                    { MakeNode("matmul", n->attrs.name + "_grad_0", {ograds[0], n->inputs[1]}, {{"transpose_a", "false"}, {"transpose_b", "false"}}), MakeNode("matmul", n->attrs.name + "_grad_1", {ograds[0], n->inputs[0]}, {{"transpose_a", "true"}, {"transpose_b", "false"}}) }; } else { // Ta = true, Tb = true // grad_x = y.T dot grad_z.T // grad_y = grad_z.T dot x.T return std::vector 
                   
                     { MakeNode("matmul", n->attrs.name + "_grad_0", {n->inputs[1], ograds[0]}, {{"transpose_a", "true"}, {"transpose_b", "true"}}), MakeNode("matmul", n->attrs.name + "_grad_1", {ograds[0], n->inputs[0]}, {{"transpose_a", "true"}, {"transpose_b", "true"}}) }; } }); 
                    
                   
                  
                 
                
               
              
             
            
           
          
         
        
       
      
     
    
  

由于nnvm后面是集成到tvm里面，可能部分代码并不是nnvm里面的实现，op 比如conv2d数学的计算，大家可以去tvm的代码去看，应该就是更高一层封装用来去注册到nnvm的graph中。

计算图的操作（pass.h/cc, pass_functions.h）

nnvm构造了一个所谓的pass（我也不知道怎么翻译）， 主要就是负责整个graph的操作，每一个操作会定义一个pass_function，举个例子SaveJSON， 将graph保存为json的格式的，会定义SaveJson实现：

Graph SaveJSON(Graph src) { std::shared_ptr 
  
    src_symbol = std::make_shared 
   
     (); src_symbol->outputs = src.outputs; JSONGraph jgraph; Symbol2JSONGraph(src_symbol, &jgraph); jgraph.attrs = src.attrs; std::ostringstream os; dmlc::JSONWriter writer(&os); jgraph.Save(&writer); Graph ret; ret.attrs["json"] = std::make_shared 
    
      (os.str()); return ret; } void Symbol2JSONGraph(std::shared_ptr 
     
       src, JSONGraph *jgraph) { std::unordered_map 
      
        node2index; jgraph->node_row_ptr.push_back(0); DFSVisit(src->outputs, [&node2index, jgraph](const NodePtr& n) { uint32_t nid = static_cast 
       
         (jgraph->nodes.size()); node2index[n.get()] = nid; if (n->is_variable()) { jgraph->arg_nodes.push_back(nid); } JSONNode jnode; jnode.node = n; jnode.inputs.reserve(n->inputs.size()); for (const NodeEntry& e : n->inputs) { jnode.inputs.emplace_back(node2index.at(e.node.get()), e.index, e.version); } for (const NodePtr& c : n->control_deps) { jnode.control_deps.push_back(node2index.at(c.get())); } jgraph->node_row_ptr.push_back(jgraph->node_row_ptr.back() + n->num_outputs()); jgraph->nodes.emplace_back(std::move(jnode)); }); for (const NodeEntry& e : src->outputs) { jgraph->heads.emplace_back(node2index.at(e.node.get()), e.index, e.version); } // recursively construct subgraphs for (JSONNode &jnode : jgraph->nodes) { // construct jnode's subgraphs const std::vector<std::shared_ptr 
        
          > &subgraphs = jnode.node->attrs.subgraphs; std::vector 
         
           &jsubgraphs = jnode.subgraphs; jsubgraphs.resize(subgraphs.size()); for (uint32_t i = 0; i < subgraphs.size(); ++i) { Symbol2JSONGraph(subgraphs[i], &jsubgraphs[i]); } } } 
          
         
        
       
      
     
    
  

然后注册pass函数：

NNVM_REGISTER_PASS(SaveJSON) .describe("Return a new empty Graph. Save graph to ret.attrs[\"json\"]") .set_body(SaveJSON) .set_change_graph(true) .provide_graph_attr("json");

如何去使用：

inline std::string SaveJSON(Graph graph) { Graph ret = ApplyPass(std::move(graph), "SaveJSON"); return ret.GetAttr 
  
    ("json"); } 
  

这样就可以对整个graph作一个SaveJson的操作，如何实现一个PassFunctionReg在pass.h中，感兴趣的可以仔细阅读下。

SaveJson是一个简单的pass function，gradient、InferShape、InferType也是类似的逻辑去理解，再gradient的代码，和SaveJson类似，都会去遍历graph，gradient会重构一个gradient computation graph来返回，如下：

Graph Gradient(Graph src) { using nnvm::FGradient; using MirrorFun = std::function 
  
    ; using AttrHintFun = std::function 
   
     ; CHECK_NE(src.attrs.count("grad_ys"), 0U) << "Gradient require grad_ys to be presented."; CHECK_NE(src.attrs.count("grad_ys_out_grad"), 0U) << "Gradient require grad_ys_out_grad to be presented."; CHECK_NE(src.attrs.count("grad_xs"), 0U) << "Gradient require grad_xs to be presented."; const std::vector 
    
      & ys = src.GetAttr<std::vector 
     
       >("grad_ys"); const std::vector 
      
        & ys_out_grad = src.GetAttr<std::vector 
       
         >("grad_ys_out_grad"); const std::vector 
        
          & xs = src.GetAttr<std::vector 
         
           >("grad_xs"); using AggFun = std::function<NodeEntry (std::vector 
          
            && inputs)>; AggFun agg_fun = DefaultAggregateGradient; if (src.attrs.count("grad_aggregate_fun") != 0) { agg_fun = src.GetAttr 
           
             ("grad_aggregate_fun"); } MirrorFun mirror_fun = nullptr; if (src.attrs.count("grad_mirror_fun") != 0) { mirror_fun = src.GetAttr 
            
              ("grad_mirror_fun"); } AttrHintFun attr_hint_fun = nullptr; if (src.attrs.count("attr_hint_fun") != 0) { attr_hint_fun = src.GetAttr 
             
               ("attr_hint_fun"); } std::vector 
              
                zero_ops; if (src.attrs.count("zero_ops") != 0) { zero_ops = src.GetAttr<std::vector 
               
                 >("zero_ops"); } const Op* copy_op = (src.attrs.count("copy_op") != 0) ? Op::Get(src.GetAttr 
                
                  ("copy_op")) : nullptr; // topo sort std::vector 
                 
                   topo_order; std::unordered_map<Node*, std::vector 
                  
                    > output_grads; DFSVisit(ys, [&](const NodePtr& node) { if (output_grads.count(node.get()) == 0) { output_grads[node.get()].resize(node->num_outputs()); } topo_order.push_back(node); }); CHECK_EQ(ys.size(), ys_out_grad.size()); for (size_t i = 0; i < ys.size(); ++i) { NodeEntry ograd = ys_out_grad[i]; output_grads[ys[i].node.get()][ys[i].index].grads = { ograd }; } // Check that all xs are reachable from ys for (size_t i = 0; i < xs.size(); ++i) { CHECK(output_grads.find(xs[i].node.get()) != output_grads.end()) << "Cannot differentiate with respect to the " << i+1 << "-th variable " << "because it is unreachable from the outputs."; } // construct mirror as memory reduction strategy if needed std::unordered_map 
                   
                     mirror_map; if (mirror_fun != nullptr) { for (const NodePtr& node_ptr : topo_order) { if (mirror_fun(*node_ptr)) { NodePtr new_node = Node::Create(); *new_node = *node_ptr; new_node->attrs.name += "_mirror"; for (auto& e : new_node->inputs) { e.node = mirror_map.at(e.node.get()); } for (auto& n : new_node->control_deps) { n = mirror_map.at(n.get()); } mirror_map[node_ptr.get()] = std::move(new_node); } else { mirror_map[node_ptr.get()] = node_ptr; } } } // traverse backward static auto& grad_fun_map = Op::GetAttr 
                    
                      ("FGradient"); static auto& finfer_shape = Op::GetAttr 
                     
                       ("FInferShape"); std::vector 
                      
                        out_agg_grads; for (auto rit = topo_order.rbegin(); rit != topo_order.rend(); ++rit) { const NodePtr& ptr = *rit; if (ptr->is_variable()) continue; out_agg_grads.clear(); auto& out_grad_vec = output_grads.at(ptr.get()); for (uint32_t i = 0; i < out_grad_vec.size i gradentry e='out_grad_vec[i];' e.sum='agg_fun(std::move(e.grads));' if e.need_attr_hint attr_hint_fun e.sum='attr_hint_fun(e.sum,' nodeentryptr 0 i out_agg_grads.push_backe.sum if rit->inputs.size() != 0) { NodePtr fwd_node = (mirror_map.size() == 0 ? ptr : mirror_map.at(ptr.get())); std::vector 
                       
                         input_grads; // Check for FGradient if (grad_fun_map.contains(ptr->op())) { input_grads = grad_fun_map[ptr->op()](fwd_node, out_agg_grads); CHECK_EQ((*rit)->inputs.size(), input_grads.size()) << "Gradient function not returning enough gradient"; } else if (CheckGradAllZero(out_agg_grads, zero_ops)) { for (size_t i = 0; i < fwd_node->num_inputs(); ++i) { std::ostringstream os; if (1 == fwd_node->num_inputs()) { os << fwd_node->attrs.name << "_backward"; } else { os << fwd_node->attrs.name << "_in" << i << _backward auto p='Node::Create();' p->attrs.op = zero_ops[0]; p->attrs.name = os.str(); p->inputs.push_back(fwd_node->inputs[i]); p->control_deps.emplace_back(fwd_node); if (p->op()->attr_parser != nullptr) { p->op()->attr_parser(&(p->attrs)); } input_grads.emplace_back(p, 0, 0); } } else { LOG(FATAL) << "Operator " << fwd_node->op()->name << " is non-differentiable " << because it didnt register fgradient attribute. for const auto nodeentry : input_grads checknodeentry.node auto git='input_grads.begin();' checkrit->inputs.size() <= input_grads.size for auto it='(*rit)-'>inputs.begin(); it != (*rit)->inputs.end(); ++it, ++git) { auto& output_grad_entry = output_grads[it->node.get()][it->index]; // if any of the backward op can do shape inference, the hint is not necessary. if (finfer_shape.contains(git->node->op())) { output_grad_entry.need_attr_hint = false; } output_grad_entry.grads.emplace_back(std::move(*git)); } } } // take out the xs' grads Graph ret; ret.outputs.resize(xs.size()); NodeEntryMap<std::pair 
                        
                          > unique_grads; size_t counter = 0; for (const NodeEntry& e : xs) { GradEntry& entry = output_grads[e.node.get()][e.index]; // aggregate sum if there haven't been if (entry.sum.node.get() == nullptr) { entry.sum = agg_fun(std::move(entry.grads)); if (entry.need_attr_hint && attr_hint_fun != nullptr) { entry.sum = attr_hint_fun(entry.sum, e); } } if (copy_op != nullptr) { auto kv = unique_grads.find(entry.sum); if (kv == unique_grads.end()) { unique_grads.emplace(std::move(entry.sum), std::make_pair(1, counter)); } else { NodePtr copy_node = Node::Create(); std::ostringstream os; os << entry.sum.node->attrs.name << "_" << kv->second.first << _copy kv->second.first++; copy_node->attrs.op = copy_op; copy_node->attrs.name = os.str(); copy_node->inputs.emplace_back(entry.sum); if (copy_node->attrs.op->attr_parser != nullptr) { copy_node->attrs.op->attr_parser(&(copy_node->attrs)); } unique_grads.emplace(NodeEntry{std::move(copy_node), 0, 0}, std::make_pair(1, counter)); } } else { ret.outputs[counter] = entry.sum; } ++counter; } if (copy_op != nullptr) { for (const auto& kv : unique_grads) { ret.outputs[kv.second.second] = kv.first; } } return ret; } // register pass NNVM_REGISTER_PASS(Gradient) .describe("Return a gradient graph of src.attrs[\"ys\"] wrt src.attrs[\"xs\"]") .set_body(Gradient) .set_change_graph(true) .depend_graph_attr("grad_ys") .depend_graph_attr("grad_xs") .depend_graph_attr("grad_ys_out_grad"); 
                         
                        
                       
                      
                     
                    
                   
                  
                 
                
               
              
             
            
           
          
         
        
       
      
     
    
  

后话

nnvm现在集合在TVM，要更加深入的理解，现在看来必须得去硬抠tvm的代码了，现在对nnvm算是理解了，但理解和做出来差了十万八千里，后续想自己写个类似nnvm的东西，而转念一想，chentianqi这些牛人开源出这些东西，就是想减少我们的工作成本，直接参与进去可能是另一种学习一种“做出来”的方法，anyway，计算图这块还是得深入下去，接下来会自己写一些测试的程序去更深入学习nnvm，大家如果有对ai system 感兴趣的小伙伴，可以私信我，我们看看能不能一起学习学习，这里面确实涉及太多东西，一个人，尤其我这种算法出身的小伙伴，真的比较难

也欢迎大家关注我的同名微信公众号 小石头的码疯窝（xiaoshitou_ml_tech），或者通过公众号加我的个人微信进行讨论