Pytorch：YOLO-v5目标检测（上）

Posted 2021-06-05 Z|Star

tags:

篇首语：本文由小常识网(cha138.com)小编为大家整理，主要介绍了Pytorch：YOLO-v5目标检测（上）相关的知识，希望对你有一定的参考价值。

在目标检测方面比较流行的是yolo和ssd两种算法。本篇内容就记录如何使用YOLO-v5进行目标检测。由于整体篇幅很长，故分为上下两篇，本篇为上篇，记录如何跑通yolo-v5代码，并对coco128数据集进行训练测试。下篇将会记录如何标注数据，训练自己的数据集。

文章目录

1.YOLO简介
2.YOLO模型
3.配置环境
4.简单测试
5.下载coco128数据集
6.开始训练
7.模型评价可视化
8.上传图片进行测试
9.参考资料

1.YOLO简介

YOLO全称为You Only Look Once（你只需看一次）。卷积神经网络处理图像时，需要将用卷积核对图像进行逐行扫描，而YOLO则是将一张图片分成无数个方格，通过机器来判断每个方格出现目标的可能性。具体的理论比较复杂，想要了解原理，我推荐看下面两个视频。
第一个是吴恩达讲解的Yolo算法（传送门)
第二个是B站讲解最清楚的Yolo-v3算法(传送门)
Yolo中的v即version，代表版本，yolo的创始人总共从v1更新到v3，而v4，v5则是另外一位作者编写，虽未得到官方认可，但在某测试集上的表现已经超越之前的v3版本。

在这里插入图片描述

2.YOLO模型

Yolo-v5总共有四个预训练模型，v5s、v5m、v5l、v5x，s为small，m为middle，l为large，比如识别一些小物体最适合用的是v5s，本篇内容也主要使用v5s作为预训练模型。
四个模型的具体性能如下所示：
在这里插入图片描述

3.配置环境

我的电脑环境：
硬件配置：GPU2060、32G内存；
软件配置：window 10、IDE: pycharm 2020、Python: 3.8、PytorchGPU版
使用git bash下载yolov5源码：

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

或者也可以复制该连接下载zip文件

将yolov5的源码克隆下来后，里面有一份requirements.txt文件，里面列举了所需安装的一些库文件
requirements.txt

# base ----------------------------------------
matplotlib>=3.2.2
numpy>=1.18.5
opencv-python>=4.1.2
Pillow
PyYAML>=5.3.1
scipy>=1.4.1
torch>=1.7.0
torchvision>=0.8.1
tqdm>=4.41.0

# logging -------------------------------------
tensorboard>=2.4.1
# wandb

# plotting ------------------------------------
seaborn>=0.11.0
pandas

# export --------------------------------------
# coremltools>=4.1
# onnx>=1.9.0
# scikit-learn==0.19.2  # for coreml quantization

可以直接在终端执行

pip install -r requirements.txt

系统会帮你把所有需要的一起安装，不过有些库比较大，容易安装出错，还是推荐手动安装。
下载完成后，文件构成如下图所示：
在这里插入图片描述
最后一份yolov5s.pt是模型权重文件，需要额外下载。
（可以在微信公众号“我有一计”中回复“yolo”，获取所有文件的下载链接）

4.简单测试

配置完环境后，就可以开始进行简单测试。
data/images文件中给了默认的两张测试图片
在这里插入图片描述

我们打开detect.py文件，看到主函数这块

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--weights', nargs='+', type=str, default='yolov5s.pt', help='model.pt path(s)')
    parser.add_argument('--source', type=str, default='data/images', help='source')  # file/folder, 0 for webcam
    parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
    parser.add_argument('--conf-thres', type=float, default=0.25, help='object confidence threshold')
    parser.add_argument('--iou-thres', type=float, default=0.45, help='IOU threshold for NMS')
    parser.add_argument('--max-det', type=int, default=1000, help='maximum number of detections per image')
    parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
    parser.add_argument('--view-img', action='store_true', help='display results')
    parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
    parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels')
    parser.add_argument('--save-crop', action='store_true', help='save cropped prediction boxes')
    parser.add_argument('--nosave', action='store_true', help='do not save images/videos')
    parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --class 0, or --class 0 2 3')
    parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')
    parser.add_argument('--augment', action='store_true', help='augmented inference')
    parser.add_argument('--update', action='store_true', help='update all models')
    parser.add_argument('--project', default='runs/detect', help='save results to project/name')
    parser.add_argument('--name', default='exp', help='save results to project/name')
    parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
    parser.add_argument('--line-thickness', default=3, type=int, help='bounding box thickness (pixels)')
    parser.add_argument('--hide-labels', default=False, action='store_true', help='hide labels')
    parser.add_argument('--hide-conf', default=False, action='store_true', help='hide confidences')
    opt = parser.parse_args()

发现它是用parser命令行传参的方式进行，每一项里面的default为默认值。
看到–souce这条内容，其选择了images文件夹，即把文件夹里面的所有内容进行测试。
运行detect.py，可以发现在runs/detect/exp文件夹下多了两张图片，这两张即检测结果
在这里插入图片描述

可以看到效果还是非常明显的。

5.下载coco128数据集

coco数据集是80分类的数据集，其内容包含飞机火车汽车等生活中常见事物，coco128即coco中的128张图片。
在下载下来的scripts文件夹下，有coco128.yaml的配置文件，打开之后，可以发现下载链接。

# download command/URL (optional)
download: https://github.com/ultralytics/yolov5/releases/download/v1.0/coco128.zip

# train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/]
train: ../coco128/images/train2017/  # 128 images
val: ../coco128/images/train2017/  # 128 images

# number of classes
nc: 80

# class names
names: [ 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
         'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
         'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
         'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard',
         'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
         'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
         'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
         'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear',
         'hair drier', 'toothbrush' ]

# with open('data/coco.yaml') as f:
#   d = yaml.safe_load(f)  # dict
#   for i, x in enumerate(d['names']):
#     print(i, x)

你也可以不用下载，直接运行train.py，它会自动帮你下载，不过网不好下载失败的可以手动下载上传，coco128数据集我已上传到微信公众号“我有一计”，回复“yolo”即可查看。
注意，下载下来的数据集需放置在和yolov5并列的文件夹下。
在这里插入图片描述
这里再对配置文件中做一点补充，train：代表训练集，val代表验证集，nc表示分类，names表示分类标签。

6.开始训练

直接运行train.py，不过其默认的batch-size比较大，可能会出现显存不足的情况，调小batch-size即可，我调整至8。其默认的epochs是300次，经实测，还是比较费时间的。初次使用可以调小一些。
可以直接调数字，也可以在终端中输入

train.py --epochs 10

这里的10会覆盖默认的300。

不过我还用直接使用默认的300次来进行训练。训练截图如下：
在这里插入图片描述
总共耗时0.7小时。
训练成功后，会有提示：

它会帮你保存两个模型，最后一次的模型last.pt，在val集上效果最好的模型best.pt。
除此之外，在runs/train/exp文件中会保存很多信息。
比如部分batch的标注和测试
标注：
在这里插入图片描述
测试：

可以看到，效果还是挺不错的。