作者:安平博,Xilinx高级工程师;来源:AI加速微信公众号
用了几章的篇幅写了一些粗读TVM代码的收获,虽然读了一点皮毛,但是还是掌握了TVM的基本架构和代码组成,算是给以后的精读打下了一点基础吧。从这章开始再从头捋一遍TVM代码,顺序是frontend-build-optimize-lower-target。
TVM前端可以适配大多数流行的深度学习框架,比如tensorflow,pytorch,onnx等。这里我选择适配tensorflow的前端代码来看。
TVM前端的输入是protoBuf格式的二进制文件,这个文件描述了由tensorflow构建的神经网络的结构。我将二进制文件转换为可视的text格式,其结构如下所示:
node {
name: "Model/Placeholder"
op: "Placeholder"
attr {
key: "_output_shapes"
value {
list {
shape {
dim {
size: 1
}
dim {
size: 1
}
}
}
}
}
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "shape"
value {
shape {
dim {
size: 1
}
dim {
size: 1
}
}
}
}
}
一个node中通常包含name,op,input,attr四种属性。Name是node的唯一标识符,如果不显示的指定名字,tensorflow会使用op+number的形式来定义name。op可以分为三类,一类是控制类op,比如switch,exit,merge,loopcond等,这些和分支跳转以及循环控制等有关,一类是计算类op,这个是最多的op,包括我们经常遇到的conv,matmul,add等。还有一类是和数据定义有关,比如placeholder,tensorarray,const等。
TVM前端就是将这些算子转换成其自定义的中间算子表示。对于计算类算子,TVM实现了一一映射,但是对于控制类算子,TVM会将其转换到其只支持的if和loop算子,merge,exit等算子会被融合进去。数据类算子会被转换为TVM的抽象表示var。
下面一行代码是实现从tensorflow算子到relay算子的转换函数:
mod, params = relay.frontend.from_tensorflow(graph_def, layout=layout, shape=shape_dict)
graph_def是从tensorflow图结构二进制文件中得到的,layout是硬件对参数排序要求,主要和conv计算相关,conv的输入和输出有4个维度信息:batch(N), height(H), width(W), channel(C)。默认是NHWC的pattern。Shape是可以由用户定义的输入数据shape信息,如果shape为空,那么TVM会从placeholder类节点中推断出输入shape。
前端代码主要在/python/tvm/relay/frontend/tensorflow.py中的_get_relay_func函数中。_get_relay_func函数会遍历graph的每个节点,并将节点转换为relay的算子表示。新的relay图会被看作一个函数并返回,这个函数就是后序build,optimize,lower的主体。
第一步检查tensorflow图中每个节点是否在TVM表示范畴之内,如果有TVM不支持的算子会报错。
missing_operators = self._parse_import_prerequisites(graph)
…
if missing_operators:
…
TVM可以支持的算子在_control_flow_nodes, _tensor_array_write_ops, _identity_list, _convert_map这个几个字典中,_control_flow_nodes包括了merge, switch, nextIteration, exit, enter, loopCond几个控制节点。_convert_map中涵盖了一些神经网络常用的计算算子,比如conv,relu等,以及一些常用的数学函数cos, sin等,还有一些有reduce性质的操作如argmax, sum等,数据线性处理的函数如split,tile等。
接下来开始遍历tensorflow的每个节点,并重新注册。
for node in graph.node:
node_name_prefix = node.name.rsplit('/', 1)[0]
self._control_flow_node_map[node_name_prefix].add(node.op)
self._tf_node_map[node.name] = node
在这里假设了name被/分割的前项是属于控制节点,认为op是和控制节点相关,所以将之加入了control_flow_node_map中(个人猜测),之后在对其进行筛查,将筛查正确的控制节点加入sorted_cf_nodes中。
if node.op == 'Placeholder' or node.op == 'PlaceholderWithDefault':
self._nodes[node.name] = [_expr.var(node.name, shape=self._input_shapes[node.name], dtype=attr['dtype'].name)]
如果节点是placeholder,首先推断出shape,然后将其转换为var类。var调用了python/tvm/relay/expr.py中的var,首先建立了tensortype,这是TVM中的数据类型,它也是一个类,其中包含了shape信息。然后调用Var类,通过__init_handle_by_constructor__实现python到C++中类的转换。可以看一下这个函数:
def __init_handle_by_constructor__(fconstructor, args):
"""Initialize handle by constructor"""
temp_args = []
values, tcodes, num_args = _make_tvm_args(args, temp_args)
ret_val = TVMValue()
ret_tcode = ctypes.c_int()
if _LIB.TVMFuncCall(
fconstructor.handle, values, tcodes, ctypes.c_int(num_args),
ctypes.byref(ret_val), ctypes.byref(ret_tcode)) != 0:
raise get_last_ffi_error()
_ = temp_args
_ = args
assert ret_tcode.value == ArgTypeCode.OBJECT_HANDLE
handle = ret_val.v_handle
return handle
_make_tvm_args将python中的变量通过ctypes转换为c格式变量,然后调用C++中的TVMFuncCall,实现了C++中对应类的建立。
elif node.op == 'Const':
tensor_value = node.attr['value'].tensor
…
self._parse_param(key, value, node.name, self._in_shape)
…
如果op为const,表明这是神经网络中的参数,数值在protobuf属性tensor中,parse_param完成了tensor到numpy array的转换,同时创建了对应的var对象。
对于control类节点,则将其注册到control_flow_nodes中,这里还假设了Exit节点一定是处于loop节点之后的。
elif node.op in _control_flow_nodes:
# We assume that the direct parent node of Exit is a while loop block
if node.op == "Exit":
self._while_loop_name_set.add(node_name_prefix)
control_flow_nodes.append(node)
在tensorflow中,exit节点通常都用在循环性质节点之后,表示退出循环。Enter节点用于循环节点开始,表示开始进入循环。比如一个介绍控制节点的例子:
对于tensorarray的节点,因为这是一个动态可变的tensor,TVM会推断出其静态的shape,动态维度会声明为Any类型,表示这个维度的大小未知。
接下来还需要对检测到的所有control节点进行一定的处理:将所有exit节点移到循环模块的最后,这里具体的处理算法还不是很清楚,主要还是对tensorflow的图结构理解不够深入,比如循环节点和exit节点的关系等。
经过以上的步骤,graph节点都已经被遍历过,并依据其不同类型进行了记录以及shape推断,同时一些自由和约束参数都进行了转换(var对象的建立)。之后就是tensorflow的节点到relay的转换了,主函数是_backtrack_construct,通过递归调用这个函数完成图的再一次搜索,在遍历过程中完成了节点到relay算子的转换。
转换过程很复杂,没有看的太懂,其基本流程是:首先将control节点完成到loop和if的转换(其实主要就是这两种控制,exit,enter等可能已经隐含于这两个控制节点之中了),然后处理tensorarray相关的变量,这块处理也比较复杂,没有看懂,基本上是推断其shape,然后再进行convert。对于ndarray变量(由const节点转化来)则建立var对象。
最后一步是获得转换后relay图的终结点,如果用户提供了outputs,那么就使用outputs作为out,如果没有提供,就使用exit节点作为out。这个时候out就可以用于遍历整个relay树结构的起点,通过遍历来寻找出自由变量,然后创建一个function。Function就是一个具有输入(自由变量)和输出的表示。这样就完成了tensorflow到relay function的一个转变。