人工智能模型無(wú)處不在，形式包括聊天機(jī)器人、分類(lèi)和摘要工具、用于分割和檢測(cè)的圖像模型、推薦模型等。人工智能機(jī)器學(xué)習(xí)（ ML ）模型有助于實(shí)現(xiàn)許多業(yè)務(wù)流程的自動(dòng)化，從數(shù)據(jù)中生成見(jiàn)解，并提供新的體驗(yàn)。

Python 是 AI/ML 開(kāi)發(fā)中最受歡迎的語(yǔ)言之一。本文將教您如何使用 NVIDIA Triton Inference Server，并利用新的 PyTriton 接口。

更具體地說(shuō)，您將學(xué)習(xí)如何在 Python 開(kāi)發(fā)環(huán)境中使用生產(chǎn)類(lèi)工具對(duì)人工智能模型進(jìn)行原型化和測(cè)試推理，以及如何使用 PyTriton 接口進(jìn)行生產(chǎn)。與 FastAPI 或 Flask 等通用 web 框架相比，您還將了解使用 PyTriton 的優(yōu)勢(shì)。這篇文章包括幾個(gè)代碼示例，說(shuō)明如何激活高性能的批處理、預(yù)處理和多節(jié)點(diǎn)推理；并實(shí)施在線學(xué)習(xí)。

什么是 PyTriton ？

PyTriton 是一個(gè)簡(jiǎn)單的接口，可讓 Python 開(kāi)發(fā)人員使用 Triton 推理服務(wù)器為 Python 代碼中的人工智能模型、簡(jiǎn)單處理功能或整個(gè)推理管道提供服務(wù)。Triton 推理服務(wù)器是一款開(kāi)源的多框架推理服務(wù)軟件，在 CPU 和 GPU 上具有較高的性能。

PyTriton 可以實(shí)現(xiàn)快速原型設(shè)計(jì)和測(cè)試 ML 模型，同時(shí)實(shí)現(xiàn)高 GPU 利用率的性能和效率。只需一行代碼，就可以調(diào)出 Triton 推理服務(wù)器，提供 動(dòng)態(tài)批處理、并發(fā)模型執(zhí)行以及從 GPU 代碼中支持 GPU 和 Python 的優(yōu)勢(shì)。

PyTriton 消除了建立模型存儲(chǔ)庫(kù)和將模型從開(kāi)發(fā)環(huán)境移植到生產(chǎn)環(huán)境的需要?，F(xiàn)有的推理管道代碼也可以在不進(jìn)行修改的情況下使用。這對(duì)于較新類(lèi)型的框架（如 JAX ）或復(fù)雜的管道（它們是 Triton 推理服務(wù)器中沒(méi)有專(zhuān)用后端的應(yīng)用程序代碼的一部分）尤其有用。

Flask 的簡(jiǎn)單性

Flask 和FastAPI 是通用的 Python Web 框架，用于部署各種各樣的 Python 應(yīng)用程序。由于它們的簡(jiǎn)單性和廣泛采用，許多開(kāi)發(fā)人員在生產(chǎn)中使用它們來(lái)部署和運(yùn)行人工智能模型。然而，這種方法存在顯著的缺點(diǎn)，包括：

通用網(wǎng)絡(luò)服務(wù)器缺乏對(duì)人工智能推理功能的支持。沒(méi)有現(xiàn)成的支持來(lái)利用像 GPU 這樣的加速器，或者打開(kāi)動(dòng)態(tài)批處理或多節(jié)點(diǎn)推理。

用戶(hù)需要構(gòu)建邏輯來(lái)滿(mǎn)足特定用例的需求，如音頻/視頻流輸入、有狀態(tài)處理或預(yù)處理輸入數(shù)據(jù)以適應(yīng)模型。

關(guān)于計(jì)算和內(nèi)存利用率或推理延遲的指標(biāo)不容易用于監(jiān)控應(yīng)用程序的性能和規(guī)模。

Triton Inference Server包含對(duì)上述功能以及更多功能的內(nèi)置支持。PyTriton 提供了 Flask 的簡(jiǎn)單性和 Python 中 Triton 的示例部署。HuggingFace 文本分類(lèi)管道使用 PyTriton 如下所示：

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 用戶(hù)提供了一個(gè)熟悉的界面，便于安裝和設(shè)置，并提供了以下好處：

