YOLO系列 --- YOLOV7算法（三）：YOLO V7算法train.py代码解析

YOLO系列 — YOLOV7算法（三）：YOLO V7算法train.py代码解析

先介绍下每个参数的含义（直接在代码上写吧）

    parser = argparse.ArgumentParser()
    parser.add_argument('--weights',type=str, default='',help='initial weights path')#初始化权重文件，如果有预训练模型，可以直接在此加载
    parser.add_argument('--cfg',type=str, default=r'E:\work\People_Detect\yolov7-main\cfg\training\yolov7x.yaml',help='model.yaml path')#网络结构配置文件
    parser.add_argument('--data',type=str, default='data/custom_data.yaml',help='data.yaml path')#训练数据集配置文件
    parser.add_argument('--hyp',type=str, default='data/hyp.scratch.p5.yaml',help='hyperparameters path')#超参数配置文件
    parser.add_argument('--epochs',type=int, default=20)#训练迭代次数
    parser.add_argument('--batch-size',type=int, default=2,help='total batch size for all GPUs')#训练批次大小
    parser.add_argument('--img-size', nargs='+',type=int, default=[640,640],help='[train, test] image sizes')#训练图片大小
    parser.add_argument('--rect', action='store_true',help='rectangular training')#是否采用矩形训练，默认False
    parser.add_argument('--resume', nargs='?', const=True, default=False,help='resume most recent training')#是否继续进行训练，如果设置成True，那么会自动寻找最近训练权重文件
    parser.add_argument('--nosave', action='store_true',help='only save final checkpoint')#不保存权重文件，默认False
    parser.add_argument('--notest', action='store_true',help='only test final epoch')#不进行test，默认False
    parser.add_argument('--noautoanchor', action='store_true',help='disable autoanchor check')#不自动调整anchor，默认False
    parser.add_argument('--evolve', action='store_true',help='evolve hyperparameters')#是否进行超参数优化，默认是False，开启该选项，会加大训练时间，一般不需要
    parser.add_argument('--bucket',type=str, default='',help='gsutil bucket')#谷歌云盘bucket，一般不会用到
    parser.add_argument('--cache-images', action='store_true',help='cache images for faster training')#是否提前将训练数据进行缓存，默认是False
    parser.add_argument('--image-weights', action='store_true',help='use weighted image selection for training')#训练的时候是否选择图片权重进行训练
    parser.add_argument('--device', default='cpu',help='cuda device, i.e. 0 or 0,1,2,3 or cpu')#训练所使用的设备
    parser.add_argument('--multi-scale', action='store_true',help='vary img-size +/- 50%%')#是否进行多尺度训练，默认False
    parser.add_argument('--single-cls', action='store_true',help='train multi-class data as single-class')#训练数据集是否只有一类
    parser.add_argument('--adam', action='store_true',help='use torch.optim.Adam() optimizer')#是否使用adam优化器
    parser.add_argument('--sync-bn', action='store_true',help='use SyncBatchNorm, only available in DDP mode')#是否使用跨卡同步BN,在DDP模式使用
    parser.add_argument('--local_rank',type=int, default=-1,help='DDP parameter, do not modify')#gpu编号
    parser.add_argument('--workers',type=int, default=8,help='maximum number of dataloader workers')#dataloader的最大worker数量
    parser.add_argument('--project', default='runs/train',help='save to project/name')#训练结果保存路径
    parser.add_argument('--entity', default=None,help='W&B entity')#wandb库对应的东西，一般不用管
    parser.add_argument('--name', default='exp',help='save to project/name')#训练结果保存文件夹名称
    parser.add_argument('--exist-ok', action='store_true',help='existing project/name ok, do not increment')#判断下训练结果保存路径是否存在，如果存在的话，就不用重新创建
    parser.add_argument('--quad', action='store_true',help='quad dataloader')#作用是兼顾速度和精度，选择折中的方案
    parser.add_argument('--linear-lr', action='store_true',help='linear LR')#用于对学习速率进行调整，默认为 false，含义是通过余弦函数来降低学习率。
    parser.add_argument('--label-smoothing',type=float, default=0.0,help='Label smoothing epsilon')#是否做标签平滑，防止出现过拟合
    parser.add_argument('--upload_dataset', action='store_true',help='Upload dataset as W&B artifact table')#wandb库对应的东西
    parser.add_argument('--bbox_interval',type=int, default=-1,help='Set bounding-box image logging interval for W&B')#wandb 库对应的东西
    parser.add_argument('--save_period',type=int, default=-1,help='Log model after every "save_period" epoch')#用于记录训练日志信息
    parser.add_argument('--artifact_alias',type=str, default="latest",help='version of dataset artifact to be used')#这一行参数表达的是想实现但还未实现的一个内容，忽略即可
    opt = parser.parse_args()

