轻松调用TensorFlow、PyTorch等多框架，Uber开源深度学习推理引擎Neuropod

近日，Uber ATG开源了深度学习推理引擎Neuropod，对外只提供一个通用接口用于运行深度学习模型，而底层可以调用不同的深度学习框架，如TensorFlow、PyTorch、Caffe2、Keras等。这将有助于企业在异构环境中大规模部署多个DL/ML模型。

在优步先进技术团队（Uber Advanced Technologies Group，ATG），我们利用深度学习来提供安全可靠的自动驾驶技术。利用深度学习，我们可以建立并训练模型来处理诸如传感器输入、识别对象以及预测这些对象可能去的地方等任务。

随着自动驾驶软件的发展，我们一直在寻找新的方法来改进模型。有时候，这意味着要试验不同的深度学习框架。随着新的深度学习框架的发布，以及 TensorFlow 和 PyTorch 等现有框架取得的进步，我们希望确保工程师和科学家能够灵活地使用最适合他们正在研究的问题的工具。

然而不幸的是，要在整个机器学习栈中添加对一个新的深度学习框架的支持，既费资源，又费时间。在 ATG，我们已经花了大量时间来简化这个过程。今天，我们通过引入 Neuropod 将这些工作成果予以开源。Neuropod 是现有深度学习框架之上的一个抽象层，它提供了一个统一的接口来运行深度学习模型。Neuropod 让研究人员可以轻松地在自己选择的框架中构建模型，同时也简化了这些模型的生产化过程。

使用多种深度学习框架

深度学习的发展非常迅速，不同的深度学习框架在不同的任务上有效。因此，在过去的几年里，Uber ATG 使用过多种深度学习框架。2016 年，Caffe2 是我们主要的深度学习框架；2017 年初，我们投入了大量工作来集成 TensorFlow。这涉及到与 CUDA 和 cuDNN 的主要集成障碍、Caffe2 和 TensorFlow 的依赖冲突、库加载问题等。2017 年底，我们开始在 PyTorch 上开发更多的模型。生产这些模型本身就需要大量的工作，而在与 TensorFlow 同时运行时，我们还发现了内存损坏问题，以及其他几个非常难以调试的问题。然后，TorchScript 发布了，我们又一次经历了类似的过程。

图 1. 多年来，Uber ATG 利用不同的流行深度学习框架，发展了其机器学习方法。

自 2018 年以来，有许多深度学习框架已经开源，包括 Ludwig（由 Uber 创建）、JAX、Trax、Flax、Haiku 和 RLax，其中许多框架都是过去九个月内发布的。

即便研究人员可以很容易地试验新框架，但在我们所有的系统和流程中添加对新深度学习框架的生产支持，也是一项艰巨的任务。其中有部分是因为，它需要对基础设施和工具的每个部分逐一进行集成和优化。

图 2. 添加对新框架的支持很困难，因为运行模型的每个基础设施都需要支持所有框架。这些基础设施组件可以是度量管道、模型服务或其他生产和测试环境。

在 2018 年末，Uber ATG 开始构建多种模型，尝试用不同的方式解决同一问题。例如，利用光学雷达进行 3D 目标检测，可以采用范围视图或鸟瞰视图。这两种方法都是有效的，但各有优缺点。不同团队构建的模型有时也在不同的框架中实现。

为了使生产化更容易，我们希望能够轻松地替换解决相同问题的模型，即使它们是在不同的框架中实现的。

我们也会遇到一些其他的情况，比如新的研究将基于PyTorch编写代码，但我们想要快速地与 TensorFlow 中的现有模型进行比较。因为每个模型都有特定于框架的、模型级的度量管道，所以比较就很难做到了。

为了在新框架中部署模型，我们需要重建模型级的度量管道，在我们所有的系统和流程中对框架进行集成，然后进行额外的优化，以确保我们能在延迟预算范围内有效地运行模型。

虽然这些步骤看上去很简单，但诸如上述问题（如内存损坏、依赖冲突等）导致我们耗费大量精力来解决，而不能专注于模型开发。

我们需要一种方法来最大程度地提高研究过程中的灵活性，而不必在过程的其他部分重复工作。

Neuropod 介绍

