作者：昇腾实战派

小模型在NPU上的推理部署： 【知识地图】

引言

Triton Inference Server 是 NVIDIA 开源的高性能推理服务框架，支持多种后端和模型格式，广泛应用于生产环境。本文将详细介绍如何在昇腾（Ascend）AI处理器上使用 Triton Server 部署 YOLOv11 目标检测模型（同样适用于 YOLOv5、YOLOv8 系列）。我们将以 YOLOv11 为例，通过 Python 后端调用昇腾推理引擎 ais_bench 加载离线模型（.om）并提供服务。文章涵盖模型仓库组织、配置文件编写、服务端代码实现、客户端调用以及服务启动命令，帮助开发者快速上手昇腾设备上的 Triton 部署。

环境准备

• 硬件：昇腾 AI 处理器（如 Atlas 300 系列）
• 软件：
  • 驱动与固件：根据昇腾社区指引安装对应版本
  • 镜像推荐：使用昇腾社区提供的 Triton Server 镜像

模型仓库目录结构

Triton Server 要求模型按指定目录结构组织。以下是一个 YOLOv11 模型的典型布局：

models/
└── yolo11                     # 模型名称，与 config.pbtxt 中的 name 一致
    ├── 1                       # 模型版本号（必须为数字）
    │   └── model.py             # Python 后端核心代码
    ├── client_yolo11.py         # 客户端测试脚本（非必需）
    ├── config.pbtxt             # 模型配置文件
    ├── yolo11s_bs8.om           # 昇腾离线模型（batch=8）
    └── yolov8s_bs8.om           # 可存放多个模型文件

• 1/ 目录表示版本号，内部必须包含 model.py（Python 后端的入口文件）。
• config.pbtxt 描述了模型的输入输出、后端类型、参数等。
• .om 文件是经过昇腾 ATC 工具转换后的离线模型，可在不同 batch size 下使用。

配置文件 config.pbtxt 详解

# 模型名称，通常与存放此配置文件的目录名保持一致
name: "yolo11"

# 指定模型运行的平台/后端，这里是 ONNX Runtime
backend: "python"

# 模型支持的最大批处理大小，根据模型和显存调整
max_batch_size: 8

# 输入节点配置
input [
  {
    name: "images"            # 输入张量的名称，必须与ONNX模型中的名称一致
    data_type: TYPE_FP32       # 数据类型，通常是 FP32（根据模型实际输入信息填写）
    dims: [ 3, 640, 640 ]      # 输入维度: [通道数, 高度, 宽度]，注意这里不包含批次维度(N)  （根据模型实际输入信息填写）
  }
]

# 输出节点配置
output [
  {
    name: "output0"          
    data_type: TYPE_FP32        # 数据类型，通常是 FP32（根据模型实际输出信息填写）
    dims: [ 84, 8400 ]          # 输出维度: [预测信息(坐标+置信度+类别), 检测框数量]
  }
]


# 实例组配置 (可选，定义使用多少个GPU实例并行)
instance_group [
  {
    count: 1                      # 实例数量
  }
]

parameters: [
  {
    key: "batch_size",
    value: { string_value: "1" }
  },
  {
    key: "model_path",
    value: { string_value: "/home/users/models/yolo11/yolo11s_bs8.om" }
  },
  {
    key: "device_id",
    value: { string_value: "0" }
  }
]

说明：

• 输入输出维度必须与转换后的 OM 模型严格匹配。
• parameters 部分可以传递任意自定义参数，在 Python 后端的 initialize 方法中通过 model_config.get('parameters', {}) 读取。
• max_batch_size 应与 OM 模型支持的 batch 一致（例如 yolo11s_bs8.om 支持 batch 8）。

服务端代码（Python 后端）

文件路径：models/yolo11/1/model.py

import json
import numpy as np
import torch
import triton_python_backend_utils as pb_utils
from ais_bench.infer.interface import InferSession

class TritonPythonModel:
    """Triton Python 后端模型类"""
    def load_model(self, model_path, device_id):
        """Load OM model using ais_bench InferSession."""
        return InferSession(int(device_id), model_path)

    def initialize(self, args):
        """模型初始化，只调用一次"""

        model_config = json.loads(args['model_config'])
        self.input_config = model_config['input'][0]  
        self.output_config = model_config['output'][0] 
        print(f"Model initialized with input: {self.input_config}, output: {self.output_config}")

        self.input_dtype = pb_utils.triton_string_to_numpy(self.input_config['data_type'])
        self.output_dtype = pb_utils.triton_string_to_numpy(self.output_config['data_type'])

        params = model_config.get('parameters', {})    
        self.batch_size = int(params['batch_size']['string_value'])  
        self.model_path = params['model_path']['string_value'] 
        self.device_id = int(params['device_id']['string_value'])    

        self.model = self.load_model(self.model_path, self.device_id)
        inputs_shape = self.model.get_inputs()[0].shape
        print(f"model inputs_shape: {inputs_shape}")
        self.input_name = model_config['input'][0]['name']  
        self.output_name = model_config['output'][0]['name']


    def execute(self, requests):
        responses = []
        for request in requests:
            input_tensor = pb_utils.get_input_tensor_by_name(request, self.input_name)
            input_data = input_tensor.as_numpy() 

            outputs = self.model.infer([input_data])[0]

            output_tensor = pb_utils.Tensor(self.output_name, outputs.astype(self.output_dtype))
            response = pb_utils.InferenceResponse(output_tensors=[output_tensor])
            responses.append(response)
        return responses

    def finalize(self):
        """清理资源（可选）"""
        print("Cleaning up resources...")

