一红选、前言
針對(duì)神經(jīng)網(wǎng)絡(luò)模型的編譯，TVM封裝了非常簡潔的python接口，如下：
# keras前端導(dǎo)入弦讽，使用llvm作為target編譯
mod, params = relay.frontend.from_keras(keras_resnet50, shape_dict)
# compile the model
target = "llvm"
dev = tvm.cpu(0)
with tvm.transform.PassContext(opt_level=0):
model = relay.build_module.create_executor("graph", mod, dev, target, params).evaluate()
print(model)
tvm_out = model(tvm.nd.array(data.astype(dtype)))
在上一篇文章中介紹了模型的算子轉(zhuǎn)換與Relay IR Module的流程，當(dāng)TVM將Relay IR Module模型編譯為runtime module時(shí)膀哲，可以通過下面的函數(shù)完成：
model = relay.build_module.create_executor("graph", mod, dev, target, params).evaluate()
它返回的是一個(gè)函數(shù)入口往产，從打印的輸出可以看到它所指向的函數(shù)：
<function graphexecutor._make_executor.._graph_wrapper at 0x7fb634b3ed30>
下面這一句相當(dāng)于將輸入傳入這個(gè)函數(shù)去運(yùn)行：
tvm_out = model(tvm.nd.array(data.astype(dtype)))
這兩個(gè)步驟在TVM中的運(yùn)行過程是怎么樣的呢被碗？這篇文章將圍繞這個(gè)問題進(jìn)行相關(guān)的介紹。
二仿村、模型編譯
首先來看一下TVM是如何調(diào)用到編譯函數(shù)的：
create_executor(...)函數(shù)會(huì)根據(jù)executor的類型返回相應(yīng)的執(zhí)行器對(duì)象蛮放，這里使用的是"graph"，所以返回的是GraphExecutor(...)對(duì)象奠宜，類class GraphExecutor()是_interpreter.Executor的子類包颁。
類class Executor(object)是一個(gè)接口類，它的成員函數(shù)_make_executor()是一個(gè)接口函數(shù)压真，繼承它的子類需要實(shí)現(xiàn)娩嚼，它的另一個(gè)成員函數(shù)evaluate()會(huì)調(diào)用_make_executor()，也即調(diào)用子類class GraphExecutor實(shí)現(xiàn)的_make_executor()成員函數(shù)滴肿。
-
_make_executor()主要的工作是調(diào)用build(...)函數(shù)對(duì)Relay IR Module進(jìn)行編譯岳悟，并且提供graph執(zhí)行器的運(yùn)行函數(shù)，也就是前言小節(jié)中提到的：
<function graphexecutor._make_executor.._graph_wrapper at 0x7fb634b3ed30>
下面介紹一下build(...)函數(shù)的過程：
首先實(shí)例化一個(gè)build_module對(duì)象bld_mod泼差，對(duì)象的實(shí)例化流程是通過tvm._ffi._init_api("relay.build_module", name)調(diào)用C++端的接口RelayBuildCreate()贵少，它會(huì)創(chuàng)建RelayBuildModule對(duì)象。
然后通過bld_mod["build"]查找到編譯函數(shù)的PackedFunc堆缘，這個(gè)主要是通過類class RelayBuildModule中的GetFunction(...)實(shí)現(xiàn)滔灶。它會(huì)調(diào)用類class RelayBuildModule的成員函數(shù)Build(...)，在對(duì)一些成員變量進(jìn)行賦值后吼肥，調(diào)用最終的BuilRelay(...)將Relay IR Module編譯為runtime module录平。
最后會(huì)返回一個(gè)GraphExecutorFactoryModule(...)對(duì)象，這個(gè)對(duì)象在初始化的時(shí)候會(huì)調(diào)用C++端的接口創(chuàng)建類class GraphExecutorFactory的對(duì)象缀皱。
其中BuildRelay的調(diào)用過程如下圖所示斗这，它主要有兩個(gè)步驟，一個(gè)是對(duì)relay_module做了OptimizeImpl(...)的優(yōu)化啤斗，對(duì)模型進(jìn)行算子的融合表箭、fold constants以及其它的一些優(yōu)化；第二個(gè)是創(chuàng)建codegen對(duì)象并生成lowered_funcs并調(diào)用tvm::build(...)進(jìn)行編譯钮莲。
整個(gè)模型編譯的過程可以總結(jié)為：將Relay IR Module編譯為runtime module并為其構(gòu)造好執(zhí)行器的對(duì)象免钻。
三、模型運(yùn)行
模型運(yùn)行時(shí)會(huì)調(diào)用_make_executor(...)里定義的_graph_wrapper(...)函數(shù)：
def _graph_wrapper(*args, **kwargs):
args = self._convert_args(self.mod["main"], args, kwargs)
# Create map of inputs.
for i, arg in enumerate(args):
gmodule.set_input(i, arg)
# Run the module, and fetch the output.
gmodule.run()
flattened = []
for i in range(gmodule.get_num_outputs()):
flattened.append(gmodule.get_output(i).copyto(_nd.cpu(0)))
unflattened = _unflatten(iter(flattened), ret_type)
return unflattened
該函數(shù)首先會(huì)遍歷輸入的數(shù)據(jù)臂痕，并調(diào)用set_input(...)為module設(shè)置輸入?yún)?shù)伯襟，然后調(diào)用run()進(jìn)行模型推理，最后獲取輸出的結(jié)果握童。
在類class GraphExecutorFactory的成員函數(shù)GetFunction(...)中股淡，如果輸入的name是模型名稱痢畜，則通過ExecutorCreate(...)創(chuàng)建GraphExecutor對(duì)象匪煌。
在GraphExecutor對(duì)象中的GetFunction(...)會(huì)根據(jù)名稱"set_input"與"run"返回相應(yīng)的PackedFunc對(duì)象。
四俺附、總結(jié)
本文主要介紹了TVM模型編譯與運(yùn)行過程中的代碼流程。
