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

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

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

什么是 PyTriton ？

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

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

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

Flask 的簡單性

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

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

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

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

Triton Inference Server包含對上述功能以及更多功能的內(nèi)置支持。PyTriton 提供了 Flask 的簡單性和 Python 中 Triton 的示例部署。HuggingFace 文本分類管道使用 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 用戶提供了一個熟悉的界面，便于安裝和設(shè)置，并提供了以下好處：

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

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

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

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

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

PyTriton 代碼示例

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

動態(tài)配料支持

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

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

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

大型語言模型的多節(jié)點推理

太大而無法放入單個 GPU 內(nèi)存的大型語言模型（LLM）需要將模型劃分為多個 GPU，在某些情況下，還需要跨多個節(jié)點進行推理。要查看示例，請訪問 Hugging Face OPT 模型在 JAX 中的多節(jié)點推理。

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

穩(wěn)定擴散

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

@batch

@group_by_values("img_size")

@first_value("img_size")

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

總結(jié)

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