TVM如何訓(xùn)練TinyML
機(jī)器學(xué)習(xí)研究人員和從業(yè)人員對“裸機(jī)”（低功耗，通常沒有操作系統(tǒng)）設(shè)備產(chǎn)生了廣泛的興趣。盡管專家已經(jīng)有可能在某些裸機(jī)設(shè)備上運(yùn)行某些模型，但是為各種設(shè)備優(yōu)化模型的挑戰(zhàn)非常艱巨，通常需要手動優(yōu)化設(shè)備特定的庫。對于那些沒有Linux支持的平臺，不存在用于部署模型的可擴(kuò)展解決方案。因此，為了定位新設(shè)備，開發(fā)人員必須實(shí)現(xiàn)一次性的定制軟件堆棧，以管理系統(tǒng)資源和調(diào)度模型執(zhí)行。
機(jī)器學(xué)習(xí)軟件的手動優(yōu)化不是裸機(jī)設(shè)備領(lǐng)域獨(dú)有的。實(shí)際上，對于使用其它硬件后端（例如GPU和FPGA）的開發(fā)人員來說，這已成為一個(gè)共同的主題。TVM已被證明可以抵御新硬件目標(biāo)的沖擊，但直到現(xiàn)在，仍無法解決微控制器的獨(dú)特特性。為了解決這一領(lǐng)域的問題，擴(kuò)展了TVM以提供稱為μTVM（腳注：發(fā)音為“ MicroTVM”）的微控制器后端。μTVM促進(jìn)了主機(jī)驅(qū)動的裸機(jī)設(shè)備上張量程序的執(zhí)行，并通過TVM內(nèi)置的張量程序優(yōu)化器AutoTVM實(shí)現(xiàn)了這些程序的自動優(yōu)化。下圖顯示了μTVM + AutoTVM基礎(chǔ)架構(gòu)的鳥瞰圖：

功能
在討論什么是TVM / MicroTVM或它如何工作之前，讓看一下實(shí)際使用示例。

標(biāo)準(zhǔn)μTVM設(shè)置，主機(jī)通過JTAG與設(shè)備通信。
上面，有一塊STM32F746ZG板，其中裝有ARM Cortex-M7處理器，考慮到在低功耗封裝中的強(qiáng)大性能，邊緣AI的理想部件。使用其USB-JTAG端口將其連接到臺式機(jī)。在桌面上，運(yùn)行OpenOCD來打開與設(shè)備的JTAG連接。反過來，OpenOCD允許μTVM使用與設(shè)備無關(guān)的TCP套接字控制M7處理器。完成此設(shè)置后，可以使用如下所示的TVM代碼運(yùn)行CIFAR-10分類器（此處為完整腳本）：
OPENOCD_SERVER_ADDR = ‘127.0.0.1’
OPENOCD_SERVER_PORT = 6666
TARGET = tvm.target.create(‘c -device=micro_dev’)
DEV_CONFIG = stm32f746xx.default_config(OPENOCD_SERVER_ADDR, OPENOCD_SERVER_PORT)

module, params = get_cifar10_cnn()
with micro.Session(device_config) as sess:
graph, c_module, params = relay.build(module[‘main’], target=TARGET, params=params)
micro_mod = micro.create_micro_mod(c_module, DEV_CONFIG)
graph_mod = graph_runtime.create(graph, micro_mod, ctx=tvm.micro_dev(0))
graph_mod.run(data=data_np)
prediction = CIFAR10_CLASSES[np.argmax(graph_mod.get_output(0).asnumpy())]
print(f’prediction was {prediction}’)
與CMSIS-NN版本5.7.0（commit a65b7c9a）（一種手動優(yōu)化的ML內(nèi)核庫）相比，以下是MicroTVM的性能結(jié)果。