这里，简单说下

argparse

这个库的用法吧，因为在detect.py和train.py中都是出现过的。

步骤：
1.parser = argparse.ArgumentParser()2.parser.add_argument（需要传入的参数）
3.opt = parser.parse_args()4.可以用**opt.**来调用对应的参数

# Set DDP variables
    opt.world_size =int(os.environ['WORLD_SIZE'])if'WORLD_SIZE'in os.environ else1
    opt.global_rank =int(os.environ['RANK'])if'RANK'in os.environ else-1
    set_logging(opt.global_rank)

跨卡多GPU训练，一般我们都是单机本地训练，所以用不到的，就不多讲了。

# Resume
    wandb_run = check_wandb_resume(opt)if opt.resume andnot wandb_run:# resume an interrupted run
        ckpt = opt.resume ifisinstance(opt.resume,str)else get_latest_run()# specified or most recent pathassert os.path.isfile(ckpt),'ERROR: --resume checkpoint does not exist'
        apriori = opt.global_rank, opt.local_rank
        withopen(Path(ckpt).parent.parent /'opt.yaml')as f:
            opt = argparse.Namespace(**yaml.load(f, Loader=yaml.SafeLoader))# replace
        opt.cfg, opt.weights, opt.resume, opt.batch_size, opt.global_rank, opt.local_rank ='', ckpt,True, opt.total_batch_size,*apriori  # reinstate
        logger.info('Resuming training from %s'% ckpt)else:# opt.hyp = opt.hyp or ('hyp.finetune.yaml' if opt.weights else 'hyp.scratch.yaml')
        opt.data, opt.cfg, opt.hyp = check_file(opt.data), check_file(opt.cfg), check_file(opt.hyp)# check filesassertlen(opt.cfg)orlen(opt.weights),'either --cfg or --weights must be specified'
        opt.img_size.extend([opt.img_size[-1]]*(2-len(opt.img_size)))# extend to 2 sizes (train, test)
        opt.name ='evolve'if opt.evolve else opt.name
        opt.save_dir = increment_path(Path(opt.project)/ opt.name, exist_ok=opt.exist_ok | opt.evolve)# increment run

首先判断下在训练超参数中是否设置了继续训练，要是设置了然后判断下传入的是不是字符串，如果是字符串就认为传入的就是继续训练的权重文件，要是不是字符串那么就调用

get_latest_run()

函数来寻找

last.pt

，然后重新编写opt.yaml文件。要是没有设置继续训练，那么直接检测一些yaml文件是否存在，扩充下测试时候图片大小，最后创建下保存训练路径。

# DDP mode
    opt.total_batch_size = opt.batch_size
    device = select_device(opt.device, batch_size=opt.batch_size)if opt.local_rank !=-1:assert torch.cuda.device_count()> opt.local_rank
        torch.cuda.set_device(opt.local_rank)
        device = torch.device('cuda', opt.local_rank)
        dist.init_process_group(backend='nccl', init_method='env://')# distributed backendassert opt.batch_size % opt.world_size ==0,'--batch-size must be multiple of CUDA device count'
        opt.batch_size = opt.total_batch_size // opt.world_size

这里是采用DDP模式，一般用不上，就不管了~

# Hyperparameterswithopen(opt.hyp)as f:
        hyp = yaml.load(f, Loader=yaml.SafeLoader)# load hyps

解析超参数文件，将yaml文件格式转换成python内置字典格式，如下图所示。

# Train
    logger.info(opt)ifnot opt.evolve:
        tb_writer =None# init loggersif opt.global_rank in[-1,0]:
            prefix = colorstr('tensorboard: ')
            logger.info(f"{prefix}Start with 'tensorboard --logdir {opt.project}', view at http://localhost:6006/")
            tb_writer = SummaryWriter(opt.save_dir)# Tensorboard
        train(hyp, opt, device, tb_writer)

