Lyft 如何使用 PyTorch 来驱动无人驾驶汽车
Lyft 如何使用 PyTorch 来驱动无人驾驶汽车
Lyft Level 5 无人驾驶团队分享 PyTorch 生产部署经验。
作者：Sammy Sidhu, Qiangui (Jack) Huang, Ray Gao 编译：McGL
Lyft 的使命是用世界上最好的交通工具改善人们的生活。我们相信，在未来，无人驾驶汽车将使交通更加安全，人人都更加方便。这就是为什么 Lyft 的无人驾驶部门 Level 5正在为 Lyft 网络开发一个完整的无人驾驶自治系统，给乘客提供这种技术的便利。然而，这是一项极其复杂的任务。
在开发中，我们使用各种各样的机器学习算法来驱动我们的无人驾驶汽车，解决地图、感知、预测和规划方面的问题。为了开发这些模型，我们对数百万张图像和密集的 LiDAR/RADAR 点云以及很多其他类型的输入(如代理轨迹或视频序列)进行了训练和验证。然后这些模型被部署在 Lyft 的无人车（autonomous vehicle / AV）平台上，在上面他们需要以毫秒级的速度推理边界框，交通灯状态，以及车辆轨迹等。
我们早期构建的框架是为了在项目的第一年快速提高机器学习的效率，使我们的产品能够上路。但是，这些解决方案并没有提供我们所需的开发速度和规模来解决随着项目增长而出现的问题。我们的机器学习框架需要满足我们的团队在几个小时内(而不是几天)建立和训练复杂的模型，验证指标，并将深度学习模型部署到无人车队中。我们现在有了一个框架，它不仅满足了我们快速迭代和规模化的所有要求，而且统一了 Lyft Level 5所有工程师和研究人员的机器学习开发。请继续往下读，了解我们是如何实现这个目标的。
构建正确的机器学习无人驾驶框架
我们相信快速的迭代和适应是 Level 5 成功的关键。这一原则同样适用于我们的机器学习(ML)模型和 ML 工具。2017年创立 Lyft Level 5时，我们在桌面上训练了一些基本的计算机视觉模型。仅仅几个月后，我们建立了第一个内部训练框架，使我们能够扩大规模。有了这个框架，我们在无人车上部署了近12个模型，但是很快我们意识到我们需要一个范式的转变，重点是以下几个关键原则:
- 以小时为单位迭代模型: 由于数据和模型大小的增大，我们的第一个生产模型需要几天的时间来训练。我们希望我们的工程师和研究人员能够从一个新的想法开始，实现模型，并在几个小时内看到生产质量的结果，哪怕我们的数据集大小在继续增长。
- 无缝模型部署: 我们的初始部署过程非常复杂，模型开发人员必须经过很多环节才能将一个训练好的模型转化为一个“ AV-ready”可部署模型，这意味着它可以在我们的汽车上运行，以 c + + 运行，具有低延迟和低误差（jitter）。我们希望优先考虑推理速度和可部署性问题。我们还需要让工程师能够轻松地在训练和推理时添加新的自定义层。
- 统一的实验和研究: 在工程师和研究人员的理想开发环境经常存在分歧。在 Lyft，我们看到了消除这种分歧的机会，认为这是快速推进我们的机器学习系统的关键之一; 我们不想要两个技术栈。诸如构建环境、云计算资源、日志等应该是全盘统一的且适用于每个人。
- 硬件资源优化: 从较低的 GPU 和 CPU 利用率看出，我们最初的训练框架不能完全利用硬件。随着我们数据集的大小和团队的增长，这导致了模型开发人员更长的训练时间和潜在的成本节约。我们希望让我们的框架能够感知硬件，以便尽可能地压榨出所有性能，从而实现更快的训练和更低的成本。
在考虑了这些原则之后，我们决定创建一个解决方案，将 PyTorch 纳入我们下一代机器学习框架的核心。在6个月的时间里，我们建立了一个原型系统，从上一个框架中移植了12个产品化的无人车模型，40多个机器学习工程师参与，并为 Level 5所有人建立了一个统一的机器学习框架。
概念验证的构建
在我们all in PyTorch 之前，我们想要验证它是否能够适应我们的场景用例。为了实现这一点，我们需要在 PyTorch 中实现我们的一个模型，对我们的数据进行训练，并将其部署到我们的 C++ 无人驾驶技术栈中。
图1—— LiDAR Segmenter模型结构图及Lidar semseg输出例子
我们选择实现的端到端的候选建模任务是 LiDAR 语义分割，这是一个接收 3D LiDAR 点云并将每个点分类的任务，类别比如地面、汽车、人等(参见图1)。
我们首先写了一个 PyTorch DataSet 类，它对带标注的点云的二值文件进行操作。然后，我们必须编写一个自定义的函数，将不同的数据项聚合到一个batch中，以便使用 PyTorch DataLoader。
鉴于我们现在有了一个可以工作的数据加载管道，我们开始实现我们的模型。我们当时的模型有以下结构:
- 输入稀疏点云 + 每点元数据
- 对每一点进行特征化
- 体素化（Voxelize） 特征化点云
- 在体素空间中运行密集卷积
- 将体素映射回每个点(deVoxelization)
- 把每一点分类
其中一些阶段，如DeVoxelization在我们先前的实现中是用 CUDA 手写，花费数周的工程时间。我们发现，在原生 PyTorch 中使用诸如 scatter_add 和 index_select 这样的原语实现这些功能，可以让我们无需手写内核就能获得类似的性能，从而可以在几天内生成相同的模型。
对于模型的其余部分，我们能够利用 torch "nn"包中的模块，利用卷积和各种损失函数等算子。在我们有了一个模型实现之后，我们编写了一些用于训练的标准样板代码，我们的模型可以在数据集上收敛。
现在我们有了在数据集上产生训练好的模型的方法，剩下的唯一事情就是让它在我们的 C++ 无人驾驶技术栈中工作。在这个概念验证的时候(PyTorch 1.3) ，我们发现了两个可以导出并在 C++ runtime 环境中运行“冻结模型”的选项:
- 将我们的 Torch 模型导出到 ONNX，并在 TensorRT 或 ONNX Runtime中运行 ONNX 模型。
- 使用TorchScript构建我们的模型，并在 LibTorch 中运行保存的序列化模型。
从我们之前的框架中学到的是: 尽管 ONNX 和 TensorRT 的历史更长，并且在某些情况下对推理速度进行了更优化，但我们重视快速部署模型的能力，以及更轻量级的管道(更少的外部库和依赖性) ，这使我们能够快速进行实验，而不必受制于 ONNX 的各种限制和编写自定义 TensorRT 插件。我们意识到我们总是需要编写自定义内核和操作，但我们宁愿编写 LibTorch 扩展(更多细节见下面的部分) ，而不是添加更多的黑盒外部层。由于这些原因，我们决定对 TorchScript 推理进行评估。
评估的最后一步是将我们的 TorchScript 模型集成到我们的无人驾驶技术栈中。首先，我们使用 PyTorch 提供的 LibTorch 共享库构建，并将其集成到构建中。然后我们能够利用 LibTorch C++ API 将模型集成到我们的 LiDAR 栈中。我们发现这个 API 虽然是C++，但对 python PyTorch API 用户很友好。
model = torch::jit::load(model_file, torch::kCUDA);
model.eval();最后，我们必须验证模型是否在延迟预算之内，我们发现实际上实现了比当前生产模型更低的延迟和误差。
总的来说，我们认为我们的概念验证进行得非常好，我们决定继续在 Lyft Level 5 使用 PyTorch 进行生产训练和部署。
创建 PyTorch 生产框架
图2—— 基于 PyTorch 的完全分布式训练框架的高层草图，称为“ Jadoo”(来源于印地语中的术语“ Magic” ，是对我们以前的机器学习框架“ Magician”的致敬)。Level 5 机器学习的研究人员和工程师都在这个框架上开发。在本地运行作业和将作业分发到云上的许多节点之间没有任何额外的步骤。实际上，我们的计算基础设施团队已经无缝地集成了许多必要的资源，以便在云（比如 AWS SageMaker 执行引擎）上安排我们的分布式作业。
作为 PyTorch 生产化过程的一部分，为了增强我们的机器学习工程师的能力，我们创建了一个名为 Jadoo 的 Lyft 内部框架(参见图2)。与某些框架不同的是，我们的目标是提供一个强大的生态系统，能够简化 runtime 环境和计算基础设施，允许机器学习工程师和研究人员快速迭代和部署代码，而不是试图让机器学习对非专家来说“更容易” ，并抽象出 PyTorch 的所有优点。
Jadoo 的一些核心特性包括:
- 所有作业都是分布式的。Jadoo 从一开始就是分布式的; 所有作业都是原生分布式作业，基础 case 是一个节点一个 GPU。工程师本地构建模型后可以使用数百个 GPU 在云中训练作业，只需修改一个命令行参数。无论一个实验是在一个 GPU、一个节点还是几十个节点上运行，我们都保持相同的代码路径和相同的 Docker image，这使我们能够避免任何本地/云训练的意外。我们还为用户提供工具，以便能够快速发现分布式设置中的问题，如 GPU 利用率、网络瓶颈和多节点日志记录
- 推理优先。在 Jadoo，所有模型都会进行 runtime 性能测试，并且可以部署到无人车系统中。我们获取所有层的所有操作计数和测试推理延迟，并存储这些信息，以便用户执行帕累托最优速度-准确性权衡。我们还确保每个经过训练的模型都可以用 C++ 或 python 在 TorchScript 中部署，并且不需要特殊的管道将其转换为推理模型。
- 将研究迭代速度与生产质量结合起来。Jadoo 的目标是给研究人员实验和迭代的自由，但同时强制实施工程最佳实践。接入 Jadoo 的所有代码都要经过基于 GPU 的持续集成(Continuous Integration，CI) ，以及单元测试、静态类型检查、样式检查、文档检查、确定性和性能分析等检查。为了实现快速的实验，我们使用代码作为配置，消除大型和笨重的 json/yaml 文件，并自动从用户代码中提取实验参数。我们通过版本控制和记录数据、实验代码和 artifacts 对所有工作执行严格的可重复性。最后，我们提供了一个用于日志记录、可视化和作业跟踪的工具生态系统，允许用户快速调试和反思他们的工作。
分布式训练
我们设计了分布式训练环境来模拟本地环境，这样用户可以在本地和分布式云训练之间无缝切换。实现这一目标的第一步是确保本地的开发环境得到良好的控制和容器化。然后，我们在环境中使用相同的容器，用于本地开发、分布式云训练和持续集成的 Jadoo/用户代码。对于分布式训练，我们可以非常依赖 PyTorch 中的分布式包。我们遵循了这样的模式，即使每个 GPU 都有自己的流程，并将我们的模型包装在 DistributedDataParallel 中。
def __init__(self) -> None:
super().__init__()
self.model = self._load_model(pretrained_backbone=True).cuda()
self.model = DistributedDataParallel(self.model, device_ids=[torch.cuda.current_device()],
output_device=torch.cuda.current_device())Jadoo 负责在节点和工作者之间进行数据共享，因此用户只需以一种易于分区的方式创建数据集。我们还发现 Nvidia NCCL 后端对于训练和其他操作(如分布式 all-reduce、scatter 和gather)都有很高的性能。计算节点的协调和供应由我们的底层基础设施来处理。我们还控制模型状态的 checkpoint，以允许节点抢占和中断，以节省成本，如热点实例的训练。
使用 LibTorch 和 TorchScript 进行推理
使用 Jadoo，我们希望优先构建能够在 AV 中通过 C++ runtime 高效运行的模型。我们发现 LibTorch 允许我们通过TorchScript轻松地部署训练好的模型，而 C++ API 使其非常容易使用。当我们部署的模型需要进行预处理或后处理时，C++ API 特别有用，因为 API 遵循熟悉的 PyTorch。
# Python
output = torch.where(score_over_threshold, label, unknown_labels)
// C++
const auto output = torch::where(score_over_threshold, label, unknown_labels);需要注意的一点是，尽管我们从 PyTorch 提供的LibTorch开始构建概念验证，但我们发现大型静态链接库很难管理。为了解决这个问题，我们使用自己的依赖项从源代码编译 LibTorch，通过共享库链接 LibTorch 的依赖项。这样我们将 LibTorch 的二进制文件大小减少一个数量级。
为了确保用户可以轻松地部署他们训练过的模型，Jadoo 检查模型是否可以在训练期间转换成 TorchScript，如果可以，定期保存包含 TorchScript 模型的 checkpoints 以及任何允许追溯模型起源的附加元数据。这些元数据包括训练运行、 GitSHA、用户名和用户选择跟踪的任何其他元数据。此外，Jadoo 自动分析这个 TorchScript 模型的延迟以及它的 MACs (乘-加)和参数计数。
当用户准备部署模型时，他们只需指向他们想要的从训练运行得到的模型，然后它就可以在我们的构建在 LibTorch 上的 C++ runtime 运行推理。
我们发现这是最佳实践，因为用户正在构建他们的模型时会记着 TorchScript 。这避免了复杂模型带来的麻烦，只需在尝试部署模型时发现，由于 TorchScript 不兼容的语法，需要更改大量模型 APIs。我们还发现，单元测试是确保模型和模块兼容并保持 TorchScript 兼容性的好方法。例如，通常一个开发人员可能会更改另一个开发人员使用的公共模块，并添加不受 TorchScript 支持的语法，这会在 CI (持续集成)的早期就被发现，并且永远不会进入生产环境。
结合研究与工程
在工程组织中，研究人员和工程师的工作方式往往有所不同; 研究人员希望自由地探索和快速迭代各种想法，而工程师则希望构建经过良好测试和可靠的产品。通过 Jadoo，我们建立了一个范式，在这个范式中，研究人员和工程师可以使用相同的框架进行共同开发，从而允许我们快速地创建一个从想法到生产的迭代循环。我们通过以下方式实现这一目标:
大量优化机器学习开发人员的迭代周期
- 用户可以在5秒钟内启动作业。
- 使用数百个 GPU 的作业几分钟内就可以启动。
- 我们对数据加载管道进行了重大修改并将硬件考虑在内，以确保分布式作业不会受到 I/O、 CPU 和网络带宽的限制。
- 对于部署在 AV 上的典型图像/lidar检测模型(见图4)，中位数作业训练时间约为1小时。
使实验易于复现和可比较
- 我们将模型和实验配置设置为 python 代码，这些代码都录入到 git 和我们的实验跟踪工具中，从而消除了管理无数模型 json 和 yaml 文件的需要。
- 我们还利用代码自省（code introspection）来跟踪和记录代码中的常量，使消融比较实验输入和结果非常容易。
- 我们自动记录实验中使用的数据集，远程 git 的 SHAs，以及本地代码的变化，使任何实验100% 可重复。这对于严格的实验是至关重要的。
保持高编码标准
- 我们打破了严格的编码实践会减慢实验周期的神话，使得诸如文档、格式、类型检查、单元测试等工作尽可能自动化，并且不会增加开发人员超过几分钟的合并代码的时间。
结果和展望
图3 —— Dashboard 显示 Lyft Level 5过去6个月所有生产训练作业的每周平均时间(P50)的大约是几个小时(以前为几天)。为了达到这个目的，每个作业都要在云端的数十到数百个 GPU 上运行分布式训练，通过我们的内部模型框架 Jadoo 和 PyTorch 的帮助，对 worker 数量进行优化，以实现非常高效的扩展。
我们相信，通过采用 PyTorch 和构建 Jadoo，我们已经在机器学习方面取得了巨大的进步。在短短几个月的工作中，我们已经能够将十几个模型从旧框架移动到 PyTorch，并使用 TorchScript + LibTorch 部署。我们还能够将2D 和3D 检测器以及语义分割器等大型生产作业的中位数训练时间从几天缩短到大约1小时(见图3) ，允许用户指定他们需要多少计算资源和时间。这让我们的用户的模型开发每天进行很多次迭代，而这在之前是不可能的。我们相信我们也已经建立了一个独特的框架，真正使我们的机器学习工程师和研究人员能够在更短的时间内做更多的事情，使他们能够快速迭代各种想法，并将它们部署到交通工具平台上，而不必尝试任何机器学习简化。
展望未来，尽管我们已经训练了数百万的数据样本，但是我们的数据集规模正在呈指数增长。我们需要能够维持对越来越多的数据的训练，还需要单个作业能够扩展到数千个 GPU 且可容错。为此，我们正在研究 PyTorch Elastic 之类的容错技术。我们还计划围绕推理性能分析和优化，以及模型和数据自省和可视化来扩展工具库。
原文：https://medium.com/pytorch/how-lyft-uses-pytorch-to-power-machine-learning-for-their-self-driving-cars-80642bc2d0ae