如所見，開箱即用的性能并不好，但這就是AutoTVM表現(xiàn)的地方。可以為設(shè)備編寫調(diào)度模板，進(jìn)行一輪自動調(diào)整，然后獲得明顯更好的結(jié)果。要插入自動調(diào)整的結(jié)果，只需要替換以下行：
graph, c_module, params = relay.build(module[‘main’], target=TARGET, params=params)
這些行：
with TARGET, autotvm.apply_history_best(TUNING_RESULTS_FILE):
graph, c_module, params = relay.build(module[‘main’], target=TARGET, params=params)
現(xiàn)在，結(jié)果如下所示：

性能提高了約2倍，現(xiàn)在，離CMSIS-NN更近了。盡管MicroTVM CIFAR10的實(shí)現(xiàn)與類似的TFLite / CMSIS-NN模型相比具有競爭優(yōu)勢，但這項(xiàng)工作才剛剛開始利用TVM的優(yōu)化功能。通過加速其它運(yùn)營商（如密集/完全連接）并利用TVM特定于模型的量化和運(yùn)營商融合功能，還有進(jìn)一步優(yōu)化的空間。帶有μTVM的TVM使能夠充分發(fā)揮作用。怎樣工作的？幕后發(fā)生了什么事？
設(shè)計(jì)

μTVM設(shè)備在RAM中的存儲器布局
μTVM旨在通過最大限度的減少必須滿足的一組要求，支持設(shè)備的最小公分母。特別是，用戶只需要提供：

1. 設(shè)備的C交叉編譯器工具鏈
2. 一種用于讀取/寫入設(shè)備內(nèi)存并在設(shè)備上執(zhí)行代碼的方法
3. 包含設(shè)備的內(nèi)存布局和一般體系結(jié)構(gòu)特征的規(guī)范
4. 一個(gè)代碼片段，為設(shè)備執(zhí)行功能做準(zhǔn)備
   大多數(shù)裸機(jī)設(shè)備都支持C和JTAG（調(diào)試協(xié)議），因此（1）和（2）通常是免費(fèi)提供的！此外，（3）和（4）通常是很小的要求。以下是STM32F746系列板卡的（3）和（4）的示例。
   device_config = {
   ‘device_id’: ‘a(chǎn)rm.stm32f746xx’, # unique identifier for the device
   ‘toolchain_prefix’: ‘a(chǎn)rm-none-eabi-’, # prefix of each binary in the cross-compilation toolchain (e.g., arm-none-eabi-gcc)
   ‘base_addr’: 0x20000000, # first address of RAM
   ‘section_sizes’: { # dictionary of desired section sizes in bytes
   ‘text’: 18000,
   ‘rodata’: 100,
   ‘data’: 100,
   …
   },
   ‘word_size’: 4, # device word size
   ‘thumb_mode’: True, # whether to use ARM’s thumb ISA
   ‘comms_method’: ‘openocd’, # method of communication with the device
   ‘server_addr’: ‘127.0.0.1’, # OpenOCD server address (if ‘comms_method’ is ‘openocd’)
   ‘server_port’: 6666, # OpenOCD server port (if ‘comms_method’ is ‘openocd’)
   }
   .syntax unified
   .cpu cortex-m7
   .fpu softvfp
   .thumb

