train.py
import argparse
import math
import os
import random
import sys
import time
from copy import deepcopy
from datetime import datetime
from pathlib import Path
import numpy as np
import torch
import torch.distributed as dist
import torch.nn as nn
import yaml
from torch.optim import lr_scheduler
from tqdm import tqdm
os.environ["GIT_PYTHON_REFRESH"] = "quiet"
os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"
FILE = Path(__file__).resolve()
ROOT = FILE.parents[0] # YOLOv5 root directory
if str(ROOT) not in sys.path:
sys.path.append(str(ROOT)) # add ROOT to PATH
ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative
import val as validate # for end-of-epoch mAP
from models.experimental import attempt_load
from models.yoloxiugai import Model
from utils.autoanchor import check_anchors
from utils.autobatch import check_train_batch_size
from utils.callbacks import Callbacks
from utils.dataloaders import create_dataloader
from utils.downloads import attempt_download, is_url
from utils.general import (LOGGER, TQDM_BAR_FORMAT, check_amp, check_dataset, check_file, check_git_info,
check_git_status, check_img_size, check_requirements, check_suffix, check_yaml, colorstr,
get_latest_run, increment_path, init_seeds, intersect_dicts, labels_to_class_weights,
labels_to_image_weights, methods, one_cycle, print_args, print_mutation, strip_optimizer,
yaml_save)
from utils.loggers import Loggers
from utils.loggers.comet.comet_utils import check_comet_resume
from utils.loss import ComputeLoss
from utils.metrics import fitness
from utils.plots import plot_evolve
from utils.torch_utils import (EarlyStopping, ModelEMA, de_parallel, select_device, smart_DDP, smart_optimizer,
smart_resume, torch_distributed_zero_first)
LOCAL_RANK = int(os.getenv('LOCAL_RANK', -1))
RANK = int(os.getenv('RANK', -1))
WORLD_SIZE = int(os.getenv('WORLD_SIZE', 1))
GIT_INFO = check_git_info()
#定义一个名为train的函数，接收四个参数：hyp（超参数），opt（选项），device（设备）和callbacks（回调函数）
def train(hyp, opt, device, callbacks):
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) # 如果evolve为真，则创建权重文件所在的目录；否则，直接创建权重文件。parents=True表示创建所有必要的父目录，exist_ok=True表示如果目录已存在，则不抛出异常。
last, best = w / 'last.pt', w / 'best.pt'#创建最后保存的权重文件和最佳权重文件的路径。
# 判断hyp是否为字符串类型
if isinstance(hyp, str):
with open(hyp, errors='ignore') as f: #以只读模式打开hyp文件，忽略错误
hyp = yaml.safe_load(f) # 使用yaml.safe_load()函数从文件中加载超参数字典
LOGGER.info(colorstr('hyperparameters: ') + ', '.join(f'{k}={v}' for k, v in hyp.items())) #打印超参数信息。
opt.hyp = hyp.copy() # 将加载的超参数复制到opt对象中，以便在后续的训练过程中使用。
# 判断是否进行模型进化
if not evolve:
yaml_save(save_dir / 'hyp.yaml', hyp) #将超参数保存到hyp.yaml文件中
yaml_save(save_dir / 'opt.yaml', vars(opt)) #将选项保存到opt.yaml文件中
# 初始化一个空的数据字典，用于存储数据集信息。
data_dict = None
if RANK in {-1, 0}: #判断当前进程是否为主进程（RANK为-1表示主进程）
loggers = Loggers(save_dir, weights, opt, hyp, LOGGER) # 创建一个Loggers实例，用于记录训练过程中的信息
# 遍历Loggers实例的所有方法
for k in methods(loggers):
callbacks.register_action(k, callback=getattr(loggers, k)) #将Loggers实例的方法注册到回调函数中
# 获取远程数据集的链接
data_dict = loggers.remote_dataset
if resume: # 如果需要从远程数据集恢复训练，则执行以下操作
weights, epochs, hyp, batch_size = opt.weights, opt.epochs, opt.hyp, opt.batch_size #从选项中获取权重、轮数、超参数和批量大小。
#配置部分开始。
plots = not evolve and not opt.noplots #创建图表
cuda = device.type != 'cpu' #检查是否有CUDA设备可用,如果不可用选CPU
init_seeds(opt.seed + 1 + RANK, deterministic=True) #初始化随机种子
with torch_distributed_zero_first(LOCAL_RANK): #确保在分布式训练中，每个进程都从第一个进程开始执行
data_dict = data_dict or check_dataset(data) # 检查数据集是否存在，如果不存在则创建一个新的数据集。
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') #检查权重文件的后缀是否为.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() # 将模型的权重转换为浮点数并获取其状态字典
csd = intersect_dicts(csd, model.state_dict(), exclude=exclude) #获取两个字典的交集
model.load_state_dict(csd, strict=False) # 加载模型的状态字典。
LOGGER.info(f'Transferred {len(csd)}/{len(model.state_dict())} items from {weights}') # 报告转移的项目数量和模型的状态字典长度。
# 如果不是预训练模型，则执行以下操作。
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 # 训练所有层
#将其设置为 False 表示该层不需要计算梯度，即该层被冻结，不会在训练过程中更新权重。
if any(x in k for x in freeze):
LOGGER.info(f'freezing {k}')
v.requires_grad = False
# 这段代码定义了图像的大小。首先计算网格大小（最大步幅），然后使用check_img_size函数验证图像大小是否为网格大小的倍数。这样可以确保输入图像与模型的步幅相匹配。
gs = max(int(model.stride.max()), 32)
imgsz = check_img_size(opt.imgsz, gs, floor=gs * 2)
# 这段代码估计最佳的批量大小。如果当前设备只有一个GPU并且没有指定批量大小，则使用check_train_batch_size函数来估算最佳批量大小，并将其记录到日志中。
if RANK == -1 and batch_size == -1:
batch_size = check_train_batch_size(model, imgsz, amp)
loggers.on_params_update({"batch_size": batch_size})
# 定义了优化器。首先设置名义上的批量大小为64，然后计算累积损失的数量。接下来，根据批量大小和累积损失的数量调整权重衰减系数。最后，使用smart_optimizer函数创建一个优化器实例。
nbs = 64
accumulate = max(round(nbs / batch_size), 1)
hyp['weight_decay'] *= batch_size * accumulate / nbs
optimizer = smart_optimizer(model, opt.optimizer, hyp['lr0'], hyp['momentum'], hyp['weight_decay'])
# 定义了一个学习率调度器。根据opt.cos_lr的值选择使用余弦退火或线性衰减策略。然后使用lr_scheduler.LambdaLR创建一个学习率调度器实例。
if opt.cos_lr:
lf = one_cycle(1, hyp['lrf'], epochs) # cosine 1->hyp['lrf']
else:
lf = lambda x: (1 - x / epochs) * (1.0 - hyp['lrf']) + hyp['lrf'] # linear
scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf) # plot_lr_scheduler(optimizer, scheduler, epochs)
# 创建了一个模型指数移动平均（EMA）实例。如果当前设备是主设备（RANK为-1），则使用ModelEMA类创建一个EMA实例；否则，将EMA实例设置为None。
ema = ModelEMA(model) if RANK in {-1, 0} else None
#用于恢复训练。如果提供了预训练模型，并且设置了Resume恢复选项，则使用smart_resume函数从检查点文件中恢复训练状态。最后，删除不再需要的检查点文件。
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
# 启用了分布式数据并行模式（DP）。如果当前设备是CUDA设备，并且不是主设备（RANK不为-1），并且有多个CUDA设备可用，则使用torch.nn.DataParallel包装模型以实现多GPU训练。
if cuda and RANK == -1 and torch.cuda.device_count() > 1:
LOGGER.warning('WARNING ⚠️ DP not recommended, use torch.distributed.run for best DDP Multi-GPU results.\n'
'See Multi-GPU Tutorial at https://github.com/ultralytics/yolov5/issues/475 to get started.')#显示一条警告信息，建议不要使用DP模式，而是使用torch.distributed.run进行多GPU训练。
model = torch.nn.DataParallel(model)
# 启用了同步批量归一化（SyncBatchNorm）
if opt.sync_bn and cuda and RANK != -1:
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device)
LOGGER.info('Using SyncBatchNorm()')
# 这一行创建了一个数据加载器（train_loader），用于从训练数据集（train_path）中加载图像和标签
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) #imgsz表示图像的大小，batch_size表示每个批次的样本数量，gs表示网格大小，single_cls表示是否使用单类别分类，hyp表示超参数，augment表示是否进行数据增强，cache表示缓存选项，rect表示是否使用矩形框，rank表示进程的等级，workers表示工作线程数，image_weights表示图像权重，quad表示是否使用四倍分辨率，prefix表示前缀字符串，shuffle表示是否打乱数据顺序，seed表示随机种子
labels = np.concatenate(dataset.labels, 0) #将数据集中的标签合并为一个NumPy数组（labels）
mlc = int(labels[:, 0].max()) #并计算最大标签类别（mlc）
assert mlc < nc, f'Label class {mlc} exceeds nc={nc} in {data}. Possible class labels are 0-{nc - 1}' #使用断言语句检查最大标签类别是否小于预设的最大类别数（nc），如果不满足条件，则抛出异常并显示错误信息。
#检查当前的进程（由RANK表示）是否在集合{-1, 0}中
if RANK in {-1, 0}:
#调用create_dataloader函数来创建一个数据加载器，用于加载验证数据集。这个函数的参数包括验证数据集的路径、图像大小、批量大小、网格大小、是否只使用单一类别、超参数等
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:
#如果选项noautoanchor为False，那么执行接下来的代码
if not opt.noautoanchor:
check_anchors(dataset, model=model, thr=hyp['anchor_t'], imgsz=imgsz) #调用check_anchors函数来检查模型的锚点是否合适。这个函数的参数包括数据集、模型、阈值和图像大小
model.half().float() # 这一行将模型的权重转换为半精度浮点数，以减少计算需求。
callbacks.run('on_pretrain_routine_end', labels, names) #运行名为'on_pretrain_routine_end'的回调函数，该函数通常用于在预训练阶段结束时执行一些操作
# 如果CUDA可用且当前进程不是主进程，那么执行接下来的代码块
if cuda and RANK != -1:
model = smart_DDP(model)
#模型属性设置包括检测层的数量、超参数等，并将它们附加到模型上
nl = de_parallel(model).model[-1].nl # 检测层的数量（用于调整超参数）
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个epochs，100个迭代）
# nw = min(nw, (epochs - start_epoch) / 2 * nb) # 限制热身迭代次数 < 训练的一半
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) # init loss class
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 # class weights
iw = labels_to_image_weights(dataset.labels, nc=nc, class_weights=cw) # image weights
dataset.indices = random.choices(range(dataset.n), weights=iw, k=dataset.n) # rand weighted 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: #遍历训练数据集中的每个批次。pbar是一个进度条对象，用于显示训练进度。
callbacks.run('on_train_batch_start') #每个批次开始时运行回调函数，可以在这里执行一些自定义操作，例如记录日志、更新学习率等
ni = i + nb * epoch #自训练开始以来的积分批次数。计算当前批次在整个训练过程中的索引值。ni表示自训练开始以来的积分批次数
imgs = imgs.to(device, non_blocking=True).float() / 255 #将输入图像数据转换为张量，并将其移动到指定的设备（如GPU）上。然后将图像数据归一化到0-1之间。
#判断当前批次是否在warmup阶段
if ni <= nw:
xi = [0, nw] #定义一个列表，表示warmup阶段的边界
#根据当前批次的索引值，计算累积梯度的数量。这里使用了numpy的interp函数进行插值计算。
accumulate = max(1, np.interp(ni, xi, [1, nbs / batch_size]).round())
for j, x in enumerate(optimizer.param_groups): #遍历优化器的参数组
# 根据当前批次的索引值，计算学习率。这里使用了numpy的interp函数进行插值计算。
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']]) #根据当前批次的索引值，计算动量值。这里使用了numpy的interp函数进行插值计算。
# 判断是否使用多尺度训练
if opt.multi_scale:
sz = random.randrange(imgsz * 0.5, 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:]] #计算缩放后的新尺寸
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)) #计算损失值和损失项
if RANK != -1: #判断是否使用分布式训练。
loss *= WORLD_SIZE #在分布式训练模式下，将损失值乘以进程数量，以实现梯度的平均
if opt.quad: #判断是否使用四倍损失
loss *= 4. #将损失值乘以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) # 将模型参数的梯度裁剪到最大范数为10.0
scaler.step(optimizer) #使用自动混合精度训练器执行一步优化
scaler.update() #更新自动混合精度训练器的状态
optimizer.zero_grad() #将优化器的梯度清零
if ema: #如果存在指数移动平均（EMA），则更新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' #获取当前GPU内存占用情况。
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 ------------------------------------------------------------------------------------------------
# 接下来的几行代码是关于学习率调度和早停机制的处理。
lr = [x['lr'] for x in optimizer.param_groups]
scheduler.step() #执行学习率调度器的一个步骤
#再次检查当前进程是否是主进程
if RANK in {-1, 0}:
#运行回调函数，处理训练周期结束后的操作
callbacks.run('on_train_epoch_end', epoch=epoch)
ema.update_attr(model, include=['yaml', 'nc', 'hyp', 'names', 'stride', 'class_weights']) #更新EMA的属性
final_epoch = (epoch + 1 == epochs) or stopper.possible_stop #判断是否到达最后一个周期或早停器可能停止训练。
if not noval or final_epoch: # 如果不需要验证集或者已经到达最后一个周期，则执行以下操作
#运行验证过程，并获取验证结果
results, maps, _ = validate.run(data_dict,
batch_size=batch_size // WORLD_SIZE * 2,
imgsz=imgsz,
half=amp,
model=ema.ema,
single_cls=single_cls,
dataloader=val_loader,
save_dir=save_dir,
plots=False,
callbacks=callbacks,
compute_loss=compute_loss)
#计算适应度值
fi = fitness(np.array(results).reshape(1, -1)) # weighted combination of [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, # 构建模型检查点字典
'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:
broadcast_list = [stop if RANK == 0 else None]
dist.broadcast_object_list(broadcast_list, 0) #在所有进程中广播广播列表
if RANK != 0: #如果不是主进程，则从广播列表中获取早停状态
stop = broadcast_list[0]
if stop:
break # #跳出所有DDP进程的训练循环
# end 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: #循环遍历last和best两个文件
if f.exists():
strip_optimizer(f) #调用了一个名为strip_optimizer的函数，用于移除优化器的信息
if f is best:
LOGGER.info(f'\nValidating {f}...') #这行代码打印出正在验证的文件名
results, _, _ = validate.run( #这行代码调用了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, # best pycocotools at iou 0.65
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) # val best model with plots
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() #清空了CUDA缓存，释放GPU内存
return results
def parse_opt(known=False):
parser = argparse.ArgumentParser()
parser.add_argument('--weights', type=str, default=ROOT / 'yolov5s.pt', help='初始化权重路径')
parser.add_argument('--cfg', type=str, default='models/yolov5S.yaml', help='模型配置文件路径')
parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='数据集配置文件路径')
parser.add_argument('--hyp', type=str, default=ROOT / 'data/hyps/hyp.scratch-low.yaml', help=' 超参数文件路径')
parser.add_argument('--epochs', type=int, default=60, help='总训练轮数，一般100-300轮')
parser.add_argument('--batch-size', type=int, default=4, help='所有GPU的总批量大小，如果为-1则自动计算，一般为8以内')
parser.add_argument('--imgsz', '--img', '--img-size', type=int, default=640, help='训练和验证图像尺寸（像素），一般不动')
parser.add_argument('--rect', action='store_true', help='是否使用矩形训练')
parser.add_argument('--resume', nargs='?', const=True, default=False, help='是否恢复最近的训练')
parser.add_argument('--nosave', action='store_true', help='是否仅保存最终检查点')
parser.add_argument('--noval', action='store_true', help='是否仅验证最终轮次')
parser.add_argument('--noautoanchor', action='store_true', help='是否禁用AutoAnchor')
parser.add_argument('--noplots', action='store_true', help='是否保存无图表文件')
parser.add_argument('--evolve', type=int, nargs='?', const=300, help='进化超参数进行x代')
parser.add_argument('--bucket', type=str, default='', help='gsutil存储桶，默认为空字符串')
parser.add_argument('--cache', type=str, nargs='?', const='ram', help='图像缓存类型，有ram/disk两种类型')
parser.add_argument('--image-weights', action='store_true', help='使用加权图像选择进行训练')
parser.add_argument('--device', default='', help='CUDA设备，例如0或0,1,2,3或cpu')
parser.add_argument('--multi-scale', action='store_true', help='是否使用多尺度训练')
parser.add_argument('--single-cls', action='store_true', help='是否将多类数据作为单类训练')
parser.add_argument('--optimizer', type=str, choices=['SGD', 'Adam', 'AdamW'], default='SGD', help='优化器类型，可选值为SGD、Adam和AdamW')
parser.add_argument('--sync-bn', action='store_true', help='是否使用SyncBatchNorm，仅在DDP模式下可用')
parser.add_argument('--workers', type=int, default=10, help='最大数据加载器工作线程数（每个RANK在DDP模式下），默认为10')
parser.add_argument('--project', default=ROOT / 'runs/train', help='保存到项目/名称，默认为runs/train')
parser.add_argument('--name', default='exp', help='保存到项目/名称，默认为exp')
parser.add_argument('--exist-ok', action='store_true', help='如果存在项目/名称，则不增加')
parser.add_argument('--quad', action='store_true', help='是否使用四倍速数据加载器')
parser.add_argument('--cos-lr', action='store_true', help='是否使用余弦学习率调度器')
parser.add_argument('--label-smoothing', type=float, default=0.0, help=' 标签平滑epsilon值')
parser.add_argument('--patience', type=int, default=100, help=' EarlyStopping耐心值（没有改进的轮次）')
parser.add_argument('--freeze', nargs='+', type=int, default=[0], help='冻结层：backbone=10, first3=0 1 2，默认为[0]')
parser.add_argument('--save-period', type=int, default=-1, help='每x个轮次保存检查点（如果<1则禁用），默认为-1')
parser.add_argument('--seed', type=int, default=0, help='全局训练种子，默认为0')
parser.add_argument('--local_rank', type=int, default=-1, help='自动DDP多GPU参数，不要修改，默认为-1')
# 这些参数用于配置日志记录和数据集上传的相关设置
parser.add_argument('--entity', default=None, help='Entity')
parser.add_argument('--upload_dataset', nargs='?', const=True, default=False, help='Upload data, "val" option')
parser.add_argument('--bbox_interval', type=int, default=-1, help='Set bounding-box image logging interval')
parser.add_argument('--artifact_alias', type=str, default='latest', help='Version of dataset artifact to use')
return parser.parse_known_args()[0] if known else parser.parse_args()
def main(opt, callbacks=Callbacks()):
# 检查是否处于主进程（RANK为-1或0），如果是，则打印opt中的参数，检查Git状态和依赖项
if RANK in {-1, 0}:
print_args(vars(opt))
check_git_status()
check_requirements()
# 代码尝试从指定的文件或最近的.pt文件中恢复训练。如果成功恢复训练，它会加载训练选项、数据集等信息，并更新opt对象。
if opt.resume and not check_comet_resume(opt) and not opt.evolve:
last = Path(check_file(opt.resume) if isinstance(opt.resume, str) else get_latest_run())
opt_yaml = last.parent.parent / 'opt.yaml' # train options yaml
opt_data = opt.data # original dataset
if opt_yaml.is_file():
with open(opt_yaml, errors='ignore') as f:
d = yaml.safe_load(f)
else:
d = torch.load(last, map_location='cpu')['opt']
opt = argparse.Namespace(**d) # replace
opt.cfg, opt.weights, opt.resume = '', str(last), True # reinstate
if is_url(opt_data):
opt.data = check_file(opt_data) # avoid HUB resume auth timeout
else:
opt.data, opt.cfg, opt.hyp, opt.weights, opt.project = \
check_file(opt.data), check_yaml(opt.cfg), check_yaml(opt.hyp), str(opt.weights), str(opt.project) # checks
assert len(opt.cfg) or len(opt.weights), 'either --cfg or --weights must be specified'
if opt.evolve:
if opt.project == str(ROOT / 'runs/train'): # if default project name, rename to runs/evolve
opt.project = str(ROOT / 'runs/evolve')
opt.exist_ok, opt.resume = opt.resume, False # pass resume to exist_ok and disable resume
if opt.name == 'cfg':
opt.name = Path(opt.cfg).stem # use model.yaml as name
opt.save_dir = str(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok))
#检查设备兼容性。如果使用多GPU DDP训练，它会确保设备兼容，并初始化进程组
device = select_device(opt.device, batch_size=opt.batch_size)
if LOCAL_RANK != -1:
msg = 'is not compatible with YOLOv5 Multi-GPU DDP training'
assert not opt.image_weights, f'--image-weights {msg}'
assert not opt.evolve, f'--evolve {msg}'
assert opt.batch_size != -1, f'AutoBatch with --batch-size -1 {msg}, please pass a valid --batch-size'
assert opt.batch_size % WORLD_SIZE == 0, f'--batch-size {opt.batch_size} must be multiple of WORLD_SIZE'
assert torch.cuda.device_count() > LOCAL_RANK, 'insufficient CUDA devices for DDP command'
torch.cuda.set_device(LOCAL_RANK)
device = torch.device('cuda', LOCAL_RANK)
dist.init_process_group(backend="nccl" if dist.is_nccl_available() else "gloo")
# 如果不是进行模型进化训练，代码会调用train函数开始训练
if not opt.evolve:
train(opt.hyp, opt, device, callbacks)
# 这段代码是用于超参数进化的。超参数进化是一种优化技术，通过在训练过程中调整超参数来找到最佳的模型性能
else:
meta = {
'lr0': (1, 1e-5, 1e-1), # initial learning rate (SGD=1E-2, Adam=1E-3)
'lrf': (1, 0.01, 1.0), # final OneCycleLR learning rate (lr0 * lrf)
'momentum': (0.3, 0.6, 0.98), # SGD momentum/Adam beta1
'weight_decay': (1, 0.0, 0.001), # optimizer weight decay
'warmup_epochs': (1, 0.0, 5.0), # warmup epochs (fractions ok)
'warmup_momentum': (1, 0.0, 0.95), # warmup initial momentum
'warmup_bias_lr': (1, 0.0, 0.2), # warmup initial bias lr
'box': (1, 0.02, 0.2), # box loss gain
'cls': (1, 0.2, 4.0), # cls loss gain
'cls_pw': (1, 0.5, 2.0), # cls BCELoss positive_weight
'obj': (1, 0.2, 4.0), # obj loss gain (scale with pixels)
'obj_pw': (1, 0.5, 2.0), # obj BCELoss positive_weight
'iou_t': (0, 0.1, 0.7), # IoU training threshold
'anchor_t': (1, 2.0, 8.0), # anchor-multiple threshold
'anchors': (2, 2.0, 10.0), # anchors per output grid (0 to ignore)
'fl_gamma': (0, 0.0, 2.0), # focal loss gamma (efficientDet default gamma=1.5)
'hsv_h': (1, 0.0, 0.1), # image HSV-Hue augmentation (fraction)
'hsv_s': (1, 0.0, 0.9), # image HSV-Saturation augmentation (fraction)
'hsv_v': (1, 0.0, 0.9), # image HSV-Value augmentation (fraction)
'degrees': (1, 0.0, 45.0), # image rotation (+/- deg)
'translate': (1, 0.0, 0.9), # image translation (+/- fraction)
'scale': (1, 0.0, 0.9), # image scale (+/- gain)
'shear': (1, 0.0, 10.0), # image shear (+/- deg)
'perspective': (0, 0.0, 0.001), # image perspective (+/- fraction), range 0-0.001
'flipud': (1, 0.0, 1.0), # image flip up-down (probability)
'fliplr': (0, 0.0, 1.0), # image flip left-right (probability)
'mosaic': (1, 0.0, 1.0), # image mixup (probability)
'mixup': (1, 0.0, 1.0), # image mixup (probability)
'copy_paste': (1, 0.0, 1.0)} # segment copy-paste (probability)
#这行代码打开一个文件，文件名由opt.hyp指定，忽略可能出现的错误。
with open(opt.hyp, errors='ignore') as f:
hyp = yaml.safe_load(f) #从文件中加载超参数字典
if 'anchors' not in hyp: #检查超参数字典中是否包含anchors这个键
hyp['anchors'] = 3
if opt.noautoanchor: #检查opt.noautoanchor是否为真
del hyp['anchors'], meta['anchors'] #删除hyp和meta字典中的anchors键
opt.noval, opt.nosave, save_dir = True, True, Path(opt.save_dir) # 将opt.noval、opt.nosave设置为True，并将opt.save_dir转换为Path对象
#创建两个文件路径对象，分别指向进化后的超参数YAML文件和CSV文件。
evolve_yaml, evolve_csv = save_dir / 'hyp_evolve.yaml', save_dir / 'evolve.csv'
if opt.bucket:
os.system(f'gsutil cp gs://{opt.bucket}/evolve.csv {evolve_csv}') #使用gsutil命令从Google Cloud Storage下载进化后的CSV文件。
for _ in range(opt.evolve): # generations to evolve
if evolve_csv.exists(): # if evolve.csv exists: select best hyps and mutate
# Select parent(s)
parent = 'single' # parent selection method: 'single' or 'weighted'
x = np.loadtxt(evolve_csv, ndmin=2, delimiter=',', skiprows=1)
n = min(5, len(x)) # number of previous results to consider
x = x[np.argsort(-fitness(x))][:n] # top n mutations
w = fitness(x) - fitness(x).min() + 1E-6 # weights (sum > 0)
if parent == 'single' or len(x) == 1:
# x = x[random.randint(0, n - 1)] # random selection
x = x[random.choices(range(n), weights=w)[0]] # weighted selection
elif parent == 'weighted':
x = (x * w.reshape(n, 1)).sum(0) / w.sum() # weighted combination
# Mutate
mp, s = 0.8, 0.2 # mutation probability, sigma
npr = np.random
npr.seed(int(time.time()))
g = np.array([meta[k][0] for k in hyp.keys()]) # gains 0-1
ng = len(meta)
v = np.ones(ng)
while all(v == 1): # mutate until a change occurs (prevent duplicates)
v = (g * (npr.random(ng) < mp) * npr.randn(ng) * npr.random() * s + 1).clip(0.3, 3.0)
for i, k in enumerate(hyp.keys()): # plt.hist(v.ravel(), 300)
hyp[k] = float(x[i + 7] * v[i]) # mutate
# 遍历所有超参数的元数据
for k, v in meta.items():
hyp[k] = max(hyp[k], v[1]) # 确保超参数的值不低于其下限
hyp[k] = min(hyp[k], v[2]) # 确保超参数的值不超过其上限
hyp[k] = round(hyp[k], 5) # 将超参数的值保留到小数点后5位
# Train mutation
results = train(hyp.copy(), opt, device, callbacks)
callbacks = Callbacks()
# Write mutation results
keys = ('metrics/precision', 'metrics/recall', 'metrics/mAP_0.5', 'metrics/mAP_0.5:0.95', 'val/box_loss',
'val/obj_loss', 'val/cls_loss')
print_mutation(keys, results, hyp.copy(), save_dir, opt.bucket)
# Plot results
plot_evolve(evolve_csv)
LOGGER.info(f'Hyperparameter evolution finished {opt.evolve} generations\n'
f"Results saved to {colorstr('bold', save_dir)}\n"
f'Usage example: $ python train.py --hyp {evolve_yaml}')
def run(**kwargs):
# Usage: import train; train.run(data='coco128.yaml', imgsz=320, weights='yolov5m.pt')
opt = parse_opt(True)
for k, v in kwargs.items():
setattr(opt, k, v)
main(opt)
return opt
if __name__ == "__main__":
opt = parse_opt()
main(opt)
detect.py
import argparse
import os
import platform
import sys
from pathlib import Path
import torch
FILE = Path(__file__).resolve()
ROOT = FILE.parents[0] # YOLOv5 root directory
if str(ROOT) not in sys.path:
sys.path.append(str(ROOT)) # add ROOT to PATH
ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative
from models.common import DetectMultiBackend
from utils.dataloaders import IMG_FORMATS, VID_FORMATS, LoadImages, LoadScreenshots, LoadStreams
from utils.general import (LOGGER, Profile, check_file, check_img_size, check_imshow, check_requirements, colorstr, cv2,
increment_path, non_max_suppression, print_args, scale_boxes, strip_optimizer, xyxy2xywh)
from utils.plots import Annotator, colors, save_one_box
from utils.torch_utils import select_device, smart_inference_mode
@smart_inference_mode()
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
):
#将输入的源数据转换为字符串类型
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) #判断源数据是否为网络摄像头流或数字，或者是否为URL且不是文件。
screenshot = source.lower().startswith('screen') #判断源数据是否为屏幕截图
if is_url and is_file: #如果源数据是URL且为文件，则下载文件
source = check_file(source) # download
# 生成保存结果的目录路径
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 #获取模型的步长、类别名称和是否使用PyTorch后端
imgsz = check_img_size(imgsz, s=stride) # 检查图像大小是否符合模型要求
# Dataloader
bs = 1 # 设置批量大小为1
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() # uint8 to fp16/32
im /= 255 # 0 - 255 to 0.0 - 1.0
if len(im.shape) == 3:
im = im[None] # expand for batch dim
#进行模型推理
with dt[1]:
visualize = increment_path(save_dir / Path(path).stem, mkdir=True) if visualize else False
pred = model(im, augment=augment, visualize=visualize)
# 进行非极大值抑制
with dt[2]:
pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)
# Second-stage classifier (optional)
# pred = utils.general.apply_classifier(pred, classifier_model, im, im0s)
# 处理预测结果，遍历每个图像的预测结果
for i, det in enumerate(pred):
seen += 1 #seen变量用于记录处理过的图像数量
if webcam: # 根据是否使用摄像头，设置图像路径、保存路径和标签路径
p, im0, frame = path[i], im0s[i].copy(), dataset.count #设置了图像路径、复制的图像和帧数
s += f'{i}: ' #用于存储字符串信息
else: #如果使用摄像头作为输入源，则执行以下代码块
p, im0, frame = path, im0s.copy(), getattr(dataset, 'frame', 0) #设置了图像路径、复制的图像和帧数
p = Path(p) # 将图片路径转换为Path对象
save_path = str(save_dir / p.name) # 保存检测结果的图片路径
txt_path = str(save_dir / 'labels' / p.stem) + ('' if dataset.mode == 'image' else f'_{frame}') # 保存检测结果的文本文件路径
s += '%gx%g ' % im.shape[2:] # 打印了图像的尺寸信息
gn = torch.tensor(im0.shape)[[1, 0, 1, 0]] # 计算了归一化增益
imc = im0.copy() if save_crop else im0 # 根据是否需要保存裁剪图像来设置imc变量
annotator = Annotator(im0, line_width=line_thickness, example=str(names)) #创建了一个Annotator对象，用于在图像上绘制边界框和标签。
if len(det): #判断是否有检测到目标
# 将边界框从原始图像尺寸缩放到当前图像尺寸
det[:, :4] = scale_boxes(im.shape[2:], det[:, :4], im0.shape).round()
#遍历所有唯一的类别
for c in det[:, 5].unique():
n = (det[:, 5] == c).sum() # 计算了每个类别的检测次数
s += f"{n} {names[int(c)]}{'s' * (n > 1)}, " #将类别名称和检测次数添加到字符串中
#历所有的检测结果
for *xyxy, conf, cls in reversed(det):
if save_txt: #判断是否需要将检测结果写入文件
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() #计算了归一化的边界框坐标。
line = (cls, *xywh, conf) if save_conf else (cls, *xywh) # 设置了标签格式
with open(f'{txt_path}.txt', 'a') as f: #打开一个文本文件以追加模式写入检测结果
f.write(('%g ' * len(line)).rstrip() % line + '\n') #将检测结果写入文件
if save_img or save_crop or view_img: # 判断是否需要将边界框添加到图像上
c = int(cls) # 将类别转换为整数
label = None if hide_labels else (names[c] if hide_conf else f'{names[c]} {conf:.2f}') #置了标签名
annotator.box_label(xyxy, label, color=colors(c, True)) #在图像上绘制边界框和标签
if save_crop: #判断是否需要保存裁剪图像
save_one_box(xyxy, imc, file=save_dir / 'crops' / names[c] / f'{p.stem}.jpg', BGR=True) #保存裁剪图像。
# 获取带有边界框和标签的图像
im0 = annotator.result()
if view_img: #判断是否需要显示图像
if platform.system() == 'Linux' and p not in windows: #判断操作系统是否为Linux，并且当前图像不在窗口列表中。
windows.append(p) #将当前图像添加到窗口列表中
cv2.namedWindow(str(p), cv2.WINDOW_NORMAL | cv2.WINDOW_KEEPRATIO) #创建一个可调整大小的窗口。
cv2.resizeWindow(str(p), im0.shape[1], im0.shape[0]) #调整窗口大小以适应图像
cv2.imshow(str(p), im0) #显示图像
cv2.waitKey(1) # 窗口等待1毫秒
# 保存带有检测结果的图像
if save_img:
if dataset.mode == 'image': #判断数据集模式是否为图像
cv2.imwrite(save_path, im0) #将带有检测结果的图像保存到指定路径
else: # 判断数据集模式是否为视频或流媒体
if vid_path[i] != save_path: #判断当前视频是否与保存路径相同
vid_path[i] = save_path #更新视频路径
if isinstance(vid_writer[i], cv2.VideoWriter): #判断当前视频写入器是否为OpenCV的VideoWriter对象。
vid_writer[i].release() #释放之前的视频写入器
if vid_cap: # 判断是否存在视频文件
fps = vid_cap.get(cv2.CAP_PROP_FPS) #获取视频的帧率
w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))#获取视频的宽度
h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))#获取视频的高度
else: #判断是否为流媒体
fps, w, h = 30, im0.shape[1], im0.shape[0] #设置流媒体的帧率为30，宽度和高度与图像相同
save_path = str(Path(save_path).with_suffix('.mp4')) # 强制在结果视频上添加.mp4后缀。
vid_writer[i] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h)) #创建一个新的视频写入器对象。
vid_writer[i].write(im0) #将带有检测结果的图像写入视频文件
# 打印输出检测的时间
LOGGER.info(f"{s}{'' if len(det) else '(no detections), '}{dt[1].dt * 1E3:.1f}ms")
# 打印输出检测结果到控制台
t = tuple(x.t / seen * 1E3 for x in dt) # speeds per image
LOGGER.info(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {(1, 3, *imgsz)}' % t)
if save_txt or save_img:
s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}{s}")
if update:
strip_optimizer(weights[0]) # update model (to fix SourceChangeWarning)
def parse_opt():
parser = argparse.ArgumentParser()
parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'yolov5s.pt', help='预训练模型的路径')
parser.add_argument('--source', type=str, default=ROOT / 'data/images', help='数据源路径，也就是要测试的资源。如果是检测摄像头改成对应的序号，一般是0')
parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='指定数据集的YAML文件路径')
parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[640], help='指定推断图像的大小，程序会自动缩放')
parser.add_argument('--conf-thres', type=float, default=0.25, help='置信度阈值')
parser.add_argument('--iou-thres', type=float, default=0.45, help='用于指定非极大值抑制（NMS）的IoU阈值')
parser.add_argument('--max-det', type=int, default=1000, help='指定每张图像的最大检测次数')
parser.add_argument('--device', default='', help='指定检测的CUDA设备，用0、1、2、3数字表示。如果是CPU可以填写CPU。可以保留空自动选择')
parser.add_argument('--view-img', action='store_true', help='是否显示结果')
parser.add_argument('--save-txt', action='store_true', help='是否将结果保存为文本文件')
parser.add_argument('--save-conf', action='store_true', help='是否在保存的文本文件中包含置信度信息')
parser.add_argument('--save-crop', action='store_true', help='是否保存裁剪后的预测框')
parser.add_argument('--nosave', action='store_true', help='是否不保存图像或视频')
parser.add_argument('--classes', nargs='+', type=int, help='用于指定要过滤的类别，如果填写对应的类别序号如0，那么显示对应的标签')
parser.add_argument('--agnostic-nms', action='store_true', help='是否进行类无关的NMS')
parser.add_argument('--augment', action='store_true', help='是否进行增强推理')
parser.add_argument('--visualize', action='store_true', help='是否可视化特征')
parser.add_argument('--update', action='store_true', help='是否更新所有模型')
parser.add_argument('--project', default=ROOT / 'runs/detect', help='保存结果的项目名称')
parser.add_argument('--name', default='exp', help='指定保存结果的名称')
parser.add_argument('--exist-ok', action='store_true', help='指定现有的项目名称是否允许覆盖')
parser.add_argument('--line-thickness', default=3, type=int, help='指定边界框的线条粗细')
parser.add_argument('--hide-labels', default=False, action='store_true', help='是否隐藏标签')
parser.add_argument('--hide-conf', default=False, action='store_true', help='是否隐藏置信度')
parser.add_argument('--half', action='store_true', help='是否使用半精度浮点数进行推理')
parser.add_argument('--dnn', action='store_true', help='是否使用OpenCV的DNN模块进行ONNX推理')
parser.add_argument('--vid-stride', type=int, default=1, help='指定视频帧率的步长')
opt = parser.parse_args() #解析命令行参数并将结果存储在变量opt中
opt.imgsz *= 2 if len(opt.imgsz) == 1 else 1 # 如果opt.imgsz的长度为1，则将其乘以2，否则保持不变
print_args(vars(opt)) #打印解析后的命令行参数
return opt
def main(opt):
check_requirements(exclude=('tensorboard', 'thop'))
run(**vars(opt))
if __name__ == "__main__":
opt = parse_opt()
main(opt)
Q 2500050191
B站 熬夜冲浪冠军
闲鱼 认真专注开卷
TB 深度学习地表最强小店
提供付费 咨询 辅导 代劳等服务
个人精力有限，提供知识付费，免费部分为爱发电不被绑架。
1万+
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
