AI模型无处不在，以聊天机器人、分类和摘要工具、图像分割和检测模型、推荐模型等形式存在。AI机器学习（ML）模型帮助自动化许多业务流程，从数据中生成洞察力，并提供新的体验。

Python是AI/ML开发中使用最广泛的语言之一。在这篇文章中，你将学习如何使用NVIDIA Triton Inference Server在你的Python代码和环境中使用新的PyTriton接口来服务模型。



更具体地说，你将学习如何使用一个生产级的工具在Python开发环境中原型和测试AI模型的推理，以及如何使用PyTriton接口进行生产。你还将了解使用PyTriton相比于使用通用的Web框架，如FastAPI或Flask，有哪些优势。文章中包含了几个代码示例，说明了你如何激活高性能批处理、预处理和多节点推理，以及如何实现在线学习。

什么是PyTriton？

PyTriton是一个简单的接口，使得Python开发者可以使用Triton Inference Server在Python代码中服务AI模型、简单的处理函数或整个推理管道。Triton Inference Server是一个开源的多框架推理服务软件，在CPU和GPU上具有高性能。

PyTriton使得ML模型的快速创建原型和测试原型成为可能，同时实现了性能和效率的提高，例如高GPU利用率。一行代码就可以启动Triton Inference Server，在Python代码中提供诸如动态批处理、并发模型执行和支持GPU和CPU等功能。

PyTriton消除了设置模型仓库和将模型从开发环境迁移到生产环境的需要。现有的推理管道代码也可以不经修改地使用。这对于JAX等新型框架或属于应用程序代码一部分的复杂管道尤其有用，而Triton推理服务器中没有专用后端。

Flask的简单性

Flask和FastAPI是通用的Python Web框架，用于部署各种Python应用程序。由于它们的简单性和广泛的采用，许多开发者使用它们来部署和运行AI模型。然而，这种方法有一些显著的缺点，包括以下几点：

1、通用的Web服务器缺乏支持AI推理功能。没有现成的支持来利用GPU等加速器，或者打开动态批处理或多节点推理。

2、用户需要构建逻辑来满足特定用例的需求，比如音频/视频流输入、有状态处理或预处理输入数据以适应模型。

3、计算和内存利用率或推理延迟等指标不容易获取，难以监控应用程序性能和规模。

Triton Inference Server内置了对上述功能以及更多功能的支持。PyTriton提供了Flask的简单性和Triton在Python中的优势。下面展示了一个使用PyTriton部署HuggingFace文本分类管道的示例：

import logging

 

import numpy as np

from transformers import BertTokenizer, FlaxBertModel  # pytype: disable=import-error

 

from pytriton.decorators import batch

from pytriton.model_config import ModelConfig, Tensor

from pytriton.triton import Triton

 

logger = logging.getLogger("examples.huggingface_bert_jax.server")

logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(name)s: %(message)s")

 

tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")

model = FlaxBertModel.from_pretrained("bert-base-uncased")

 

 

@batch

def _infer_fn(**inputs: np.ndarray):

	(sequence_batch,) = inputs.values()

 

	# need to convert dtype=object to bytes first

	# end decode unicode bytes

	sequence_batch = np.char.decode(sequence_batch.astype("bytes"), "utf-8")

 

	last_hidden_states = []

	for sequence_item in sequence_batch:

    	tokenized_sequence = tokenizer(sequence_item.item(), return_tensors="jax")

    	results = model(**tokenized_sequence)

    	last_hidden_states.append(results.last_hidden_state)

	last_hidden_states = np.array(last_hidden_states, dtype=np.float32)

	return [last_hidden_states]

 

 

with Triton() as triton:

	logger.info("Loading BERT model.")

	triton.bind(

    	model_name="BERT",

    	infer_func=_infer_fn,

    	inputs=[

        	Tensor(name="sequence", dtype=np.bytes_, shape=(1,)),

    	],

    	outputs=[

        	Tensor(name="last_hidden_state", dtype=np.float32, shape=(-1,)),

    	],

PyTriton提供了一个类似于Flask用户熟悉的接口，方便安装和设置，并提供了以下优势：

1、用一行代码就可以启动NVIDIA Triton

2、无需设置模型仓库和模型格式转换（对于使用Triton Inference Server进行高性能实现很重要）

3、无需修改现有的推理管道代码

4、支持多种装饰器来适配模型输入

无论你是在做一个生成式AI应用还是其他模型，PyTriton都可以让你在自己的开发环境中获得Triton Inference Server的好处。它可以帮助你利用GPU在很短的时间内（毫秒或秒，取决于用例）产生推理响应。它还可以帮助你让GPU以高负载运行，并同时服务多个推理请求，从而降低基础设施成本。

PyTriton代码示例

本节提供了一些你可以用来开始使用PyTriton的代码示例。它们从本地机器开始，这是测试和原型的理想选择，并提供了Kubernetes配置来进行扩展部署。

动态批处理支持

Flask/FastAPI和PyTriton之间的一个关键区别是，动态批处理可以将来自多个调用应用程序的推理请求批量处理给模型，同时保持延迟要求。两个示例分别是HuggingFace BART PyTorch和HuggingFace ResNET PyTorch。

在线学习

在线学习是指在生产中不断地从新数据中学习。使用PyTriton，你可以控制支持你的推理服务器的不同模型实例的数量。这使得你可以同时从两个不同的端点训练和服务同一个模型。

大语言模型的多节点推理

大语言模型（LLM）由于太大而无法放入单个GPU内存，需要将模型分割到多个GPU上，甚至在某些情况下需要跨多个节点进行推理。

Stable Diffusion

使用PyTriton，你可以使用预处理装饰器来执行高级批处理操作，比如使用简单的定义将相同大小的图像批量处理在一起：

@batch

@group_by_values("img_size")

@first_value("img_size")

总结：

PyTriton提供了一个简单的接口，使得Python开发者可以使用NVIDIA Triton Inference Server来服务一个模型、一个简单的处理函数或一个完整的推理管道。这种在Python中对Triton Inference Server的原生支持使得ML模型的快速创建原型和测试原型成为可能，同时实现了性能和效率的提高。一行代码就可以启动Triton Inference Server。动态批处理、并发模型执行和在Python代码中支持GPU和CPU等都是它的优势。PyTriton提供了Flask的简单性和Triton Inference Server在Python中的优势。

 

来源：https://developer.nvidia.com/blog/how-to-deploy-an-ai-model-in-python-with-pytriton/