TASO

TASO: O SuperOptimizer de Álgebra Tensor para Aprendizado Profundo

A TASO otimiza os gráficos de cálculo dos modelos DNN usando transformações gráficas geradas e verificadas automaticamente. Para um modelo DNN arbitrário, a TASO usa as transformações gráficas geradas automaticamente para construir um grande espaço de busca de gráficos de computação potenciais que são equivalentes ao modelo DNN original. A TASO emprega um algoritmo de busca baseado em custos para explorar o espaço e descobre automaticamente gráficos de computação altamente otimizados. A TASO supera os otimizadores gráficos em estruturas de aprendizagem profunda existentes em até 3x.

Comparação de desempenho de inferência ponta a ponta em uma GPU NVIDIA V100.

Instale TASO

Veja as instruções para instalar o TASO a partir do código-fonte. Também fornecemos imagens docker pré-construídas com todas as dependências pré-instaladas.

Usar TASO

A TASO pode otimizar diretamente qualquer modelo DNN pré-treinado nos formatos de gráfico ONNX, TensorFlow e PyTorch. A TASO também fornece uma interface Python para otimizar arquiteturas DNN arbitrárias. A TASO suporta a exportação de gráficos de computação otimizados para ONNX, que podem ser usados ​​diretamente como entradas pela maioria das estruturas de aprendizagem profunda existentes.

Otimize modelos ONNX

A TASO pode otimizar diretamente modelos ONNX pré-treinados, e isso pode ser feito em apenas algumas linhas de código Python. O trecho de código a seguir mostra como carregar um modelo DNN pré-treinado do ONNX, otimizar o modelo e salvar o modelo otimizado em um arquivo 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" )

O modelo otimizado tem a mesma precisão do original e pode ser usado diretamente pelas estruturas de aprendizagem profunda existentes. Alguns arquivos ONNX originais e otimizados para TASO estão disponíveis na pasta onnx .

Otimize modelos do TensorFlow

A TASO pode otimizar modelos TensorFlow convertendo o modelo para ONNX usando tf2onnx.

• Primeiro, instale tf2onnx do PyPi como segue ou da fonte.

 pip install -U tf2onnx

• Em segundo lugar, converta um modelo TensorFlow em ONNX usando tf2onnx .

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

• Terceiro, use o TASO para otimizar o modelo no ONNX seguindo as instruções acima.

Otimize modelos PyTorch

PyTorch tem suporte integrado para ONNX como parte do pacote torch.onnx. A TASO pode otimizar diretamente os modelos PyTorch no formato ONNX.

Otimize modelos DNN arbitrários usando a interface Python

A TASO também pode otimizar arquiteturas DNN arbitrárias usando a interface TASO Python. O trecho de código a seguir cria o gráfico DNN mais à esquerda representado na figura. A TASO executa automaticamente uma série de transformações não triviais e, eventualmente, descobre o gráfico DNN mais à direita, que é 1,3x mais rápido em uma GPU V100. Mais exemplos de DNN estão disponíveis na pasta 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" )

Publicação

• Zhihao Jia, Oded Padon, James Thomas, Todd Warszawski, Matei Zaharia e Alex Aiken. TASO: Otimizando Computação de Aprendizado Profundo com Geração Automatizada de Substituições de Gráficos. Em Anais do Simpósio sobre Princípios de Sistemas Operacionais (SOSP), Ontário, Canadá, outubro de 2019.

• Zhihao Jia, James Thomas, Todd Warszawski, Mingyu Gao, Matei Zaharia e Alex Aiken. Otimizando a computação DNN com substituições relaxadas de gráficos. Em Anais da Conferência sobre Sistemas e Aprendizado de Máquina (SysML), Palo Alto, CA, abril de 2019.