TASO: супероптимизатор тензорной алгебры для глубокого обучения

TASO оптимизирует графы вычислений моделей DNN, используя автоматически генерируемые и проверенные преобразования графов. Для произвольной модели DNN TASO использует автоматически сгенерированные преобразования графов для построения большого пространства поиска потенциальных графов вычислений, которые эквивалентны исходной модели DNN. TASO использует алгоритм поиска на основе затрат для исследования пространства и автоматически обнаруживает высокооптимизированные графы вычислений. TASO превосходит оптимизаторы графов в существующих средах глубокого обучения почти в 3 раза.

Комплексное сравнение производительности вывода на графическом процессоре NVIDIA V100.

Установить ТАСО

См. инструкции по установке TASO из исходного кода. Мы также предоставляем готовые образы Docker со всеми предустановленными зависимостями.

Используйте ТАСО

TASO может напрямую оптимизировать любые предварительно обученные модели DNN в графических форматах ONNX, TensorFlow и PyTorch. TASO также предоставляет интерфейс Python для оптимизации произвольных архитектур DNN. TASO поддерживает экспорт оптимизированных графиков вычислений в ONNX, которые могут напрямую использоваться в качестве входных данных для большинства существующих платформ глубокого обучения.

Оптимизация моделей 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" )

Оптимизированная модель имеет ту же точность, что и исходная, и может напрямую использоваться существующими платформами глубокого обучения. Некоторые оригинальные и оптимизированные для TASO файлы ONNX доступны в папке onnx .

Оптимизация моделей TensorFlow

TASO может оптимизировать модели TensorFlow, преобразовав их в ONNX с помощью tf2onnx.

Сначала установите tf2onnx из PyPi следующим образом или из исходного кода.

 pip install -U tf2onnx

Во-вторых, преобразуйте модель TensorFlow в ONNX с помощью tf2onnx .

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

В-третьих, используйте TASO для оптимизации модели в ONNX, следуя приведенным выше инструкциям.

Оптимизация моделей PyTorch

PyTorch имеет встроенную поддержку ONNX как часть пакета torch.onnx. TASO может напрямую оптимизировать модели PyTorch в формате ONNX.

Оптимизируйте произвольные модели DNN с помощью интерфейса Python

TASO также может оптимизировать произвольные архитектуры DNN, используя интерфейс TASO Python. Следующий фрагмент кода строит крайний левый график DNN, изображенный на рисунке. TASO автоматически выполняет серию нетривиальных преобразований и в конечном итоге обнаруживает самый правый график DNN, который на графическом процессоре V100 работает в 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 г.
Чжихао Цзя, Джеймс Томас, Тодд Варшавски, Мингю Гао, Матей Захария и Алекс Эйкен. Оптимизация вычислений DNN с помощью расслабленных замен графов. В материалах конференции по системам и машинному обучению (SysML), Пало-Альто, Калифорния, апрель 2019 г.

Расширять

Дополнительная информация