这里判断是否进行超参数优化，其实一般我们都是默认False的，常规训练不需要该操作，效果已经很好了。而且如果进行超参数优化，训练时间会明显加长，所以这里我就没有把else模块的代码放出来了。下面直接进入

train

开始训练。

    logger.info(colorstr('hyperparameters: ')+', '.join(f'{k}={v}'for k, v in hyp.items()))
    save_dir, epochs, batch_size, total_batch_size, weights, rank = \
        Path(opt.save_dir), opt.epochs, opt.batch_size, opt.total_batch_size, opt.weights, opt.global_rank

进行简单的赋值

# Directories
    wdir = save_dir /'weights'
    wdir.mkdir(parents=True, exist_ok=True)# make dir
    last = wdir /'last.pt'
    best = wdir /'best.pt'
    results_file = save_dir /'results.txt'

创建权重文件保存路径，权重名字和训练日志txt文件

# Save run settingswithopen(save_dir /'hyp.yaml','w')as f:
        yaml.dump(hyp, f, sort_keys=False)withopen(save_dir /'opt.yaml','w')as f:
        yaml.dump(vars(opt), f, sort_keys=False)

将训练的相关参数全部写入

# Configure
    plots =not opt.evolve  # create plots
    cuda = device.type!='cpu'
    init_seeds(2+ rank)withopen(opt.data)as f:
        data_dict = yaml.load(f, Loader=yaml.SafeLoader)# data dict
    is_coco = opt.data.endswith('coco.yaml')

主要作用是解析data.yaml文件，判断训练集是否是coco数据集格式

    loggers ={'wandb':None}# loggers dictif rank in[-1,0]:
        opt.hyp = hyp  # add hyperparameters
        run_id = torch.load(weights).get('wandb_id')if weights.endswith('.pt')and os.path.isfile(weights)elseNone
        wandb_logger = WandbLogger(opt, Path(opt.save_dir).stem, run_id, data_dict)
        loggers['wandb']= wandb_logger.wandb
        data_dict = wandb_logger.data_dict
        if wandb_logger.wandb:
            weights, epochs, hyp = opt.weights, opt.epochs, opt.hyp  # WandbLogger might update weights, epochs if resuming

    nc =1if opt.single_cls elseint(data_dict['nc'])# number of classes
    names =['item']if opt.single_cls andlen(data_dict['names'])!=1else data_dict['names']# class namesassertlen(names)== nc,'%g names found for nc=%g dataset in %s'%(len(names), nc, opt.data)# check

这里主要是进行训练类别的计算。如果自定义数据集不止一个类别，但是又不小心将

single_cls 

设置为Ture的话，其实代码是不会报错的，但是这样就会在测试的时候不会正确的显示类别，所有的类别都变成了"items"。

# Model
    pretrained = weights.endswith('.pt')if pretrained:with torch_distributed_zero_first(rank):
            attempt_download(weights)# download if not found locally
        ckpt = torch.load(weights, map_location=device)# load checkpoint
        model = Model(opt.cfg or ckpt['model'].yaml, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device)# create
        exclude =['anchor']if(opt.cfg or hyp.get('anchors'))andnot opt.resume else[]# exclude keys
        state_dict = ckpt['model'].float().state_dict()# to FP32
        state_dict = intersect_dicts(state_dict, model.state_dict(), exclude=exclude)# intersect
        model.load_state_dict(state_dict, strict=False)# load
        logger.info('Transferred %g/%g items from %s'%(len(state_dict),len(model.state_dict()), weights))# reportelse:
        model = Model(opt.cfg, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device)# createwith torch_distributed_zero_first(rank):
        check_dataset(data_dict)# check
    train_path = data_dict['train']
    test_path = data_dict['val']

这部分代码首先是判断有没有传入预训练权重文件，分两类：
1.如果传入权重文件，直接

model.load_state_dict

加载模型
2.如果没有传入权重文件，创建模型实例，从头开始训练
这里，暂时先不讲网络结构相关的代码，等下篇blog单独来说说该部分~
最后，获取的

train_path

和

test_path

分别表示在data.yaml中训练数据集和测试数据集的地址。

