TASO

TASO: Der Tensor-Algebra-SuperOptimizer für Deep Learning

TASO optimiert die Berechnungsgraphen von DNN-Modellen mithilfe automatisch generierter und verifizierter Graphtransformationen. Für ein beliebiges DNN-Modell verwendet TASO die automatisch generierten Graphtransformationen, um einen großen Suchraum potenzieller Berechnungsgraphen zu erstellen, die dem ursprünglichen DNN-Modell äquivalent sind. TASO verwendet einen kostenbasierten Suchalgorithmus, um den Raum zu erkunden, und entdeckt automatisch hochoptimierte Berechnungsdiagramme. TASO übertrifft die Graphoptimierer in vorhandenen Deep-Learning-Frameworks um das Dreifache.

End-to-End-Inferenzleistungsvergleich auf einer NVIDIA V100-GPU.

Installieren Sie TASO

Siehe Anweisungen zur Installation von TASO aus der Quelle. Wir stellen auch vorgefertigte Docker-Images mit allen vorinstallierten Abhängigkeiten bereit.

Verwenden Sie TASO

TASO kann alle vorab trainierten DNN-Modelle in den Diagrammformaten ONNX, TensorFlow und PyTorch direkt optimieren. TASO bietet außerdem eine Python-Schnittstelle zur Optimierung beliebiger DNN-Architekturen. TASO unterstützt den Export der optimierten Berechnungsdiagramme nach ONNX, die von den meisten vorhandenen Deep-Learning-Frameworks direkt als Eingaben verwendet werden können.

Optimieren Sie ONNX-Modelle

TASO kann vorab trainierte ONNX-Modelle direkt optimieren, und dies kann in nur wenigen Zeilen Python-Code erfolgen. Der folgende Codeausschnitt zeigt, wie man ein vorab trainiertes DNN-Modell von ONNX lädt, das Modell optimiert und das optimierte Modell in einer ONNX-Datei speichert.

 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" )

Das optimierte Modell weist die gleiche Genauigkeit wie das Original auf und kann direkt von vorhandenen Deep-Learning-Frameworks verwendet werden. Einige originale und TASO-optimierte ONNX-Dateien sind im Ordner onnx verfügbar.

Optimieren Sie TensorFlow-Modelle

TASO kann TensorFlow-Modelle optimieren, indem es das Modell mithilfe von tf2onnx in ONNX konvertiert.

• Installieren Sie zunächst tf2onnx von PyPi wie folgt oder von der Quelle.

 pip install -U tf2onnx

• Zweitens konvertieren Sie ein TensorFlow-Modell mit tf2onnx in ONNX.

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

• Drittens verwenden Sie TASO, um das Modell in ONNX zu optimieren, indem Sie die obigen Anweisungen befolgen.

Optimieren Sie PyTorch-Modelle

PyTorch verfügt über eine integrierte Unterstützung für ONNX als Teil des Torch.onnx-Pakets. TASO kann PyTorch-Modelle im ONNX-Format direkt optimieren.

Optimieren Sie beliebige DNN-Modelle mithilfe der Python-Schnittstelle

TASO kann mithilfe der TASO-Python-Schnittstelle auch beliebige DNN-Architekturen optimieren. Der folgende Codeausschnitt erstellt das in der Abbildung dargestellte DNN-Diagramm ganz links. TASO führt automatisch eine Reihe nicht trivialer Transformationen durch und erkennt schließlich das DNN-Diagramm ganz rechts, das auf einer V100-GPU 1,3-mal schneller ist. Weitere DNN-Beispiele finden Sie im 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" )

Veröffentlichung

• Zhihao Jia, Oded Padon, James Thomas, Todd Warszawski, Matei Zaharia und Alex Aiken. TASO: Optimierung der Deep-Learning-Berechnung durch automatisierte Generierung von Graphsubstitutionen. In Proceedings of the Symposium on Operating Systems Principles (SOSP), Ontario, Kanada, Oktober 2019.

• Zhihao Jia, James Thomas, Todd Warszawski, Mingyu Gao, Matei Zaharia und Alex Aiken. Optimierung der DNN-Berechnung mit entspannten Graphsubstitutionen. In Proceedings of the Conference on Systems and Machine Learning (SysML), Palo Alto, CA, April 2019.