TVM：通過Python接口（AutoTVM）來編譯和優化模型
上次我們已經介紹了如何從源碼編譯安裝 tvm，本文我們將介紹在本機中使用 tvm Python 接口來編譯優化模型的一個demo。
TVM 是一個深度學習編譯器框架，有許多不同的模塊可用于處理深度學習模型和運算符。 在本教程中，我們將學習如何使用 Python API 加載、編譯和優化模型。
在本文中，我們將使用 Python 接口的 tvm 完成以下任務：
- 為 tvm runtime 編譯一個預訓練好的 ResNet50-v2 模型
- 在編譯好的模型上運行一張真實的圖像，并得到正確的結果
- 使用 tvm 在 CPU 上 tune 模型
- 使用 tvm 收集的數據重新編譯并優化模型
- 再次運行一張真實的圖像，對比優化前后模型的輸出和性能
導入必要的包
- onnx：用于模型的加載和轉換
- PIL：用于處理圖像數據的 Python 圖像庫
- numpy：用于圖像數據預處理和后處理的
- 用于下載測試數據的輔助程序
- TVM relay 框架和 TVM Graph Executor
import onnx
from tvm.contrib.download import download_testdata
from PIL import Image
import numpy as np
import tvm.relay as relay
import tvm
from tvm.contrib import graph_executor
下載并加載onnx模型
在本文中，我們將使用 ResNet-50 v2。
TVM 提供了一個幫助庫來下載預先訓練的模型。 通過模塊提供模型 URL、文件名和模型類型，TVM 將下載模型并將其保存到磁盤。 對于 ONNX 模型的實例，我們可以使用 ONNX runtime 將其加載到內存中。另外提一下一個很方便的查看 onnx 模型的工具：netron。
model_url = "".join(["https://github.com/onnx/models/raw/","master/vision/classification/resnet/model/","resnet50-v2-7.onnx",]
)model_path = download_testdata(model_url, "resnet50-v2-7.onnx", module="onnx")
onnx_model = onnx.load(model_path)
下載、預處理并加載測試圖像
我們從網絡上下載一只小貓的圖像作為測試圖像。
img_url = "https://s3.amazonaws.com/model-server/inputs/kitten.jpg"
img_path = download_testdata(img_url, "imagenet_cat.png", module="data")# 將圖像尺寸調整為 (224, 224)
resized_image = Image.open(img_path).resize((224, 224))
img_data = np.asarray(resized_image).astype("float32")# 此時我們圖像的數據排布是 HWC，但是 onnx 需要的是 CHW，所以要轉換以下
img_data = np.transpose(img_data, (2, 0, 1))# 根據 ImageNet 數據集的標準進行歸一化
imagenet_mean = np.array([0.485, 0.456, 0.406]).reshape((3, 1, 1))
imagenet_stddev = np.array([0.229, 0.224, 0.225]).reshape((3, 1, 1))
norm_img_data = (img_data / 255 - imagenet_mean) / imagenet_stddev# 增加通道維，此時我們輸入的數據排布為 NCHW
img_data = np.expand_dims(norm_img_data, axis=0)
注意，以上的模型和測試圖像完全可以替換成自己的，只要按要求轉換為指定的格式即可
通過relay編譯模型
target = 'llvm'
注意：請定義正確的定義正確的target
指定正確的 target 會對編譯模塊的性能產生巨大影響，因為它可以利用 target 上可用的硬件功能。 有關更多信息，請參閱自動調整 x86 CPU 的卷積網絡。 我們建議確定您正在運行的 CPU 以及可選功能，并適當設置 target 。 例如，對于某些處理器 target = “llvm -mcpu=skylake”，或 target = “llvm -mcpu=skylake-avx512” 用于具有 AVX-512 矢量指令集的處理器。
# 注意這里 'input_name' 可能會根據模型不同而不同，大家可以使用上面提到 netron 工具來查看輸入名稱
input_name = "data"
shape_dict = {input_name: img_data.shape}mod, params = relay.frontend.from_onnx(onnx_model, shape_dict)with tvm.transform.PassContext(opt_level=3):lib = relay.build(mod, target=target, params=params)dev = tvm.device(str(target), 0)
module = graph_executor.GraphModule(lib["default"](dev))
在 TVM Runtime 上執行
在模型編譯完成之后，我們可以使用 TVM Runtime 來用模型來輸出預測結果。要運行 TVM Runtime 來完成預測，我們需要：
- 編譯好的模型，就是我們剛剛做的
- 有效地模型輸入
dtype = "float32"
module.set_input(input_name, img_data)
module.run()
output_shape = (1, 1000)
tvm_output = module.get_output(0, tvm.nd.empty(output_shape)).numpy()
收集基本性能數據
我們這里要收集這個未優化模型相關的一些基本性能數據，然后將其與 tune 后的模型進行比較。 為了消除 CPU 噪聲的影響，我們以多次重復的方式在多個批次中運行計算，然后收集一些關于平均值、中值和標準偏差的基礎統計數據。
import timeittiming_number = 10
timing_repeat = 10
unoptimized = (np.array(timeit.Timer(lambda: module.run()).repeat(repeat=timing_repeat, number=timing_number))* 1000/ timing_number
)
unoptimized = {"mean": np.mean(unoptimized),"median": np.median(unoptimized),"std": np.std(unoptimized),
}print(unoptimized)
此處輸出：
{'mean': 229.1864895541221, 'median': 228.7280524149537, 'std': 1.0664440211813757}
對結果進行后處理
如前所述，不同的模型輸出張量的方式可能不同。
在我們的例子中，我們需要進行一些后處理，使用為模型提供的查找表將 ResNet-50-V2 的輸出呈現為更易讀的形式。
from scipy.special import softmax# 下載標簽列表
labels_url = "https://s3.amazonaws.com/onnx-model-zoo/synset.txt"
labels_path = download_testdata(labels_url, "synset.txt", module="data")with open(labels_path, "r") as f:labels = [l.rstrip() for l in f]# 打開并讀取輸出張量
scores = softmax(tvm_output)
scores = np.squeeze(scores)
ranks = np.argsort(scores)[::-1]
for rank in ranks[0:5]:print("class='%s' with probability=%f" % (labels[rank], scores[rank]))
此處輸出：
class='n02123045 tabby, tabby cat' with probability=0.610551
class='n02123159 tiger cat' with probability=0.367180
class='n02124075 Egyptian cat' with probability=0.019365
class='n02129604 tiger, Panthera tigris' with probability=0.001273
class='n04040759 radiator' with probability=0.000261
調整 （tune）模型
之前編譯的模型工作在 TVM Runtime 上，但是并未提供任何針對特定硬件平臺的優化。這里我們來演示如何構建一個針對特定硬件平臺的優化模型。
在某些情況下，使用我們自己編譯的模塊運行推理時，性能可能無法達到預期。 在這種情況下，我們可以利用自動調諧器（Auto-tuner）為模型找到更好的配置并提高性能。 TVM 中的調優是指優化模型以在給定目標上運行得更快的過程。 這與訓練（training）和微調（fine-tuning）的不同之處在于它不會影響模型的準確性，而只會影響運行時性能。 作為調優過程的一部分，TVM 將嘗試運行許多不同的算子實現的可能，以查看哪個性能最佳。 并將這些運行的結果存儲在調整記錄文件中。
在最簡單的形式下，tuning 需要我們指定三項：
- 我們想要運行該模型的目標設備的規格
- 存儲調整記錄輸出文件的路徑
- 要調整的模型的路徑
首先我們導入一些需要的庫：
import tvm.auto_scheduler as auto_scheduler
from tvm.autotvm.tuner import XGBTuner
from tvm import autotvm
為運行器（runner）設置一些基本的參數，運行其會根據這組特定的參數來生成編譯代碼并測試其性能。
number指定我們將要測試的不同配置的數目repeat指定我們對每種配置測試多少次min_repeat_ms執行運行每次配置測試的多長時間，如果重復次數低于此值，則會增加。該選項對于 GPU tuning 時必須的，對于 CPU tuning 則不需要。將其設為 0 即禁用它。timeout指定了每次配置測試的運行時間上限。
number = 10
repeat = 1
min_repeat_ms = 0 # 由于我們是 CPU tuning，故不需要該參數
timeout = 10 # 秒# 創建 TVM runner
runner = autotvm.LocalRunner(number=number,repeat=repeat,timeout=timeout,min_repeat_ms=min_repeat_ms,enable_cpu_cache_flush=True,
)
創建一個簡單的結構來保存調整選項。
tunner：我們使用 XGBoost 算法來指導搜索。 在實際中可能需要根據模型復雜度、時間限制等因素選擇其他算法。tirals：對于實際項目，您需要將試驗次數設置為大于此處使用的值 10。 CPU 推薦 1500，GPU 3000-4000。 所需的試驗次數可能取決于特定模型和處理器，因此值得花一些時間評估一系列值的性能，以找到調整時間和模型優化之間的最佳平衡。early_stopping參數是在應用提前停止搜索的條件之前要運行的最小 trial 數。measure_option指定將在何處構建試用代碼以及將在何處運行。 在本例中，我們使用我們剛剛創建的 LocalRunner 和一個 LocalBuilder。tuning_records選項指定一個文件來寫入調整數據。
tuning_option = {"tuner": "xgb","trials": 10,"early_stopping": 100,"measure_option": autotvm.measure_option(builder=autotvm.LocalBuilder(build_func="default"), runner=runner),"tuning_records": "resnet-50-v2-autotuning.json",
}
注意：在此示例中，為了節省時間，我們將試驗次數和提前停止次數設置為 10。如果將這些值設置得更高，我們可能會看到更多的性能改進，但這是以花費調優時間為代價的。 收斂所需的試驗次數將根據模型和目標平臺的具體情況而有所不同。
# 開始從 onnx 模型中提取 tasks
tasks = autotvm.task.extract_from_program(mod["main"], target=target, params=params)# 一次 tune 提取到的 tasks
for i, task in enumerate(tasks):prefix = "[Task %2d/%2d] " % (i + 1, len(tasks))tuner_obj = XGBTuner(task, loss_type="rank")tuner_obj.tune(n_trial=min(tuning_option["trials"], len(task.config_space)),early_stopping=tuning_option["early_stopping"],measure_option=tuning_option["measure_option"],callbacks=[autotvm.callback.progress_bar(tuning_option["trials"], prefix=prefix),autotvm.callback.log_to_file(tuning_option["tuning_records"]),],)
此處輸出：
[Task 1/25] Current/Best: 0.00/ 0.00 GFLOPS | Progress: (0/10) | 0.00 s
[Task 1/25] Current/Best: 33.79/ 49.04 GFLOPS | Progress: (4/10) | 5.66 s...Done.[Task 25/25] Current/Best: 3.13/ 3.13 GFLOPS | Progress: (4/10) | 3.17 s
[Task 25/25] Current/Best: 2.48/ 3.13 GFLOPS | Progress: (8/10) | 15.43 s
[Task 25/25] Current/Best: 0.00/ 3.13 GFLOPS | Progress: (10/10) | 45.72 s
使用 tuning data 編譯優化過的模型
作為上述調優過程的輸出，我們獲得了存儲在 resnet-50-v2-autotuning.json 中的調優記錄。 編譯器將根據該結果為指定 target 上的模型生成高性能代碼。
現在已經收集了模型的調整數據，我們可以使用優化的算子重新編譯模型以加快計算速度。
with autotvm.apply_history_best(tuning_option["tuning_records"]):with tvm.transform.PassContext(opt_level=3, config={}):lib = relay.build(mod, target=target, params=params)dev = tvm.device(str(target), 0)
module = graph_executor.GraphModule(lib["default"](dev))
驗證優化過后的模型的運行后的輸出結果與之前的相同：
dtype = "float32"
module.set_input(input_name, img_data)
module.run()
output_shape = (1, 1000)
tvm_output = module.get_output(0, tvm.nd.empty(output_shape)).numpy()scores = softmax(tvm_output)
scores = np.squeeze(scores)
ranks = np.argsort(scores)[::-1]
for rank in ranks[0:5]:print("class='%s' with probability=%f" % (labels[rank], scores[rank]))
此處輸出：
class='n02123045 tabby, tabby cat' with probability=0.610552
class='n02123159 tiger cat' with probability=0.367180
class='n02124075 Egyptian cat' with probability=0.019365
class='n02129604 tiger, Panthera tigris' with probability=0.001273
class='n04040759 radiator' with probability=0.000261
確是是相同的。
比較調整過的和未調整過的模型
這里我們同樣收集與此優化模型相關的一些基本性能數據，以將其與未優化模型進行比較。 根據底層硬件、迭代次數和其他因素，在將優化模型與未優化模型進行比較時，我們能看到性能改進。
import timeittiming_number = 10
timing_repeat = 10
optimized = (np.array(timeit.Timer(lambda: module.run()).repeat(repeat=timing_repeat, number=timing_number))* 1000/ timing_number
)
optimized = {"mean": np.mean(optimized), "median": np.median(optimized), "std": np.std(optimized)}print("optimized: %s" % (optimized))
print("unoptimized: %s" % (unoptimized))
此處輸出：
optimized: {'mean': 211.9480087934062, 'median': 211.2688914872706, 'std': 1.1843122740378864}
unoptimized: {'mean': 229.1864895541221, 'median': 228.7280524149537, 'std': 1.0664440211813757}
在本教程中，我們給出了一個簡短示例，說明如何使用 TVM Python API 編譯、運行和調整模型。 我們還討論了對輸入和輸出進行預處理和后處理的必要性。 在調整過程之后，我們演示了如何比較未優化和優化模型的性能。
這里我們展示了一個在本地使用 ResNet 50 V2 的簡單示例。 但是，TVM 支持更多功能，包括交叉編譯、遠程執行和分析/基準測試。這將會在以后的教程中介紹。
Ref：
https://tvm.apache.org/docs/tutorial/autotvm_relay_x86.html#sphx-glr-tutorial-autotvm-relay-x86-py
來處理算子)