我们针对这一问题的解决方案是 Neuropod，这是一个开源库，它能使所有深度学习框架在运行模型时看起来都是一样的。有了Neuropod，在所有工具和基础设施中添加对新框架的支持，就像将其添加到 Neuropod 中一样简单。

图3. Neuropod 是一个抽象层，提供了一个统一的接口，可以基于多个框架运行深度学习模型。

目前，Neuropod 支持的框架包括：TensorFlow、PyTorch、Keras 和 TorchScript，同时也可以轻松地添加新的框架。

自 2019 年初在内部发布以来，Neuropod 在 Uber 快速部署新模型方面发挥了重要作用。在过去的一年里，我们已经在 Uber ATG、Uber AI 和 Uber 的核心业务部署了数百个 Neuropod 模型。这些模型包括需求预测模型、预计到达时间（Estimated time of arrival，ETA）预测、Uber Eats 的菜单转录以及用于自动驾驶汽车的目标检测模型。随着 Neuropod 的开源发布，我们希望机器学习社区的其他人也会发现它能帮上大忙！

概述

Neuropod 从问题定义的概念开始：模型要解决的“问题”的形式化描述。在这种情况下，问题可能是图像的语义分割或文本的语言翻译。通过形式化地定义问题，我们可以将其视为一个接口，并抽象出具体的实现。每个 Neuropod 模型都实现了一个问题定义。因此，解决相同问题的任何模型都是可互换的，即使它们使用不同的框架。

Neuropod 的工作原理是将现有模型包装在 Neuropod 包（或简称为“Neuropod”）中。此包包含原始模型以及元数据、测试数据和自定义操作（如果有的话）。

使用 Neuropod，任何模型都能以任何支持的语言执行。例如，如果用户想基于 C++运行 PyTorch 模型，Neuropod 将在后台启动一个 Python 解释器，并与其通信来运行模型。这么做是必要的，因为 PyTorch 模型需要 Python 解释器才能运行。这一功能允许我们在将 PyTorch 模型转换为 TorchScript 之前进行快速测试，而 TorchScript 可以在 C++上直接运行。

Neuropod 目前支持基于Python 和 C++运行模型。但是，为库编写额外的语言绑定也很简单。例如，Uber 的机器学习平台 Michelangelo，使用 Neuropod 作为其核心深度学习模型格式，并实现了 Go 绑定来从 Go 运行其生产模型。

图 4. Neuropod 提供了特定于框架的封装 API 和与框架无关的推理 API。

推理概述

在我们深入探讨 Neuropod 的工作原理之前，让我们先看看过去用传统方法是如何将深度学习模型集成到应用程序中的：

图 5. 通常，应用程序在整个推理过程中直接与深度学习框架交互。

在上图中，应用程序在推理过程中的所有部分都是直接与 TensorFlow API 进行交互的。

通过使用 Neuropod，应用程序只与与框架无关的 API 进行交互（下面的所有紫色部分），并且 Neuropod 将这些与框架无关的调用转换为对底层框架的调用。我们尽可能使用零拷贝操作来高效地实现这一点。更多细节请参阅下面的“优化”部分。

图 6. 使用 Neuropod，应用程序可以与框架无关的 API 进行交互，而 Neuropod 与底层框架进行交互。

Neuropod 有一个可插拔的后端层，每个支持的框架都有自己的实现。这使得向 Neuropod 添加新框架变得非常简单。

使用 Neuropod 进行深度学习

让我们看一下使用 Neuropod 时的整个深度学习过程，以了解它是如何帮助简化实验、部署和迭代的。

问题定义

要封装 Neuropod，我们必须首先创建一个问题定义。如上所述，这是对我们试图解决的问题的输入和输出的规范描述。这个定义包括所有输入和输出张量的名称、数据类型和形状。例如，2D 目标检测的问题定义可能是类似下面这样的：

注意，上面的定义在形状定义中使用了“符号”（num_classes 和 num_detections）。符号的每个实例都必须在运行时解析为相同的值。与仅将形状元素设置为 None 相比，这提供了一种更可靠的约束形状的方式。在上面的例子中，num_detections 跨 boxes 和 object_class_probability 必须是相同的。

为简单起见，我们将在本文中使用一个更简单的问题：加法。

上面的代码段还定义了测试输入和输出数据，我们将在下面讨论。

生成占位符模型

一旦定义了问题，我们就可以使用 Neuropod 自动生成一个占位符模型来实现问题规范。这允许我们在没有实际模型的情况下开始集成。