关键点：

• ​必须实现 Triton Python 后端标准接口​：在 model.py 中，需要定义 TritonPythonModel 类，并实现 initialize、execute 和 finalize 三个方法。这是 Triton Server 调用 Python 后端的约定，确保模型能被正确加载和执行。
• ​模型加载与初始化​：在 initialize 方法中，通过 ais_bench.infer.interface.InferSession 加载昇腾 OM 模型，并读取配置文件 config.pbtxt 中 parameters 传递的自定义参数（如模型路径、设备 ID 等）。同时获取输入输出的名称和数据类型，供后续推理使用。
• ​核心推理逻辑​：execute 方法处理每个推理请求。它遍历请求列表，从每个请求中提取输入张量（转换为 NumPy 数组），调用 model.infer 执行昇腾模型推理，然后将输出结果封装为 pb_utils.Tensor，最后构造 pb_utils.InferenceResponse 返回。
• ​注意输出格式​：model.infer 返回的是一个列表，因为 OM 模型可能有多个输出。本例中取第一个输出 [0]，实际使用时需根据模型输出数量调整索引。同时，输出数据需转换为配置文件中声明的数据类型（self.output_dtype）。
• ​资源清理​：finalize 方法在模型卸载时调用，可用于释放模型句柄或其他资源（虽然不是必须的，但建议实现以避免资源泄漏）。

客户端代码示例

文件路径：models/yolo11/client_yolo11.py

import numpy as np
import tritonclient.http as httpclient

# 创建随机输入数据（模拟 batch=8 的图像）
input_data = np.random.randn(8, 3, 640, 640).astype(np.float32)

# 创建 HTTP 客户端（Triton 默认端口 8000，此处使用 9000）
client = httpclient.InferenceServerClient(url='localhost:9000')

# 定义输入输出
inputs = [httpclient.InferInput('images', input_data.shape, 'FP32')]
inputs[0].set_data_from_numpy(input_data)
outputs = [httpclient.InferRequestedOutput('output0')]

# 发送推理请求
response = client.infer('yolo11', inputs=inputs, outputs=outputs)
result = response.as_numpy('output0')

print(f"推理成功！输出形状: {result.shape}")

# 提示：前后处理需自行实现
def preprocess():
    pass

def postprocess():
    pass

说明：

• 客户端使用 HTTP 协议，端口需与 Triton 启动时指定的 --http-port 一致。
• 输入数据是模拟的随机数组，实际使用时需替换为真实图像的预处理结果。
• 输出 result 的形状应与 OM 模型输出一致，例如 (8, 84, 8400)。

启动 Triton Server

在昇腾镜像上启动 Triton Server 服务，执行以下命令：

/opt/tritonserver/bin/tritonserver \
    --model-repository=/path/to/your/models \
    --http-port=9000 \
    --grpc-port=9002

• --model-repository：指定模型仓库的根目录（绝对路径），例如 /home/user/models。
• --http-port：HTTP 服务端口，客户端通过该端口发送 HTTP 请求（默认 8000，此处改为 9000）。
• --grpc-port：gRPC 服务端口（默认 8001，此处改为 9002）。

启动后，Triton Server 会加载 models/ 下所有有效模型，并打印日志。您可以通过 curl 或客户端验证服务是否正常。

总结与后续工作

本文提供了一个在昇腾设备上使用 Triton Server 部署 YOLOv11 模型的完整示例，包括：

• 模型仓库的目录结构
• 配置文件 config.pbtxt 的详细解释
• Python 后端服务端代码（基于 ais_bench 推理接口）
• 客户端调用示例
• Triton Server 启动命令

注意：本文仅聚焦于模型的服务化部署，未包含预处理（如图像缩放、归一化）和后处理（如 NMS、结果解析）逻辑。开发者需根据实际业务需求自行实现这些部分。

前后处理的实现可参考昇腾官方仓库示例：

🔗 Ascend/modelzoo-GPL - Yolov8 for PyTorch 前后处理参考