项目(一)
数据加载核心组件
1. 图像格式支持
YOLOv5 支持多种图像格式:
img_formats = ['.bmp', '.jpg', '.jpeg', '.png', '.tif', '.tiff', '.dng'] vid_formats = ['.mov', '.avi', '.mp4', '.mpg', '.mpeg', '.m4v', '.wmv', '.mkv']
2. EXIF 方向校正
def exif_size(img):# 获取图像 EXIF 信息中的方向标签,自动校正图像方向s = img.size # (width, height)try:rotation = dict(img._getexif().items())[orientation]if rotation == 6 or rotation == 8: # 旋转 270 或 90 度s = (s[1], s[0])except:passreturn s
3. 数据加载器创建函数
def create_dataloader(path, imgsz, batch_size, stride, opt, hyp=None, augment=False, cache=False, pad=0.0, rect=False,rank=-1, world_size=1, workers=8):
关键参数说明:
path: 数据集路径imgsz: 输入图像尺寸batch_size: 批次大小stride: 模型步长augment: 是否进行数据增强cache: 是否缓存图像到内存rect: 是否使用矩形训练rank: 分布式训练中的进程排名
分布式训练支持
# 确保 DDP 中只有第一个进程处理数据集,其他进程使用缓存 with torch_distributed_zero_first(rank):dataset = LoadImagesAndLabels(path, imgsz, batch_size,augment=augment,hyp=hyp,rect=rect,cache_images=cache,stride=int(stride),pad=pad,rank=rank)
命令行参数配置
YOLOv5 提供了丰富的命令行参数用于推理和训练:
模型相关参数
--weights yolov5s.pt # 模型权重文件路径 --img-size 640 # 推理尺寸(像素)
推理设置
--source inference/images # 输入源:文件/文件夹/摄像头 --output inference/output # 输出文件夹 --conf-thres 0.4 # 目标置信度阈值 --iou-thres 0.5 # NMS 的 IOU 阈值
功能开关
--view-img # 显示推理结果 --save-txt # 保存结果到 txt 文件 --agnostic-nms # 类别无关的 NMS --augment # 增强推理 --update # 更新所有模型
类别过滤
--classes 0 2 3 # 按类别过滤:只检测类别 0,2,3
数据加载优化
1. 工作进程数计算
nw = min([os.cpu_count() // world_size, batch_size if batch_size > 1 else 0, workers])
自动计算最优的工作进程数量,考虑 CPU 核心数、批次大小和用户设置。
2. 分布式采样器
sampler = torch.utils.data.distributed.DistributedSampler(dataset) if rank != -1 else None
在分布式训练中确保数据正确分发到各个进程。
3. 无限数据加载器
使用 InfiniteDataLoader 确保在长时间训练中数据不会耗尽。
