天天看點
傳回

【從零開始學深度學習編譯器】一，深度學習編譯器及TVM 介紹

0x0. 介紹

大家好呀，在過去的半年到一年時間裡，我分享了一些算法解讀，算法優化，模型轉換相關的一些文章。這篇文章是自己開啟學習深度學習編譯器的第一篇文章，後續也會努力更新這個系列。這篇文章是開篇，是以我不會太深入講解TVM的知識，更多的是介紹一下深度學習編譯器和TVM是什麼？以及為什麼我要選擇學習TVM，最後我也會給出一個讓讀者快速體驗TVM效果的一個開發環境搭建的簡要教程以及一個簡單例子。

0x1. 為什麼需要深度學習編譯器？

深度學習編譯器這個詞語，我們可以先拆成兩個部分來看。

首先談談深度學習領域。從訓練架構角度來看，Google的TensorFlow和FaceBook的Pytorch是全球主流的深度學習架構，另外亞馬遜的MxNet，百度的Paddle，曠視的MegEngine，華為的Mindspore以及一流科技的OneFlow也逐漸在被更多人接受和使用。這麼多訓練架構，我們究竟應該選擇哪個？如果追求易用性，可能你會選擇Pytorch，如果追求項目部署落地，可能你會選擇TensorFlow，如果追求分布式訓練最快可能你會體驗OneFlow。是以這個選擇題沒有确定答案，在于你自己的喜好。從推理架構角度來看，無論我們選擇何種訓練架構訓練模型，我們最終都是要将訓練好的模型部署到實際場景的，在模型部署的時候我們會發現我們要部署的裝置可能是五花八門的，例如Intel CPU/Nvidia GPU/Intel GPU/Arm CPU/Arm GPU/FPGA/NPU(華為海思)/BPU(地平線)/MLU(寒武紀)，如果我們要手寫一個用于推理的架構在所有可能部署的裝置上都達到良好的性能并且易于使用是一件非常困難的事。

一般要部署模型到一個指定裝置上，我們一般會使用硬體廠商自己推出的一些前向推理架構，例如在Intel的CPU/GPU上就使用OpenVINO，在Arm的CPU/GPU上使用NCNN/MNN等，在Nvidia GPU上使用TensorRT。雖然針對不同的硬體裝置我們使用特定的推理架構進行部署是最優的，但這也同時存在問題，比如一個開發者訓練了一個模型需要在多個不同類型的裝置上進行部署，那麼開發者需要将訓練的模型分别轉換到特定架構可以讀取的格式，并且還要考慮各個推理架構OP實作是否完全對齊的問題，然後在不同平台部署時還容易出現的問題是開發者訓練的模型在一個硬體上可以高效推理，部署到另外一個硬體上性能驟降。并且從之前幾篇探索ONNX的文章來看，不同架構間模型轉換工作也是阻礙各種訓練架構模型快速落地的一大原因。

接下來，我們要簡單描述一下編譯器。實際上在編譯器發展的早期也和要将各種深度學習訓練架構的模型部署到各種硬體面臨的情況一下，曆史上出現了非常多的程式設計語言，比如C/C++/Java等等，然後每一種硬體對應了一門特定的程式設計語言，再通過特定的編譯器去進行編譯産生機器碼，可以想象随着硬體和語言的增多，編譯器的維護難度是多麼困難。還好現代的編譯器已經解決了這個問題，那麼這個問題編譯器具體是怎麼解決的呢？

為了解決上面的問題，科學家為編譯器抽象出了編譯器前端，編譯器中端，編譯器後端等概念，并引入IR (Intermediate Representation)的機率。解釋如下：

  • 編譯器前端：接收C/C++/Java等不同語言，進行代碼生成，吐出IR
  • 編譯器中端：接收IR，進行不同編譯器後端可以共享的優化，如常量替換，死代碼消除，循環優化等，吐出優化後的IR
  • 編譯器後端：接收優化後的IR，進行不同硬體的平台相關優化與硬體指令生成，吐出目标檔案

以LLVM編譯器為例子，借用藍色（知乎ID）大佬的圖：

【從零開始學深度學習編譯器】一，深度學習編譯器及TVM 介紹

受到編譯器解決方法的啟發，深度學習編譯器被提出，我們可以将各個訓練架構訓練出來的模型看作各種程式設計語言，然後将這些模型傳入深度學習編譯器之後吐出IR，由于深度學習的IR其實就是計算圖，是以可以直接叫作Graph IR。針對這些Graph IR可以做一些計算圖優化再吐出IR分發給各種硬體使用。這樣，深度學習編譯器的過程就和傳統的編譯器類似，可以解決上面提到的很多繁瑣的問題。仍然引用藍色大佬的圖來表示這個思想。

