保姆级教程:在Ubuntu 20.04上用YOLOv5s训练自己的人脸检测模型(附数据集)
从零构建人脸检测模型:YOLOv5在Ubuntu 20.04的实战指南
当计算机视觉遇上边缘计算,人脸检测技术正从实验室走向日常生活。无论是智能门锁的身份识别,还是社交媒体中的AR滤镜,背后都离不开高效的目标检测算法。本文将带您用YOLOv5这一前沿框架,在普通消费级显卡上打造专属的人脸检测系统。
1. 开发环境搭建
在开始模型训练前,我们需要构建稳定的深度学习环境。推荐使用Anaconda管理Python环境,它能有效解决依赖冲突问题。以下是具体步骤:
# 创建专用环境(Python3.8与PyTorch1.7.1最佳适配) conda create -n yolov5_face python=3.8 conda activate yolov5_face # 安装PyTorch与CUDA工具包(注意选择与显卡驱动兼容的版本) conda install pytorch==1.7.1 torchvision==0.8.2 torchaudio==0.7.2 cudatoolkit=11.0 -c pytorch环境验证环节不可忽视,执行以下命令确认GPU可用性:
import torch print(torch.cuda.is_available()) # 应返回True print(torch.cuda.get_device_name(0)) # 显示显卡型号提示:若出现CUDA版本不兼容,建议通过NVIDIA官网更新显卡驱动至最新版
2. 数据集准备与标注
优质数据是模型性能的基石。我们采用CASIA-FaceV5数据集,包含多角度、多光照条件下的人脸图像。数据集组织结构应遵循YOLO标准格式:
CASIA-FaceV5/ ├── images │ ├── train # 训练集图片 │ ├── val # 验证集图片 │ └── test # 测试集图片 └── labels ├── train # 对应标注文件 ├── val └── test标注文件为.txt格式,每行代表一个目标:
<class_id> <x_center> <y_center> <width> <height>使用LabelImg工具标注时,注意以下要点:
- 标注框应紧贴人脸边缘
- 避免包含过多背景区域
- 对遮挡人脸采用合理推测标注
3. 模型配置与训练
YOLOv5通过.yaml文件定义模型结构。我们复制yolov5s.yaml并修改关键参数:
# yolov5_face.yaml nc: 1 # 仅人脸类别 names: ['face'] # 类别名称数据配置文件定义数据集路径:
# face_data.yaml train: ../CASIA-FaceV5/images/train val: ../CASIA-FaceV5/images/val nc: 1 names: ['face']启动训练的命令行参数解析:
python train.py \ --data face_data.yaml \ --cfg yolov5_face.yaml \ --weights yolov5s.pt \ --epochs 100 \ --batch-size 16 \ --img-size 640 \ --device 0 # 使用第一块GPU训练过程监控指标包括:
- mAP@0.5:交并比阈值为0.5时的平均精度
- Precision:预测为正样本中真实正样本比例
- Recall:真实正样本中被正确预测的比例
4. 模型验证与优化
训练完成后,使用验证集评估模型表现:
python val.py \ --data face_data.yaml \ --weights runs/train/exp/weights/best.pt \ --batch-size 8 \ --task test # 在测试集上评估常见性能提升技巧:
- 数据增强:调整hyp.scratch.yaml中的参数
- 随机旋转:degrees: 10
- 色彩抖动:hsv_h: 0.015
- 迁移学习:冻结部分层参数
python train.py --freeze 10 # 冻结前10层 - 超参数搜索:使用进化算法优化
python train.py --evolve
5. 部署与实时检测
将训练好的模型转换为ONNX格式便于跨平台部署:
python export.py \ --weights runs/train/exp/weights/best.pt \ --img-size 640 \ --include onnx实时摄像头检测命令:
python detect.py \ --weights runs/train/exp/weights/best.pt \ --source 0 \ # 摄像头设备号 --conf-thres 0.5 \ # 置信度阈值 --view-img # 实时显示窗口性能优化建议:
- 使用TensorRT加速推理
- 调整img-size平衡速度与精度
- 对视频流采用跳帧检测策略
在GTX 1660 Ti上的典型性能:
| 分辨率 | FPS | 内存占用 |
|---|---|---|
| 640x640 | 45 | 2.1GB |
| 320x320 | 120 | 1.4GB |
