TASO

TASO : le superoptimiseur d'algèbre tensorielle pour l'apprentissage profond

TASO optimise les graphiques de calcul des modèles DNN à l'aide de transformations graphiques générées et vérifiées automatiquement. Pour un modèle DNN arbitraire, TASO utilise les transformations graphiques générées automatiquement pour créer un grand espace de recherche de graphiques de calcul potentiels équivalents au modèle DNN d'origine. TASO utilise un algorithme de recherche basé sur les coûts pour explorer l'espace et découvre automatiquement des graphiques de calcul hautement optimisés. TASO surpasse jusqu'à 3 fois les optimiseurs de graphiques dans les cadres d'apprentissage en profondeur existants.

Comparaison des performances d'inférence de bout en bout sur un GPU NVIDIA V100.

Installer TASO

Voir les instructions pour installer TASO à partir des sources. Nous fournissons également des images Docker prédéfinies avec toutes les dépendances préinstallées.

Utiliser TASO

TASO peut optimiser directement tous les modèles DNN pré-entraînés aux formats graphiques ONNX, TensorFlow et PyTorch. TASO fournit également une interface Python pour optimiser les architectures DNN arbitraires. TASO prend en charge l'exportation des graphiques de calcul optimisés vers ONNX, qui peuvent être directement utilisés comme entrées par la plupart des frameworks d'apprentissage en profondeur existants.

Optimiser les modèles ONNX

TASO peut optimiser directement les modèles ONNX pré-entraînés, et cela peut être fait en quelques lignes de code Python seulement. L'extrait de code suivant montre comment charger un modèle DNN pré-entraîné à partir d'ONNX, optimiser le modèle et enregistrer le modèle optimisé dans un fichier 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" )

Le modèle optimisé a la même précision que l'original et peut être directement utilisé par les frameworks d'apprentissage profond existants. Certains fichiers ONNX originaux et optimisés pour TASO sont disponibles dans le dossier onnx .

Optimiser les modèles TensorFlow

TASO peut optimiser les modèles TensorFlow en convertissant le modèle en ONNX à l'aide de tf2onnx.

• Tout d’abord, installez tf2onnx depuis PyPi comme suit ou depuis les sources.

 pip install -U tf2onnx

• Deuxièmement, convertissez un modèle TensorFlow en ONNX à l'aide de tf2onnx .

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

• Troisièmement, utilisez TASO pour optimiser le modèle dans ONNX en suivant les instructions ci-dessus.

Optimiser les modèles PyTorch

PyTorch prend en charge ONNX dans le cadre du package torch.onnx. TASO peut directement optimiser les modèles PyTorch au format ONNX.

Optimiser les modèles DNN arbitraires à l'aide de l'interface Python

TASO peut également optimiser les architectures DNN arbitraires à l'aide de l'interface TASO Python. L'extrait de code suivant crée le graphique DNN le plus à gauche représenté dans la figure. TASO effectue automatiquement une série de transformations non triviales et découvre finalement le graphe DNN le plus à droite, qui est 1,3 fois plus rapide sur un GPU V100. D'autres exemples DNN sont disponibles dans le dossier 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" )

Publication

• Zhihao Jia, Oded Padon, James Thomas, Todd Warszawski, Matei Zaharia et Alex Aiken. TASO : Optimisation du calcul d'apprentissage en profondeur avec la génération automatisée de substitutions de graphiques. Dans Actes du Symposium sur les principes des systèmes d'exploitation (SOSP), Ontario, Canada, octobre 2019.

• Zhihao Jia, James Thomas, Todd Warszawski, Mingyu Gao, Matei Zaharia et Alex Aiken. Optimisation du calcul DNN avec des substitutions de graphiques détendues. Dans Actes de la conférence sur les systèmes et l'apprentissage automatique (SysML), Palo Alto, Californie, avril 2019.