使用 NVIDIA TensorRT 和 Google Dataflow 扩展 ML 推理

Google Cloud 和 NVIDIA 之间的合作让 Apache Beam 的用户能够使用 NVIDIA TensorRT 和 NVIDIA GPU 与新的 Apache Beam TensorRTEngineHandler 结合，以最大限度地提高其数据处理流水线中 ML 模型的性能。

NVIDIA TensorRT SDK 提供高性能神经网络推理，使开发人员能够在 NVIDIA GPU 上优化和部署经过训练的 ML 模型，在实现最高吞吐量和最低延迟的同时保持模型预测准确性。TensorRT 经过专门设计，用于支持多类深度学习模型，包括卷积神经网络 (CNN)、循环神经网络 (RNN) 和基于 Transformer 的模型。

如何部署和管理端到端 ML 推理流水线，同时最大限度地提高基础设施利用率并最大限度地降低总成本一直是一项难题。在生产数据处理流水线中集成机器 学习模型以提取见解，需要解决与三个主要工作流段相关的挑战：

1. 对来自多个数据源的大量原始数据进行预处理，用其训练 ML 模型以“推断/预测”结果的输入，然后利用下游的 ML 模型输出以纳入业务流程。

2. 在数据处理流水线中调用 ML 模型，同时支持不同的推理用例：批处理、流式处理、集成模型、远程推理或本地推理。流水线不限于单一模型，通常需要一组模型才能产生所需的业务成果。

3. 优化 ML 模型的性能，以在应用程序的准确性、吞吐量和延迟限制范围内交付结果。对于为 NLP 等用例使用复杂的计算密集型模型或同时需要多个 ML 模型的流水线，这些模型的响应时间通常成为性能瓶颈。这可能会导致硬件利用率低下，并且需要更多计算资源才能在生产中部署您的流水线，从而会导致更高的运营成本。

Google Cloud Dataflow 是一个完全托管的运行器，适用于使用 Apache Beam 编写的流或批处理流水线。为了使开发人员能够轻松地将 ML 模型整合到数据处理流水线中，Dataflow 最近宣布支持 Apache Beam 的通用 ML 预测和推理转换 RunInference。RunInference 转换使得开发人员无需大量样板代码就可以在生产流水线中使用模型，从而简化了 ML 流水线创建过程。

您可以在以下代码示例中查看它与 Apache Beam 的用法示例。 请注意，engine_handler 作为配置传递给 RunInference 转换，它将用户从运行模型的实现细节中抽象出来。

    engine_handler = TensorRTEngineHandlerNumPy(

              min_batch_size=4,

              max_batch_size=4,

              engine_path=

              'gs://gcp_bucket/single_tensor_features_engine.trt')

    pcoll = pipeline | beam.Create(SINGLE_FEATURE_EXAMPLES)

    predictions = pcoll | RunInference(engine_handler)

将 Apache Beam 与 Dataflow runner 和 TensorRT 结合在一起，可以帮助用户应对三个主要挑战。Dataflow runner 负责大规模预处理数据，准备数据以用作模型输入。Apache Beam 用于批处理和流处理流水线的单一 API，意味着 RunInference 自动可用于这两种用例。Apache Beam 定义复杂多路径流水线的能力也让创建具有多个模型的流水线变得更加容易。借助 TensorRT 支持，Dataflow 现在还能够优化 NVIDIA GPU 上模型的推理性能。