ONNX:开放神经网络交换格式
一场关于人工智能灵魂的开源战争正在进行。 工业巨头、大学和全球机器学习研究人员社区都在参与这场战争。 本文记录了这场战斗中的一场小规模冲突:神经网络的标准化文件格式。 关键在于在众多工具之间进行开放数据交换,而不是相互竞争的单一框架。
好消息是,战场是自由和开放的。 没有一家大型厂商在推动闭源解决方案。 无论是 Google 支持的 Keras 和 Tensorflow、Amazon 认可的 Apache MXNet,还是 Facebook 支持的 Caffe2 或 PyTorch,所有解决方案都是开源软件。
不幸的是,虽然这些项目是开放的,但它们不是可互操作的。 每个框架都构成一个完整的堆栈,直到最近都无法以任何方式与其他框架进行交互。 一个新的行业支持的标准,开放神经网络交换格式,可能会改变这种情况。
现在,想象一下这样一个世界:您可以在 Keras 中训练一个神经网络,通过 NNVM 优化编译器运行训练好的模型,并将其部署到 MXNet 的生产环境中。 想象一下,这只是无数种可互操作的深度学习工具组合中的一种,包括可视化工具、性能分析器和优化器。 研究人员和 DevOps 不再需要在提供平庸建模环境和一般部署性能的单一工具链上妥协。
所需的是一种标准化格式,它可以表达任何机器学习模型并存储训练后的参数和权重,并且可以被一套独立开发的软件读取和写入。
这就是 开放神经网络交换格式 (ONNX)。
愿景要理解像 ONNX 这样的标准对互操作性的迫切需求,我们首先必须理解我们对现有单一框架的荒谬要求。
深度学习框架的普通用户可能会认为它是一种用于指定神经网络的语言。 例如,我想要 100 个输入神经元,三个全连接层,每层有 50 个 ReLU 输出,以及输出端的 softmax。 我选择的框架有一种领域语言来指定这一点(如 Caffe),或者绑定到像 Python 这样的语言,并具有清晰的 API。
然而,网络架构的规范仅仅是冰山一角。 一旦定义了网络结构,框架仍然需要做大量复杂的工作才能使其在您的 CPU 或 GPU 集群上运行。
显然,Python 无法在 GPU 上运行。 为了使您的网络定义在 GPU 上运行,它需要被编译成 CUDA (NVIDIA) 或 OpenCL (AMD 和 Intel) API 的代码,或者在 CPU 上运行时以高效的方式进行处理。 这种编译很复杂,这也是为什么大多数框架不支持 NVIDIA 和 AMD GPU 后端的原因。
但工作尚未完成。 您的框架还必须平衡您正在使用的硬件的资源分配和并行性。 您是在拥有 3,000 多个计算核心的 Titan X 卡上运行,还是在核心数量不到一半的 GTX 1060 上运行? 您的显卡有 16GB 内存还是只有 4GB? 所有这些都会影响计算的优化和运行方式。
更糟糕的是,您是否有 50 台多 GPU 机器的集群来训练您的网络? 您的框架也需要处理这个问题。 网络协议、高效分配、参数共享——您对单个框架的要求有多高?
现在您说您想部署到生产环境? 您希望自动扩展您的集群? 您想要一种具有安全 API 的可靠语言?
当您把所有这些加起来时,要求一个单一的整体项目来处理所有这些需求似乎绝对是疯狂的。 您不能期望编写完美网络定义语言的作者与在 Kubernetes 中集成部署系统或编写最佳 CUDA 编译器的作者是同一批人。
ONNX 的目标是打破单一框架。 让贡献者生态系统开发这些组件中的每一个,并通过通用的规范格式将它们粘合在一起。
生态系统(和政治)互操作性是开放生态系统的健康标志。 不幸的是,直到最近,深度学习领域还不存在互操作性。 每个框架都有自己的格式来存储计算图和训练好的模型。
去年年底,这种情况开始改变。 开放神经网络交换格式倡议由 Facebook、Amazon 和 Microsoft 发起,并得到了 AMD、ARM、IBM、Intel、华为、NVIDIA 和 Qualcomm 的支持。 让我换句话说,除了 Google 之外的所有人。 该格式已包含在大多数知名的框架中,除了 Google 的 TensorFlow(为此存在第三方转换器)。
这似乎是典型的场景,即明显的市场领导者 Google 对于为了开放性而颠覆其主导地位几乎没有兴趣。 较小的参与者正在联合起来对抗这只 500 磅重的大猩猩。
Google 致力于其自己的 TensorFlow 模型和权重文件格式 SavedModel,它与 ONNX 共享许多功能。 Google 正在围绕该格式构建自己的生态系统,包括 TensorFlow Server、Estimator 和 Tensor2Tensor 等。
ONNX 解决方案构建一个可以表达所有深度学习框架的所有功能的单一文件格式并非易事。 您如何描述卷积或具有内存的循环网络? 注意力机制? Dropout 层? 快速文本或 StarSpace 中发现的嵌入和最近邻算法呢?
ONNX 从 TensorFlow 中借鉴了一个想法,并声明一切都是张量运算图。 然而,仅凭这一声明是不够的。 必须支持数十甚至数百种操作,并非所有操作都将得到所有其他工具和框架的支持。 某些框架也可能以不同于其同类框架的方式实现操作。
ONNX 社区就张量运算应建模在哪个级别上进行了相当多的辩论。 ONNX 应该是一个可以支持具有正弦和乘法等原语的任意方程的数学工具箱,还是应该支持更高级别的结构,例如集成的 GRU 单元或 Layer Normalization 作为单一的整体操作?
就目前而言,ONNX 目前定义了大约 100 种操作。 它们的复杂程度从算术加法到完整的长短期记忆实现不等。 并非所有工具都支持所有操作,因此仅仅因为您可以生成模型的 ONNX 文件并不意味着它可以在任何地方运行。
在某些框架中,生成 ONNX 模型文件也可能很麻烦,因为它依赖于图结构中操作顺序的严格定义。 例如,PyTorch 在定义模型时以非常 Pythonic 的命令式体验而自豪。 您可以使用 Python 逻辑来布置模型的流程,但您不会像 TensorFlow 等其他框架那样定义严格的图结构。 因此,没有操作图可以保存; 您实际上必须运行模型并跟踪操作。 操作的跟踪会保存到 ONNX 文件中。
结论深度学习互操作性还处于早期阶段。 大多数用户仍然选择一个框架并坚持使用它。 越来越多的用户正在使用 TensorFlow。 Google 为此投入了大量资源和实际生产经验——这很难抗拒。
所有框架在某些领域都很强大,而在另一些领域则很薄弱。 每个新框架都必须重新实现功能的完整“堆栈”。 打破堆栈,您就可以发挥各个工具的优势。 这将带来更健康的生态系统。
ONNX 是朝着正确方向迈出的一步。
注意:ONNX GitHub 页面是 这里。
免责声明Braddock Gaskill 是 eBay Inc. 的一名研究科学家。 他以个人身份为本文做出了贡献。 所表达的观点是他自己的观点,不一定代表 eBay Inc. 的观点。
关于作者Braddock Gaskill 在 AI 和算法软件开发方面拥有 25 年的经验。 他还共同创立了 Internet-in-a-Box 开源项目,并开发了 libre Humane Wikipedia Reader,用于将内容传递给发展中国家的学生。