【從零開始學(xué)深度學(xué)習(xí)編譯器】一，深度學(xué)習(xí)編譯器及TVM 介紹
0x0. 介紹
大家好呀，在過去的半年到一年時間里，我分享了一些算法解讀，算法優(yōu)化，模型轉(zhuǎn)換相關(guān)的一些文章。這篇文章是自己開啟學(xué)習(xí)深度學(xué)習(xí)編譯器的第一篇文章，后續(xù)也會努力更新這個系列。這篇文章是開篇，所以我不會太深入講解TVM的知識，更多的是介紹一下深度學(xué)習(xí)編譯器和TVM是什么？以及為什么我要選擇學(xué)習(xí)TVM，最后我也會給出一個讓讀者快速體驗TVM效果的一個開發(fā)環(huán)境搭建的簡要教程以及一個簡單例子。
0x1. 為什么需要深度學(xué)習(xí)編譯器？
深度學(xué)習(xí)編譯器這個詞語，我們可以先拆成兩個部分來看。
首先談?wù)勆疃葘W(xué)習(xí)領(lǐng)域。從訓(xùn)練框架角度來看，Google的TensorFlow和FaceBook的Pytorch是全球主流的深度學(xué)習(xí)框架，另外亞馬遜的MxNet，百度的Paddle，曠視的MegEngine，華為的Mindspore以及一流科技的OneFlow也逐漸在被更多人接受和使用。這么多訓(xùn)練框架，我們究竟應(yīng)該選擇哪個？如果追求易用性，可能你會選擇Pytorch，如果追求項目部署落地，可能你會選擇TensorFlow，如果追求分布式訓(xùn)練最快可能你會體驗OneFlow。所以這個選擇題沒有確定答案，在于你自己的喜好。從推理框架角度來看，無論我們選擇何種訓(xùn)練框架訓(xùn)練模型，我們最終都是要將訓(xùn)練好的模型部署到實際場景的，在模型部署的時候我們會發(fā)現(xiàn)我們要部署的設(shè)備可能是五花八門的，例如Intel CPU/Nvidia GPU/Intel GPU/Arm CPU/Arm GPU/FPGA/NPU(華為海思)/BPU(地平線)/MLU(寒武紀(jì))，如果我們要手寫一個用于推理的框架在所有可能部署的設(shè)備上都達(dá)到良好的性能并且易于使用是一件非常困難的事。
一般要部署模型到一個指定設(shè)備上，我們一般會使用硬件廠商自己推出的一些前向推理框架，例如在Intel的CPU/GPU上就使用OpenVINO，在Arm的CPU/GPU上使用NCNN/MNN等，在Nvidia GPU上使用TensorRT。雖然針對不同的硬件設(shè)備我們使用特定的推理框架進(jìn)行部署是最優(yōu)的，但這也同時存在問題，比如一個開發(fā)者訓(xùn)練了一個模型需要在多個不同類型的設(shè)備上進(jìn)行部署，那么開發(fā)者需要將訓(xùn)練的模型分別轉(zhuǎn)換到特定框架可以讀取的格式，并且還要考慮各個推理框架OP實現(xiàn)是否完全對齊的問題，然后在不同平臺部署時還容易出現(xiàn)的問題是開發(fā)者訓(xùn)練的模型在一個硬件上可以高效推理，部署到另外一個硬件上性能驟降。并且從之前幾篇探索ONNX的文章來看，不同框架間模型轉(zhuǎn)換工作也是阻礙各種訓(xùn)練框架模型快速落地的一大原因。
接下來，我們要簡單描述一下編譯器。實際上在編譯器發(fā)展的早期也和要將各種深度學(xué)習(xí)訓(xùn)練框架的模型部署到各種硬件面臨的情況一下，歷史上出現(xiàn)了非常多的編程語言，比如C/C++/Java等等，然后每一種硬件對應(yīng)了一門特定的編程語言，再通過特定的編譯器去進(jìn)行編譯產(chǎn)生機(jī)器碼，可以想象隨著硬件和語言的增多，編譯器的維護(hù)難度是多么困難。還好現(xiàn)代的編譯器已經(jīng)解決了這個問題，那么這個問題編譯器具體是怎么解決的呢？
為了解決上面的問題，科學(xué)家為編譯器抽象出了編譯器前端，編譯器中端，編譯器后端等概念，并引入IR (Intermediate Representation)的概率。解釋如下：
編譯器前端：接收C/C++/Java等不同語言，進(jìn)行代碼生成，吐出IR 編譯器中端：接收IR，進(jìn)行不同編譯器后端可以共享的優(yōu)化，如常量替換，死代碼消除，循環(huán)優(yōu)化等，吐出優(yōu)化后的IR 編譯器后端：接收優(yōu)化后的IR，進(jìn)行不同硬件的平臺相關(guān)優(yōu)化與硬件指令生成，吐出目標(biāo)文件
以LLVM編譯器為例子，借用藍(lán)色（知乎ID）大佬的圖：
受到編譯器解決方法的啟發(fā)，深度學(xué)習(xí)編譯器被提出，我們可以將各個訓(xùn)練框架訓(xùn)練出來的模型看作各種編程語言，然后將這些模型傳入深度學(xué)習(xí)編譯器之后吐出IR，由于深度學(xué)習(xí)的IR其實就是計算圖，所以可以直接叫作Graph IR。針對這些Graph IR可以做一些計算圖優(yōu)化再吐出IR分發(fā)給各種硬件使用。這樣，深度學(xué)習(xí)編譯器的過程就和傳統(tǒng)的編譯器類似，可以解決上面提到的很多繁瑣的問題。仍然引用藍(lán)色大佬的圖來表示這個思想。
0x02. TVM
基于上面深度學(xué)習(xí)編譯器的思想，陳天奇領(lǐng)銜的TVM橫空出世。TVM就是一個基于編譯優(yōu)化的深度學(xué)習(xí)推理框架（暫且說是推理吧，訓(xùn)練功能似乎也開始探索和接入了），我們來看一下TVM的架構(gòu)圖。來自于：https://tvm.apache.org/2017/10/06/nnvm-compiler-announcement
從這個圖中我們可以看到，TVM架構(gòu)的核心部分就是NNVM編譯器（注意一下最新的TVM已經(jīng)將NNVM升級為了Realy，所以后面提到的Relay也可以看作是NNVM）。NNVM編譯器支持直接接收深度學(xué)習(xí)框架的模型，如TensorFlow/Pytorch/Caffe/MxNet等，同時也支持一些模型的中間格式如ONNX、CoreML。這些模型被NNVM直接編譯成Graph IR，然后這些Graph IR被再次優(yōu)化，吐出優(yōu)化后的Graph IR，最后對于不同的后端這些Graph IR都會被編譯為特定后端可以識別的機(jī)器碼完成模型推理。比如對于CPU，NNVM就吐出LLVM可以識別的IR，再通過LLVM編譯器編譯為機(jī)器碼到CPU上執(zhí)行。
0x03. 環(huán)境配置
工欲善其事，必先利其器，再繼續(xù)探索TVM之前我們先了解一下TVM的安裝流程。這里參考著官方的安裝文檔提供兩種方法。
0x03.1 基于Docker的方式
我們可以直接拉安裝配置好TVM的docker，在docker中使用TVM，這是最快捷最方便的。例如拉取一個編譯了cuda后端支持的TVM鏡像，并啟動容器的示例如下：
docker pull tvmai/demo-gpu
nvidia-docker run --rm -it tvmai/demo-gpu bash
這樣就可以成功進(jìn)入配置好tvm的容器并且使用TVM了。
0x03.2 本地編譯以Ubuntu為例
如果有修改TVM源碼或者給TVM貢獻(xiàn)的需求，可以本地編譯TVM，以Ubuntu為例編譯和配置的流程如下：
git clone --recursive https://github.com/apache/tvm tvm
cd tvm
mkdir build
cp cmake/config.cmake build
cd build
cmake ..
make -j4
export TVM_HOME=/path/to/tvm
export PYTHONPATH=$TVM_HOME/python:${PYTHONPATH}
這樣我們就配置好了TVM，可以進(jìn)行開發(fā)和測試了。
我的建議是本地開發(fā)和調(diào)試使用后面的方式，工業(yè)部署使用Docker的方式。
0x04. 樣例展示
在展示樣例前說一下我的環(huán)境配置，pytorch1.7.0 && TVM 0.8.dev0
這里以Pytorch模型為例，展示一下TVM是如何將Pytorch模型通過Relay（可以理解為NNVM的升級版，）構(gòu)建TVM中的計算圖并進(jìn)行圖優(yōu)化，最后再通過LLVM編譯到Intel CPU上進(jìn)行執(zhí)行。最后我們還對比了一下基于TVM優(yōu)化后的Relay Graph推理速度和直接使用Pytorch模型進(jìn)行推理的速度。這里是以torchvision中的ResNet18為例子，結(jié)果如下：
Relay top-1 id: 282, class name: tiger cat
Torch top-1 id: 282, class name: tiger cat
Relay time: 1.1846002000000027 seconds
Torch time: 2.4181047000000007 seconds
可以看到在預(yù)測結(jié)果完全一致的情況下，TVM能帶來2倍左右的加速。這里簡單介紹一下代碼的流程。這個代碼可以在這里(https://github.com/BBuf/tvm_learn)找到。
0x04.1 導(dǎo)入TVM和Pytorch并加載ResNet18模型
import time
import tvm
from tvm import relay
import numpy as np
from tvm.contrib.download import download_testdata
# PyTorch imports
import torch
import torchvision
######################################################################
# Load a pretrained PyTorch model
# -------------------------------
model_name = "resnet18"
model = getattr(torchvision.models, model_name)(pretrained=True)
model = model.eval()
# We grab the TorchScripted model via tracing
input_shape = [1, 3, 224, 224]
input_data = torch.randn(input_shape)
scripted_model = torch.jit.trace(model, input_data).eval()
需要注意的是Relay在解析Pytorch模型的時候是解析TorchScript格式的模型，所以這里使用torch.jit.trace跑一遍原始的Pytorch模型并導(dǎo)出TorchScript模型。
0x04.2 載入測試圖片
加載一張測試圖片，并執(zhí)行一些后處理過程。
from PIL import Image
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))
# Preprocess the image and convert to tensor
from torchvision import transforms
my_preprocess = transforms.Compose(
[
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
]
)
img = my_preprocess(img)
# 新增Batch維度
img = np.expand_dims(img, 0)
0x04.3 Relay導(dǎo)入TorchScript模型并編譯到LLVM后端
接下來我們將PyTorch的graph導(dǎo)入到Relay成為Relay Graph，這里輸入層的名字可以任意指定。然后將Gpath使用給定的配置編譯到LLVM目標(biāo)硬件上。
######################################################################
# Import the graph to Relay
# -------------------------
# Convert PyTorch graph to Relay graph. The input name can be arbitrary.
input_name = "input0"
shape_list = [(input_name, img.shape)]
mod, params = relay.frontend.from_pytorch(scripted_model, shape_list)
######################################################################
# Relay Build
# -----------
# Compile the graph to llvm target with given input specification.
target = "llvm"
target_host = "llvm"
ctx = tvm.cpu(0)
with tvm.transform.PassContext(opt_level=3):
lib = relay.build(mod, target=target, target_host=target_host, params=params)
0x04.4 在目標(biāo)硬件上進(jìn)行推理并輸出分類結(jié)果
這里加了一個計時函數(shù)用來記錄推理的耗時情況。
######################################################################
# Execute the portable graph on TVM
# ---------------------------------
# Now we can try deploying the compiled model on target.
from tvm.contrib import graph_runtime
tvm_t0 = time.clock()
for i in range(10):
dtype = "float32"
m = graph_runtime.GraphModule(lib["default"](ctx))
# Set inputs
m.set_input(input_name, tvm.nd.array(img.astype(dtype)))
# Execute
m.run()
# Get outputs
tvm_output = m.get_output(0)
tvm_t1 = time.clock()
接下來我們在1000類的字典里面查詢一下Top1概率對應(yīng)的類別并輸出，同時也用Pytorch跑一下原始模型看看兩者的結(jié)果是否一致和推理耗時情況。
#####################################################################
# Look up synset name
# -------------------
# Look up prediction top 1 index in 1000 class synset.
synset_url = "".join(
[
"https://raw.githubusercontent.com/Cadene/",
"pretrained-models.pytorch/master/data/",
"imagenet_synsets.txt",
]
)
synset_name = "imagenet_synsets.txt"
synset_path = download_testdata(synset_url, synset_name, module="data")
with open(synset_path) as f:
synsets = f.readlines()
synsets = [x.strip() for x in synsets]
splits = [line.split(" ") for line in synsets]
key_to_classname = {spl[0]: " ".join(spl[1:]) for spl in splits}
class_url = "".join(
[
"https://raw.githubusercontent.com/Cadene/",
"pretrained-models.pytorch/master/data/",
"imagenet_classes.txt",
]
)
class_name = "imagenet_classes.txt"
class_path = download_testdata(class_url, class_name, module="data")
with open(class_path) as f:
class_id_to_key = f.readlines()
class_id_to_key = [x.strip() for x in class_id_to_key]
# Get top-1 result for TVM
top1_tvm = np.argmax(tvm_output.asnumpy()[0])
tvm_class_key = class_id_to_key[top1_tvm]
# Convert input to PyTorch variable and get PyTorch result for comparison
torch_t0 = time.clock()
for i in range(10):
with torch.no_grad():
torch_img = torch.from_numpy(img)
output = model(torch_img)
# Get top-1 result for PyTorch
top1_torch = np.argmax(output.numpy())
torch_class_key = class_id_to_key[top1_torch]
torch_t1 = time.clock()
tvm_time = tvm_t1 - tvm_t0
torch_time = torch_t1 - torch_t0
print("Relay top-1 id: {}, class name: {}".format(top1_tvm, key_to_classname[tvm_class_key]))
print("Torch top-1 id: {}, class name: {}".format(top1_torch, key_to_classname[torch_class_key]))
print('Relay time: ', tvm_time / 10.0, 'seconds')
print('Torch time: ', torch_time / 10.0, 'seconds')
0x05. 小節(jié)
這一節(jié)是對TVM的初步介紹，暫時講到這里，后面的文章會繼續(xù)深度了解和介紹深度學(xué)習(xí)編譯器相關(guān)的知識。
0x06. 參考
http://tvm.apache.org/docs/tutorials/frontend/from_pytorch.html#sphx-glr-tutorials-frontend-from-pytorch-py https://zhuanlan.zhihu.com/p/50529704
歡迎關(guān)注GiantPandaCV, 在這里你將看到獨(dú)家的深度學(xué)習(xí)分享，堅持原創(chuàng)，每天分享我們學(xué)習(xí)到的新鮮知識。( ? ?ω?? )?
有對文章相關(guān)的問題，或者想要加入交流群，歡迎添加BBuf微信：
為了方便讀者獲取資料以及我們公眾號的作者發(fā)布一些Github工程的更新，我們成立了一個QQ群，二維碼如下，感興趣可以加入。