TinyNN：輕鬆讓深度神經網路加速數倍的開源工具
深度學習在產業界與學術界最大的差別，就是學術界只要訓練出可用模型並且證明理論的可行性就好了；但是產業界必須要把訓練出的模型部署到機器上，並且對性能要求也比較嚴苛。
阿里巴巴內部也碰到了很多模型效能的問題。舉例來說，2 年前我們團隊曾經嘗試要在低端平板電腦上部署 Google 的超輕量人臉檢測模型 BlazeFace。根據論文，BlazeFace 在各種手機上的運行時間不超過 6ms (166 FPS)，但是在我們的低端平板電腦卻要運行 200ms (5 FPS)，幾乎是不可用的程度。
為了解決模型效能的問題，我們團隊設計了一套方便的神經網路加速工具 TinyNeuralNetwork，並已經在本月開源：
這篇文章會簡單介紹一下 TinyNeuralNetwork 的主要功能，介紹背後的機制原理，並且提供簡單的範例。
TinyNeuralNetwork
TinyNeuralNetwork, TinyNN 是一套基於 PyTorch 的模型的加速工具。實務中，你可以把 TinyNN 當作模型轉換工具，輕鬆把 Torch 模型轉換成 TFLite 模型。更進一步，你可以用 TinyNN 來設計更輕量化的網路結構，達到神經網路加速的效果。
TinyNN 主要從 3方面來實現神經網路加速：
- Compute Graph Optimization, 計算圖最佳化
- Quantization Aware Training, 量化訓練
- Structured Pruning, 結構化剪枝
Compute Graph Optimization 計算圖最佳化
計算圖最佳化，指的是透過 OP 融合的方式，儘量減少計算圖裡面的冗余步驟。冗余的計算圖通常有 2 個來源：一個是訓練產生的，一個是模型轉換產生的。
訓練產生的計算圖冗余，是因為有些 OP 在訓練與推理的特性不一致導致。舉個例子，相鄰的 Conv OP與 BatchNorm OP 在訓練階段必須拆開計算，因為 Conv 的權重是透過 back propagation 更新，而 BatchNorm 是透過 forward propagation 更新的。但是在推理階段，這兩個 OP 都是屬於線性代數運算，其實是可以完全可以合併成一個 Conv OP；如果不做融合，就會造成冗余的計算成本。
模型轉換產生的計算圖冗余，主要來自於 PyTorch 與 TensorFlow 之間模型轉換時，會碰到 Tensor 格式不一樣的問題。PyTorch 預設是使用 NCHW 的格式，而 TFLite 採用 NWHC，因此轉換後的模型會包含大量無用的格式互換，會造成大量的計算成本。
TinyNN 透過計算圖分析與最佳化，能夠成功捕捉上述的 2 種計算圖冗余並且消除。因此，使用 TinyNN 轉換的 TFLite 模型，會比自己轉換的快上數倍。
下面是 TinyNN 轉換模型 (Torch -> TFLite) 的基本範例：
import io
import os
import sys
import torch
from examples.models.cifar10.mobilenet import DEFAULT_STATE_DICT, Mobilenet
from tinynn.converter import TFLiteConverter
CURRENT_PATH = os.path.abspath(os.path.dirname(__file__))
def main_worker():
model = Mobilenet()
model.load_state_dict(torch.load(DEFAULT_STATE_DICT))
model.cpu()
model.eval() # 注意，這裡是還是 Torch 的 NCHW 排列
dummy_input = torch.rand((1, 3, 224, 224))
output_path = os.path.join(CURRENT_PATH, 'out', 'mbv1_224.tflite')
# 如果是量化模型，這裡必須要設置
torch.backends.quantized.engine = 'qnnpack'
# 轉換成 TFLite
converter = TFLiteConverter(model, dummy_input, output_path)
converter.convert()
if __name__ == '__main__':
main_worker()
Quantization Aware Training 量化訓練
在之前的文章《神經網路的推理速度》裡面曾經提到，量化 (Quantization) 是影響模型速度的一個關鍵因子。經過網路量化將 float32 運算替換成 uint8/int8 運算，可以提升網路的理論速度到 3、4 倍。其中，量化訓練 又比 後量化 更能降低精度的損失。
PyTorch 雖然支援量化訓練，但是使用上特別麻煩，而且很多手工更改的部分。整體來說，你需要做以下幾件事情：
- 在量化的起點跟終點，手動插入
QuantStub、DeQuantStub - 手動把所有
add,mul,cat都換成torch.nn.quantized.FloatFunctional裡面的方法。 - 手動實作
fuse_model方法來融合 Conv 與 BatchNorm 的計算圖冗余。 - 使用
torch.quantization.prepare_qat產生量化訓練用的計算圖。 - 訓練完後，使用
torch.quantization.convert生成 TorchScript。 - 如果你需要 TFLite 或 ONNX，需要再做額外轉換。
為了自動化的實現上述的所有步驟，TinyNN 使用動態圖分析以及程式碼生成的方法，可以自動化的幫你產生適合量化訓練的模型程式碼，等價於上述的步驟 1, 2。而步驟 3, 4 , 6，則會在計算圖最佳化的部分幫你自動完成。
一個簡化的量化訓練 TinyNN Sample Code 如下所示：
with model_tracer():
model = Mobilenet()
model.load_state_dict(torch.load(DEFAULT_STATE_DICT))
# 注意，這裡是還是 Torch 的 NCHW 排列
dummy_input = torch.rand((1, 3, 224, 224))
# QATQuantizer 會自動分析模型，並且生成量化模型程式碼到 work_dir
quantizer = QATQuantizer(model, dummy_input, work_dir='out')
qat_model = quantizer.quantize()
# print 出最後的模型結構
print(qat_model)# 訓練模型...# 產生量化模型的計算圖
qat_model = torch.quantization.convert(qat_model)
torch.backends.quantized.engine = 'qnnpack'
# 轉換成 TFLite
converter = TFLiteConverter(qat_model, dummy_input, tflite_path='out/qat_model.tflite')
converter.convert()
當然，TinyNN 也可以支援後量化。雖然後量化的精度不如量化訓練，但是當你想在 data 有限的情況下量化現成模型的時候就特別好用。
with model_tracer():
model = Mobilenet()
model.load_state_dict(torch.load(DEFAULT_STATE_DICT))
# 注意，這裡是還是 Torch 的 NCHW 排列
dummy_input = torch.rand((1, 3, 224, 224))
# PostQuantizer 會自動分析模型，並且生成量化模型程式碼到 work_dir
quantizer = PostQuantizer(model, dummy_input, work_dir='out')
qat_model = quantizer.quantize()
# print 出最後的模型結構
print(qat_model)# 透过多次推理 qat_model，校正後量化模型權重
device = torch.device('cuda')
qat_model.to(device=device)
qat_model.eval()
for image in dataset:
qat_model(image.to(device=device))# Export 量化模型的設置
qat_model = torch.quantization.convert(qat_model)
torch.backends.quantized.engine = 'qnnpack'
# 轉換成 TFLite
converter = TFLiteConverter(qat_model, dummy_input, tflite_path='out/qat_model.tflite')
converter.convert()
Structured Pruning 結構化剪枝
Pruning，神經網路剪枝，是一種直接改變模型結構來達到網路壓縮的加速方式。Pruning 是透過刪除神經網路中不必要的神經元來達到減少計算量的目的。根據目前的研究，經過 Pruning 裁減過的模型，在裁減量不大的情況下甚至能夠提升模型效果。
Pruning 可以分為 結構化 與 非結構化 剪枝。
- Structured Pruning 結構化剪枝：指的是透過某種規則裁減網路，例如 Channel 裁減、Kernal 裁減，裁減後的網路仍然可以用一般的 Conv, FC 運算來表示，具備較好的部署通用性。
- Unstructured Pruning 非結構化剪枝：指的是沒有任何限制的任意剪枝。非結構化剪枝雖然可以達到比結構化剪枝更高的理論壓縮率，但是它破壞了網路結構，導致裁剪後的模型無法以正常的 Conv, FC 運算表示，所以需要特定的硬體才能達到加速效果。
TinyNN 支持的是結構化剪枝，因此剪枝出來的模型可以在任意的硬體上運行，例如 CPU、GPU、NPU。目前 TinyNN 支持 3 種剪枝演算法：NetAdapt、ADMM、OneShot。剪枝算法的優劣很難評測，因為根據裁剪模型與目標任務的不同會有不同的剪枝效果。
下面是 OneShot 的剪枝範例程式碼：
from tinynn.prune import OneShotChannelPruner
from examples.models.cifar10 import mobilenet# 原始的模型
model = mobilenet.Mobilenet()
model.load_state_dict(torch.load(mobilenet.DEFAULT_STATE_DICT))
model.to(device=device)# config_yaml 是用來設定 Pruner 參數的設定檔路徑
pruner = OneShotChannelPruner(model, torch.ones(1, 3, 224, 224), config_yaml)
# 這一步會更新設定檔，會把每個 op 的剪枝設定更詳細的寫出來。你可以手動更新設定檔。
pruner.generate_config(config_yaml)
# 開始剪枝 model
pruner.prune()
Prunner 設定檔是 YAML 格式，範例如下。Sparsity 0.25 代表減去 0.25 的通道，理論上能降低約 (1-0.25)² = 56% 的推理時間。
sparsity: 0.25
metrics: l2_norm # 可以參考 tinynn/graph/modifier.py 得到可用的設置結語
TinyNN 是我在阿里參與的第一個開源碼。雖然如此，我其實在裡面的 coding 量不是很多，而是試著基於 TinyNN 的架構去嘗試更多的剪枝演算法，並反饋給 TinyNN 的核心開發成員。
目前我的工作之一是試著在 TinyNN 裡面建立更好的剪枝演算法，包括對 RNN 與 Transformer 系列結構更有效的 Pruner。如果大家有任何問題可以加入 Github ReadMe 裡面的釘釘群組，當然也可以直接找我交流。
