TASO

TASO：用於深度學習的張量代數超級優化器

TASO 使用自動產生和驗證的圖轉換來最佳化 DNN 模型的計算圖。對於任意 DNN 模型，TASO 使用自動產生的圖轉換來建立與原始 DNN 模型等效的潛在計算圖的大型搜尋空間。 TASO 採用基於成本的搜尋演算法來探索空間，並自動發現高度最佳化的計算圖。 TASO 的效能比現有深度學習框架中的圖形優化器高出 3 倍。

NVIDIA V100 GPU 上的端對端推理效能比較。

安裝TASO

請參閱從來源安裝 TASO 的說明。我們還提供預先建置的 docker 映像，並預先安裝了所有相依性。

使用TASO

TASO 可以直接優化 ONNX、TensorFlow 和 PyTorch 圖形格式的任何預訓練 DNN 模型。 TASO 也提供了一個用於優化任意 DNN 架構的 Python 介面。 TASO 支援將最佳化後的計算圖匯出到 ONNX，可以直接用作大多數現有深度學習框架的輸入。

優化 ONNX 模型

TASO 可以直接優化預先訓練的 ONNX 模型，並且只需幾行 Python 程式碼即可完成。以下程式碼片段展示如何從 ONNX 載入預先訓練的 DNN 模型、優化模型並將優化後的模型儲存到 ONNX 檔案中。

 import taso
import onnx

old_model = taso . load_onnx ( "/path/to/load/onnx/model" )
taso_graph = taso . optimize ( old_model )
new_model = taso . export_onnx ( taso_graph )
onnx . save ( new_model , "/path/to/save/new/onnx/model" )

優化後的模型與原始模型具有相同的精度，並且可以直接被現有的深度學習框架使用。 onnx資料夾中提供了一些原始和 TASO 優化的 ONNX 檔案。

優化 TensorFlow 模型

TASO 可以透過使用 tf2onnx 將模型轉換為 ONNX 來優化 TensorFlow 模型。

• 首先，從 PyPi 安裝tf2onnx ，如下所示或從原始碼安裝。

 pip install -U tf2onnx

• 其次，使用tf2onnx將 TensorFlow 模型轉換為 ONNX。

 python -m tf2onnx.convert 
       --saved-model /path/to/tensorflow/saved/model 
       --output /path/to/onnx/model/file

• 第三，請依照上述說明使用 TASO 在 ONNX 中最佳化模型。

優化 PyTorch 模型

PyTorch 作為 torch.onnx 套件的一部分內建了對 ONNX 的支援。 TASO 可以直接最佳化 ONNX 格式的 PyTorch 模型。

使用 Python 介面優化任意 DNN 模型

TASO 也可以使用 TASO Python 介面優化任意 DNN 架構。以下程式碼片段建構了圖中最左邊的 DNN 圖。 TASO 自動執行一系列重要的轉換，並最終發現最右側的 DNN 圖，該圖在 V100 GPU 上速度提高了 1.3 倍。 examples夾中提供了更多 DNN 範例。

 import taso
import onnx

#Build DNN model
graph = taso . new_graph ()
input = graph . new_input ( dims = ( 1 , 128 , 56 , 56 ))
w1 = graph . new_weight ( dims = ( 128 , 128 , 3 , 3 ))
w2 = graph . new_weight ( dims = ( 128 , 128 , 1 , 1 ))
w3 = graph . new_weight ( dims = ( 128 , 128 , 3 , 3 ))
left = graph . conv2d ( input = input , weight = w1 , strides = ( 1 , 1 ), padding = "SAME" , activation = "RELU" )
left = graph . conv2d ( input = left , weight = w3 , strides = ( 1 , 1 ), padding = "SAME" )
right = graph . conv2d ( input = input , weight = w2 , strides = ( 1 , 1 ), padding = "SAME" , activation = "RELU" )
output = graph . add ( left , right )
output = graph . relu ( output )

#Optimize DNN model
new_graph = taso . optimize ( graph )
onnx_model = taso . export_onnx ( new_graph )
onnx . save ( onnx_model , "/path/to/save/new/onnx/model" )

出版品

• 賈志浩、奧德·帕登、詹姆斯·托馬斯、陶德·華沙斯基、馬泰·扎哈里亞和亞歷克斯·艾肯。 TASO：透過自動產生圖替換來優化深度學習計算。作業系統原理研討會 (SOSP) 論文集，加拿大安大略省，2019 年 10 月。

• 賈志浩、詹姆斯·托馬斯、陶德·華沙斯基、高明宇、馬泰·扎哈里亞和亞歷克斯·艾肯。使用寬鬆圖替換優化 DNN 計算。系統與機器學習會議 (SysML) 會議記錄，加州帕洛阿爾托，2019 年 4 月。