yolov5使用TorchServer部署

参考（抄袭）文章

• 参考大佬文章–pytorch部署新利器TorchServer
• U神大佬的yolov5源码
• TorchServer官方源码及教程
• 版本信息

  torch                  1.7.1+cu101
  torch-model-archiver   0.3.0
  torchaudio             0.7.2
  torchvision            0.8.2+cu101
  Python                 3.7.9
  

步骤过程

一、容器环境准备

• 先安装docker容器，我也没有装过，找it运维小伙子装的，版本为Docker version 19.03.8
• 拉取server容器，由于本机CUDA版本是10.0，好像不支持，加上没得权限升级，拉下的是cpu版本TorchServer容器，找对应的TorchServer容器，如下图使用命令拉取容器
  docker pull pytorch/torchserve:latest

二、模型准备

• 自己先跑yolov5得到自己的模型文件，best.pt文件；
  

• 将训练好的best.pt模型文件导出为best.torchscript.pt

  # 这里有一个小坑，需要修改 export.py 文件 将下面一行修改为False，本身该文件为cpu导出
  # model.model[-1].export = False  # set Detect() layer export=True
  python models/export.py --weights ./runs/train/exp6/weights/best.pt --img 640 --batch 1
  

• 运行上述命令后得到如下文件，其中best.torchscript.pt是我们需要的
  

• 验证模型是否正确输出

  model = torch.jit.load("./assets/best.torchscript.pt")
  img = Image.open("./images/people(47).jpg").convert("RGB")
  img_ = img.resize((640, 640))
  inp = tfrm.ToTensor()(img_).unsqueeze(0)
  print('input shape', inp.shape)
  
  pred = model(inp)
  
  pred = pred[0]
  print('pred shape', pred.shape)
  pred = non_max_suppression(pred)[0].cpu().numpy().tolist()
  print(pred)
  
  draw = ImageDraw.Draw(img, "RGBA")
  
  for p in pred:
  	p = adjust_ratio_xyxy(p, img.size, 640)
  	x1, y1, x2, y2 = p[:4]
  	draw.rectangle(((x1, y1), (x2, y2)), fill=(200, 100, 0, 127))
  	draw.rectangle(((x1, y1), (x2, y2)), outline=(0, 0, 0, 127), width=3)
  
  img.show()
  

• 将上一步导出的best.torchscript.pt模型及环境打包

  torch-model-archiver --model-name fallDown --version 1.0 --serialized-file /data/share/imageAlgorithm/zhangcheng/code/detect_torchserve/inference-endpoint/assets/best.torchscript.pt --extra-files /data/share/imageAlgorithm/zhangcheng/code/detect_torchserve/inference-endpoint/assets/index_to_name.json,/data/share/imageAlgorithm/zhangcheng/code/detect_torchserve/inference-endpoint/objectDetectionHandler.py,/data/share/imageAlgorithm/zhangcheng/code/detect_torchserve/inference-endpoint/utils.py,/data/share/imageAlgorithm/zhangcheng/code/detect_torchserve/inference-endpoint/config.py --handler /data/share/imageAlgorithm/zhangcheng/code/detect_torchserve/inference-endpoint/endpointHandler.py --export-path /data/share/imageAlgorithm/zhangcheng/code/detect_torchserve/inference-endpoint/model-store -f
  

• index_to_name.json文件，字典对应的类别

  {
  "0": "fall"
  }
  
  

• config.py 配置文件

  scaledSize = 640
  idx_to_class = {
  			    "0": "fall"
  				}
  

• endpointHandler.py

  # coding=utf-8
  """
  @Time: 2021/3/17 17:49
  @IDE: PyCharm
  @Author: chengzhang
  @File: endpointHandler.py
  """
  from objectDetectionHandler import ObjectDetectionHandler
  
  _service = ObjectDetectionHandler()
  
  
  def handle(data, context):
   	if not _service.initialized:
     	 	_service.initialize(context)
  
  	if data is None:
      	return None
  
  	data = _service.preprocess(data)
  	data = _service.inference(data)
  	data = _service.postprocess(data)
  	return data
  

  参数意思来源–pytorch部署新利器TorchServer

    ```
    # 参数意思  
    --model-name: 模型的名称，后来的接口名称和管理的模型名称都是这个
    --serialized-file: 模型环境及代码及参数的打包文件
    --export-path: 本次打包文件存放位置
    --extra-files: handle.py中需要使用到的其他文件
    --handler: 指定handler函数。（模型名:函数名）
    -f 覆盖之前导出的同名打包文件
    ```
  

准备下班，不写了