如何使用tritonclient Python包为Triton模型传入输入并执行推理

前置准备 #

先安装对应协议的tritonclient Python包：

# HTTP协议客户端
pip install tritonclient[http]
# 若使用gRPC协议，安装这个
pip install tritonclient[grpc]

推理代码示例 #

以下代码完全匹配你提供的config.pbtxt配置，直接替换服务地址即可运行：

import numpy as np
import tritonclient.http as httpclient

# 初始化HTTP客户端，替换为你的Triton服务地址和端口
client = httpclient.InferenceServerClient(url="localhost:8000")

# 构造符合要求的输入数据：生成2个样本的随机数据，维度为(2, 32, 32, 3)，数据类型float32
batch_size = 2
input_data = np.random.rand(batch_size, 32, 32, 3).astype(np.float32)

# 创建输入张量，严格对应config里的input名称、维度和数据类型
inputs = [
    httpclient.InferInput(
        name="input_1",
        shape=(batch_size, 32, 32, 3),
        datatype="FP32"
    )
]
# 给输入张量赋值
inputs[0].set_data_from_numpy(input_data)

# 指定要获取的输出，对应config里的output名称
outputs = [
    httpclient.InferRequestedOutput(name="fc10")
]

# 发送推理请求，指定模型名"cifar10"
response = client.infer(model_name="cifar10", inputs=inputs, outputs=outputs)

# 提取并处理结果
result = response.as_numpy("fc10")
print(f"推理结果形状: {result.shape}")
print(f"第一个样本的分类概率输出: {result[0]}")

关键注意事项 #

• 维度匹配：输入张量必须加上batch维度（比如单个样本的话shape是(1,32,32,3)），与config.pbtxt里的dims配合使用。
• 协议选择：如果你的Triton服务用gRPC启动，只需把导入的tritonclient.http换成tritonclient.grpc，对应类名同步替换即可。
• 批量推理：利用max_batch_size:10000的配置，可将batch_size设置为1到10000之间的任意值，批量提交提升效率。

内容的提问来源于stack exchange，提问作者Mahesh