生成的模型接受问题规范中描述的输入，并返回与输出规范匹配的随机数据。

构建模型

在建立了问题定义（并可选地生成占位符模型）之后，我们就可以构建模型了。我们完成了构建和训练模型的所有正常步骤，但现在我们在该过程的最后添加了一个 Neuropod 导出步骤：

在上面的代码段中，我们将模型导出为 Neuropod，同时还提供了可选的测试数据。如果提供了测试数据，库将在导出后立即对模型进行自检。

Create_TensorFlow_Neuropod和所有其他封装程序的选项都有很好的文档说明。

构建度量管道

现在我们有了自己的模型，就可以用 Python 为这个问题构建度量管道。与没有 Neuropod 的情况下做这件事的唯一区别是，我们现在使用 Neuropod Python 库来运行模型，而不是使用特定于框架的 API。

Neuropod 的文档包含有关 load_neuropod 和 infer 的更多详细信息。

集成

现在，我们可以将模型集成到生产 C++系统中。下面的示例显示了 Neuropod C++ API 的使用非常简单，但该库也支持更复杂的使用，支持高效的零拷贝操作和包装现有内存。更多细节，请参考 Neuropod 文档。

与没有 Neuropod 的集成过程不同的是，这个步骤对于每个问题只需执行一次，而不是每个框架执行一次。用户无需理解 TensorFlow 或者 Torch C++ API 的复杂性，但当研究人员决定要使用哪种框架时，仍然可以提供很大的灵活性。

此外，由于核心 Neuropod 库是用 C++ 编写的，因此，我们可以为其他各种编程语言（包括 Go、Java 等）编写绑定。

优化

在 Uber ATG，我们对延迟的要求相当严格，因此对许多关键操作都有零拷贝路径。我们在分析和优化方面投入了大量工作，现在，Neuropod 可以成为实现适用于所有模型的推理优化的核心位置。

作为这项工作的一部分，每个 Neuropod 提交都在持续集成（CI）管道中的以下平台上进行测试：

• Mac、Linux、Linux（GPU）
• Python 的五个版本
• 每个支持的深度学习框架的五个版本

要了解支持的平台和框架的最新列表，请查看文档。

Neuropod 还提供了一种使用高性能共享内存桥在工作进程中运行模型的方法，这使得我们可以在不引入显著延迟损失的情况下，将模型彼此隔离开来。我们将在本文末尾更详细地讨论这一点。

迭代

一旦我们构建并集成了模型的第一个版本，就可以迭代并改进解决方案。

作为这个过程的一部分，如果我们想尝试用 TorchScript 模型来代替上面创建的 TensorFlow 模型，那就可以直接替换。

如果没有 Neuropod，就将需要重新执行之前的许多步骤。有了 Neuropod，任何实现相同问题规范的模型都是可互换的。我们可以重用前面创建的度量管道以及之前做的所有集成工作。

所有运行模型的系统和流程都与框架无关，从而在构建模型时提供了更多的灵活性。Neuropod 让用户专注于他们试图解决的问题，而不是他们用来解决问题的技术。

与问题无关的工具

尽管 Neuropod 是从关注一个“问题”开始的，比如给定图像的 2D 目标检测，文本的情感分析等等，但我们可以在 Neuropod 的基础上构建一个与框架和问题均无关的工具。这让我们能够构建通用的基础设施，可以用于任何模型。

规范输入构建管道

Neuropod 的一个有趣的、与问题无关的用例是规范输入构建管道。在 ATG，我们有一个已定义的格式/规范，用于说明如何以张量表示输入数据。这涵盖了传感器数据，包括光学雷达、雷达、相机图像以及其他信息，如高分辨率地图等。定义这种标准格式，使得在训练方面管理超大型数据集变得轻松容易。这也使我们能够快速构建新模型，因为我们的许多模型都使用这些输入数据的子集（例如，只对相机和光学雷达进行操作的模型）。

通过将这种通用输入格式与 Neuropod 结合起来，我们可以构建一个单一的优化输入构建管道，供我们所有的模型使用，而不管它们是在什么框架中实现的。

只要每个模型都使用同一组特性的子集，输入构建器就与问题无关。输入构建器中的任何优化都有助于改进我们所有的模型。

