TASO は、自動的に生成され検証されたグラフ変換を使用して、DNN モデルの計算グラフを最適化します。任意の DNN モデルの場合、TASO は自動生成されたグラフ変換を使用して、元の DNN モデルと同等の潜在的な計算グラフの大規模な検索空間を構築します。 TASO はコストベースの検索アルゴリズムを採用して空間を探索し、高度に最適化された計算グラフを自動的に発見します。 TASO は、既存の深層学習フレームワークのグラフ オプティマイザーよりも最大 3 倍優れています。
TASO をソースからインストールする手順を参照してください。また、すべての依存関係がプリインストールされた、ビルド済みの Docker イメージも提供します。
TASO は、ONNX、TensorFlow、および PyTorch グラフ形式の事前トレーニング済み DNN モデルを直接最適化できます。 TASO は、任意の DNN アーキテクチャを最適化するための Python インターフェイスも提供します。 TASO は、最適化された計算グラフの ONNX へのエクスポートをサポートしており、既存の深層学習フレームワークのほとんどで入力として直接使用できます。
TASO は、事前トレーニングされた ONNX モデルを直接最適化できます。これは、わずか数行の Python コードで実行できます。次のコード スニペットは、事前トレーニングされた DNN モデルを ONNX から読み込み、モデルを最適化し、最適化されたモデルを 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 ファイルの一部は、 onnx
フォルダーにあります。
TASO は、tf2onnx を使用してモデルを ONNX に変換することで、TensorFlow モデルを最適化できます。
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
PyTorch には、torch.onnx パッケージの一部として ONNX のサポートが組み込まれています。 TASO は、ONNX 形式の PyTorch モデルを直接最適化できます。
TASO は、TASO Python インターフェイスを使用して任意の DNN アーキテクチャを最適化することもできます。次のコード スニペットは、図に示されている一番左の DNN グラフを構築します。 TASO は一連の重要な変換を自動的に実行し、最終的に右端の DNN グラフを検出します。これは V100 GPU で 1.3 倍高速です。さらに多くの DNN サンプルがexamples
フォルダーにあります。
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 月。