TASO

TASO: El superoptimizador de álgebra tensorial para el aprendizaje profundo

TASO optimiza los gráficos de cálculo de los modelos DNN utilizando transformaciones de gráficos generadas y verificadas automáticamente. Para un modelo DNN arbitrario, TASO utiliza las transformaciones de gráficos generadas automáticamente para construir un gran espacio de búsqueda de gráficos de cálculo potenciales que son equivalentes al modelo DNN original. TASO emplea un algoritmo de búsqueda basado en costos para explorar el espacio y descubre automáticamente gráficos de cálculo altamente optimizados. TASO supera a los optimizadores de gráficos en los marcos de aprendizaje profundo existentes hasta 3 veces.

Comparación del rendimiento de inferencia de un extremo a otro en una GPU NVIDIA V100.

Instalar TASO

Consulte las instrucciones para instalar TASO desde la fuente. También proporcionamos imágenes de Docker prediseñadas con todas las dependencias preinstaladas.

Utilice TASO

TASO puede optimizar directamente cualquier modelo DNN previamente entrenado en formatos de gráficos ONNX, TensorFlow y PyTorch. TASO también proporciona una interfaz Python para optimizar arquitecturas DNN arbitrarias. TASO admite la exportación de gráficos de cálculo optimizados a ONNX, que pueden usarse directamente como entradas para la mayoría de los marcos de aprendizaje profundo existentes.

Optimizar modelos ONNX

TASO puede optimizar directamente los modelos ONNX previamente entrenados, y esto se puede hacer con solo unas pocas líneas de código Python. El siguiente fragmento de código muestra cómo cargar un modelo DNN previamente entrenado desde ONNX, optimizar el modelo y guardar el modelo optimizado en un archivo 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" )

El modelo optimizado tiene la misma precisión que el original y puede ser utilizado directamente por los marcos de aprendizaje profundo existentes. Algunos archivos ONNX originales y optimizados para TASO están disponibles en la carpeta onnx .

Optimizar los modelos de TensorFlow

TASO puede optimizar los modelos de TensorFlow convirtiendo el modelo a ONNX usando tf2onnx.

• Primero, instale tf2onnx desde PyPi de la siguiente manera o desde la fuente.

 pip install -U tf2onnx

• En segundo lugar, convierta un modelo de TensorFlow a ONNX usando tf2onnx .

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

• En tercer lugar, utilice TASO para optimizar el modelo en ONNX siguiendo las instrucciones anteriores.

Optimizar modelos de PyTorch

PyTorch tiene soporte integrado para ONNX como parte del paquete torch.onnx. TASO puede optimizar directamente los modelos PyTorch en formato ONNX.

Optimice los modelos DNN arbitrarios utilizando la interfaz Python

TASO también puede optimizar arquitecturas DNN arbitrarias utilizando la interfaz TASO Python. El siguiente fragmento de código crea el gráfico DNN situado más a la izquierda que se muestra en la figura. TASO realiza automáticamente una serie de transformaciones no triviales y, finalmente, descubre el gráfico DNN más a la derecha, que es 1,3 veces más rápido en una GPU V100. Hay más ejemplos de DNN disponibles en la carpeta de 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" )

Publicación

• Zhihao Jia, Oded Padon, James Thomas, Todd Warszawski, Matei Zaharia y Alex Aiken. TASO: Optimización de la computación de aprendizaje profundo con generación automatizada de sustituciones de gráficos. En Actas del Simposio sobre principios de sistemas operativos (SOSP), Ontario, Canadá, octubre de 2019.

• Zhihao Jia, James Thomas, Todd Warszawski, Mingyu Gao, Matei Zaharia y Alex Aiken. Optimización del cálculo DNN con sustituciones de gráficos relajados. En Actas de la Conferencia sobre Sistemas y Aprendizaje Automático (SysML), Palo Alto, CA, abril de 2019.