图 7. Neuropod 允许我们构建单一的、优化的输入构建管道，该管道可以与许多模型一起工作，而不必为每个模型或每个框架构建一个单独的管道。

模型服务

另一个与问题无关的基础设施是模型服务。在 Uber ATG，我们有一些模型需要大量的输入，而且模型本身也很大。对于部分模型，在诸如离线模拟之类的任务中在 CPU 上运行是不可行的。从资源效率的角度来看，在所有集群机器中包含 GPU 是没有意义的，因此，我们提供了一种服务，可以让用户在远程 GPU 上运行模型。这与基于 gRPC 的模型服务非常相似。

如果没有 Neuropod，那么模型服务平台就需要擅长远程运行 Keras、远程运行 TensorFlow、远程运行 PyTorch、远程运行 TorchScript 等。它可能需要实现序列化、反序列化、与 C++ API 的交互，以及对每个框架的优化。所有的应用还需要擅长在本地运行所有这些框架的模型。

因为在本地和远程运行有不同的实现，所以整个系统需要关注 2 * # of frameworks 情况。

但是，通过使用 Neuropod，模型服务可以很好地远程运行 Neuropod，而 Neuropod 可以很好地在多个框架中运行模型。

通过将关注点分开，系统必须关注的案例数量以相机啊的方式 (2 + # framworks) ，而不是以乘积的方式增加。随着越来越多的框架得到支持，这种差异就变得更加明显。

不过，我们可以让它变得更强大。如果我们添加模型服务作为 Neuropod 后端，任何基础设施都可以轻松地远程运行模型。

图 8. 添加模型服务作为Neuropod后端，允许任何使用 Neuropod 远程运行模型的应用程序，而无需进行重大更改。

在后台，它可能看起来像这样：

图 9. 应用程序可以使用 Neuropod 将模型执行代理到远程机器。

这个解决方案不是特定于问题或特定于框架的，而是为使用 Neuropod 的应用程序提供了更大的灵活性。

进程外执行

图 10. Neuropod 支持使用低延迟共享内存桥在独立的工作进程中运行模型。

由于 Uber ATG 的输入通常很大，并且我们对延迟非常敏感，所以本地的 gRPC 并不是分离模型的理想方法。相反，我们可以使用优化的基于进程外执行（Out-of-process Execution，OPE）的共享内存实现来避免复制数据。这让我们可以将模型彼此隔离开来，而不会造成严重的延迟损失。

这种隔离很重要，因为在过去，我们已经看到过在统一进程中运行的框架之间存在冲突，包括微妙的 CUDA 错误和内存损坏。这些问题都很难追查。

在运行 Python 模型时，这种隔离也有助于提高性能，因为它能让用户避免在所有模型之间共享 GIL。

在单独的工作进程中运行每个模型还可以启用“下一步”章节中提到的一些附加功能。

下一步

Neuropod 使 Uber 能够快速构建和部署新的深度学习模型，但这仅仅是开始。

我们正在积极开展的工作包括：

1. 版本选择：该功能使用户能够在导出模型时，指定框架所需的版本范围。例如，一个模型可能需要 TensorFlow 1.13.1，而 Neuropod 将使用正确版本的框架的 OPE 自动运行该模型。这使得用户能够在一个应用程序中使用多个框架和每个框架的多个版本。
2. 封装操作：这个功能使应用程序能够指定使用张量“完成”的时间。一旦张量被密封，Neuropod 可以在推理运行之前将数据异步传输到正确的目的地（例如本地的 GPU 或辅助进程等）。这有助于用户并行化数据传输和计算。
3. Dockerized 工作进程：这样可以在模型之间提供更多的隔离。例如，使用此功能，即使需要不同 CUDA 版本的模型也可以在同一应用程序中运行。

随着我们继续通过增强当前功能并引入新功能来扩展 Neuropod 时，我们期待着与开源社区合作来改进这个库。

Neuropod 在 Uber 的各种深度学习团队中都非常有用，我们希望它也能对其他人有所帮助。请试用一下吧！

作者介绍：

Vivek Panyam，Uber ATG 高级自动驾驶工程师。在感知团队工作，领导 Neuropod 的开发。此前，曾带领团队将深度学习和机器学习模型集成到自动驾驶车队中。

原文链接：

https://eng.uber.com/introducing-neuropod/