MXNet之NNVM代码简析
Op
class Op: 每个OP都有name和desc,还有argument和attr attr都保存在全局的OpManager里。 有同一个name的Op不一定是同一个Op,比如ElementWiseSum可能在不同的Node里会有不同数量的input,每个不同node里的Op都会生成一个Op实例,虽然这些Op名字一样,但attr不一样。不同的Op其实是用Op*或Op.index来区分,Op.index就是这个Op加入OpManager的序号 静态属性是在注册Op时确定的;动态属性(输入个数等)是在创建Node时由node的attr来确定的,创建时会调用op.parse(node.attr)。
Symbol
Symbol里只有vector
Graph
Symbol提供了许多图的接口,便于前端访问,而Graph里面没几个接口,主要就是有个indexed_graph,便于底层训练时快速访问。Symbol非常灵活,以后有可能支持动态图,但每次动态变化后都要先转成Graph,底层不太支持动态度,因为都是vector用index来索引node,不太适合中间插入一个node。
里面有vector
每个graph有以下attr: 网络结构的json string
每个NodeEntry的TShape的vector
IndexedGraph
里面一堆vector,把Graph里的node映射成vector的下标,方便快速访问。提供了一些类似Symbol里的接口,不过主要根据node的index来访问vector,而Symbol更像通过node组成的DAG来访问。
GraphExecutor
InitFullGraph
做一件事,生成带backward的Graph,调用了nnvm::pass::Gradient
copy_op,消除重复,比如c=a+b
,dc/da
和dc/db
其实一样,只用一个NodeEntry表示就行了。
AssignContext
只做一件事,设置graph的两个属性:g.attr["context"]和g.attr["device"],里面会调用nnvm::pass::PlaceDevice。 设置了in_arg、auc_arg和arg_grad的context,中间节点的context在哪里设置的?猜测可能在PlaceDevice里面
InferShape&InferType&InferStorageType
调用nnvm::pass::InferXXX推断,做的事情很类似,主要是设置g.attr["shape"]、g.attr["dtype"]、g.attr["storage_type"]属性
InitArguments
使用用户提供的参数初始化g里面的数据,g.data_entry_、g.grad_store_等 有一个版本是带shared_buffer的
FinishInitGraph
设置g.attrs["dispatch_stypes"]和g.attrs["storage"]等属性
###AttachOpExecs&AttachOpResources 遍历g,为每个节点生成一个OpExecutor,设置g.attrs["op_execs"]属性。 有三种类型Op:1)带状态的FStatefulCompute 2)backward节点,使用forward的State 3)普通节点FCompute
InitDataEntryMemory
为data_entry_里的每个数据生成一个NDArray
InitCachedOps
使用g.attrs["op_execs"],初始化op_nodes_,里面的每个元素都是一个Engine可执行的Op,包括了use_vars_vars等
InitOpSegs
把上一步的op_nodes_分段隔成多个大的exec,每个大exec里其实有许多小exec,内部按序执行小exec,大exec的use_vars和mutute_vars是小exec的汇总,这样在engine里大exec作为一个执行单元