10.2 基于YOLO的目标检测
YOLO(You Only Look Once)是一种基于深度学习的目标检测算法,它具有实时性和高准确性的特点。YOLO算法通过单个神经网络同时完成目标定位和分类,以极高的速度在图像或视频中检测多个目标物体。YOLO v5是YOLO算法的一个版本,它在原始的YOLO算法基础上进行了一些关键性的改进,以提升检测性能和准确性。在本节将讲解使用YOLO v5实现目标检测的方法。
10.2.1 YOLO v5的改进
YOLO v5是YOLO算法的改进版本,YOLO v5相比于原始的YOLO算法,在检测准确性和速度上都有显著的提升,使得它成为一个广泛使用的目标检测算法。然后,YOLO算法继续进行改进和优化,发展出了后续版本的YOLO,如YOLO v3、YOLO v4和YOLO v5,以进一步提升目标检测的性能。这些改进版本在网络结构、特征提取、预测机制等方面进行了创新和改进,以满足不同应用场景下的需求。
在Python中使用YOLO v5前需要先通过如下命令安装:
pip install ultralytics
作为初学者来说,可以直接使用YOLO v5提供的预训练模型。预训练的YOLO v5模型可以从多个来源下载,以下是一些常用的下载来源:
- Darknet官方网站:YOLO v5是由Joseph Redmon开发的Darknet框架的一部分。您可以从Darknet官方网站下载YOLO v5的预训练模型。访问链接:https://pjreddie.com/darknet/yolo/
- YOLO官方GitHub页面:YOLO官方GitHub仓库也提供了YOLO v5的预训练模型。您可以在https://github.com/pjreddie/darknet/releases 页面找到预训练模型的下载链接。
- 第三方资源:除了官方来源,还有一些第三方资源库和社区提供了YOLO v5的预训练模型的下载。一些常见的资源库包括Model Zoo、PyTorch Hub、Hugging Face Model Hub等。
在下载预训练的YOLO v5模型时,请确保使用可信赖的来源,并查看模型的许可和使用条款。预训练模型的下载可能需要注册、登录或同意特定的使用条件,具体取决于下载来源。
注意:由于YOLO v2模型在不同的深度学习框架中可能具有不同的实现方式和权重文件格式,因此需要确保下载的模型与您使用的深度学习框架兼容。一般来说,官方提供的预训练模型会与官方支持的框架兼容性较好。
10.2.2 基于YOLO v5的训练、验证和预测
请看下面的例子,功能是使用YOLO v5实现训练、验证和预测功能,让大家初步了解YOLO v5实现图像目标检测的基本功能。
实例10-1:使用YOLO v5实现模型训练、验证和预测
源码路径:daima\10\yolov5-master\
1. 目标检测
目标检测是一种计算机视觉任务,旨在从图像或视频中定位和识别出特定对象的位置。通常使用深度学习模型来进行目标检测,例如使用卷积神经网络(CNN)或相关的模型,如YOLO(You Only Look Once)和Faster R-CNN(Region-based Convolutional Neural Networks)。在本实例中,编写文件detect.py实现目标检测任务,这是一个运行YOLOv5推理的脚本,可以在各种来源上进行推理,自动从YOLO v5发布中下载模型,并将结果保存到“runs/detect”目录下。文件detect.py的具体实现流程如下所示:
(1)定义了一个名为run的函数,该函数用于运行目标检测任务。具体实现代码如下所示。
def run(
weights=ROOT / 'yolov5s.pt', # model path or triton URL
source=ROOT / 'data/images', # file/dir/URL/glob/screen/0(webcam)
data=ROOT / 'data/coco128.yaml', # dataset.yaml path
imgsz=(640, 640), # inference size (height, width)
conf_thres=0.25, # confidence threshold
iou_thres=0.45, # NMS IOU threshold
max_det=1000, # maximum detections per image
device='', # cuda device, i.e. 0 or 0,1,2,3 or cpu
view_img=False, # show results
save_txt=False, # save results to *.txt
save_conf=False, # save confidences in --save-txt labels
save_crop=False, # save cropped prediction boxes
nosave=False, # do not save images/videos
classes=None, # filter by class: --class 0, or --class 0 2 3
agnostic_nms=False, # class-agnostic NMS
augment=False, # augmented inference
visualize=False, # visualize features
update=False, # update all models
project=ROOT / 'runs/detect', # save results to project/name
name='exp', # save results to project/name
exist_ok=False, # existing project/name ok, do not increment
line_thickness=3, # bounding box thickness (pixels)
hide_labels=False, # hide labels
hide_conf=False, # hide confidences
half=False, # use FP16 half-precision inference
dnn=False, # use OpenCV DNN for ONNX inference
vid_stride=1, # video frame-rate stride
):
在上述代码中,各个参数的具体说明如下:
- weights:模型的路径或Triton URL,默认值为yolov5s.pt,表示模型的权重文件路径。
- source:推理的来源,可以是文件、目录、URL、通配符、屏幕截图或者摄像头。默认值为data/images,表示推理的来源为data/images目录。
- data:数据集的配置文件路径,默认值为data/coco128.yaml,表示使用COCO128数据集的配置文件。
- imgsz:推理时的图像尺寸,默认为(640, 640),表示推理时将图像调整为高度和宽度都为640的尺寸。
- conf_thres:置信度阈值,默认值为0.25,表示只保留置信度大于该阈值的检测结果。
- iou_thres:NMS(非极大值抑制)的IoU(交并比)阈值,默认值为0.45,用于去除重叠度较高的重复检测结果。
- max_det:每张图像的最大检测数量,默认值为1000,表示每张图像最多保留1000个检测结果。
- device:设备类型,默认为空字符串,表示使用默认设备(GPU或CPU)进行推理。
- view_img:是否显示结果图像,默认值为False,表示不显示结果图像。
- save_txt:是否将结果保存为文本文件,默认值为False。
- save_conf:是否将置信度保存在保存的文本标签中,默认值为False。
- save_crop:是否保存裁剪的预测框,默认值为False。
- nosave:是否禁止保存图像或视频,默认值为False。
- classes:根据类别进行过滤,默认为None,表示不进行类别过滤。
- agnostic_nms:是否使用类别不可知的NMS,默认值为False。
- augment:是否进行增强推理,默认值为False。
- visualize:是否可视化特征,默认值为False。
- update:是否更新所有模型,默认值为False。
- project:保存结果的项目路径,默认为runs/detect。
- name:保存结果的名称,默认为exp。
- exist_ok:是否允许存在的项目/名称,如果为True,则不递增项目/名称,默认值为False。
- line_thickness:边界框线条的粗细,默认为3个像素。
- hide_labels:是否隐藏标签,默认值为False。
- hide_conf:是否隐藏置信度,默认值为False。
- half:是否使用FP16的半精度推理,默认值为False。
- dnn:是否使用OpenCV DNN进行ONNX推理,默认值为False。
- vid_stride:视频帧率步长,默认值为1。
这些参数可以根据需要进行调整,以适应不同的目标检测场景和要求。
(2)运行目标检测推理,并根据不同的输入来源进行相应的处理和操作。具体实现代码如下所示。
source = str(source) # 将输入的来源转换为字符串类型
save_img = not nosave and not source.endswith('.txt') # 是否保存推理图像
is_file = Path(source).suffix[1:] in (IMG_FORMATS + VID_FORMATS) # 是否为文件路径
is_url = source.lower().startswith(('rtsp://', 'rtmp://', 'http://', 'https://')) # 是否为URL
webcam = source.isnumeric() or source.endswith('.streams') or (is_url and not is_file) # 是否为摄像头输入
screenshot = source.lower().startswith('screen') # 是否为屏幕截图输入
if is_url and is_file:
source = check_file(source) # 下载文件
save_dir = increment_path(Path(project) / name, exist_ok=exist_ok) # 生成保存结果的目录
(save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True) # 创建目录
# 加载模型
device = select_device(device) # 选择设备
model = DetectMultiBackend(weights, device=device, dnn=dnn, data=data, fp16=half) # 实例化模型
stride, names, pt = model.stride, model.names, model.pt # 获取模型参数
imgsz = check_img_size(imgsz, s=stride) # 检查图像尺寸
# 数据加载器
bs = 1 # batch_size
if webcam:
view_img = check_imshow(warn=True) # 检查是否显示图像
dataset = LoadStreams(source, img_size=imgsz, stride=stride, auto=pt, vid_stride=vid_stride) # 加载数据流
bs = len(dataset)
elif screenshot:
dataset = LoadScreenshots(source, img_size=imgsz, stride=stride, auto=pt) # 加载截图
else:
dataset = LoadImages(source, img_size=imgsz, stride=stride, auto=pt, vid_stride=vid_stride) # 加载图像/视频
vid_path, vid_writer = [None] * bs, [None] * bs
# 执行推理
model.warmup(imgsz=(1 if pt or model.triton else bs, 3, *imgsz)) # 模型预热
seen, windows, dt = 0, [], (Profile(), Profile(), Profile())
for path, im, im0s, vid_cap, s in dataset:
with dt[0]:
im = torch.from_numpy(im).to(model.device)
im = im.half() if model.fp16 else im.float() # 转换数据类型
im /= 255 # 范围调整为[0, 1]
if len(im.shape) == 3:
im = im[None] # 添加批处理维度
# 推理
with dt[1]:
visualize = increment_path(save_dir / Path(path).stem, mkdir=True) if visualize else False # 可视化结果的路径
pred = model(im, augment=augment, visualize=visualize) # 推理
# NMS(非最大值抑制)
with dt[2]:
pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_ # NMS操作_
对上述代码的具体说明如下:
- 首先,对输入的source进行了一系列判断和处理。将source转换为字符串类型,并根据条件判断是否保存推理图像。接着,判断source是文件路径还是URL,以及是否为摄像头或屏幕截图。如果source是URL且为文件路径,则会进行文件下载操作。
- 接下来,创建保存结果的目录,根据project和name参数生成保存结果的路径,并在指定路径下创建目录。如果save_txt为True,则在目录下创建一个名为'labels'的子目录,否则直接创建主目录。
- 然后,加载模型并选择设备。根据指定的设备类型,选择相应的设备进行推理。同时,实例化DetectMultiBackend类的对象model,并获取其步长(stride)、类别名称(names)和模型(pt)。
- 接下来,检查图像尺寸并创建数据加载器。根据不同的输入来源,选择相应的数据加载器对象:LoadStreams用于摄像头输入,LoadScreenshots用于屏幕截图输入,LoadImages用于图像或视频输入。同时,根据是否为摄像头输入,确定批处理大小(bs)。
- 接下来,进行推理过程。首先,通过调用model.warmup()方法进行模型预热,其中输入图像尺寸根据模型类型进行调整。然后,使用迭代器遍历数据加载器,获取输入图像及相关信息。将图像转换为PyTorch张量,并根据模型的精度要求进行数据类型和范围的调整。如果输入图像维度为3维,则添加一个批处理维度。然后根据是否需要可视化结果,确定是否保存推理结果的路径。在得到推理结果后,根据设定的置信度阈值、iou阈值和其他参数,进行非最大值抑制(NMS)操作。
- 最后,对推理结果进行处理。包括计算推理的时间、保存推理结果图像和输出结果信息。其中,推理时间分为3个阶段,即数据准备、模型推理和NMS操作。保存推理结果图像的路径根据是否需要可视化和输入图像的文件名进行生成。然后,根据设置的参数决定是否将推理结果保存为文本文件。
2. 训练
在深度学习中,"train" 表示训练模型的过程。训练是指通过给定的输入数据和相应的标签,使用梯度下降等优化算法来调整模型的参数,使其逐渐适应给定的任务。训练过程通常包括前向传播、计算损失函数、反向传播和参数更新等步骤。通过反复迭代训练数据集,模型可以逐渐学习到数据中的模式和特征,从而提高在给定任务上的性能。在本实例中,编写文件train.py实现训练功能。具体实现流程如下所示:
(1)编写函数train()用于训练模型,首先通过下面的代码进行训练前的准备工作,包括解析参数、创建目录、加载超参数、保存运行设置和创建日志记录器等。它为后续的训练过程提供了必要的信息和设置。
def train(hyp, opt, device, callbacks): # hyp is path/to/hyp.yaml or hyp dictionary
save_dir, epochs, batch_size, weights, single_cls, evolve, data, cfg, resume, noval, nosave, workers, freeze = \
Path(opt.save_dir), opt.epochs, opt.batch_size, opt.weights, opt.single_cls, opt.evolve, opt.data, opt.cfg, \
opt.resume, opt.noval, opt.nosave, opt.workers, opt.freeze
callbacks.run('on_pretrain_routine_start')
# 目录
w = save_dir / 'weights' # 权重目录
(w.parent if evolve else w).mkdir(parents=True, exist_ok=True) # 创建目录
last, best = w / 'last.pt', w / 'best.pt'
# 超参数
if isinstance(hyp, str):
with open(hyp, errors='ignore') as f:
hyp = yaml.safe_load(f) # 加载超参数字典
LOGGER.info(colorstr(' hyperparameters: ') + ', '.join(f'{k}={v}' for k, v in hyp.items()))
opt.hyp = hyp.copy() # 用于保存检查点的超参数
# 保存运行设置
if not evolve:
yaml_save(save_dir / 'hyp.yaml', hyp)
yaml_save(save_dir / 'opt.yaml', vars(opt))
# 记录器
data_dict = None
if RANK in {-1, 0}:
loggers = Loggers(save_dir, weights, opt, hyp, LOGGER) # 记录器实例
# 注册操作
for k in methods(loggers):
callbacks.register_action(k, callback=getattr(loggers, k))
# 处理自定义数据集工件链接
data_dict = loggers.remote_dataset
if resume: # 如果从远程工件恢复运行
weights, epochs, hyp, batch_size = opt.weights, opt.epochs, opt.hyp, opt.batch_size
对上述代码的具体说明如下:
- 首先,将函数的参数和选项进行解析和赋值。包括hyp(超参数路径或字典)、opt(选项参数对象)、device(设备)、callbacks(回调函数集合)等。
- 创建保存目录,并设置权重目录w。
- 加载超参数。如果hyp是一个字符串,则从文件中加载超参数字典。然后将超参数保存在opt.hyp中,以便在训练过程中保存到检查点。
- 如果不是进化训练(evolve=False),则保存运行设置。将超参数保存为hyp.yaml文件,将选项参数保存为opt.yaml文件。
- 创建日志记录器(Loggers)实例。如果当前进程是主进程(RANK=-1或RANK=0),则创建日志记录器并注册回调函数。
- 如果是从远程artifact恢复运行,则更新权重、轮数、超参数和批大小。
(2)继续训练模型前的准备工作,包括配置设置、加载模型、冻结层、设置图像尺寸、设置批大小、创建优化器、实现学习率调度器、设置EMA指数滑动平均、恢复训练等。具体实现代码如下所示。
plots = not evolve and not opt.noplots # 创建绘图
cuda = device.type != 'cpu'
init_seeds(opt.seed + 1 + RANK, deterministic=True)
with torch_distributed_zero_first(LOCAL_RANK):
data_dict = data_dict or check_dataset(data) # 检查是否为None
train_path, val_path = data_dict['train'], data_dict['val']
nc = 1 if single_cls else int(data_dict['nc']) # 类别数
names = {0: 'item'} if single_cls and len(data_dict['names']) != 1 else data_dict['names'] # 类别名称
is_coco = isinstance(val_path, str) and val_path.endswith('coco/val2017.txt') # 是否为COCO数据集
# 模型
check_suffix(weights, '.pt') # 检查权重文件
pretrained = weights.endswith('.pt')
if pretrained:
with torch_distributed_zero_first(LOCAL_RANK):
weights = attempt_download(weights) # 如果本地不存在,则下载
ckpt = torch.load(weights, map_location='cpu') # 加载检查点到CPU以避免CUDA内存泄漏
model = Model(cfg or ckpt['model'].yaml, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device) # 创建模型
exclude = ['anchor'] if (cfg or hyp.get('anchors')) and not resume else [] # 排除的键
csd = ckpt['model'].float().state_dict() # 将检查点状态字典转换为FP32
csd = intersect_dicts(csd, model.state_dict(), exclude=exclude) # 取交集
model.load_state_dict(csd, strict=False) # 加载模型状态
LOGGER.info(f'从{weights}转移了{len(csd)}/{len(model.state_dict())}个项目') # 报告
else:
model = Model(cfg, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device) # 创建模型
amp = check_amp(model) # 检查AMP
# 冻结层
freeze = [f'model.{x}.' for x in (freeze if len(freeze) > 1 else range(freeze[0]))] # 要冻结的层
for k, v in model.named_parameters():
v.requires_grad = True # 训练所有层
if any(x in k for x in freeze):
LOGGER.info(f'冻结{k}')
v.requires_grad = False
# 图像尺寸
gs = max(int(model.stride.max()), 32) # 网格大小(最大步长)
imgsz = check_img_size(opt.imgsz, gs, floor=gs * 2) # 确保imgsz是gs的倍数
# 批处理大小
if RANK == -1 and batch_size == -1: # 仅在单GPU情况下,估算最佳批处理大小
batch_size = check_train_batch_size(model, imgsz, amp)
loggers.on_params_update({'batch_size': batch_size})
# 优化器
nbs = 64 # 名义批处理大小
accumulate = max(round(nbs / batch_size), 1) # 在优化之前积累损失
hyp['weight_decay'] *= batch_size * accumulate / nbs # 缩放weight_decay
optimizer = smart_optimizer(model, opt.optimizer, hyp['lr0'], hyp['momentum'], hyp['weight_decay'])
# 调度器
if opt.cos_lr:
lf = one_cycle(1, hyp['lrf'], epochs) # 余弦1->hyp['lrf']
else:
lf = lambda x: (1 - x / epochs) * (1.0 - hyp['lrf']) + hyp['lrf'] # 线性
scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf) # 绘制lr_scheduler(optimizer, scheduler, epochs)
# EMA
ema = ModelEMA(model) if RANK in {-1, 0} else None
# 恢复
best_fitness, start_epoch = 0.0, 0
if pretrained:
if resume:
best_fitness, start_epoch, epochs = smart_resume(ckpt, optimizer, ema, weights, epochs, resume)
del ckpt, csd
(3)使用深度学习框架PyTorch训练目标检测模型,具体实现代码如下所示。
train_loader, dataset = create_dataloader(train_path,
imgsz,
batch_size // WORLD_SIZE,
gs,
single_cls,
hyp=hyp,
augment=True,
cache=None if opt.cache == 'val' else opt.cache,
rect=opt.rect,
rank=LOCAL_RANK,
workers=workers,
image_weights=opt.image_weights,
quad=opt.quad,
prefix=colorstr('train: '),
shuffle=True,
seed=opt.seed)
labels = np.concatenate(dataset.labels, 0)
mlc = int(labels[:, 0].max()) # 最大标签类别
assert mlc < nc, f'Label class {mlc} exceeds nc={nc} in {data}. Possible class labels are 0-{nc - 1}'
# 进程0
if RANK in {-1, 0}:
val_loader = create_dataloader(val_path,
imgsz,
batch_size // WORLD_SIZE * 2,
gs,
single_cls,
hyp=hyp,
cache=None if noval else opt.cache,
rect=True,
rank=-1,
workers=workers * 2,
pad=0.5,
prefix=colorstr('val: '))[0]
if not resume:
if not opt.noautoanchor:
check_anchors(dataset, model=model, thr=hyp['anchor_t'], imgsz=imgsz) # 运行AutoAnchor
model.half().float() # 预先减少锚精度
callbacks.run('on_pretrain_routine_end', labels, names)
# DDP模式
if cuda and RANK != -1:
model = smart_DDP(model)
# 模型属性
nl = de_parallel(model).model[-1].nl # 检测层数量(用于缩放hyps)
hyp['box'] *= 3 / nl # 缩放到层
hyp['cls'] *= nc / 80 * 3 / nl # 缩放到类别和层
hyp['obj'] *= (imgsz / 640) ** 2 * 3 / nl # 缩放到图像尺寸和层
hyp['label_smoothing'] = opt.label_smoothing
model.nc = nc # 将类别数附加到模型
model.hyp = hyp # 将超参数附加到模型
model.class_weights = labels_to_class_weights(dataset.labels, nc).to(device) * nc # 附加类别权重
model.names = names
# 开始训练
t0 = time.time()
nb = len(train_loader) # 批次数量
nw = max(round(hyp['warmup_epochs'] * nb), 100) # 温暖迭代次数,最大(3个epoch,100个迭代)
# nw = min(nw, (epochs - start_epoch) / 2 * nb) # 限制温暖到小于1/2的训练
last_opt_step = -1
maps = np.zeros(nc) # 每个类别的mAP
results = (0, 0, 0, 0, 0, 0, 0) # P,R,mAP@.5,mAP@.5-.95,val_loss(box,obj,cls)
scheduler.last_epoch = start_epoch - 1 # 不要移动
scaler = torch.cuda.amp.GradScaler(enabled=amp)
stopper, stop = EarlyStopping(patience=opt.patience), False
compute_loss = ComputeLoss(model) # 初始化损失类
callbacks.run('on_train_start')
LOGGER.info(f'Image sizes {imgsz} train, {imgsz} val\n'
f'Using {train_loader.num_workers * WORLD_SIZE} dataloader workers\n'
f"Logging results to {colorstr('bold', save_dir)}\n"
f'Starting training for {epochs} epochs...')
for epoch in range(start_epoch, epochs): # epoch ------------------------------------------------------------------
callbacks.run('on_train_epoch_start')
model.train()
# 更新图像权重(可选,仅单GPU)
if opt.image_weights:
cw = model.class_weights.cpu().numpy() * (1 - maps) ** 2 / nc # 类别权重
iw = labels_to_image_weights(dataset.labels, nc=nc, class_weights=cw) # 图像权重
dataset.indices = random.choices(range(dataset.n), weights=iw, k=dataset.n) # 随机加权idx
# 更新马赛克边框(可选)
# b = int(random.uniform(0.25 * imgsz, 0.75 * imgsz + gs) // gs * gs)
# dataset.mosaic_border = [b - imgsz, -b] # 高度,宽度边界
mloss = torch.zeros(3, device=device) # 平均损失
if RANK != -1:
train_loader.sampler.set_epoch(epoch)
pbar = enumerate(train_loader)
LOGGER.info(('\n' + '%11s' * 7) % ('Epoch', 'GPU_mem', 'box_loss', 'obj_loss', 'cls_loss', 'Instances', 'Size'))
if RANK in {-1, 0}:
pbar = tqdm(pbar, total=nb, bar_format=TQDM_BAR_FORMAT) # 进度条
optimizer.zero_grad()
for i, (imgs, targets, paths, _) in pbar: # batch -------------------------------------------------------------
callbacks.run('on_train_batch_start')
ni = i + nb * epoch # 自训练开始以来集成的批次数
imgs = imgs.to(device, non_blocking=True).float() / 255 # uint8转换为float32,0-255转换为0.0-1.0
# 温暖
if ni <= nw:
xi = [0, nw] # x interp
# compute_loss.gr = np.interp(ni, xi, [0.0, 1.0]) # iou损失比率(obj_loss = 1.0或iou)
accumulate = max(1, np.interp(ni, xi, [1, nbs / batch_size]).round())
for j, x in enumerate(optimizer.param_groups):
# 偏差lr从0.1下降到lr0,所有其他lr从0.0上升到lr0
x['lr'] = np.interp(ni, xi, [hyp['warmup_bias_lr'] if j == 0 else 0.0, x['initial_lr'] * lf(epoch)])
if 'momentum' in x:
x['momentum'] = np.interp(ni, xi, [hyp['warmup_momentum'], hyp['momentum']])
# 多尺度
if opt.multi_scale:
sz = random.randrange(int(imgsz * 0.5), int(imgsz * 1.5) + gs) // gs * gs # 尺寸
sf = sz / max(imgs.shape[2:]) # 缩放因子
if sf != 1:
ns = [math.ceil(x * sf / gs) * gs for x in imgs.shape[2:]] # 新形状(拉伸到gs-multiple)
imgs = nn.functional.interpolate(imgs, size=ns, mode='bilinear', align_corners=False)
# 前向
with torch.cuda.amp.autocast(amp):
pred = model(imgs) # 前向
loss, loss_items = compute_loss(pred, targets.to(device)) # 损失按batch_size缩放
if RANK != -1:
loss *= WORLD_SIZE # DDP模式下在设备之间平均梯度
if opt.quad:
loss *= 4.
# 后向
scaler.scale(loss).backward()
# 优化 - https://pytorch.org/docs/master/notes/amp_examples.html
if ni - last_opt_step >= accumulate:
scaler.unscale_(optimizer) # 取消缩放梯度
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=10.0) # 截断梯度
scaler.step(optimizer) # 优化器步骤
scaler.update()
optimizer.zero_grad()
if ema:
ema.update(model)
last_opt_step = ni
# 日志
if RANK in {-1, 0}:
mloss = (mloss * i + loss_items) / (i + 1) # 更新平均损失
mem = f'{torch.cuda.memory_reserved() / 1E9 if torch.cuda.is_available() else 0:.3g}G' # (GB)
pbar.set_description(('%11s' * 2 + '%11.4g' * 5) %
(f'{epoch}/{epochs - 1}', mem, *mloss, targets.shape[0], imgs.shape[-1]))
callbacks.run('on_train_batch_end', model, ni, imgs, targets, paths, list(mloss))
if callbacks.stop_training:
return
# end batch -------------------------------
对上述代码的具体说明如下:
- 首先,创建了训练数据集的数据加载器,并获取了数据集中的标签信息。
- 接下来,根据训练模式进行不同的操作。如果是分布式训练模式且使用了多个GPU,则使用torch.nn.DataParallel将模型包装起来,以实现多GPU并行训练。如果开启了同步BatchNorm,并且不是分布式训练模式,则将模型中的BatchNorm层转换为SyncBatchNorm,以实现跨GPU同步。
- 然后,对模型的一些属性进行设置,包括调整损失函数的权重,附加类别权重、类别名称等信息到模型上。
- 在开始训练之前,首先需要设置一些参数,比如学习率调整策略、Early Stopping 的条件等。然后创建优化器和损失函数,为训练做好准备。在训练过程中,会将整个数据集分成多个批次,然后在每个 epoch 中对这些批次进行遍历和训练。在训练过程中,根据设置的参数进行一些预处理操作,如数据增强、图片尺寸调整等。然后将数据输入模型进行前向传播,得到预测结果,并计算损失函数。
- 接下来进行反向传播和优化器更新操作。如果达到一定的条件(如累计一定数量的批次),则进行一次优化器的更新操作。
在训练过程中,会输出一些训练信息,如当前的epoch、GPU内存占用、损失值等。同时会调用一些回调函数,如在每个训练批次开始前和结束后运行的函数。整个训练过程会持续进行多个epoch,直到达到指定的训练轮数为止。
(4)训练模型,使用循环训练目标检测模型。具体实现代码如下所示。
# 计算适应度
fi = fitness(np.array(results).reshape(1, -1)) # 权重组合[P, R, mAP@.5, mAP@.5-.95]
stop = stopper(epoch=epoch, fitness=fi) # 提前停止检查
if fi > best_fitness:
best_fitness = fi
log_vals = list(mloss) + list(results) + lr
callbacks.run('on_fit_epoch_end', log_vals, epoch, best_fitness, fi)
# 保存模型
if (not nosave) or (final_epoch and not evolve): # 如果需要保存
ckpt = {
'epoch': epoch,
'best_fitness': best_fitness,
'model': deepcopy(de_parallel(model)).half(),
'ema': deepcopy(ema.ema).half(),
'updates': ema.updates,
'optimizer': optimizer.state_dict(),
'opt': vars(opt),
'git': GIT_INFO, # 如果是git仓库,包括{remote, branch, commit}
'date': datetime.now().isoformat()}
# 保存最后、最佳并删除
torch.save(ckpt, last)
if best_fitness == fi:
torch.save(ckpt, best)
if opt.save_period > 0 and epoch % opt.save_period == 0:
torch.save(ckpt, w / f'epoch{epoch}.pt')
del ckpt
callbacks.run('on_model_save', last, epoch, final_epoch, best_fitness, fi)
# 提前停止
if RANK != -1: # 如果是DDP训练
broadcast_list = [stop if RANK == 0 else None]
dist.broadcast_object_list(broadcast_list, 0) # 将'stop'广播给所有rank
if RANK != 0:
stop = broadcast_list[0]
if stop:
break # 必须终止所有DDP rank
# 结束epoch
# 结束训练
if RANK in {-1, 0}:
LOGGER.info(f'\n{epoch - start_epoch + 1} epochs completed in {(time.time() - t0) / 3600:.3f} hours.')
for f in last, best:
if f.exists():
strip_optimizer(f) # 去除优化器
if f is best:
LOGGER.info(f'\nValidating {f}...')
results, _, _ = validate.run(
data_dict,
batch_size=batch_size // WORLD_SIZE * 2,
imgsz=imgsz,
model=attempt_load(f, device).half(),
iou_thres=0.65 if is_coco else 0.60, # 在iou 0.65时使用最佳pycocotools
single_cls=single_cls,
dataloader=val_loader,
save_dir=save_dir,
save_json=is_coco,
verbose=True,
plots=plots,
callbacks=callbacks,
compute_loss=compute_loss) # 用plot验证最佳模型
if is_coco:
callbacks.run('on_fit_epoch_end', list(mloss) + list(results) + lr, epoch, best_fitness, fi)
callbacks.run('on_train_end', last, best, epoch, results)
torch.cuda.empty_cache()
return results
对上述代码的具体说明如下:
- fi = fitness(np.array(results).reshape(1, -1)): 计算模型在验证集上的综合指标。results是一个包含模型在验证集上表现的元组,通过调用fitness函数计算综合指标。
- stop = stopper(epoch=epoch, fitness=fi): 判断是否满足停止训练的条件。stopper是一个用于判断是否进行Early Stopping的对象,根据当前的训练轮数epoch和综合指标fitness来判断是否停止训练。
- if fi > best_fitness: best_fitness = fi: 更新最佳的综合指标best_fitness,如果当前的综合指标大于最佳指标。
- log_vals = list(mloss) + list(results) + lr: 将当前的损失值、验证集上的结果和学习率组成一个列表log_vals,用于记录训练过程中的日志。
- ckpt = {...}: 创建一个字典ckpt,保存训练过程中的相关信息,包括当前的轮数epoch、最佳综合指标best_fitness、模型参数、优化器状态等。
- torch.save(ckpt, last): 将ckpt保存到文件last,这里保存的是最后一轮的模型。
- if best_fitness == fi: torch.save(ckpt, best): 如果当前的综合指标等于最佳指标,将ckpt保存到文件best,这里保存的是最佳的模型。
- if opt.save_period > 0 and epoch % opt.save_period == 0: torch.save(ckpt, w / f'epoch{epoch}.pt'): 如果设置了保存周期save_period且当前轮数是保存周期的倍数,将ckpt保存到文件epoch{epoch}.pt,用于定期保存模型。
- callbacks.run('on_model_save', last, epoch, final_epoch, best_fitness, fi): 运行回调函数on_model_save,将保存的模型文件路径、当前轮数、是否是最后一轮、最佳综合指标和当前综合指标等参数传递给回调函数。
- if stop: break: 如果满足停止训练的条件,跳出训练循环,结束训练。
- if RANK in {-1, 0}: LOGGER.info(f'\n{epoch - start_epoch + 1} epochs completed in {(time.time() - t0) / 3600:.3f} hours.'): 如果是主进程(RANK为-1或0),打印训练完成的信息,包括训练轮数和所花费的时间。
- for f in last, best: ...: 遍历最后一轮的模型和最佳模型。
- if f.exists(): ...: 如果模型文件存在。
- strip_optimizer(f): 去除模型文件中的优化器信息。
- if f is best: ...: 如果是最佳模型,运行验证函数对模型进行评估。
- callbacks.run('on_fit_epoch_end', list(mloss) + list(results) + lr, epoch, best_fitness, fi): 运行回调函数on_fit_epoch_end,将损失值、验证集结果和学习率等参数传递给回调函数。
- callbacks.run('on_train_end', last, best, epoch, results): 运行回调函数on_train_end,将最后一轮模型、最佳模型、当前轮数和验证集结果等参数传递给回调函数。
- torch.cuda.empty_cache(): 清空GPU缓存。
3. val模型验证
在深度学习中,"val" 通常指的是验证数据集。验证数据集是在训练模型过程中用于评估模型性能和调整超参数的数据集。训练过程通常分为训练集、验证集和测试集三部分。验证集用于在训练过程中评估模型的性能,并根据验证结果调整模型的超参数,以优化模型的泛化能力。与训练集用于训练模型不同,验证集的目的是评估模型在未见过的数据上的性能,以避免过拟合和选择合适的模型。在本项目中,编写文件val.py实现模型验证功能,具体实现流程如下所示。
(1)编写函数save_one_txt(),功能是将目标检测模型的预测结果保存到文本文件中,具体实现流程如下:通过计算归一化增益,将预测框的坐标从xyxy格式转换为归一化的xywh格式。
- 遍历每个预测框,并根据是否保存置信度确定输出格式。
- 将结果写入文本文件中,每个预测结果占一行。
函数save_one_txt()的具体实现代码如下所示。
def save_one_txt(predn, save_conf, shape, file):
# Save one txt result
gn = torch.tensor(shape)[[1, 0, 1, 0]] # normalization gain whwh
for *xyxy, conf, cls in predn.tolist():
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh
line = (cls, *xywh, conf) if save_conf else (cls, *xywh) # label format
with open(file, 'a') as f:
f.write(('%g ' * len(line)).rstrip() % line + '\n')
(2)编写函数save_one_json()的功能是将目标检测模型的预测结果保存为JSON格式的文件,具体实现流程如下:
- 根据输入文件路径提取图像ID。
- 将预测框的坐标从xyxy格式转换为xywh格式,并将中心坐标转换为左上角坐标。
- 遍历每个预测框,将预测结果以字典的形式添加到列表中。
- 字典包括图像ID、类别ID、边界框坐标和置信度。
函数save_one_json()的具体实现代码如下所示。
def save_one_json(predn, jdict, path, class_map):
# Save one JSON result {"image_id": 42, "category_id": 18, "bbox": [258.15, 41.29, 348.26, 243.78], "score": 0.236}
image_id = int(path.stem) if path.stem.isnumeric() else path.stem
box = xyxy2xywh(predn[:, :4]) # xywh
box[:, :2] -= box[:, 2:] / 2 # xy center to top-left corner
for p, b in zip(predn.tolist(), box.tolist()):
jdict.append({
'image_id': image_id,
'category_id': class_map[int(p[5])],
'bbox': [round(x, 3) for x in b],
'score': round(p[4], 5)})
(2)编写函数process_batch(),功能是根据预测框和标签框计算正确的预测矩阵,具体实现流程如下:
- 创建一个全零矩阵correct,用于存储预测框和标签框之间的匹配情况。
- 计算预测框和标签框之间的IoU(交并比)。
- 对于每个IoU阈值,筛选出IoU大于阈值且类别匹配的预测框。
- 将匹配的结果记录在correct矩阵中,用布尔值表示匹配情况。
- 返回一个correct张量,其中每一行表示一个预测框,每一列表示一个IoU阈值,值为True表示匹配正确。
函数process_batch()的具体实现代码如下所示。
def process_batch(detections, labels, iouv):
# 初始化正确矩阵
correct = np.zeros((detections.shape[0], iouv.shape[0])).astype(bool)
# 计算IoU
iou = box_iou(labels[:, 1:], detections[:, :4])
# 检查类别是否匹配
correct_class = labels[:, 0:1] == detections[:, 5]
# 遍历每个IoU阈值
for i in range(len(iouv)):
# 寻找匹配项
x = torch.where((iou >= iouv[i]) & correct_class) # IoU > 阈值且类别匹配
if x[0].shape[0]:
# 组合匹配项
matches = torch.cat((torch.stack(x, 1), iou[x[0], x[1]][:, None]), 1).cpu().numpy() # [标签,检测,IoU]
# 对匹配项按IoU降序排序并去重
if x[0].shape[0] > 1:
matches = matches[matches[:, 2].argsort()[::-1]]
matches = matches[np.unique(matches[:, 1], return_index=True)[1]]
# matches = matches[matches[:, 2].argsort()[::-1]]
matches = matches[np.unique(matches[:, 0], return_index=True)[1]]
# 标记正确检测
correct[matches[:, 1].astype(int), i] = True
return torch.tensor(correct, dtype=torch.bool, device=iouv.device)
通过运行下面的命令可以展示识别结果:
python detect.py --weights yolov5s.pt --img 640 --conf 0.25 --source data/images
标签:opt,10,检测,模型,YOLO,hyp,epoch,model,save
From: https://blog.csdn.net/asd343442/article/details/137246315