.section .text.UTVMInit
.type UTVMInit, %function
UTVMInit:
/* enable fpu /
ldr r0, =0xE000ED88
ldr r1, [r0]
ldr r2, =0xF00000
orr r1, r2
str r1, [r0]
dsb
isb
/ set stack pointer /
ldr sp, =_utvm_stack_pointer_init
bl UTVMMain
.size UTVMInit, .-UTVMInit
μTVM基礎(chǔ)架構(gòu)和設(shè)備Runtime僅用于滿足這些要求，正在努力通過支持通用的開源Runtime平臺（例如mBED OS）來處理編譯和鏈接過程，以降低這些要求。
設(shè)備會話
鑒于微控制器交互的網(wǎng)絡(luò)性質(zhì)，通過引入的概念略微偏離了標(biāo)準(zhǔn)TVM代碼MicroSession。
μTVM中的每個(gè)功能，都依賴于與目標(biāo)設(shè)備的開放會話。如果熟悉TVM，可能會注意到有一行代碼與第一個(gè)代碼段中的規(guī)范有所不同-即，這是一個(gè)代碼：
…
with micro.Session(device_config) as sess:
…
該with塊內(nèi)的每一行都可以調(diào)用μTVM中的函數(shù)，上下文是由所指定的設(shè)備device_config。這條線在做很多事情，拆開包裝。
首先，使用指定的任何一種通信方法（通常是OpenOCD）來初始化與設(shè)備的連接。然后，使用指定的交叉編譯器交叉編譯μTVM設(shè)備的Runtime。最后，由主機(jī)分配用于已編譯二進(jìn)制文件的空間，并使用打開的連接將二進(jìn)制文件加載到設(shè)備上。
有了設(shè)備上的Runtime，自然會希望一些功能通過運(yùn)行。
模塊加載
TVM中的核心抽象之一是模塊的抽象。模塊存儲用于特定設(shè)備/Runtime目標(biāo)的一組相關(guān)功能。鑒于微控制器通常沒有操作系統(tǒng)，因此μTVM需要做很多額外的工作來維持這種高級抽象。要查看發(fā)生了什么，將跟蹤創(chuàng)建和加載與μTVM兼容的模塊的過程。
假設(shè)有一個(gè)micro.Session開放的設(shè)備和一個(gè)實(shí)現(xiàn)2D卷積的TVM調(diào)度。如果想將其加載到微控制器上，需要發(fā)出C代碼。為此，只需要target在tvm.build或中設(shè)置即可relay.build。例子：
graph, c_module, params = relay.build(module[‘main’], target=‘c -device=micro_dev’, params=params)
通過這樣設(shè)置目標(biāo)，構(gòu)建過程將貫穿C代碼生成后端。但是，生成的C模塊仍駐留在主機(jī)上。為了將其加載到設(shè)備上，通過μTVM基礎(chǔ)架構(gòu)中的核心功能之一運(yùn)行create_micro_mod。例子：
micro_mod = micro.create_micro_mod(c_module, DEV_CONFIG)
上面的行交叉編譯模塊中的C源代碼，為所得的二進(jìn)制文件分配空間（以便可以與Runtime在設(shè)備內(nèi)存中共存），然后將二進(jìn)制文件的每個(gè)部分發(fā)送到其在設(shè)備上分配的插槽中。一旦模塊二進(jìn)制文件貼緊在設(shè)備內(nèi)存中，便會修補(bǔ)二進(jìn)制文件中的功能指針，以使模塊可以在設(shè)備Runtime訪問輔助功能（例如，用于分配暫存器）。
現(xiàn)在，在將內(nèi)核加載到設(shè)備上之后，可以獲取卷積函數(shù)的遠(yuǎn)程句柄，如下所示：
micro_func = micro_mod[‘conv2d’]
張量加載
如果要調(diào)用算子，首先需要一些張量作為參數(shù)：
data_np, kernel_np = get_conv_inputs()
ctx = tvm.micro_dev(0)
data = tvm.nd.array(data_np, ctx=ctx)
kernel = tvm.nd.array(kernel_np, ctx=ctx)
根據(jù)它的數(shù)據(jù)類型（例如，int8，float32等）和形狀，各張量的字節(jié)大小被計(jì)算，并在主機(jī)分配所述設(shè)備的堆存儲器的區(qū)域中。然后將張量的數(shù)據(jù)加載到分配的區(qū)域中。
函數(shù)調(diào)用
算子執(zhí)行可能是該系統(tǒng)中最棘手的部分。為了簡化其表示，將首先介紹嚴(yán)格執(zhí)行（在調(diào)用操作符后立即執(zhí)行操作），然后是延遲執(zhí)行（僅在需要其結(jié)果后才執(zhí)行操作符）–后者是系統(tǒng)的實(shí)際運(yùn)行方式。
嚴(yán)格執(zhí)行
調(diào)用函數(shù)時(shí)，輸入張量和輸出張量均作為參數(shù)傳遞，即所謂的目標(biāo)傳遞風(fēng)格：
conv2D(data, kernel, output)
鑒于這些張量已在設(shè)備上分配，只需要將元數(shù)據(jù)發(fā)送到設(shè)備（設(shè)備地址，形狀和數(shù)據(jù)類型），知道要使用哪個(gè)駐留張量。函數(shù)調(diào)用的Runtime表示形式包括此元數(shù)據(jù)以及被調(diào)用函數(shù)的地址（如下所示）。在構(gòu)造此表示形式之前，需要將元數(shù)據(jù)序列化到為此目的明確存在的設(shè)備上的arguments部分中。
/

• task struct for uTVM
  /
  typedef struct {
  / pointer to function to call for this task /
  int32_t (func)(void, void, int32_t);
  /* array of argument tensors /
  TVMValue arg_values;
  /* array of datatype codes for each argument /
  int arg_type_codes;
  /* number of arguments */
  int32_t num_args;
  } UTVMTask;
  在嚴(yán)格的設(shè)置中，只有一個(gè)全局UTVMTask實(shí)例，從主機(jī)端將其寫入其中。一旦寫入任務(wù)，Runtime就具有執(zhí)行該功能所需的一切，并且可以在Runtime的入口點(diǎn)開始執(zhí)行。Runtime將執(zhí)行一些輕量級的初始化，運(yùn)行算子，然后將控制權(quán)返回給主機(jī)。
  執(zhí)行
  在實(shí)踐中，由于通信開銷開始占主導(dǎo)地位，一旦用戶要求執(zhí)行算子就變得非常耗資源。可以通過延遲評估，直到用戶希望獲得調(diào)用結(jié)果的方式來提高系統(tǒng)的吞吐量。
  從實(shí)現(xiàn)的角度來看，UTVMTask現(xiàn)在不急于序列化參數(shù)元數(shù)據(jù)和數(shù)據(jù)，而是需要在主機(jī)端累積函數(shù)調(diào)用元數(shù)據(jù)，然后再將其刷新到設(shè)備中。設(shè)備Runtime還需要進(jìn)行一些更改：（1）現(xiàn)在必須具有的全局?jǐn)?shù)組，UTVMTask并且（2）需要依次遍歷并執(zhí)行每個(gè)任務(wù)。
  帶MicroTVM的AutoTVM
  到目前為止，描述的Runtime對于模型部署似乎并不是很有用，因?yàn)榉浅Ｒ蕾囍鳈C(jī)。這是有意為之的，實(shí)際上，Runtime是為實(shí)現(xiàn)另一個(gè)目標(biāo)而設(shè)計(jì)的：AutoTVM支持。
  通常，AutoTVM會提出候選內(nèi)核，并使用隨機(jī)輸入在目標(biāo)后端上運(yùn)行，然后使用計(jì)時(shí)結(jié)果來改善其搜索過程。鑒于AutoTVM只關(guān)心單個(gè)算子的執(zhí)行，將Runtime設(shè)計(jì)為面向算子，而不是面向模型。但是對于μTVM，與設(shè)備的通信通常會占據(jù)執(zhí)行時(shí)間。惰性執(zhí)行使可以多次運(yùn)行同一算子，而無需將控制權(quán)交還給主機(jī)，因此，通信成本在每次Runtime均攤銷，可以更好地了解性能概況。
  由于AutoTVM需要在大量候選內(nèi)核上進(jìn)行快速迭代，因此μTVM基礎(chǔ)架構(gòu)目前僅使用RAM。但是，對于自托管Runtime，肯定需要同時(shí)使用閃存和RAM。
  托管圖Runtime
  盡管托管的Runtime是為AutoTVM設(shè)計(jì)的，但仍然可以運(yùn)行完整的模型（只要沒有任何控制流）。僅通過使用TVM的圖形Runtime即可免費(fèi)使用此功能，但具有μTVM上下文。實(shí)際上，圖Runtime對主機(jī)的唯一依賴是張量分配和算子調(diào)度（這只是依賴圖的一種拓?fù)漕愋?#xff09;。
  評估
  有了這個(gè)基礎(chǔ)架構(gòu)，試圖回答以下問題：

