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

TASO 使用自动生成并经过验证的图变换来优化深度神经网络（DNN）模型的计算图。 对于任意 DNN 模型，TASO 利用自动生成的图变换构建一个庞大的搜索空间，其中包含与原始 DNN 模型等价的潜在计算图。 TASO 采用基于代价的搜索算法来探索这一空间，并自动发现高度优化的计算图。 TASO 的性能优于现有深度学习框架中的图优化器，最高可提升至 3 倍。

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

安装 TASO

请参阅 安装说明，以从源代码安装 TASO。 我们还提供了预构建的 Docker 镜像，其中已预先安装所有依赖项。

使用 TASO

TASO 可直接优化以 ONNX、TensorFlow 和 PyTorch 图格式表示的任何预训练 DNN 模型。 TASO 还提供了一个 Python 接口，用于优化任意 DNN 架构。 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 将 TensorFlow 模型转换为 ONNX 格式来进行优化。

• 首先，按照以下方式或从 源代码 安装 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 内置了对 ONNX 的支持，可通过 torch.onnx 包实现。 TASO 可直接优化以 ONNX 格式表示的 PyTorch 模型。

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

TASO 还可以通过其 Python 接口优化任意 DNN 架构。 以下代码片段构建了图中所示的最左侧 DNN 计算图。TASO 自动执行一系列复杂的变换，最终发现最右侧的 DNN 计算图，其在 V100 GPU 上的速度提升了 1.3 倍。更多 DNN 示例可在 examples 文件夹中找到。

import taso
import onnx

# 构建 DNN 模型
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)

# 优化 DNN 模型
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 月。

TASO 快速上手指南

TASO (The Tensor Algebra SuperOptimizer) 是一款专为深度学习设计的计算图超级优化器。它通过自动生成和验证的图变换技术，自动探索等效的计算图空间，从而发现比现有框架（如 TensorFlow、PyTorch 内置优化器）性能高出多达 3 倍的高效计算图。

环境准备

• 操作系统: Linux (推荐 Ubuntu 18.04 或更高版本)
• 硬件要求: NVIDIA GPU (测试环境为 V100，支持 CUDA)
• 前置依赖:
  • Python 3.x
  • CMake, GCC/G++
  • ONNX, Protobuf
  • (可选) Docker：如果你希望避免手动配置依赖，官方提供了预装所有依赖的 Docker 镜像。

提示：国内用户建议使用清华源或阿里源加速 pip 和 apt 包的安装。

安装步骤

你可以选择从源码编译或使用预构建的 Docker 镜像。

方式一：使用 Docker（推荐，最快捷）

直接拉取包含所有依赖的镜像，无需手动配置环境：

docker pull jiazhihao/taso:latest
docker run --gpus all -it jiazhihao/taso:latest /bin/bash

方式二：从源码安装

如果需要在本地环境部署，请参照以下步骤：

1. 克隆仓库：

git clone https://github.com/jiazhihao/TASO.git
cd TASO

2. 安装依赖并编译（详细步骤请参考项目根目录下的 INSTALL.md）：

# 示例：安装 Python 依赖
pip install -r requirements.txt

# 构建项目
mkdir build && cd build
cmake ..
make -j$(nproc)

基本使用

TASO 支持直接优化 ONNX、TensorFlow 和 PyTorch 模型。核心工作流是：加载模型 -> 执行优化 -> 导出为 ONNX 格式（可直接被主流框架加载推理）。

场景 1：优化现有的 ONNX 模型

这是最简单的用法，只需几行代码即可完成优化。

import taso
import onnx

# 1. 加载预训练的 ONNX 模型
old_model = taso.load_onnx("/path/to/load/onnx/model")

# 2. 执行 TASO 优化
taso_graph = taso.optimize(old_model)

# 3. 导出优化后的模型为 ONNX 格式
new_model = taso.export_onnx(taso_graph)
onnx.save(new_model, "/path/to/save/new/onnx/model")

注：优化后的模型精度与原模型保持一致，可直接用于推理。

场景 2：优化 TensorFlow 或 PyTorch 模型

TASO 原生支持 ONNX 格式。对于 TF 或 PyTorch 模型，需先转换为 ONNX，再按上述步骤优化。

TensorFlow 转换步骤：

1. 安装转换工具：

pip install -U tf2onnx

2. 转换为 ONNX：

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

3. 使用 TASO 优化生成的 ONNX 文件（参考场景 1）。

PyTorch 转换步骤： 利用 torch.onnx 将模型导出为 ONNX，随后同样参考场景 1 进行优化。

场景 3：使用 Python API 构建并优化自定义架构

你也可以直接使用 TASO 的 Python 接口定义网络结构，让 TASO 自动寻找最优算子组合。

import taso
import onnx

# 构建 DNN 模型
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)

# 优化模型
new_graph = taso.optimize(graph)

# 导出结果
onnx_model = taso.export_onnx(new_graph)
onnx.save(onnx_model, "/path/to/save/new/onnx/model")