【triton】语意分割（Deep Labv3 ）基于triton ensemble部署推理服务

deeplabv3 triton ensemble摘要triton ensemble 部署 onnx 模型1、torch转onnx2、项目代码1deeplabv3_preprocess (预处理)config.pbtxtpreprocess.py2deeplabv3_inference (推理)config.pbtxt3deeplabv3_postprocess (后处理)config.pbtxtpostprocess.py4deeplabv3_ensemble (集成调度)config.pbtxt3、模型部署triton ensemble 部署 tensorrt 模型1、onnx转tensorrt2、项目代码3、模型部署摘要本博客详细地描述如何把语意分割deeplabv3模型的预处理、推理、后处理部署在同一个triton服务里面其中使用预处理和后处理采用python为backend实现推理使用onnxruntime/tensorrt为backend最后通过triton的ensemble定制预处理、推理和后处理的流水线。完整的项目代码可参考triton_ensemble_model_zootriton ensemble 部署 onnx 模型1、torch转onnximporttorchimporttorchvisionimportonnx# 1. 加载预训练模型 modeltorchvision.models.segmentation.deeplabv3_resnet50(weightsNone,# 不使用内置预训练权重aux_lossTrue,# 启用 aux_classifier以匹配权重文件num_classes21# COCO 标准是 21 类含背景请按需修改)devicecudaiftorch.cuda.is_available()elsecpuweight_pathrweights/deeplabv3_resnet50_coco.pthstate_dicttorch.load(weight_path,map_locationdevice)model.load_state_dict(state_dict)modelmodel.to(device).eval()# 2. 准备示例输入 # ONNX 导出需要提供一个示例输入张量用于追踪计算图。# 这里 batch_size 设为 1但稍后会通过 dynamic_axes 声明为动态。batch_size1num_channels3height,width512,512# DeepLabV3 期望的典型输入尺寸dummy_inputtorch.randn(batch_size,num_channels,height,width)# 3. 定义动态轴 # 在导出时指定哪些维度是动态的可变长度。# 我们通常希望 batch 维度动态变化也可以让高度/宽度动态但需要注意模型内部可能对尺寸有隐含约束。dynamic_axes{input:{0:batch_size},# 输入张量的第 0 维是 batchoutput:{0:batch_size},# 输出字典中的 out 张量的第 0 维是 batch# 如果还需要动态图像大小可以取消注释下面的行但可能需要对预处理做额外处理# input: {2: height, 3: width},# output: {2: height, 3: width},}# 4. 导出 ONNX 模型 onnx_file_pathweights/deeplabv3_resnet50_coco.onnxtorch.onnx.export(model,dummy_input,# 示例输入onnx_file_path,# 保存路径input_names[input],# 输入名字output_names[output],# 输出名字dynamic_axesdynamic_axes,# 动态轴配置opset_version11,# ONNX opset 版本建议 11do_constant_foldingTrue,# 折叠常量优化dynamoFalse,# 是否打印导出日志)print(fONNX 模型已保存至:{onnx_file_path})# 5. 验证导出的 ONNX 模型可选 # 加载 ONNX 模型进行结构验证onnx_modelonnx.load(onnx_file_path)onnx.checker.check_model(onnx_model)print(ONNX 模型验证通过)# 6. 使用 ONNX Runtime 进行简单推理测试可选try:importonnxruntimeasortimportnumpyasnp# 创建 ONNX Runtime 推理会话ort_sessionort.InferenceSession(onnx_file_path)# 准备不同 batch 大小的输入测试动态 batch 是否正常工作fortest_batchin[1,2,4]:test_inputnp.random.randn(test_batch,3,height,width).astype(np.float32)outputsort_session.run([output],{input:test_input})print(fBatch size{test_batch}- 输出形状:{outputs[0].shape})exceptImportError:print(未安装 onnxruntime跳过动态 batch 测试。可运行 pip install onnxruntime 安装)2、项目代码triton服务对文件夹目录有严格要求文件夹如下存放models/ ├── deeplabv3_ensemble/ │ ├── 1/ │ └── config.pbtxt ├── deeplabv3_inference/ │ ├── 1/ │ │ └── deeplabv3_resnet50_coco.onnx │ └── config.pbtxt ├── deeplabv3_postprocess/ │ ├── 1/ │ │ └── postprocess.py │ └── config.pbtxt └── deeplabv3_preprocess/ ├── 1/ │ └── preprocess.py └── config.pbtxtdeeplabv3_preprocess (预处理)作用负责接收原始图像数据进行解码、Resize调整大小至模型输入尺寸如 512x512、Normalize归一化等操作将其转换为模型所需的 Tensor 格式。deeplabv3_inference (推理)存放模型和配置文件。deeplabv3_postprocess (后处理)作用负责接收模型的原始输出通常是 Logits 或 Argmax 后的掩码图执行颜色映射Color Map、调整回原图尺寸等操作最终生成可视化的分割结果。deeplabv3_ensemble (集成调度)这是一个特殊的“虚拟”模型。它不包含实际代码或权重而是通过 config.pbtxt 定义上述三个步骤的执行顺序和数据流向Preprocess - Inference - Postprocess对外提供一个统一的 API 接口。1deeplabv3_preprocess (预处理)config.pbtxtname:deeplabv3_preprocessbackend:pythonmax_batch_size: 128 default_model_filename:preprocess.pyinput[{name:RAW_IMAGEdata_type: TYPE_STRING dims:[1]}]output[{name:PREPROCESSED_IMAGEdata_type: TYPE_FP32 dims:[3,512,512]},{name:IMAGE_SHAPEdata_type: TYPE_INT64 dims:[2]}]instance_group[{count: 32 kind: KIND_CPU}]preprocess.pyimporttriton_python_backend_utilsaspb_utilsimportnumpyasnpimportcv2importbase64importlogging logging.basicConfig(levellogging.INFO)loggerlogging.getLogger(__name__)classTritonPythonModel:definitialize(self,args):self.input_height512self.input_width512self.meannp.array([0.485,0.456,0.406],dtypenp.float32)self.stdnp.array([0.229,0.224,0.225],dtypenp.float32)defimg_preprocess(self,image_bgr)-np.ndarray: 使用 OpenCV 和 NumPy 预处理图像 Args: image_bgr: BGR 格式的图像 (H, W, 3)dtypeuint8 Returns: input_tensor: (1, 3, H, W) float32 numpy array已归一化 # 1. BGR - RGBimage_rgbcv2.cvtColor(image_bgr,cv2.COLOR_BGR2RGB)# 2. Resize 到固定尺寸 (520, 520)使用双线性插值resizedcv2.resize(image_rgb,(self.input_width,self.input_height),interpolationcv2.INTER_LINEAR)# 3. 归一化uint8 [0,255] - float [0,1]img_floatresized.astype(np.float32)/255.0# 4. 标准化 (ImageNet 统计)img_norm(img_float-self.mean)/self.std# 5. HWC - CHW 并添加 batch 维度input_tensornp.transpose(img_norm,(2,0,1))# (3, H, W)returninput_tensordefbase64_to_image(self,base64_str):try:raw_bytesbase64.b64decode(base64_str)nparrnp.frombuffer(raw_bytes,np.uint8)imgcv2.imdecode(nparr,cv2.IMREAD_COLOR)returnimgexceptExceptionaserror:logger.error(fError:{error})logger.error(ferror line:{error.__traceback__.tb_lineno})defexecute(self,requests):responses[]forrequestinrequests:# 1. 获取原始图像数据 (base64), 是一个批次in_tensorpb_utils.get_input_tensor_by_name(request,RAW_IMAGE)raw_batchin_tensor.as_numpy()batch_sizeraw_batch.shape[0]imgs_resized[]origins_shape[]# 2. 处理批次中的每个图像foriinrange(batch_size):base64_strraw_batch[i][0]# 获取第i个图像的字节# 3. 把 base64 转成 imageimgself.base64_to_image(base64_str)ifimgisNone:# 如果解码失败可以插入一个黑色图像或者报错这里我们插入一个黑色图像logger.info(Base64 converted to image failed!)imgnp.zeros((self.img_size,self.img_size,3),dtypenp.uint8)origins_shape.append([self.input_width,self.input_height])else:logger.info(Base64 converted to image Successfully!)orig_h,orig_wimg.shape[:2]origins_shape.append([orig_h,orig_w])# 4. 图像预处理img_resizedself.img_preprocess(img)imgs_resized.append(img_resized)iflen(imgs_resized)0:# 如果没有图像创建一个0批次batch_imgsnp.zeros((batch_size,3,self.input_width,self.input_height),dtypenp.float32)batch_shapesnp.zeros((batch_size,2),dtypenp.int64)else:batch_imgsnp.stack(imgs_resized,axis0)# (batch, 3, 520, 520)batch_shapesnp.array(origins_shape,dtypenp.int64)# (batch_size, 2)# 5. 构建输出张量out_tensorpb_utils.Tensor(PREPROCESSED_IMAGE,batch_imgs)out_shapepb_utils.Tensor(IMAGE_SHAPE,batch_shapes)responsepb_utils.InferenceResponse(output_tensors[out_tensor,out_shape])responses.append(response)returnresponses2deeplabv3_inference (推理)config.pbtxtname:deeplabv3_inferencebackend:onnxruntimedefault_model_filename:deeplabv3_resnet50_coco.onnxmax_batch_size: 128 dynamic_batching{max_queue_delay_microseconds: 100000 preferred_batch_size:[4,16,32,64,128]}input[{name:inputdata_type: TYPE_FP32 dims:[3,512,512]}]output[{name:outputdata_type: TYPE_FP32 dims:[-1,-1,-1]},{name:621data_type: TYPE_FP32 dims:[-1,-1,-1]}]instance_group[{count: 1 kind: KIND_GPU gpus:[0]# 使用第 0 号 GPU}]3deeplabv3_postprocess (后处理)config.pbtxtname:deeplabv3_postprocessbackend:pythonmax_batch_size: 128 default_model_filename:postprocess.pyinput[{name:MASKSdata_type: TYPE_FP32 dims:[-1,-1,-1]},{name:IMAGE_SHAPEdata_type: TYPE_INT64 dims:[2]}]output[{name:OUTPUT_RESULTSdata_type: TYPE_STRING dims:[1]}]instance_group[{count: 32 kind: KIND_CPU}]postprocess.pyimporttriton_python_backend_utilsaspb_utilsimportnumpyasnpimportbase64importjsonimportcv2classTritonPythonModel:definitialize(self,args):passdefmask_to_base64(self,mask): input: mask: [h, w], 像素值在[0, 21] output: base64_encodingbase64 mask_imagemask[:,:,np.newaxis]success,buffercv2.imencode(.png,mask_image)ifnotsuccess:raiseValueError(图像编码失败)base64_maskbase64.b64encode(buffer).decode(utf-8)returnbase64_maskdefpostprocess_mask(self,output,original_sizeNone): 后处理取 argmax可选上采样回原始尺寸 Args: output: (num_classes, H, W) numpy array original_size: (width, height) 原始图像尺寸 Returns: mask: (H, W) 或 (orig_H, orig_W) numpy arraydtypeuint8 # 取类别索引masknp.argmax(output,axis0)# (H, W)maskmask.astype(np.uint8)# (H, W)iforiginal_sizeisnotNone:orig_h,orig_woriginal_size maskcv2.resize(mask,(orig_w,orig_h),interpolationcv2.INTER_NEAREST)returnmaskdefexecute(self,requests):responses[]forrequestinrequests:masks_tensorpb_utils.get_input_tensor_by_name(request,MASKS)masksmasks_tensor.as_numpy()shape_tensorpb_utils.get_input_tensor_by_name(request,IMAGE_SHAPE)origins_shapeshape_tensor.as_numpy()batch_sizemasks.shape[0]batch_results[]foriinrange(batch_size):maskself.postprocess_mask(masks[i],origins_shape[i])mask_item{mask:self.mask_to_base64(mask),mask_shape:list(mask.shape),min_class:int(mask.min()),max_class:int(mask.max()),# origin shape:: list(origins_shape[i])}result_json_strjson.dumps(mask_item)batch_results.append(result_json_str.encode(utf-8))output_arraynp.array(batch_results,dtypeobject)out_tensorpb_utils.Tensor(OUTPUT_RESULTS,output_array)responsepb_utils.InferenceResponse(output_tensors[out_tensor])responses.append(response)returnresponses4deeplabv3_ensemble (集成调度)config.pbtxtname:deeplabv3_ensembleplatform:ensemblemax_batch_size: 128 input[{name:RAW_IMAGEdata_type: TYPE_STRING dims:[1]}]output[{name:OUTPUT_RESULTSdata_type: TYPE_STRING dims:[1]}]ensemble_scheduling{step[{model_name:deeplabv3_preprocessmodel_version:-1 input_map{key:RAW_IMAGEvalue:RAW_IMAGE}output_map[{key:PREPROCESSED_IMAGEvalue:preprocessed_image},{key:IMAGE_SHAPEvalue:IMAGE_SHAPE}]},{model_name:deeplabv3_inferencemodel_version:-1 input_map{key:inputvalue:preprocessed_image}output_map{key:outputvalue:output}},{model_name:deeplabv3_postprocessmodel_version:-1 input_map[{key:MASKSvalue:output},{key:IMAGE_SHAPEvalue:IMAGE_SHAPE}]output_map{key:OUTPUT_RESULTSvalue:OUTPUT_RESULTS}}]}3、模型部署docker run-d \--gpus 1 \--name tritonserver \-p 127.0.0.1:8000:8000 \-v deeplabv3/models:/models \ nvcr.io/nvidia/tritonserver:23.01-py3-v0.0.1 \ CUDA_VISIBLE_DEVICES0 tritonserver--model-repository/models--strict-model-configfalse--log-verbose1注意tritonserver镜像需要下载opencv依赖库triton ensemble 部署 tensorrt 模型1、onnx转tensorrtnerdctl run--gpus 1-v $(pwd):/workspace-it nvcr.io/nvidia/tensorrt:23.01-py3 \ bash-c \cd /workspace \ trtexec \ --onnxdeeplabv3_resnet50_coco.onnx \ --minShapesinput:1x3x512x512 \ --optShapesinput:64x3x512x512 \ --maxShapesinput:128x3x512x512 \ --workspace8192 \ --saveEnginedeeplabv3_resnet50_coco_fp16.plan \ --explicitBatch \ --fp162、项目代码项目代码和部署onnx一致只需要修改deeplabv3_inference部分把tēnsorrt模型替换models/deeplabv3_inference/1/中的onnx模型或者在文件夹1的同目录下创建文件2把tensorrt模型放在文件2里面。修改models/deeplabv3_inference/config.pbtxtname:deeplabv3_inferencebackend:tensorrt## 把onnxruntime改成tensorrtdefault_model_filename:deeplabv3_resnet50_coco_fp16.plan# 修改为tensorrt的名称max_batch_size: 128 dynamic_batching{max_queue_delay_microseconds: 100000 preferred_batch_size:[4,16,32,64,128]}input[{name:inputdata_type: TYPE_FP32 dims:[3,512,512]}]output[{name:outputdata_type: TYPE_FP32 dims:[-1,-1,-1]},{name:621data_type: TYPE_FP32 dims:[-1,-1,-1]}]instance_group[{count: 1 kind: KIND_GPU gpus:[0]# 使用第 0 号 GPU}]3、模型部署和onnx部署一致。