# Freeze
    freeze =[]# parameter names to freeze (full or partial)for k, vin model.named_parameters():
        v.requires_grad = True  # train all layersif any(x in k forxin freeze):
            print('freezing %s' % k)
            v.requires_grad = False

这部分代码是设置冻结层的。简单来说，就是冻结指定的层的权重，让他们在反向传播的时候不更新权重大小。作者这里列出来这部分代码的目的其实并不是鼓励使用冻结指定层，因为作者认为这样效果其实并不是很好。

# Optimizer
    nbs =64# nominal batch size
    accumulate = max(round(nbs / total_batch_size), 1)# accumulate loss before optimizing
    hyp['weight_decay'] *= total_batch_size * accumulate / nbs  # scale weight_decay
    logger.info(f"Scaled weight_decay = {hyp['weight_decay']}")

nbs为模拟的batch_size; 就比如默认的话上面设置的opt.batch_size为16,这个nbs就为64，也就是模型梯度累积了64/16=4(accumulate)次之后再更新一次模型，变相的扩大了batch_size。

for k, vin model.named_modules():
        if hasattr(v, 'bias') and isinstance(v.bias, nn.Parameter):
            pg2.append(v.bias)# biasesif isinstance(v, nn.BatchNorm2d):
            pg0.append(v.weight)# no decayelif hasattr(v, 'weight') and isinstance(v.weight, nn.Parameter):
            pg1.append(v.weight)# apply decay

这里，先不着急解析代码，先简单来说下

model.named_modules()

：
先说model.modules()迭代遍历模型的所有子层，而model.named_modules()不但返回模型的所有子层，还会返回这些层的名字。还有一个 model.parameters()，它的作用是迭代地返回模型的所有参数。
然后，用

hasattr

函数来判断遍历的每个层对象是否拥有相对应的属性，将所有参数分成三类：weight、bn, bias。

if opt.adam:
        optimizer = optim.Adam(pg0, lr=hyp['lr0'], betas=(hyp['momentum'], 0.999))# adjust beta1 to momentum
    else:
        optimizer = optim.SGD(pg0, lr=hyp['lr0'], momentum=hyp['momentum'], nesterov=True)
    
    optimizer.add_param_group({'params': pg1, 'weight_decay': hyp['weight_decay']})# add pg1 with weight_decay
    optimizer.add_param_group({'params': pg2})# add pg2 (biases)
    logger.info('Optimizer groups: %g .bias, %g conv.weight, %g other' % (len(pg2), len(pg1), len(pg0)))
    del pg0, pg1, pg2

选用优化器，并设置weights的优化方式。然后将其他的参数喂入优化器中。

if opt.linear_lr:
        lf = lambda x: (1 - x / (epochs - 1)) * (1.0 - hyp['lrf']) + hyp['lrf']# linear
    else:
        lf = one_cycle(1, hyp['lrf'], epochs)# cosine 1->hyp['lrf']
    scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)

这里是设置学习率衰减方式，在训练参数那块opt.linear_lr默认的是False，所以最后还是使用的余弦退火方式进行衰减，相对应的函数是

one_cycle

。

# Image sizes
    gs = max(int(model.stride.max()), 32)# grid size (max stride)nl= model.model[-1].nl  # number of detection layers (used for scaling hyp['obj'])
    imgsz, imgsz_test =[check_img_size(x, gs)forxin opt.img_size]# verify imgsz are gs-multiples

获取网络的最大步长，预测头数量和训练及测试图片长宽大小。

# DP modeif cuda and rank == -1 and torch.cuda.device_count()>1:
        model = torch.nn.DataParallel(model)# SyncBatchNormif opt.sync_bn and cuda and rank != -1:
        model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device)
        logger.info('Using SyncBatchNorm()')

这块代码是跨卡训练，一般不考虑。

# Trainloader
    dataloader, dataset = create_dataloader(train_path, imgsz, batch_size, gs, opt,
                                            hyp=hyp, augment=True, cache=opt.cache_images, rect=opt.rect, rank=rank,
                                            world_size=opt.world_size, workers=opt.workers,
                                            image_weights=opt.image_weights, quad=opt.quad, prefix=colorstr('train: '))
    mlc = np.concatenate(dataset.labels, 0)[:, 0].max()# max label class
    nb = len(dataloader)# number of batches
    assert mlc < nc, 'Label class %g exceeds nc=%g in %s. Possible class labels are 0-%g' % (mlc, nc, opt.data, nc - 1)

