本文记录了YOLOv5的部署过程，包含以下代码，下载方式在文末。
（1）onnxruntime部署代码；
（3）ncnn部署代码；
Setp1：安装环境，前面我们已经安装好了yolov5需要的环境pip install -r requirements.txt现在我们还需要再安装onnx和coremltools。
|
1 2 3 4 |
pip install onnx>=1.7.0 # for ONNX export pip installcoremltools==4.0 # for CoreMLexport |
Setp2：将预先训练的 YOLOv5s 模型导出为 ONNX、TorchScript 和 CoreML 格式
|
1 |
export PYTHONPATH="$PWD" && python models/export.py --weights ./weights/yolov5s.pt --img 640 --batch 1 |
转换过程会有如下日志输出：
导出的3种模型与原始 PyTorch 模型一起保存
|
1 2 3 |
model.model[-1].export = True # set Detect() layer export=True |
Setp3：检测
首先，先加载模型：
|
1 2 3 4 5 |
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] |
然后，对图像做好预处理后，进行模型的推理：
|
1 2 3 4 5 6 7 8 |
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，非常简单，使用下面的命令进行安装。
|
1 |
pip install flask |
每次启动模型，加载模型是一件非常费时间的事情，所以最好的办法就是load一次模型。我们直接使用YOLOv5里面的attempt_load函数进行加载。
|
1 |
model = attempt_load(weights, map_location=device) # load FP32 model |
接着，就是 YOLOv5 中标准的预处理流程，首先将图片resize到固定的大小，然后转换成 tensor，接着做标准化。
|
1 2 3 4 5 6 |
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 函数来进行模型的前向传播。
|
1 2 3 4 |
def detect(): save_img = False form_data = request.json print(form_data) ... |
最后我们在 main 函数中调用
|
1 2 |
if __name__ == '__main__': app.run() |
测试时，需要先启动flask服务：
|
1 2 |
python flask_detect.py <path-to-your-model> python flask_detect.py <path-to-your-model> <GPU-device-number(s)> |
然后，使用curl命令或者提供的python脚本去访问服务器：
|
1 2 |
python client_flask.py curl -X POST -H "Content-Type: application/json" -d @api_data.json http://localhost:5000/detect |
测试结果如下
3 NCNN部署
|
1 |
python3 -m onnxsim best.onnx best-sim.onnx |
Step2：将简化后的onnx模型转化为ncnn的模型：
|
1 |
onnx2ncnn best-sim.onnx best.param best.bin |
Setp3：推理
加载放在example/yolov5目录下的模型：
|
1 |
yolov5.load_param("../examples/yolov5/yolov5b.param");yolov5.load_model("../examples/yolov5/yolov5b.bin"); |
然后进行YOLOv5 中标准的预处理流程：
|
1 2 3 4 5 6 7 8 9 10 11 12 13 |
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操作：
|
1 2 3 4 5 6 7 8 9 |
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文件简单的修改网络结构，代码的可读性更高，能够更方便修改代码。
数据科学与编程 » YOLOv5从入门到部署之：模型部署