?用一行代碼調(diào)出 NVIDIA Triton

無(wú)需設(shè)置模型存儲(chǔ)庫(kù)和模型格式轉(zhuǎn)換（對(duì)于使用 Triton 推理服務(wù)器的高性能實(shí)現(xiàn)非常重要）

使用現(xiàn)有推理管道代碼而不進(jìn)行修改

支持許多裝飾器來(lái)調(diào)整模型輸入

無(wú)論是在generative AI應(yīng)用程序還是其他模型中，PyTriton 可以讓您在自己的開(kāi)發(fā)環(huán)境中獲得 Triton InferenceServer 的好處。它可以幫助利用 GPU 在很短的時(shí)間內(nèi)（毫秒或秒，取決于用例）生成推理響應(yīng)。它還有助于以高容量運(yùn)行 GPU ，并且可以同時(shí)為許多推理請(qǐng)求提供服務(wù)，且基礎(chǔ)設(shè)施成本低。

PyTriton 代碼示例

本節(jié)提供了一些可以用來(lái)開(kāi)始 PyTriton 的代碼示例。它們從本地機(jī)器開(kāi)始，這是測(cè)試和原型的理想選擇，并為大規(guī)模部署提供 Kubernetes 配置。

動(dòng)態(tài)配料支持

Flask/FastAPI 和 PyTriton 之間的一個(gè)關(guān)鍵區(qū)別是，動(dòng)態(tài)批處理允許對(duì)來(lái)自模型的多個(gè)調(diào)用應(yīng)用程序的推理請(qǐng)求進(jìn)行批處理，同時(shí)保留延遲要求。兩個(gè)示例是HuggingFace BART PyTorch和HuggingFace ResNET PyTorch。

在線學(xué)習(xí)

在線學(xué)習(xí)是指在生產(chǎn)中不斷從新數(shù)據(jù)中學(xué)習(xí)。使用 PyTriton，您可以控制支持推理服務(wù)器的不同模型實(shí)例的數(shù)量，從而使您能夠同時(shí)訓(xùn)練和服務(wù)同一個(gè)模型。想要了解更多關(guān)于如何使用 PyTriton 在 MNIST 數(shù)據(jù)集上同時(shí)訓(xùn)練和推斷模型的信息，請(qǐng)?jiān)L問(wèn) PyTriton 的示例。

大型語(yǔ)言模型的多節(jié)點(diǎn)推理

太大而無(wú)法放入單個(gè) GPU 內(nèi)存的大型語(yǔ)言模型（LLM）需要將模型劃分為多個(gè) GPU，在某些情況下，還需要跨多個(gè)節(jié)點(diǎn)進(jìn)行推理。要查看示例，請(qǐng)?jiān)L問(wèn) Hugging Face OPT 模型在 JAX 中的多節(jié)點(diǎn)推理。

想要查看NeMo Megatron GPT 模型部署，使用NVIDIA NeMo 1.3B 參數(shù)模型。使用 Slurm 和 Kubernetes 展示了多節(jié)點(diǎn)推理部署編排。

穩(wěn)定擴(kuò)散

使用 PyTriton ，您可以使用預(yù)處理裝飾器來(lái)執(zhí)行高級(jí)批處理操作，例如使用簡(jiǎn)單的定義將相同大小的圖像批處理在一起：

@batch

@group_by_values("img_size")

@first_value("img_size")

想了解更多信息，請(qǐng)查看此示例，該示例使用 Hugging Face 的 Stable Diffusion 1.5 圖像生成管道。

總結(jié)

PyTriton 提供了一個(gè)簡(jiǎn)單的接口，使 GPU 開(kāi)發(fā)人員能夠使用 NVIDIA Triton InferenceServer 為模型、簡(jiǎn)單的處理功能或整個(gè)推理管道提供服務(wù)。這種對(duì) Python 中的 Triton 推理服務(wù)器的本地支持使 ML 模型的快速原型設(shè)計(jì)和測(cè)試具有性能和效率。一行代碼就會(huì)顯示 Triton 推理服務(wù)器。動(dòng)態(tài)批處理、并發(fā)模型執(zhí)行以及 Python 代碼中對(duì) GPU 和 Python 的支持都是其中的好處。 PyTriton 提供了 Flask 的簡(jiǎn)單性和 Python 中 Triton InferenceServer 的優(yōu)點(diǎn)。