【Triton Inference Server 多输入|多输出|无输出】如何用 triton_client.infer 调用多输入、多输出的模型进行推理呢？

Triton Inference Server 是一个强大的推理平台，支持多种框架和部署环境。本文将介绍如何使用 tritonclient 库中的 infer 方法，通过 HTTP 客户端实现对多输入、多输出模型的推理。我们将通过一个示例代码详细讲解这一过程。

环境准备

在开始之前，请确保已安装 tritonclient 库。你可以通过以下命令进行安装：

pip install tritonclient[all]

代码示例

导入必要的库

import argparse
import sys
import gevent.ssl
import numpy as np
import tritonclient.http as httpclient
from tritonclient.utils import InferenceServerException

定义推理函数

定义两个推理函数 test_infer 和 test_infer_no_outputs，分别用于处理有输出和无输出的情况。

函数 test_infer

该函数用于执行多输入、多输出的推理：

def test_infer(
    model_name,
    input0_data,
    input1_data,
    headers=None,
    request_compression_algorithm=None,
    response_compression_algorithm=None,
):
    inputs = []
    outputs = []
    inputs.append(httpclient.InferInput("INPUT0", [1, 16], "INT32"))
    inputs.append(httpclient.InferInput("INPUT1", [1, 16], "INT32"))

    # 初始化数据
    inputs[0].set_data_from_numpy(input0_data, binary_data=False)
    inputs[1].set_data_from_numpy(input1_data, binary_data=True)

    outputs.append(httpclient.InferRequestedOutput("OUTPUT0", binary_data=True))
    outputs.append(httpclient.InferRequestedOutput("OUTPUT1", binary_data=False))
    query_params = {"test_1": 1, "test_2": 2}
    results = triton_client.infer(
        model_name,
        inputs,
        outputs=outputs,
        query_params=query_params,
        headers=headers,
        request_compression_algorithm=request_compression_algorithm,
        response_compression_algorithm=response_compression_algorithm,
    )

    return results

1. 定义输入和输出： 创建 InferInput 和 InferRequestedOutput 对象，分别用于定义模型的输入和输出。
2. 初始化数据： 使用 set_data_from_numpy 方法将 numpy 数据赋值给输入对象。
3. 调用推理： 使用 triton_client.infer 方法进行推理，并返回结果。

函数 test_infer_no_outputs

该函数用于执行只有输入的推理：

def test_infer_no_outputs(
    model_name,
    input0_data,
    input1_data,
    headers=None,
    request_compression_algorithm=None,
    response_compression_algorithm=None,
):
    inputs = []
    inputs.append(httpclient.InferInput("INPUT0", [1, 16], "INT32"))
    inputs.append(httpclient.InferInput("INPUT1", [1, 16], "INT32"))

    # 初始化数据
    inputs[0].set_data_from_numpy(input0_data, binary_data=False)
    inputs[1].set_data_from_numpy(input1_data, binary_data=True)

    query_params = {"test_1": 1, "test_2": 2}
    results = triton_client.infer(
        model_name,
        inputs,
        outputs=None,
        query_params=query_params,
        headers=headers,
        request_compression_algorithm=request_compression_algorithm,
        response_compression_algorithm=response_compression_algorithm,
    )

    return results

1. 定义输入： 与 test_infer 函数类似，但没有定义输出。
2. 调用推理： 使用 triton_client.infer 方法进行推理，并返回结果。

使用示例

假设我们有一个模型名为 my_model，并且我们有两个输入数据 input0 和 input1，可以按如下方式调用：

model_name = "my_model"
input0 = np.random.randint(0, 100, size=(1, 16), dtype=np.int32)
input1 = np.random.randint(0, 100, size=(1, 16), dtype=np.int32)

# 有输出的推理
results_with_outputs = test_infer(model_name, input0, input1)
print("Results with outputs:", results_with_outputs)

# 无输出的推理
results_no_outputs = test_infer_no_outputs(model_name, input0, input1)
print("Results with no outputs:", results_no_outputs)

另附个人代码参考

需要注意 triton 和原算法中的数据格式不太一样，triton 一般是 numpy，原算法需要 tensor。

def create_triton_client(url='localhost:8000'):
    """
    创建并返回一个 Triton 客户端，同时进行健康检查。
    
    参数:
    url (str): Triton 服务器的地址，默认为 'localhost:8000'。
    
    返回:
    httpclient.InferenceServerClient: 已初始化并检查过的 Triton 客户端。
    """
    # 创建 Triton 客户端
    try:
        triton_client = httpclient.InferenceServerClient(url=url)
    except Exception as e:
        print("Failed to create Triton client: ", str(e))
        return None

    # 检查服务器健康状态
    try:
        if not triton_client.is_server_live():
            print("Triton server is not live")
            return None
        if not triton_client.is_server_ready():
            print("Triton server is not ready")
            return None
    except InferenceServerException as e:
        print("Server not ready: ", str(e))
        return None

    return triton_client

def triton_model_1(triton_client, pixel_seq, xy_seq, text, device, model_name='model_1', model_version='1'):
    # 检查模型是否准备好
    try:
        if not triton_client.is_model_ready(model_name, model_version):
            print(f"Model {model_name} version {model_version} is not ready")
            return None
    except InferenceServerException as e:
        print("Model not ready: ", str(e))
        return None
    
    # 设置输入
    inputs = []
    inputs.append(httpclient.InferInput('pixel_seq', pixel_seq.shape, 'INT32'))
    inputs.append(httpclient.InferInput('xy_seq', xy_seq.shape, 'INT32'))
    inputs.append(httpclient.InferInput('text', text.shape, 'INT32'))

    if pixel_seq.dtype != np.int32:
        pixel_seq = pixel_seq.astype(np.int32)
    
    inputs[0].set_data_from_numpy(pixel_seq)

    if xy_seq.dtype != np.int32:
        xy_seq = xy_seq.astype(np.int32)

    inputs[1].set_data_from_numpy(xy_seq)

    if text.dtype != np.int32:
        text = text.astype(np.int32)

    inputs[2].set_data_from_numpy(text)

    # 设置输出
    outputs = httpclient.InferRequestedOutput('output')

    # 执行推理
    try:
        results = triton_client.infer(model_name=model_name, inputs=inputs, outputs=[outputs])
    except InferenceServerException as e:
        print("Inference failed: ", str(e))
        return None

    # 获取输出结果
    triton_output = results.as_numpy('output')
    text_features = torch.tensor(triton_output, device=device)
    return text_features

通过这种方式，你可以轻松地在 Triton Inference Server 上进行多输入、多输出模型的推理。希望这篇博客能帮助你更好地理解和使用 Triton Inference Server。如果有任何问题或建议，欢迎在下方留言！

参考 triton-inference-server 官方代码库：
https://github.com/triton-inference-server/client/blob/main/src/python/examples/simple_http_infer_client.py