1. μTVM是否真的與設(shè)備無關(guān)？
2. 使用μTVM進(jìn)行優(yōu)化試驗(yàn)需要多少精力？
   為了評估（1），在兩個(gè)目標(biāo)上進(jìn)行了實(shí)驗(yàn)：
   ? 一個(gè)手臂STM32F746NG開發(fā)板，采用了的Cortex-M7處理器
   ? μTVM主機(jī)仿真設(shè)備，可在主機(jī)上創(chuàng)建一個(gè)內(nèi)存競技場，與之連接的主機(jī)就像裸機(jī)設(shè)備一樣。
   為了評估（2），探索了Arm板的優(yōu)化方案，這些方案可以最大程度地降低成本。
   作為比較，從Arm的本教程中提取了量化的CIFAR-10 CNN。CMSIS-NN（Arm專家高度優(yōu)化的內(nèi)核庫）用作算子庫，使該CNN成為完美的評估目標(biāo)，因?yàn)楝F(xiàn)在可以直接將μTVM的結(jié)果與Arm上的CMSIS-NN進(jìn)行比較木板。
   

CIFAR-10 CNN圖
方法
在實(shí)驗(yàn)中，使用HEAD的TVM（commit 9fa8341），CMSIS-NN的5.7.0版（commit a65b7c9a），STM32CubeF7的1.16.0版以及Arm的適用于Arm嵌入式處理器的GNU工具的GCC 9-2019-q4-major 9.2 .1工具鏈（修訂版277599）。實(shí)驗(yàn)中使用的主機(jī)運(yùn)行Ubuntu Linux 18.04.4 LTS，并運(yùn)行帶有62GB RAM的AMD Ryzen Threadripper 2990WX 32核處理器。
特定于手臂的優(yōu)化
使用CMSIS-NN，第一個(gè)卷積映射到其RGB卷積實(shí)現(xiàn)（專門用于輸入層），而后兩個(gè)卷積映射到其“快速”卷積實(shí)現(xiàn)。經(jīng)過較早的泛型優(yōu)化后，覺得性能對于RGB卷積已經(jīng)足夠接近了，但是對快速卷積結(jié)果卻不滿意。幸運(yùn)的是，Arm發(fā)布了一篇描述CMSIS-NN中使用的優(yōu)化的論文，發(fā)現(xiàn)正在從SIMD內(nèi)在函數(shù)中獲得巨大的加速。在本文中，提出了一種使用SIMD內(nèi)在函數(shù)的矩陣乘法微內(nèi)核（下圖）。雖然可以在TVM的代碼生成工具中添加對內(nèi)在函數(shù)的一流支持，這從長遠(yuǎn)來看可能是最好的做法，但TVM提供了張量化是支持SIMD的“快捷方法”。

