1. yolov5 train.py代码注释详解
1.1 训练器的初始化
在train.py代码中，我们首先会看到的是训练器的初始化部分，包括一些参数的设置，比如学习率、momentum、weight_decay等。
# 参数设置
parser = argparse.ArgumentParser()
parser.add_argument('--epochs', type=int, default=300) # 训练轮数
parser.add_argument('--batch-size', type=int, default=16) # batch size大小
parser.add_argument('--cfg', type=str, default='./models/yolov5s.yaml', help='模型配置文件路径') # 模型配置路径
parser.add_argument('--data', type=str, default='./data/coco128.yaml', help='数据配置文件路径') # 数据集配置路径
parser.add_argument('--weights', type=str, default='./weights/yolov5s.pt', help='初始化权重文件路径') # 模型初始权重路径
parser.add_argument('--device', default='', help='cuda设备号 eg. 0 or 0,1 or cpu') # 运行设备 cuda or cpu
parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='训练时图像大小') # 图像输入大小
parser.add_argument('--rect', action='store_true', help='矩形输入 размер')
parser.add_argument('--resume', nargs='?', const=True, default=False, help='从checkpoint恢复调整训练位置')
parser.add_argument('--nosave', action='store_true', help='只保存最终权重')
parser.add_argument('--notest', action='store_true', help='不进行测试')
parser.add_argument('--evolve', action='store_true', help='进化操作') # 进化操作，优化超参数
parser.add_argument('--bucket', type=str, default='', help='gsutil所在的S3路径（gs://bucket 或 s3://bucket）')
parser.add_argument('--cache-images', action='store_true', help='进行预加载（缓存）')
parser.add_argument('--hyp', type=str, default='data/hyp.scratch.yaml', help='超参数路径')
parser.add_argument('--project', default='runs/train', help='tensorboard保存路径')
parser.add_argument('--name', default='exp', help='tensorboard保存名称') # tensorboard记录名称
parser.add_argument('--exist-ok', action='store_true', help='文件已存在时运行报错')
parser.add_argument('--quad', action='store_true', help='quad transform')
parser.add_argument('--linear-lr', action='store_true', help='linear LR')
opt = parser.parse_args()
这些参数设置会影响模型的训练结果，比如学习率的大小决定了模型的训练速度和稳定性，而momentum和weight_decay则可以优化损失函数的下降和防止过拟合等问题。
1.2 数据读取和模型加载
刚刚设置好了训练器的参数，我们接下来要加载数据集，并且载入模型，为训练做好准备。
def train(hyp, opt, device):
# 加载模型配置
model_cfg = opt.cfg
assert model_cfg.endswith('.yaml'), '仅支持yaml模型配置文件'
with open(model_cfg) as f:
model = Model(yaml.safe_load(f), ch=3, nc=int(opt.data.split('.')[-1]))
# 加载模型的预训练权重模型
ckpt = opt.weights if opt.weights.endswith('.pt') else last_checkpoint(opt.weights)
start_epoch, best_fitness = 0, 0.0
if ckpt is not None:
# 加载权重模型的状态
ckpt = torch.load(ckpt, map_location=device)
# 从检查点中加载学习率、轮次和最佳的fitness
if opt.resume or opt.transfer:
start_epoch = ckpt['epoch'] + 1
# 忽略学习率为“None”的CKPT参数
hyp['lr0'] = hyp.get('lr0', 0) or 0.01
# 我们不要恢复best_fitness，除非我们从检查点剪切或转移
best_fitness = ckpt['best_fitness'] if 'best_fitness' in ckpt else 0.0
# 加载模型权重
state_dict = ckpt['model'].float().state_dict()
# 加载权重
try:
# 尝试加载所有模型参数
model.float().state_dict().update(state_dict)
except:
# 加载模型参数失败，则只加载最后一层的权重数据
model_info(model)
for k, v in state_dict.items():
k = k.replace('model.', '')
model.load_state_dict(state_dict, strict=False)
if opt.transfer: # 从预处理函数中进行transfer装换
print(f'Transferring {ckpt["model"].yaml["nc"]} -> {model.nc}...')
model = model.transfer(ckpt['model'], ip_size=tuple(opt.img_size))
del ckpt, state_dict
# 加载数据集
with open(opt.data) as f:
data_dict = yaml.load(f, Loader=yaml.FullLoader)
# 定义数据集以及数据增强器
train_path = data_dict['train']
val_path = data_dict['val']
nc = int(data_dict['nc'])
hyp['cls'] *= nc / 80.0
train = LoadImagesAndLabels(train_path, opt.img_size, opt.batch_size, augment=True, hyp=hyp, rect=opt.rect, cache_images=opt.cache_images)
val = LoadImagesAndLabels(val_path, opt.img_size, opt.batch_size * 2, augment=False, hyp=hyp, rect=True)
# 开启数据集缓存(可加速数据加载)
if opt.cache_images:
cache = CacheImages(train, path=os.path.join(os.getcwd(), 'train.cache'), mmap=False)
train = ConcatDataset([cache.create(i) for i in range(len(train.sources))])
val = CacheImages(val, path=os.path.join(os.getcwd(), 'val.cache'), mmap=False)
return model, train, val, start_epoch, best_fitness
这部分代码主要是加载模型的配置文件，并将其传入Model初始化方法中，以此构建模型。此外，还会根据接下来要训练的数据集设置数据集的路径，并且构建一个数据增强器，提高数据的多样性，增强模型的泛化能力。
1.3 模型训练
接下来是模型的训练部分，这部分代码的主要作用是根据前面已经设置好的参数，开始对模型进行训练并优化权重。
def train(hyp, opt, device):
# ...
# 训练器的初始化
t0 = time.time()
nw = max(round(hyp['warmup_epochs'] * len(train) / opt.batch_size), 1000) # warmup运行时间
# 完整的训练逻辑
print(f'{prefix}开始训练{epochs}个epochs，batch_size = {opt.batch_size}，img_sz = {train.shapes[-1][0]}')
for epoch in range(start_epoch, epochs):
t1 = time.time()
# 在EPOCHS末尾的超参数变化LR_POLY，MOMSCHED
adjust_learning_rate(optimizer, epoch, hyp['epochs'], hyp['lr0'])
adjust_momentum(optimizer, epoch, hyp['epochs'], hyp['momentum'])
#训练单次epoch（一遍图片）
train_losses = []
nbs = len(train) # 在不同的图片之间切换的次数（”，“数属于data.yaml中，图片的数量除以batch size，batch size是16）
for i, (imgs, targets, _, _) in enumerate(train):
# 在warmup期间，调整LR，从0逐渐达到LR0
ni = i + nbs * epoch
if ni <= nw:
xi = [0, nw] # 学习率变化的范围
# 我们在这里用imag的第二(initial_lr)参数来调整LR变化的幅度(lf)
# 这是为了避免数据依赖性超出一个EPOCH，保证了一致性集成运算
lf = lambda x: 1.0 * (1 - x / nw) ** 0.9 # \x\ 0-nw 之间进行线性降低
hyp['lr'] = opt.hyp['lr0'] * lf(ni / (nw - 1) if nw > 1 else 1) # 迭代周期t
# 更新网络学习（前向和反向）
loss, _ = model(imgs.to(device), targets.to(device)) # 运行 loss = model(imgs, targets)
# 记录训练损失
loss.backward() # 对权重进行反向、偏向、反向传播
if ni <= nw:
# 在warmup期间权重不会改变
accumulate(model_grads, model) # 累加渐变
if ni % accumulate == 0:
# 使用更新后的（已累积过的梯度）梯度更新，再将梯度重设为0
optimizer.step()
optimizer.zero_grad()
# 规范化损失
loss = loss.item()
train_losses.append(loss)
# 打印训练结果
if rank in [-1, 0]:
# 进度条日志记录logger[master].tensorboard
s = ('%8s%12s' + '%10.3g' * 6) % (
f'{epoch}/{epochs - 1}', f'{i}/{nbs - 1}', *losses.avg, hyp['lr'], imgs.shape[-1])
pbar.set_description(s)
# 保存训练结果到tensorboard
n_iter = epoch * nbs + i
tb_writer.add_scalar('loss', loss, n_iter)
tb_writer.add_scalar('lr', hyp['lr'], n_iter)
tb_writer.add_scalar('n_img', imgs.shape[0], n_iter)
# 内存管理gc
if gc and ni % 3 == 0:
#gc.collect()
torch.cuda.empty_cache()
# 进行一次验证
results, maps = val.run(data_dict, epoch, img_size=img_size, batch_size=batch_size * 2, adam=None)
if rank in [-1, 0]:
tb_writer.add_scalar('mAP', maps, epoch)
# 保存最优结果
fitness = sum([k[-1] for k in results.values()]) / len(results)
if fitness > best_fitness:
best_fitness = fitness
# 保存权重模型
if rank in [-1, 0]:
save = not opt.nosave or (epoch == epochs - 1)
if save:
with open(f'{opt.name}.pt', 'wb') as f:
f.write(get_latest_run()[-1]) # 最新的运行模型
# 保存最优参数，并显示
if best_fitness == fitness and epoch > 0:
with open(f'{opt.name}_best.pt', 'wb') as f:
f.write(get_latest_run()[-1])
print(f'已经保存最优模型的参数到 {opt.name}_best.pt...')
# CPU释放
del results, maps
#gc.collect();torch.cuda.empty_cache() # 内存管理gc
# 训练器运行时间记录
if rank in [-1, 0]:
tb_writer.close()
dist.destroy_process_group() if worlds_size > 1 else None
torch.cuda.empty_cache()
print(f"{prefix}训练完毕, 共花费{(time.time() - t0) / 3600:.2f}小时")
可以看到，在训练过程中，每个epoch都需要调整学习率和动量等超参数，以优化模型参数的更新。训练结束后，我们需要保存最优的权重，以便后续的使用。
2. yolov5 train.py使用教程
2.1 环境安装
在使用yolov5 train.py之前，需要确保计算机已经安装好了必要的环境，比如CUDA、cuDNN等，同时在Python中需要安装好torch和torchvision，可以执行以下代码尝试：
# 安装PyTorch
pip install torch torchvision
# 安装yolov5的依赖库
pip install -U -r requirements.txt
2.2 数据准备
在运行训练之前，需要准备训练数据集，需要满足以下条件：
- 数据集图片需要处理为统一的尺寸（比如640x640）。
- 数据集需要按照一定的目录结构组织，比如：
data/
├── train/
│ ├── 1.jpg
│ ├── 1.txt
│ ├── 2.jpg
│ ├── 2.txt
│ ├── ...
│ ├── classes.txt
│ └── train.txt
└── val/
├── 1.jpg
├── 1.txt
├── 2.jpg
├── 2.txt
├── ...
├── classes.txt
└── val.txt
# 其中，train.txt和val.txt保存的是图片路径和标注信息的映射
# classes.txt保存了类别信息
在上述文件夹中，train/表示训练集，val/表示验证集。每个图片都对应一个txt文件，里面保存了该图片上各个物体的位置信息与类别信息。
2.3 模型训练
准备好数据之后，我们就可以开始进行模型训练了。这里，我们以训练yolov5s模型为例，可以执行以下代码：
# 设置参数
$ python train.py --data /path/to/data --cfg ./models/yolov5s.yaml --weights '' --batch-size 16
# 进行训练：10 epochs以及正常的batch_size为16和img_size为640，使用默认数据集、模型配置和超参数
$ python train.py --epochs 10
# 继续历史训练
$ python train.py --resume
# 实时迭代超参数（超参数优化）
$ python train.py --evolve
# 多GPU训练
$ python train.py --device 0,1,2,3
# 更多参数运行方式详见README.md
2.4 模型应用
模型训练完成之后，我们可以使用训练好的权重文件进行模型应用。整个过程比较简单，只需要调用detect.py文件，并设置好相应参数即可，可以执行以下代码：
$ python detect.py --source /path/to/image/or/video --weights /path/to/weights --conf 0.25 ...
这里，source表示要进行检测的图片或视频路径，weights表示训练好的权重文件。conf参数是检测的置信度阈值，当检测结果的置信度小于该阈值时，该检测结果会被过滤掉。
3.5 结语
yolov5是一种基于深度学习的目标检测算法，它具有模型轻量化、速度快等优点