create_dataloader

这部分代码是训练数据的读取，后面也是会单独拿出来解析，这里就跳过了。然后获取标签文件中所有类别数大小，如果与设定的类别数不一致就会报错。这里值得注意，

dataset.labels

是一个列表，内部一个element就是一个图片相对应的标签，而

np.concatenate(dataset.labels, 0)

就会将所有的标签的每行标签进行保存在一起，如下图所示：


上图是dataset.labels，下图是np.concatenate(dataset.labels, 0)。

# Process 0if rank in[-1, 0]:
        testloader = create_dataloader(test_path, imgsz_test, batch_size * 2, gs, opt,  # testloaderhyp=hyp, cache=opt.cache_images and not opt.notest, rect=True, rank=-1,
                                       world_size=opt.world_size, workers=opt.workers,
                                       pad=0.5, prefix=colorstr('val: '))[0]if not opt.resume:
            labels = np.concatenate(dataset.labels, 0)
            c = torch.tensor(labels[:, 0])# classes# cf = torch.bincount(c.long(), minlength=nc) + 1.  # frequency# model._initialize_biases(cf.to(device))if plots:
                #plot_labels(labels, names, save_dir, loggers)if tb_writer:
                    tb_writer.add_histogram('classes', c, 0)# Anchorsif not opt.noautoanchor:
                check_anchors(dataset, model=model, thr=hyp['anchor_t'], imgsz=imgsz)
            model.half().float()# pre-reduce anchor precision

创建测试集dataloader（也是调用

create_dataloader

类，后面一起解析吧），然后将所有样本的标签拼接到一起shape为(total, 5)，统计后做可视化，同时获得所有样本的类别，根据上面的统计对所有样本的类别，中心点xy位置，长宽wh做可视化。

check_anchors

计算默认锚点anchor与数据集标签框的长宽比值，标签的长h宽w与anchor的长h_a宽w_a的比值, 即h/h_a, w/w_a都要在(1/hyp[‘anchor_t’], hyp[‘anchor_t’])是可以接受的，如果标签框满足上面条件的数量小于总数的99%，则根据k-mean算法聚类新的锚点anchor。

# Model parameters
    hyp['box'] *=3. / nl# scale to layers
    hyp['cls'] *=nc / 80. * 3. / nl# scale to classes and layers
    hyp['obj'] *=(imgsz / 640) ** 2 * 3. / nl# scale to image size and layers
    hyp['label_smoothing']= opt.label_smoothing
    model.nc =nc# attach number of classes to model
    model.hyp = hyp  # attach hyperparameters to model
    model.gr =1.0# iou loss ratio (obj_loss = 1.0 or iou)
    model.class_weights = labels_to_class_weights(dataset.labels, nc).to(device) * nc# attach class weights
    model.names = names

根据自己数据集的类别数设置分类损失的系数，位置损失的系数。设置类别数，超参数等操作

# Start training
    t0 = time.time()
    nw = max(round(hyp['warmup_epochs'] * nb), 1000)# number of warmup iterations, max(3 epochs, 1k iterations)# nw = min(nw, (epochs - start_epoch) / 2 * nb)  # limit warmup to < 1/2 of training
    maps = np.zeros(nc)# mAP per class
    results =(0, 0, 0, 0, 0, 0, 0)# P, R, [email protected], [email protected], val_loss(box, obj, cls)
    scheduler.last_epoch = start_epoch - 1# do not move 
    scaler = amp.GradScaler(enabled=cuda)
    compute_loss_ota = ComputeLossOTA(model)# init loss class
    compute_loss = ComputeLoss(model)# init loss class
    logger.info(f'Image sizes {imgsz} train, {imgsz_test} test\n'
                f'Using {dataloader.num_workers} dataloader workers\n'
                f'Logging results to {save_dir}\n'
                f'Starting training for {epochs} epochs...')
    torch.save(model, wdir / 'init.pt')forepochin range(start_epoch, epochs):  # epoch ------------------------------------------------------------------
        model.train()

首先计算下热身训练的迭代次数，然后

model.train()

开始训练。