开源一年多的模型交换格式ONNX，已经一统框架江湖了？
机器之心原创
作者：思源
近日，微软亚洲研究院和华为举办了 ONNX 合作伙伴研讨会，这是 ONNX 开源社区成立以来首次在中国举办的活动。在研讨会中，微软、Facebook、华为和英特尔等的开发者介绍了他们在 ONNX 上的开源贡献及思考。
在过去的一年多中，ONNX 这种「通用」的神经网络交换格式已经有了很长远的发展，用不同框架编写的模型可以在不同的平台中流通。在这次研讨会中，我们确切地感受到了这一点，因为开源社区围绕着 ONNX 介绍了很多优化工具和资源库。
微软上个月开源了 ONNX Runtime，其专为 ONNX 格式的模型设计了高性能推理引擎。Facebook 早两个月前开源了 ONNXIFI，其为 ONNX 提供了用于框架集成的接口，即一组用于加载和执行 ONNX 计算图的跨平台 API。更早一些，英特尔在今年 3 月份就开源 nGraph，它能编译 ONNX 格式的模型，并在 CPU 或 GPU 等硬件加速模型的运行。
而到了昨天，微软又开源了 ONNX.JS，它是一种在浏览器和 Node.js 上运行 ONNX 模型的 JavaScript 库。它部署的模型效率非常高，且能实现交互式的直观推理。该开源项目给出了图像分类的交互式演示，且在 Chrome 浏览器和 CPU 下比 TensorFlow.JS 快了近 8 倍，后文将详细介绍这一开源库。
当然除了这些开源工作，ONNX 社区还有更多的实践，例如如何部署 ONNX 模型到边缘设备、如何维护一个包罗万象的 ONNX Model Zoo 等。本文主要从什么是 ONNX、怎样用 ONNX，以及如何优化 ONNX 三方面看看 ONNX 是不是已经引领「框架间的江湖」了。
什么是 ONNX
很多开发者在玩 GitHub 的时候都有这样「悲痛」的经历，好不容易找到令人眼前一亮的项目，然而发现它使用我们不熟悉的框架写成。其实我们会发现很多优秀的视觉模型是用 Caffe 写的，很多新的研究论文是用 PyTorch 写的，而更多的模型用 TensorFlow 写成。因此如果我们要测试它们就必须拥有对应的框架环境，但 ONNX 交换格式令我们在同一环境下测试不同模型有了依靠。
简而言之 ONNX 就是一种框架间的转换格式，例如我们用 TensorFlow 写的模型可以转换为 ONNX 格式，并在 Caffe2 环境下运行该模型。
- 项目地址：https://github.com/onnx/onnx
ONNX 定义了一种可扩展的计算图模型、一系列内置的运算单元（OP）和标准数据类型。每一个计算流图都定义为由节点组成的列表，并构建有向无环图。其中每一个节点都有一个或多个输入与输出，每一个节点称之为一个 OP。这相当于一种通用的计算图，不同深度学习框架构建的计算图都能转化为它。
如下所示，目前 ONNX 已经支持大多数框架，使用这些框架构建的模型可以转换为通用的 ONNX 计算图和 OP。现阶段 ONNX 只支持推理，所以导入的模型都需要在原框架完成训练。
其中 Frameworks 下的框架表示它们已经内嵌了 ONNX，开发者可以直接通过这些框架的内置 API 将模型导出为 ONNX 格式，或采用它们作为推理后端。而 Converters 下的框架并不直接支持 ONNX 格式，但是可以通过转换工具导入或导出这些框架的模型。
其实并不是所有框架都支持导入和导出 ONNX 格式的模型，有一些并不支持导入 ONNX 格式的模型，例如 PyTorch 和 Chainer 等，TensorFlow 的 ONNX 导入同样也正处于实验阶段。下图展示了各框架对 ONNX 格式的支持情况：
怎样使用 ONNX
对于内建了 ONNX 的框架而言，使用非常简单，只需要调用 API 导出或导入已训练模型就可以了。例如对 PyTorch 而言，只需要几个简单的步骤就能完成模型的导出和导入。简单而言，首先加载 torch.onnx 模块，然后导出预训练模型并查看模型结构信息，最后再将导出的 ONNX 模型加载到另外的框架就能执行推理了。
from torch.autograd import Variable
import torch.onnx
import torchvision
dummy_input = Variable(torch.randn(10, 3, 224, 224)).cuda()
model = torchvision.models.alexnet(pretrained=True).cuda()
input_names = [ "actual_input_1" ] + [ "learned_%d" % i for i in range(16) ]
output_names = [ "output1" ]
torch.onnx.export(model, dummy_input, "alexnet.onnx", verbose=True, input_names=input_names, output_names=output_names)如上所示将导出 ONNX 格式的 AlexNet 模型，其中"alexnet.onnx"为保存的模型，input_names、output_names 和 verbose=True 都是为了打印出模型结构信息。同样随机产生的「图像」dummy_input 也是为了了解模型结构，因为我们可以通过它理解输入与每一层具体的参数维度。以下展示了 ONNX 输出的简要模型信息：
graph(%actual_input_1 : Float(10, 3, 224, 224)
%learned_0 : Float(64, 3, 11, 11)
%learned_1 : Float(64)
# ---- omitted for brevity ----
%learned_14 : Float(1000, 4096)
%learned_15 : Float(1000)) {
%17 : Float(10, 64, 55, 55) = onnx::Conv[dilations=[1, 1], group=1, kernel_shape=[11, 11], pads=[2, 2, 2, 2], strides=[4, 4]](%actual_input_1, %learned_0, %learned_1), scope: AlexNet/Sequential[features]/Conv2d[0]
%18 : Float(10, 64, 55, 55) = onnx::Relu(%17), scope: AlexNet/Sequential[features]/ReLU[1]
%19 : Float(10, 64, 27, 27) = onnx::MaxPool[kernel_shape=[3, 3], pads=[0, 0, 0, 0], strides=[2, 2]](%18), scope: AlexNet/Sequential[features]/MaxPool2d[2]
# ---- omitted for brevity ----
%output1 : Float(10, 1000) = onnx::Gemm[alpha=1, beta=1, broadcast=1, transB=1](%45, %learned_14, %learned_15), scope: AlexNet/Sequential[classifier]/Linear[6]
return (%output1);
}其实我们也可以借助 ONNX 检查中间表征，不过这里并不介绍。后面加载另外一个框架并执行推理同样非常简单。如下所示，我们可以从 caffe2 中加载 ONNX 的后端，并将前面保存的模型加载到该后端，从而在新框架下进行推理。这里我们能选择执行推理的硬件，并直接推理得出输出结果。
import caffe2.python.onnx.backend as backend
import numpy as np
import onnx