ResNet-50 triton ensemble摘要triton ensemble 部署onnx 模型1、模型准备1下载resnet50预训练权重2pt转onnx2、项目代码1resnet50_preprocess (预处理)config.pbtxtpreprocess.py2resnet50_inference (推理)config.pbtxt3resnet50_postprocess (后处理)config.pbtxtpostprocess.py4resnet50_ensemble (集成调度)config.pbtxt3、模型部署triton ensemble 部署tensorrt模型1、onnx转tensorrt2、项目代码3、模型部署摘要本博客详细描述如何把resent50模型的预处理、推理、后处理集合在同一个triton服务里面。其中预处理和后处理基于python为backend实现推理采用onnxruntime/tenosrr为backend最后使用triton服务的ensemble为backend把预处理、推理、后处理定制流水线。完整代码可以参考triton_ensemble_model_zootriton ensemble 部署onnx 模型1、模型准备准备好onnx模型如果已经有onnx模型可以跳过这一步如果没有可以参照我以下教程下载开源的resnet50的预训练权重并转化成onnx.1下载resnet50预训练权重importosimporttorchimporttorchvision.modelsasmodelsdefsave_resnet50_weights(save_dirweights,weight_nameresnet50_imagenet_v1.pth): 下载 ResNet-50 的 ImageNet 预训练权重并保存到指定目录。 Args: save_dir (str): 权重保存目录的路径 # 创建目录如果不存在os.makedirs(save_dir,exist_okTrue)# 定义保存路径save_pathos.path.join(save_dir,weight_name)# 检查是否已存在避免重复下载ifos.path.exists(save_path):print(f文件已存在:{save_path})returnprint(正在下载 ResNet-50 预训练权重 (IMAGENET1K_V2) ...)# 使用 torchvision 官方权重V2 版本准确率更高# weights 参数会自动下载并缓存但我们单独保存到指定目录modelmodels.resnet50(weightsmodels.ResNet50_Weights.IMAGENET1K_V1)# 提取 state_dict 并保存torch.save(model.state_dict(),save_path)print(f权重已保存至:{save_path})if__name____main__:save_resnet50_weights()2pt转onnximporttorchimporttorchvision.modelsasmodelsimportosdefconvert_to_onnx(weight_path,onnx_path,input_size(3,224,224)): 将 ResNet-50 的 PyTorch 权重转换为 ONNX 模型 Args: weight_path (str): 本地 .pth 权重文件路径 onnx_path (str): 输出的 ONNX 文件路径 input_size (tuple): 输入张量的形状 (C, H, W)默认 (3, 224, 224) # 1. 创建模型结构不加载预训练权重modelmodels.resnet50(weightsNone)devicecudaiftorch.cuda.is_available()elsecpu# 2. 加载本地权重ifnotos.path.exists(weight_path):raiseFileNotFoundError(f权重文件不存在:{weight_path})state_dicttorch.load(weight_path,map_locationdevice)model.load_state_dict(state_dict)# 3. 设置设备并切换到推理模式modelmodel.to(device)model.eval()# 4. 构造示例输入batch_size1c,h,winput_size example_inputtorch.randn(batch_size,c,h,w).to(device)# 5. 执行一次前向传播验证模型可用withtorch.no_grad():outputmodel(example_input)print(f前向验证通过输出形状:{output.shape})# 6. 导出 ONNXprint(f开始导出 ONNX 到:{onnx_path})torch.onnx.export(model,example_input,onnx_path,export_paramsTrue,# 保存模型参数opset_version11,# ONNX opset 版本常用 11/12do_constant_foldingTrue,# 常量折叠优化input_names[input],# 输入节点名output_names[output],# 输出节点名dynamic_axes{input:{0:batch_size},output:{0:batch_size}},# 动态 batch 维度dynamoFalse# 关键强制使用 TorchScript 导出器)print(ONNX 导出成功)if__name____main__:weight_pathrweights/resnet50_imagenet_v2.pthonnx_pathrweights/resnet50_imagenet_v2.onnxinput_size(3,224,224)convert_to_onnx(weight_path,onnx_path,input_size)2、项目代码triton服务对模型的目录有严格的要求需要按照以下格式进行存放models/ ├── resnet50_ensemble/ │ ├── 1/ │ └── config.pbtxt ├── resnet50_inference/ │ ├── 1/ │ │ └── resnet50_imagenet_v2.onnx │ └── config.pbtxt ├── resnet50_postprocess/ │ ├── 1/ │ │ ├── imagenet_classes.txt │ │ └── postprocess.py │ └── config.pbtxt └── resnet50_preprocess/ ├── 1/ │ └── preprocess.py └── config.pbtxtresnet50_preprocess (预处理)负责接收原始图像数据进行解码、Resize调整大小、Normalize归一化等操作将其转换为模型所需的 Tensor 格式。resnet50_inference (推理)负责模型推理resnet50_postprocess (后处理)负责接收模型的原始输出通常是 Logits 或 Probabilities执行 Softmax、Top-K 筛选等操作并利用 classes.txt 将数字 ID 映射为人类可读的分类结果如“猫”、“狗”。resnet50_ensemble (集成调度)这是一个特殊的“虚拟”模型。它不包含实际代码或权重而是通过 config.pbtxt 定义上述三个步骤的执行顺序和数据流向Preprocess - Inference - Postprocess对外提供一个统一的接口。1resnet50_preprocess (预处理)config.pbtxtname:resnet50_preprocessbackend:pythonmax_batch_size: 256 default_model_filename:preprocess.pyinput[{name:RAW_IMAGEdata_type: TYPE_STRING dims:[1]}]output[{name:PREPROCESSED_IMAGEdata_type: TYPE_FP32 dims:[3,224,224]}]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.resize_size256self.crop_size224# ImageNet 均值与标准差 (RGB 顺序)self.meannp.array([0.485,0.456,0.406],dtypenp.float32)self.stdnp.array([0.229,0.224,0.225],dtypenp.float32)defimg_preprocess(self,img)-np.ndarray:# 1. 获取原图尺寸h,wimg.shape[:2]# 2. 短边缩放至 256ifwh:new_wself.resize_size new_hint(h*self.resize_size/w)else:new_hself.resize_size new_wint(w*self.resize_size/h)img_resizedcv2.resize(img,(new_w,new_h),interpolationcv2.INTER_LINEAR)# (256, 256, 3)img_resizedcv2.cvtColor(img_resized,cv2.COLOR_BGR2RGB)# BGR to RGB# 3. 中心裁剪 224x224start_x(new_w-self.crop_size)//2start_y(new_h-self.crop_size)//2img_croppedimg_resized[start_y:start_yself.crop_size,start_x:start_xself.crop_size]# (224, 224, 3)# 4. 归一化到 [0,1] 并转为 float32img_normimg_cropped.astype(np.float32)/255.0# 5. 标准化img_resized(img_norm-self.mean)/self.std# 6. 转化通道img_resizednp.transpose(img_resized,(2,0,1))# (H,W,C) - (C,H,W)returnimg_resized.astype(np.float32)defbase64_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)ifin_tensorisNone:logger.error(Input tensor RAW_IMAGE not found!)continueraw_batchin_tensor.as_numpy()batch_sizeraw_batch.shape[0]imgs_resized[]# 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)else:logger.info(Base64 converted to image Successfully!)# 3. 图像预处理img_resizedself.img_preprocess(img)imgs_resized.append(img_resized)iflen(imgs_resized)0:# 如果没有图像创建一个0批次batch_imgsnp.zeros((batch_size,3,self.img_size,self.img_size),dtypenp.float32)else:batch_imgsnp.stack(imgs_resized,axis0)# (batch, 3, 224, 224)# 4. 构建输出张量out_tensorpb_utils.Tensor(PREPROCESSED_IMAGE,batch_imgs)responsepb_utils.InferenceResponse(output_tensors[out_tensor])responses.append(response)returnresponses2resnet50_inference (推理)config.pbtxtname:resnet50_inferencebackend:onnxruntimedefault_model_filename:resnet50_imagenet_v2.onnxmax_batch_size: 256 dynamic_batching{max_queue_delay_microseconds: 100000 preferred_batch_size:[16,32,64,128,256]}input[{name:inputdata_type: TYPE_FP32 dims:[3,224,224]# 固定输入尺寸NCHW 格式}]output[{name:outputdata_type: TYPE_FP32 dims:[1000]# 输出维度[batch, 1000类别]}]instance_group[{count: 1 kind: KIND_GPU gpus:[0]# 使用第 0 号 GPU}]3resnet50_postprocess (后处理)config.pbtxtname:resnet50_postprocessbackend:pythonmax_batch_size: 256 default_model_filename:postprocess.pyinput[{name:LOGITSdata_type: TYPE_FP32 dims:[1000]}]output[{name:OUTPUT_RESULTSdata_type: TYPE_STRING dims:[1]}]instance_group[{count: 32 kind: KIND_CPU}]postprocess.pyimporttriton_python_backend_utilsaspb_utilsimportnumpyasnpimportosimportjsonclassTritonPythonModel:definitialize(self,args):# 获取当前模型postprocess的版本目录model_repositoryargs[model_repository]# 例如: /models/resnet50_postprocessmodel_versionargs[model_version]# 例如: 1label_fileos.path.join(model_repository,model_version,imagenet_classes.txt)# 加载ImageNet的1000个标签withopen(label_file,r,encodingutf-8)asf:self.labels[line.strip()forlineinf.readlines()]defsoftmax(self,logits):exp_logitsnp.exp(logits-np.max(logits,axis1,keepdimsTrue))probsexp_logits/np.sum(exp_logits,axis1,keepdimsTrue)returnprobsdefexecute(self,requests):responses[]forrequestinrequests:# 1. 获取ONNX模型的输出logitslogits_tensorpb_utils.get_input_tensor_by_name(request,LOGITS)logitslogits_tensor.as_numpy()# shape: (batch_size, 1000)# 2. Softmax计算概率probsself.softmax(logits)# 3. 获取Top-5的索引和对应标签top5_indicesnp.argsort(probs,axis1)[:,-5:][:,::-1]batch_sizelogits.shape[0]batch_results[]foriinrange(batch_size):top5_labels[self.labels[idx]foridxintop5_indices[i]]top5_probs[float(probs[i][idx])foridxintop5_indices[i]]classify_item{top5_labels:top5_labels,top5_probs:top5_probs}result_json_strjson.dumps(classify_item)batch_results.append(result_json_str.encode(utf-8))# 4. 输出最终结果这里简单返回字符串和概率output_arraynp.array(batch_results,dtypeobject)out_tensorpb_utils.Tensor(OUTPUT_RESULTS,output_array)responsepb_utils.InferenceResponse(output_tensors[out_tensor])responses.append(response)returnresponses4resnet50_ensemble (集成调度)config.pbtxtname:resnet50_ensembleplatform:ensemblemax_batch_size: 256 input[{name:RAW_IMAGEdata_type: TYPE_STRING dims:[1]}]output[{name:OUTPUT_RESULTSdata_type: TYPE_STRING dims:[1]}]ensemble_scheduling{step[{model_name:resnet50_preprocessmodel_version:-1 input_map{key:RAW_IMAGEvalue:RAW_IMAGE}output_map{key:PREPROCESSED_IMAGEvalue:preprocessed_image}},{model_name:resnet50_inferencemodel_version:-1 input_map{key:inputvalue:preprocessed_image}output_map{key:outputvalue:logits}},{model_name:resnet50_postprocessmodel_version:-1 input_map{key:LOGITSvalue:logits}output_map{key:OUTPUT_RESULTSvalue:OUTPUT_RESULTS}}]}3、模型部署docker run-d \--gpus 1 \--name tritonserver \-p 127.0.0.1:8000:8000 \-v/path/to/resnet50/models:/models \ nvcr.io/nvidia/tritonserver:23.01-py3-v0.0.1 \ CUDA_VISIBLE_DEVICES1 tritonserver--model-repository/models--strict-model-configfalse--log-verbose1triton ensemble 部署tensorrt模型1、onnx转tensorrtdocker run--gpus 1-v $(pwd):/workspace-it nvcr.io/nvidia/tensorrt:23.01-py3 \ bash-c \cd /workspace \ trtexec \ --onnxresnet50_imagenet_v2.onnx \ --minShapesinput:1x3x224x224 \ --optShapesinput:256x3x224x224 \ --maxShapesinput:512x3x224x224 \ --workspace8192 \ --saveEngineresnet50_imagenet_v2_fp16.plan \ --explicitBatch \ --fp162、项目代码项目代码和部署onnx的一样只需要修改resnet50_inference部分内容把models/resnet50_inference/1/路径的onnx模型替换成tensorrt或者在文件夹1的同目录下创建文件夹2把tensorrt模型放在文件夹2里面修改models/resnet50_inference/config.pbtxtname:resnet50_inferencebackend:tensorrt# 把onnxruntime改成tensorrtdefault_model_filename:resnet50_imagenet_v2_fp16.plan# 改成tensorrt模型名称max_batch_size: 256 dynamic_batching{max_queue_delay_microseconds: 100000 preferred_batch_size:[16,32,64,128,256]}input[{name:inputdata_type: TYPE_FP32 dims:[3,224,224]# 固定输入尺寸NCHW 格式}]output[{name:outputdata_type: TYPE_FP32 dims:[1000]# 输出维度[batch, 1000类别]}]instance_group[{count: 1 kind: KIND_GPU gpus:[0]# 使用第 0 号 GPU}]3、模型部署部署和onnx部署一致。