【從零開始學深度學習編譯器】一，深度學習編譯器及TVM 介紹

0x02. TVM

基于上面深度學習編譯器的思想，陳天奇領銜的TVM橫空出世。TVM就是一個基于編譯優化的深度學習推理架構（暫且說是推理吧，訓練功能似乎也開始探索和接入了），我們來看一下TVM的架構圖。來自于：https://tvm.apache.org/2017/10/06/nnvm-compiler-announcement

【從零開始學深度學習編譯器】一，深度學習編譯器及TVM 介紹

從這個圖中我們可以看到，TVM架構的核心部分就是NNVM編譯器（注意一下最新的TVM已經将NNVM更新為了Realy，是以後面提到的Relay也可以看作是NNVM）。NNVM編譯器支援直接接收深度學習架構的模型，如TensorFlow/Pytorch/Caffe/MxNet等，同時也支援一些模型的中間格式如ONNX、CoreML。這些模型被NNVM直接編譯成Graph IR，然後這些Graph IR被再次優化，吐出優化後的Graph IR，最後對于不同的後端這些Graph IR都會被編譯為特定後端可以識别的機器碼完成模型推理。比如對于CPU，NNVM就吐出LLVM可以識别的IR，再通過LLVM編譯器編譯為機器碼到CPU上執行。

0x03. 環境配置

工欲善其事，必先利其器，再繼續探索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

這樣就可以成功進入配置好tvm的容器并且使用TVM了。

0x03.2 本地編譯以Ubuntu為例

如果有修改TVM源碼或者給TVM貢獻的需求，可以本地編譯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，可以進行開發和測試了。

我的建議是本地開發和調試使用後面的方式，工業部署使用Docker的方式。

0x04. 樣例展示

在展示樣例前說一下我的環境配置，pytorch1.7.0 && TVM 0.8.dev0

這裡以Pytorch模型為例，展示一下TVM是如何将Pytorch模型通過Relay（可以了解為NNVM的更新版，）建構TVM中的計算圖并進行圖優化，最後再通過LLVM編譯到Intel CPU上進行執行。最後我們還對比了一下基于TVM優化後的Relay Graph推理速度和直接使用Pytorch模型進行推理的速度。這裡是以torchvision中的ResNet18為例子，結果如下：

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

可以看到在預測結果完全一緻的情況下，TVM能帶來2倍左右的加速。這裡簡單介紹一下代碼的流程。這個代碼可以在這裡(https://github.com/BBuf/tvm_learn)找到。

0x04.1 導入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模型并導出TorchScript模型。

0x04.2 載入測試圖檔

加載一張測試圖檔，并執行一些後處理過程。

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導入TorchScript模型并編譯到LLVM後端

接下來我們将PyTorch的graph導入到Relay成為Relay Graph，這裡輸入層的名字可以任意指定。然後将Gpath使用給定的配置編譯到LLVM目标硬體上。

######################################################################
# 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 在目标硬體上進行推理并輸出分類結果

這裡加了一個計時函數用來記錄推理的耗時情況。

######################################################################
# 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機率對應的類别并輸出，同時也用Pytorch跑一下原始模型看看兩者的結果是否一緻和推理耗時情況。

#####################################################################
# 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. 小節

這一節是對TVM的初步介紹，暫時講到這裡，後面的文章會繼續深度了解和介紹深度學習編譯器相關的知識。

0x06. 參考

  • http://tvm.apache.org/docs/tutorials/frontend/from_pytorch.html#sphx-glr-tutorials-frontend-from-pytorch-py
  • https://zhuanlan.zhihu.com/p/50529704

歡迎關注GiantPandaCV, 在這裡你将看到獨家的深度學習分享，堅持原創，每天分享我們學習到的新鮮知識。( • ̀ω•́ )✧

有對文章相關的問題，或者想要加入交流群，歡迎添加BBuf微信：

為了友善讀者擷取資料以及我們公衆号的作者釋出一些Github工程的更新，我們成立了一個QQ群，二維碼如下，感興趣可以加入。

編譯器 深度學習 Docker
上一篇: vue + element 項目，表格多選狀态如何禁選
下一篇: Visual Studio 2005 and UnrealScript

繼續閱讀