TASO：用于深度学习的张量代数超级优化器

TASO 使用自动生成和验证的图转换来优化 DNN 模型的计算图。对于任意 DNN 模型，TASO 使用自动生成的图转换来构建与原始 DNN 模型等效的潜在计算图的大型搜索空间。 TASO 采用基于成本的搜索算法来探索空间，并自动发现高度优化的计算图。 TASO 的性能比现有深度学习框架中的图形优化器高出 3 倍。

NVIDIA V100 GPU 上的端到端推理性能比较。

安装TASO

请参阅从源安装 TASO 的说明。我们还提供预构建的 docker 镜像，并预安装了所有依赖项。

使用TASO

TASO 可以直接优化 ONNX、TensorFlow 和 PyTorch 图形格式的任何预训练 DNN 模型。 TASO 还提供了一个用于优化任意 DNN 架构的 Python 接口。 TASO 支持将优化后的计算图导出到 ONNX，可以直接用作大多数现有深度学习框架的输入。

优化 ONNX 模型

TASO 可以直接优化预训练的 ONNX 模型，并且只需几行 Python 代码即可完成。以下代码片段展示了如何从 ONNX 加载预训练的 DNN 模型、优化模型并将优化后的模型保存到 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" )

优化后的模型与原始模型具有相同的精度，并且可以直接被现有的深度学习框架使用。 onnx文件夹中提供了一些原始和 TASO 优化的 ONNX 文件。

优化 TensorFlow 模型

TASO 可以通过使用 tf2onnx 将模型转换为 ONNX 来优化 TensorFlow 模型。

首先，从 PyPi 安装tf2onnx ，如下所示或从源代码安装。

 pip install -U tf2onnx

其次，使用tf2onnx将 TensorFlow 模型转换为 ONNX。

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

第三，按照上述说明使用 TASO 在 ONNX 中优化模型。

优化 PyTorch 模型

PyTorch 作为 torch.onnx 包的一部分内置了对 ONNX 的支持。 TASO 可以直接优化 ONNX 格式的 PyTorch 模型。

使用 Python 接口优化任意 DNN 模型

TASO 还可以使用 TASO Python 接口优化任意 DNN 架构。以下代码片段构建了图中最左侧的 DNN 图。 TASO 自动执行一系列重要的转换，并最终发现最右侧的 DNN 图，该图在 V100 GPU 上速度提高了 1.3 倍。 examples夹中提供了更多 DNN 示例。

 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 年 10 月。
贾志浩、詹姆斯·托马斯、托德·华沙斯基、高明宇、马泰·扎哈里亚和亚历克斯·艾肯。使用宽松图替换优化 DNN 计算。系统和机器学习会议 (SysML) 会议记录，加利福尼亚州帕洛阿尔托，2019 年 4 月。

展开

附加信息