pytorch yolov5 推理和训练环境搭建

关于之前pytorch框架下yolov3推理和训练环境搭建可以参考，本文所使用的环境和这篇记录一致。

pytorch yolov3 推理和训练环境搭建_papaofdoudou的博客-CSDN博客_yolov3环境搭建

下载代码框架环境：

 git clone https://github.com/ultralytics/yolov5

安装依赖：

pip install -r requirements.txt

下载YOLOV5模型文件：

下载链接 ：Releases · ultralytics/yolov5 · GitHub以下模型均可用：

拷贝模型文件到yolov5顶层目录，之后执行命令：

python detect.py --source data/images --weights yolov5s.pt --conf 0.25

输出如下，一开始会去下载字库，猜测是用来画标签用：

根据LOG，推理结果在runs/detect/exp目录：

 python代码推理：

由于我们已经将YOLOV5的pytorch训练框架下载过来，接下来我们就可以编写代码来完成推理了，代码如下所示：

import torch

# Model
model = torch.hub.load('ultralytics/yolov5', 'yolov5s')  # or yolov5m, yolov5l, yolov5x, custom

# Images
img = './data/images/zidane.jpg'  # or file, Path, PIL, OpenCV, numpy, list

# Inference
results = model(img)

# Results
results.print()  # or .show(), .save(), .crop(), .pandas(), etc.
results.show()

 执行命令python inf.py 推理完成后：

YOLOV5X.pt

python detect.py --source data/images --weights yolov5x.pt --conf 0.25

yolov5x.pt

导出ONNX模型数据

导出ONNX模型，还需要依赖两个python安装包：

pip install onnx==1.8.0 -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install onnxruntime==1.8.0 -i https://pypi.tuna.tsinghua.edu.cn/simple

 导出模型：

python export.py --weights yolov5s.pt --img 640 --batch 1

可以看到，生成了yolov5s.onnx模型文件，我们看其结构：

修改导图尺寸：

python export.py --weights yolov5s.pt --img 320 --batch 1

 

yolov5x6大模型转换

python export.py --weights yolov5x6.pt --img 320 --batch 1

 batch = 2的时候

python export.py --weights yolov5x6.pt --img 320 --batch 2

此时，网络一次吃两张图，同时也输出两张图的结果：

训练：

 python train.py --data coco.yaml --cfg yolov5n.yaml --weights '' --batch-size 128

YOLOV5s中很多类似的结构，开始我还以为是卷积后求sigmoid之后和原卷积输出矩阵做矩阵乘法，后面看了libonnx的代码，才发现不是，矩阵对应元素做乘法。

ONNX推理

用生成的ONNX模型做推理，程序如下：

 代码：

import cv2
import numpy as np
import onnxruntime as rt

height, width = 640, 640
img0 = cv2.imread('./data/images/bus.jpg') 
img = cv2.resize(img0, (height, width))  # 尺寸变换
img = img / 255.
img = img[:, :, ::-1].transpose((2, 0, 1))  # HWC转CHW
data = np.expand_dims(img, axis=0)  # 扩展维度至[1,3,640,640]
sess = rt.InferenceSession('yolov5s.onnx')
input_name = sess.get_inputs()[0].name
label_name = sess.get_outputs()[0].name
pred_onx = sess.run([label_name], {input_name: data.astype(np.float32)})[0] 
pred = np.squeeze(pred_onx) 
print(pred)
print(pred.shape)

参考资料：

https://github.com/ultralytics/yolov5/issues/251

---

结束