CMSIS-NN論文的圖表顯示了2x2矩陣乘法微內(nèi)核
張量化通過定義可插入TVM算子最內(nèi)層循環(huán)的微內(nèi)核來工作。使用這種機(jī)制，添加對Arm板的SIMD支持就像在C中定義一個(gè)微內(nèi)核一樣簡單（可在此處找到），該微內(nèi)核反映了其論文中的實(shí)現(xiàn)。定義了一個(gè)調(diào)度，使用該微內(nèi)核（在此處找到），對其進(jìn)行自動調(diào)整，然后得到“ μTVM SIMD調(diào)整”結(jié)果。
盡管能夠使用SIMD微內(nèi)核進(jìn)行直接卷積，但是CMSIS-NN使用他們所謂的“部分im2col”作為其實(shí)現(xiàn)策略，這在性能和內(nèi)存使用之間進(jìn)行了權(quán)衡。代替一次顯示整個(gè)im2col矩陣，部分im2col一次只生成幾列。然后，對于每一批，他們可以將矩陣發(fā)送到其SIMD matmul函數(shù)。
假設(shè)是，除其它優(yōu)化外，可以通過自動調(diào)整找到最佳的批量大小。在實(shí)踐中，發(fā)現(xiàn)部分im2col比直接卷積實(shí)現(xiàn)要慢得多，因此在其余結(jié)果中不包括。
當(dāng)然，還可以從CMSIS-NN中獲得其它優(yōu)化來進(jìn)一步縮小差距：
? 將int8權(quán)重批量擴(kuò)展為int16，以減少SIMD的重復(fù)擴(kuò)展
? 將卷積拆分為3x3的圖塊以減少填充檢查
但是，目標(biāo)是展示μTVM可以完成的工作的大致范圍。即使這樣，這也不是競爭，因?yàn)镃MSIS-NN（以及任何其它手動優(yōu)化的庫）可以使用Bring Your Own Codegen框架直接插入TVM 。
端到端
CIFAR-10
在探索卷積優(yōu)化之后，著手測量其對端到端性能的影響。對于ARM板，收集了未調(diào)整的結(jié)果，這是調(diào)整的結(jié)果沒有任何使用SIMD，這是調(diào)整的結(jié)果與SIMD和結(jié)果使用CMSIS-NN。對于模擬的主機(jī)設(shè)備，僅收集未調(diào)整的結(jié)果和通用的調(diào)整結(jié)果。
https://github.com/areusch/microtvm-blogpost-eval

