YOLOv5从入门到部署之：模型部署

本文记录了 YOLOv5 的部署过程，包含以下代码，下载方式在文末。

（1）onnxruntime 部署代码；

（3）ncnn 部署代码；

Setp1：安装环境，前面我们已经安装好了 yolov5 需要的环境 pip install -r requirements.txt现在我们还需要再安装 onnx 和 coremltools。

pip install onnx>=1.7.0 
# for ONNX export
pip installcoremltools==4.0  
# for CoreMLexport

Setp2：将预先训练的 YOLOv5s 模型导出为 ONNX、TorchScript 和 CoreML 格式

export PYTHONPATH="$PWD" && python models/export.py --weights ./weights/yolov5s.pt --img 640 --batch 1

转换过程会有如下日志输出：

导出的 3 种模型与原始 PyTorch 模型一起保存

model.model[-1].export = True  
# set Detect() 
layer export=True

首先，先加载模型：

session = onnxruntime.InferenceSession('runs/train/exp1/weights/best.onnx')
print("The model expects input shape:", session.get_inputs()[0].shape)
batch_size = session.get_inputs()[0].shape[0]
img_size_h = session.get_inputs()[0].shape[2]
img_size_w = session.get_inputs()[0].shape[3]

然后，对图像做好预处理后，进行模型的推理：

image_src = Image.open(image_path)
resized = letterbox_image(image_src, (img_size_w, img_size_h))
img_in = np.transpose(resized, (2, 0, 1)).astype(np.float32)  
# HWC -> CHW
img_in = np.expand_dims(img_in, axis=0)
img_in /= 255.0
# inferenceinput_name = session.get_inputs()[0].name
outputs = session.run(None, {input_name: img_in})

最后，对包含解码的 ONNX 模型，只需要进行 MNS 操作即可，对没有后处理的 ONNX 模型，还需要先进行解码再进行 MNS。

2 Flask 部署

首先确保安装了 pytorch，因为需要使用 flask 这个 web 框架，所以当然需要安装 flask，非常简单，使用下面的命令进行安装。

pip install flask

每次启动模型，加载模型是一件非常费时间的事情，所以最好的办法就是 load 一次模型。我们直接使用 YOLOv5 里面的 attempt_load 函数进行加载。

model = attempt_load(weights, map_location=device) # load FP32 model

接着，就是 YOLOv5 中标准的预处理流程，首先将图片 resize 到固定的大小，然后转换成 tensor，接着做标准化。

if webcam:    
view_img = True   
cudnn.benchmark = True  # set True to speed up constant image size inference    
dataset = LoadStreams(source, img_size=imgsz)else:    
save_img = True    
dataset = LoadImages(source, img_size=imgsz)

然后，定义请求方式为 POST，表示向服务器传输数据，接着定义一个 predict 函数来进行模型的前向传播。

def detect(): 
   save_img = False    
form_data = request.json    
print(form_data)    ...

最后我们在 main 函数中调用

if __name__ == '__main__':    
app.run()

测试时，需要先启动 flask 服务：

python flask_detect.py <path-to-your-model>
python flask_detect.py <path-to-your-model> <GPU-device-number(s)>

然后，使用 curl 命令或者提供的 python 脚本去访问服务器：

python client_flask.py
curl -X POST -H "Content-Type: application/json" -d @api_data.json http://localhost:5000/detect

测试结果如下

3 NCNN 部署

python3 -m onnxsim best.onnx best-sim.onnx

onnx2ncnn best-sim.onnx best.param best.bin

加载放在 example/yolov5 目录下的模型：

yolov5.load_param("../examples/yolov5/yolov5b.param");yolov5.load_model("../examples/yolov5/yolov5b.bin");

然后进行 YOLOv5 中标准的 预处理流 程：

const int target_size = 640;
float r_w = input_size / (img_w*1.0);
float r_h = input_size/ (img_h*1.0);
int true_w, true_h, x, y;if (r_h > r_w) {true_w = input_size;true_h = r_w * img_h;x = 0;y = (input_size - true_h) / 2;} 
else {true_w = r_h* img_w;true_h = input_size;x = (input_size - true_w) / 2;y = 0;}
//ncnn::Mat in = ncnn::Mat::
from_pixels_resize(bgr.data, ncnn::Mat::PIXEL_BGR2RGB, bgr.cols, bgr.rows, target_size, target_size);
ncnn::Mat in = ncnn::Mat::
from_pixels(bgr.data, ncnn::Mat::PIXEL_BGR2RGB, bgr.cols, bgr.rows);
const float mean_vals[3] = {0, 0, 0};
const float norm_vals[3] = {1 / 255.f, 1 / 255.f, 1 / 255.f};
in.substract_mean_normalize(mean_vals, norm_vals);

最后进行解码和 NMS 操作：

ncnn::Extractor ex = yolov5.create_extractor();ex.set_num_threads(4);
ex.input(0, in);std::
vector<BoxInfo> result;for(auto& 
layer: layers){ncnn::Mat blob;    ex.extract(layer.name.c_str(),blob);
auto boxes = decode_infer(blob,layer.stride,{img_w, img_h},true_w, 
true_h,num_class,layer.anchors,threshold, x, y);    
result.insert(result.begin(),boxes.begin(),boxes.end());}
nms(result,nms_threshold);

Yolov5 从它出现以来就饱受争议，大家认为它的创新不足更应该叫做 yolov4 的 pytorch 版本。但是不可否认的是 yolov5 是一个优秀的目标检测框架，yolov5 大量整合了计算机视觉领域的 State-of-the-art，它相对于 yolov4 更方便我们的使用，我们可以在 yaml 文件简单的修改网络结构，代码的可读性更高，能够更方便修改代码。