Triton 模型导航器#

欢迎使用 Triton 模型导航器，这是一款推理工具包，旨在优化和部署深度学习模型，重点关注 NVIDIA GPU。 Triton 模型导航器简化了将 PyTorch、TensorFlow 和/或 ONNX 中实现的模型和管线迁移到 TensorRT 的过程。

Triton 模型导航器自动化了几个关键步骤，包括模型导出、转换、正确性测试和性能分析。通过为各种受支持的框架提供单一入口点，用户可以使用每个框架的优化功能有效地搜索最佳部署选项。由此产生的优化模型已准备好部署在 PyTriton 或 Triton 推理服务器 上。

功能概览#

Triton 模型导航器的独特功能在功能矩阵中进行了总结

功能

描述

易用性

单行代码即可直接从源代码运行所有可能的优化路径

广泛的框架支持

兼容各种机器学习框架，包括 PyTorch、TensorFlow 和 ONNX

模型优化

增强 ResNET 和 BERT 等模型的性能，以实现高效的推理部署

管线优化

使用 Inplace 优化（PyTorch 独有）简化 Stable Diffusion 和 Whisper 等模型的 Python 代码管线

模型导出和转换

自动化在各种格式之间导出和转换模型的过程，重点关注 TensorRT 和 Torch-TensorRT

正确性测试

确保转换后的模型生成正确的输出，并针对原始模型进行验证

性能分析

分析模型，以根据延迟和吞吐量等性能指标选择最佳格式，从而优化目标硬件利用率

模型部署

通过专用 API 自动化模型和管线在 PyTriton 和 Triton 推理服务器上的部署

文档#

在文档中了解有关 Triton 模型导航器功能的更多信息。

先决条件#

在继续安装 Triton 模型导航器之前，请确保您的系统满足以下条件

  • 操作系统：Linux（建议使用 Ubuntu 20.04+）

  • Python：版本 3.8 或更高版本

  • NVIDIA GPU

您可以使用 PyTorch 和 TensorFlow 的 NGC 容器，其中包含所有必要的依赖项

安装#

可以从 pypi.org 安装 Triton 模型导航器。

使用 PyTorch 额外组件安装#

要使用 PyTorch 依赖项进行安装，请使用

pip install -U --extra-index-url https://pypi.ngc.nvidia.com triton-model-navigator[torch]

使用 TensorFlow 额外组件安装#

要使用 TensorFlow 依赖项进行安装，请使用

pip install -U --extra-index-url https://pypi.ngc.nvidia.com triton-model-navigator[tensorflow]

为 CUDA 11 安装 onnxruntime-gpu#

自 1.19.0 起，ONNXRuntime 的默认 CUDA 版本为 CUDA 12。要安装 CUDA 11 支持，请使用以下额外索引 URL

.. --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-11/pypi/simple/ ..

快速入门#

快速入门部分提供了 Triton 模型导航器中提供的可能的优化和部署路径示例。

使用 Inplace 优化 Stable Diffusion#

Inplace 优化允许对模型进行无缝优化以进行部署，例如将其转换为 TensorRT，而无需对原始 Python 管线进行任何更改。

以下代码介绍了 Stable Diffusion 管线优化。但首先，在运行示例之前，请安装所需的软件包

pip install transformers diffusers torch

然后，初始化管线并使用 nav.Module 包装模型组件：

import model_navigator as nav
from transformers.modeling_outputs import BaseModelOutputWithPooling
from diffusers import DPMSolverMultistepScheduler, StableDiffusionPipeline


def get_pipeline():
    # Initialize Stable Diffusion pipeline and wrap modules for optimization
    pipe = StableDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1")
    pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
    pipe = pipe.to("cuda")
    pipe.text_encoder = nav.Module(
        pipe.text_encoder,
        name="clip",
        output_mapping=lambda output: BaseModelOutputWithPooling(**output),
    )
    pipe.unet = nav.Module(
        pipe.unet,
        name="unet",
    )
    pipe.vae.decoder = nav.Module(
        pipe.vae.decoder,
        name="vae",
    )
    return pipe

准备一个简单的数据加载器

# Please mind, the first element in tuple need to be a batch size
def get_dataloader():
    return [(1, "a photo of an astronaut riding a horse on mars")]

执行模型优化

pipe = get_pipeline()
dataloader = get_dataloader()

nav.optimize(pipe, dataloader)

管线优化完成后，您可以显式加载模块的最优性能版本，执行

nav.load_optimized()

此时，您可以简单地使用原始管线，直接在 Python 中使用优化的模型生成预测

pipe.to("cuda")

images = pipe(["a photo of an astronaut riding a horse on mars"])
image = images[0][0]

image.save("an_astronaut_riding_a_horse.png")

有关如何通过 PyTriton 服务 Stable Diffusion 管线的示例，请参见此处。

当您计划将代码放在 Python 脚本中时，请阅读运行 Python 脚本时的错误隔离。

优化 ResNET 并在 Triton 上部署#

Triton 模型导航器还支持在 Triton 上部署的优化路径。此路径支持输入为张量的 nn.Module、keras.Model 或 ONNX 文件。

要从 TorchHub 优化 ResNet50 模型，请运行以下代码

import torch
import model_navigator as nav

# Initialize the model
resnet50 = torch.hub.load('NVIDIA/DeepLearningExamples:torchhub', 'nvidia_resnet50', pretrained=True).eval()

# Wrap model in nav.Module
resnet50 = nav.Module(resnet50, name="resnet50")

# Optimize Torch model loaded from TorchHub
nav.optimize(resnet50, dataloader=[(1, [torch.randn(1, 3, 256, 256)])])

完成优化后，创建用于在 Triton 上部署的模型仓库就像以下代码一样简单

import pathlib

# Generate the model store from optimized model
resnet50.triton_model_store(
    model_repository_path=pathlib.Path("model_repository"),
)

当您计划将代码放在 Python 脚本中时，请阅读运行 Python 脚本时的错误隔离。

在 Python 中分析任何模型或可调用对象#

Triton 模型导航器增强了模型和管线，并为分析任何 Python 函数、可调用对象或模型提供了统一的方法。目前，我们的支持严格限于静态批量性能分析方案。

例如，我们将使用一个简单的函数，该函数仅休眠 50 毫秒

import time


def custom_fn(input_):
    # wait 50ms
    time.sleep(0.05)
    return input_

让我们提供一个数据加载器，我们将用于性能分析

# Tuple of batch size and data sample
dataloader = [(1, ["This is example input"])]

最后，使用准备好的数据加载器运行该函数的性能分析

nav.profile(custom_fn, dataloader)

运行 Python 脚本时的错误隔离#

重要提示：请查看以下部分，以防止在运行 optimize 时出现意外问题。

为了更好地隔离错误，某些转换和导出在单独的子进程中使用多处理在 spawn 模式下运行。这意味着全局范围内的所有内容都将在子进程中运行。当优化代码放置在 Python 脚本中并执行为以下内容时，您可能会遇到意外问题

python optimize.py

为了防止嵌套优化，您必须将优化代码放在

if __name__ == "__main__":
    # optimization goes here

或

import multiprocessing as mp
if mp.current_process().name == "MainProcess":
    # optimization goes here

如果以上方法均不适用于您，您可以在单个进程中运行所有优化，但会牺牲错误隔离，方法是设置以下环境变量

NAVIGATOR_USE_MULTIPROCESSING=False

示例#

我们提供了全面的分步指南，展示了 Triton 模型导航器各种功能的利用率。 这些指南旨在阐明使用 PyTriton 和 Triton 推理服务器 优化、分析、测试和部署模型的过程。

在本页中