int8Arm STM32F746NG進(jìn)行量化的CIFAR-10 CNN比較（從上方轉(zhuǎn)貼）

int8μTVM的仿真主機(jī)設(shè)備上對量化的CIFAR-10 CNN進(jìn)行比較
在Arm STM32系列板上，與最初的未調(diào)整算子相比，能夠?qū)⑿阅芴岣呒s2倍，并且所獲得的結(jié)果更接近CMSIS-NN。此外，能夠顯著提高主機(jī)仿真設(shè)備上的性能。盡管x86的數(shù)字意義不大，表明可以使用相同的基礎(chǔ)架構(gòu)（μTVM）來在極為不同的體系結(jié)構(gòu)上優(yōu)化性能。
隨著更廣泛地?cái)U(kuò)展此方法，在將來繼續(xù)關(guān)注更多端到端基準(zhǔn)測試。
自托管Runtime：最終領(lǐng)域

設(shè)想的μTVM優(yōu)化和部署流程
如上所述，雖然當(dāng)前Runtime已經(jīng)可以獲取端到端基準(zhǔn)測試結(jié)果，但目前仍在路線圖上以獨(dú)立能力部署這些模型。差距在于面向AutoTVM的Runtime當(dāng)前依賴于主機(jī)來分配張量并計(jì)劃函數(shù)執(zhí)行。然而，為了在邊緣是有用的，需要通過μTVM，其產(chǎn)生一個(gè)管道單一待裸機(jī)設(shè)備上運(yùn)行的二進(jìn)制。然后，用戶可以通過在邊緣應(yīng)用程序中包含此二進(jìn)制文件，輕松地將快速M(fèi)L集成到他們的應(yīng)用程序中。該管道的每個(gè)階段都已經(jīng)到位，現(xiàn)在只需將粘合在一起即可，因此期待在此方面的最新進(jìn)展。
結(jié)論
用于單內(nèi)核優(yōu)化的MicroTVM現(xiàn)已準(zhǔn)備就緒，并且是該用例的選擇。現(xiàn)在，當(dāng)建立自托管的部署支持時(shí)，希望也和使μTVM成為模型部署的選擇一樣興奮。但是，這不只是一場觀看比賽-記住：這都是開源的！μTVM仍處于起步階段，因此每個(gè)人對其軌跡都會產(chǎn)生很大的影響。

總結(jié)

以上是生活随笔為你收集整理的TVM如何训练TinyML的全部內(nèi)容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯(cuò)，歡迎將生活随笔推薦給好友。