一笆载、前言
本文將從源碼分析一個算子在TVM中是如何從前端到后端的轉(zhuǎn)換過程侵贵。首先來看一下keras模型編譯推理的一個示例：
通過pip安裝好keras和tensorflow卷拘，按照TVM官方文檔的示例進行一些修改蔼两，因為環(huán)境配置的不同匈仗，python包的導入以及編譯target進行了修改，示例使用cuda悍手，這里使用cpu帘睦。跑通的代碼如下：
import tvm
from tvm import te
import tvm.relay as relay
from tvm.contrib.download import download_testdata
from tensorflow import keras
import numpy as np
from PIL import Image
from matplotlib import pyplot as plt
from tensorflow.keras.applications.resnet50 import preprocess_input
# 下載并導入resnet50模型參數(shù)
if tuple(keras.__version__.split(".")) < ("2", "4", "0"):
weights_url = "".join(
[
"https://github.com/fchollet/deep-learning-models/releases/",
"download/v0.2/resnet50_weights_tf_dim_ordering_tf_kernels.h5",
]
)
weights_file = "resnet50_keras_old.h5"
else:
weights_url = "".join(
[
" https://storage.googleapis.com/tensorflow/keras-applications/",
"resnet/resnet50_weights_tf_dim_ordering_tf_kernels.h5",
]
)
weights_file = "resnet50_keras_new.h5"
weights_path = download_testdata(weights_url, weights_file, module="keras")
keras_resnet50 = keras.applications.resnet50.ResNet50(
include_top=True, weights=None, input_shape=(224, 224, 3), classes=1000
)
keras_resnet50.load_weights(weights_path)
# 下載一張測試圖像并做一些預處理袍患，主要是數(shù)據(jù)轉(zhuǎn)為float32和layout從NHWC轉(zhuǎn)為NCHW
img_url = "https://github.com/dmlc/mxnet.js/blob/main/data/cat.png?raw=true"
img_path = download_testdata(img_url, "cat.png", module="data")
img = Image.open(img_path).resize((224, 224))
plt.imshow(img)
plt.show()
# input preprocess
data = np.array(img)[np.newaxis, :].astype("float32")
data = preprocess_input(data).transpose([0, 3, 1, 2]) #NHWC -> NCHW
print("input_1", data.shape)
shape_dict = {"input_1": data.shape}
# keras前端導入，使用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()
# 使用編譯后的model進行推理得到結(jié)果
dtype = "float32"
tvm_out = model(tvm.nd.array(data.astype(dtype)))
tvm_out = tvm_out.numpy()[0]
top1_tvm = np.argmax(tvm_out)
synset_url = "".join(
[
"https://gist.githubusercontent.com/zhreshold/",
"4d0b62f3d01426887599d4f7ede23ee5/raw/",
"596b27d23537e5a1b5751d2b0481ef172f58b539/",
"imagenet1000_clsid_to_human.txt",
]
)
synset_name = "imagenet1000_clsid_to_human.txt"
synset_path = download_testdata(synset_url, synset_name, module="data")
with open(synset_path) as f:
synset = eval(f.read())
print("Relay top-1 id: {}, prob: {}, class name: {}".format(top1_tvm, tvm_out[top1_tvm], synset[top1_tvm]))
# confirm correctness with keras output
keras_out = keras_resnet50.predict(data.transpose([0, 2, 3, 1]))
keras_out = keras_out[0]
top1_keras = np.argmax(keras_out)
print("Keras top-1 id: {}, prob: {}, class name: {}".format(top1_keras, keras_out[top1_keras], synset[top1_keras]))
在jupyter notebook下運行的結(jié)果如下竣付，可以看到诡延，輸出的結(jié)果一致：
二、TVM代碼結(jié)構(gòu)
src/relay 的代碼主要處理神經(jīng)網(wǎng)絡的計算圖古胆，圖節(jié)點的編譯和執(zhí)行由src的其它代碼實現(xiàn)肆良。python 目錄提供了C++ API的封裝，圖節(jié)點對應的算子在src/relay/op中注冊逸绎，算子的具體實現(xiàn)在topi惹恃，通過C++或者Python實現(xiàn)。
當用戶使用relay.build(...)執(zhí)行圖編譯時桶良，TVM會對圖中的每個節(jié)點執(zhí)行下面的動作：
在算子注冊表中查找算子的實現(xiàn)
生成計算表達式和調(diào)度
將算子編譯成目標對象代碼
三座舍、 前端流程
以keras前端轉(zhuǎn)換conv2d算子為例沮翔，python部分的關鍵代碼調(diào)用關系如下圖所示：
在keras的示例中陨帆，從前端加載模型的調(diào)用為：
mod, params = relay.frontend.from_keras(keras_resnet50, shape_dict)
from_keras()函數(shù)會調(diào)用_convert_layer()，此函數(shù)會調(diào)用keras_op_to_layers()將keras前端定義的模型layers轉(zhuǎn)換成TVM的Relay表達式采蚀，主要是通過一個全局的轉(zhuǎn)換字典_convert_map={layer名稱：轉(zhuǎn)換函數(shù)}疲牵，比如其中的卷積層：{"Conv2D": _convert_convolution}，而_convert_convolution()函數(shù)會將卷積相關的參數(shù)weights榆鼠、kernel size纲爸、padding以及l(fā)ayout等傳入_op.nn.conv2d()函數(shù)，而_op.nn.conv2d()會調(diào)用_make.conv2d()來運行算子創(chuàng)建函數(shù)妆够，主要是通過tvm._ffi._init_api("relay.op.nn._make", name)函數(shù)實現(xiàn)识啦。那么此函數(shù)是如何找到對應的conv2d算子的創(chuàng)建函數(shù)呢？
tvm._ffi._init_api("relay.op.nn._make", name)函數(shù)調(diào)用_init_api_prefix("relay.op.nn._make", ...)神妹，主要是通過sys.modules["relay.op.nn._make"]查找到對應的模塊颓哮，然后通過list_global_func_names()獲取已經(jīng)注冊的全局函數(shù)列表，這里調(diào)用了C++實現(xiàn)的函數(shù)TVMFuncListGlobalNames()鸵荠，在加載FFI模塊的時候冕茅，會調(diào)用_load_lib()通過ctypes的庫函數(shù)加載tvm_runtime.so。這個表的內(nèi)容是一些已經(jīng)注冊的函數(shù)名稱蛹找，在_init_api_prefix()函數(shù)中加打印的部分輸出如下：
遍歷這個表姨伤，查找節(jié)點對應的函數(shù)名，通過_get_global_func(name,...)獲取到PackedFunc的控制句柄庸疾，調(diào)用的也是C++部分的接口_LIB.TVMFuncGetGlobal()乍楚。
四、算子relay轉(zhuǎn)換
在上面的Python代碼處理過程中届慈，調(diào)用了兩個C++的接口函數(shù)徒溪，一個是TVMFuncListGlobalNames()凌箕，它主要是返回了Registry注冊表中已經(jīng)注冊的函數(shù)名稱列表，另一個TVMFuncGetGlobal()則是根據(jù)name從注冊表中查找對應的PackedFunc對象词渤。那么算子是在哪里注冊的呢牵舱？
在TVM中，Python與C++之間是通過自己實現(xiàn)的PackedFunc來進行連接的缺虐，比如conv2d算子芜壁，主要是通過三個宏注冊相應的接口：
TVM_REGISTER_NODE_TYPE(Conv2DAttrs)：創(chuàng)建conv2d的attrs對象，在TVM中高氮，算子的計算所需要的參數(shù)都以Attribute的方式來定義慧妄。
TVM_REGISTER_GLOBAL("relay.op.nn._make.conv2d).set_body_typed()：將"relay.op.nn._make.conv2d注冊到注冊表中，并調(diào)用MakeConv()剪芍，MakeConv會返回以op和attrs為參數(shù)初始化的Call對象塞淹，Call對象的初始化過程中會創(chuàng)建CallNode對象并賦值。
RELAY_REGISTER_OR("nn.conv2d")：將con2d的OpNode注冊到relay的節(jié)點注冊表中并設置節(jié)點的屬性參數(shù)罪裹。
在TVM中一個算子就是通過上面三個宏進行注冊的饱普，但是這里只是注冊了算子并且設置了算子計算所需要的屬性參數(shù)，并不涉及算子的計算状共。從第二小節(jié)可以知道同套耕，算子計算實現(xiàn)在topi里。
五峡继、Relay到TOPI實現(xiàn)的調(diào)用過程
Relay到TOPI算子的連接是在tvm/python/tvm/relay/op/目錄下實現(xiàn)的冯袍，根據(jù)不同的算子類型歸類到幾個子目錄下，比如vision是處理視覺的算子碾牌，image處理圖像康愤，nn處理神經(jīng)網(wǎng)絡等等。一般情況下舶吗，Relay OP與TOPI OP是通過@reg.register_compute()進行連接征冷，如下：
@reg.register_compute("nn.upsampling")
def compute_upsampling(attrs, inputs, out_dtype):
...
而對于conv2d這種具有多種類型的算子，tvm會使用策略模式的方式來實現(xiàn)Relay與TOPI的連接裤翩，調(diào)用過程為：
首先通過reg.register_strategy("nn.conv2d", strategy.conv2d_strategy)注冊conv2d的處理策略资盅，在conv2d_strategy()中會根據(jù)不同的layout或者group參數(shù)選擇不同的卷積計算策略，通過wrap_compute_conv2d(topi.nn.conv2d_nchw)與topi模塊連接踊赠，在conv2d_nchw()中會調(diào)用te.compute()進行計算呵扛，此函數(shù)的第二個參數(shù)fcompute就是定義的計算規(guī)則，它會轉(zhuǎn)為body傳到C++的函數(shù)TensorComputeOp()或者ComputeOp()進行運算筐带。
六今穿、總結(jié)
本文主要從代碼層面介紹了TVM算子從keras前端到Relay，再到TOPI算子的轉(zhuǎn)換過程伦籍。
