• YOLOR剪枝【附代码】

    本文中的剪枝采用的是通道剪枝，为离线剪枝的一种，也就是可以直接对已经训练好的模型进行剪枝后再微调训练，不用稀疏训练【也就是不用边训练边剪枝】。

    剪枝参考的paper：Pruning Filters for Efficient ConvNets。

    相关剪枝原理部分可以参考我另一篇文章：YOLOv4剪枝【附代码】_爱吃肉的鹏的博客-CSDN博客_yolo剪枝

    本文主要实现：

    1.可训练自己的数据集 

    2.对单独某一个卷积的剪枝，

    3.对某些层的剪枝。

     PS：最终的效果本文并不保证，需要根据自己剪枝方案进行评估，剪的地方不同以及剪枝率的不同都会有影响，本文只是把功能进行了实现。这里需要说一下的是，我在实际测试的时候，发现YOLOR其实并没多好，而且参数量也挺大的，所以大家选这个算法要慎重。

    目录

    训练部分代码讲解：

    训练参数说明：

    训练阶段：

     相关路径代码：

    保存训练过程参数的yaml文件

     相关配置

    Model的定义

     optimizer相关配置代码

    optimizer 预权重加载：

    训练集的加载 

    验证集的加载

     绘制labels

    训练

    训练自己的数据集

    检测推理

    剪枝

    1.保存完整的权重文件

    2.网络剪枝

    1.单独卷积的剪枝

    2.卷积层(某个模块)的剪枝

    3.剪枝后的微调训练 

    4.剪枝后的推理检测

    github代码：

    所遇问题：

    1.剪枝训练期间在测mAP的时候报错：

     

     本文暂时不讲YOLOR的理论部分【有些细节我也还在研究】。

    clone代码至本地后可以先测试一下：

    git clone https://github.com/YINYIPENG-EN/Pruning_for_YOLOR_pytorch

    cd Pruning_for_YOLOR_pytorch 

    python detect.py --source inference/images/horses.jpg --cfg cfg/yolor_csp.cfg --weights yolor_csp.pt --conf 0.25 --img-size 640 --device 0

     

     

    训练部分代码讲解：

    YOLOR的训练代码和其他YOLO系列差不多。这一部分将会分块说明训练代码中各个模块的功能【如果不想看这一部分可以略去】

    训练参数说明：

    首先看一下参数部分，这里只说一些经常用到的。

    --weights 是权重路径

    --cfg 网络结构配置路径

    --data 数据集yaml配置路径

    --hyp 训练期间超参数的配置路径

    --epochs  训练总的epochs

    --batch-size batch size大小

    --img-size 输入网络的图像大小

    --nosave 如果开启该功能，只保存最后一轮的训练结构，关闭该功能则每epoch都保存

    --notest 开启后只对最后一轮进行测试，关闭该功能则轮都测试

    --device 设备id，如果只有一个GPU，设置为0

    --adam  采用adam优化器训练，默认为SGD

    --pt 剪枝后的微调训练

    --period 测试mAP的周期，默认为每两轮测试一此mAP

    1. parser = argparse.ArgumentParser()
    2.     parser.add_argument('--weights', type=str, default='yolor_csp.pt', help='initial weights path')
    3.     parser.add_argument('--cfg', type=str, default='cfg/yolor_csp.cfg', help='model.yaml path')
    4.     parser.add_argument('--data', type=str, default='data/coco.yaml', help='data.yaml path')
    5.     parser.add_argument('--hyp', type=str, default='data/hyp.scratch.640.yaml', help='hyperparameters path')
    6.     parser.add_argument('--epochs', type=int, default=300)
    7.     parser.add_argument('--batch-size', type=int, default=4, help='total batch size for all GPUs')
    8.     parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='[train, test] image sizes')
    9.     parser.add_argument('--rect', action='store_true', help='rectangular training')
    10.     parser.add_argument('--resume', nargs='?', const=True, default=False, help='resume most recent training')
    11.     parser.add_argument('--nosave', action='store_true', help='only save final checkpoint')
    12.     parser.add_argument('--notest', action='store_true', help='only test final epoch')
    13.     parser.add_argument('--noautoanchor', action='store_true', help='disable autoanchor check')
    14.     parser.add_argument('--evolve', action='store_true', help='evolve hyperparameters')
    15.     parser.add_argument('--bucket', type=str, default='', help='gsutil bucket')
    16.     parser.add_argument('--cache-images', action='store_true', help='cache images for faster training')
    17.     parser.add_argument('--image-weights', action='store_true', help='use weighted image selection for training')
    18.     parser.add_argument('--device', default='0', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
    19.     parser.add_argument('--multi-scale', action='store_true', help='vary img-size +/- 50%%')
    20.     parser.add_argument('--single-cls', action='store_true', help='train as single-class dataset')
    21.     parser.add_argument('--adam', action='store_true', help='use torch.optim.Adam() optimizer')
    22.     parser.add_argument('--sync-bn', action='store_true', help='use SyncBatchNorm, only available in DDP mode')
    23.     parser.add_argument('--local_rank', type=int, default=-1, help='DDP parameter, do not modify')
    24.     parser.add_argument('--log-imgs', type=int, default=16, help='number of images for W&B logging, max 100')
    25.     parser.add_argument('--workers', type=int, default=4, help='maximum number of dataloader workers')
    26.     parser.add_argument('--project', default='runs/train', help='save to project/name')
    27.     parser.add_argument('--name', default='exp', help='save to project/name')
    28.     parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
    29.     parser.add_argument('--pt', action='store_true', default=False, help='pruned model train')
    30.     parser.add_argument('--period', type=int, default=2, help='test period')
    31.     opt = parser.parse_args()

    训练阶段：

     训练阶段的代码会调用train()函数，主要是hyp【超参数】，opt【传入参数】，device【设备】。

    1. # Train
    2.     logger.info(opt)
    3.     if not opt.evolve:
    4.         tb_writer = None  # init loggers
    5.         if opt.global_rank in [-1, 0]: # tensorboard的写入
    6.             logger.info(f'Start Tensorboard with "tensorboard --logdir {opt.project}", view at http://localhost:6006/')
    7.             tb_writer = SummaryWriter(opt.save_dir)  # Tensorboard
    8.         train(hyp, opt, device, tb_writer, wandb)  # 训练过程

     相关路径代码：

    1.     # Directories
    2.     wdir = save_dir / 'weights'  # 保存路径
    3.     wdir.mkdir(parents=True, exist_ok=True)  # make dir
    4.     last = wdir / 'last.pt'  # 最后一次权重
    5.     best = wdir / 'best.pt'  # 最好的权重
    6.     results_file = save_dir / 'results.txt'  # 结果txt路径(记录训练过程)

    保存训练过程参数的yaml文件

    1.     # 保存训练过程中的参数文件，方便中断训练，yaml格式
    2.     with open(save_dir / 'hyp.yaml', 'w') as f:
    3.         yaml.dump(hyp, f, sort_keys=False)
    4.     with open(save_dir / 'opt.yaml', 'w') as f:
    5.         yaml.dump(vars(opt), f, sort_keys=False)

     相关配置

    相关配置包括了是否开启绘图功能、yaml文件的读取，获取训练集、验证集，类的数量。

    1.     # 相关配置
    2.     plots = not opt.evolve  # create plots
    3.     cuda = device.type != 'cpu'
    4.     init_seeds(2 + rank)  # 初始化随机种子
    5.     with open(opt.data) as f:  # 读取数据集的yaml文件
    6.         data_dict = yaml.load(f, Loader=yaml.FullLoader)  # data dict
    7.     with torch_distributed_zero_first(rank):  # 分布式的初始化
    8.         check_dataset(data_dict)  # check  检查数据集
    9.     train_path = data_dict['train']  # 读取训练集
    10.     test_path = data_dict['val']  # 读取验证集
    11.     # nc:类的数量，names:类名
    12.     nc, names = (1, ['item']) if opt.single_cls else (int(data_dict['nc']), data_dict['names'])  # number classes, names
    13.     assert len(names) == nc, '%g names found for nc=%g dataset in %s' % (len(names), nc, opt.data)  # check

    Model的定义

     包含了预权重的加载以及model的定义，其中opt.pt的含义是是否采用剪枝微调训练。

    1.     # Model的相关定义
    2.     pretrained = weights.endswith('.pt')  # 读取预训练权重
    3.     if pretrained:
    4.         if opt.pt:
    5.             ckpt = torch.load(weights)
    6.             model = ckpt['model']
    7.             model.to(device)
    8.         else:
    9.             with torch_distributed_zero_first(rank):
    10.                 attempt_download(weights)  # download if not found locally
    11.             ckpt = torch.load(weights, map_location=device)  # load checkpoint
    12.             model = Darknet(opt.cfg).to(device)  # 创建网络
    13.             state_dict = {k: v for k, v in ckpt['model'].items() if model.state_dict()[k].numel() == v.numel()}
    14.             model.load_state_dict(state_dict, strict=False)  # 加载预权重
    15.             print('Transferred %g/%g items from %s' % (len(state_dict), len(model.state_dict()), weights))  # report
    16.     else:
    17.         model = Darknet(opt.cfg).to(device) # create

     optimizer相关配置代码

    这里包括了常用的优化器以及超参数配置代码，代码默认采用的SGD【Adam发现有时候在剪枝训练的有问题，还没搞清楚什么问题】。

    1.     # Optimizer相关配置
    2.     nbs = 64  # nominal batch size
    3.     accumulate = max(round(nbs / total_batch_size), 1)  # accumulate loss before optimizing
    4.     hyp['weight_decay'] *= total_batch_size * accumulate / nbs  # scale weight_decay
    5.  
    6.     pg0, pg1, pg2 = [], [], []  # optimizer parameter groups
    7.     for k, v in dict(model.named_parameters()).items():
    8.         if '.bias' in k:
    9.             pg2.append(v)  # biases
    10.         elif 'Conv2d.weight' in k:
    11.             pg1.append(v)  # apply weight_decay
    12.         elif 'm.weight' in k:
    13.             pg1.append(v)  # apply weight_decay
    14.         elif 'w.weight' in k:
    15.             pg1.append(v)  # apply weight_decay
    16.         else:
    17.             pg0.append(v)  # all else
    18.  
    19.     if opt.adam: # 采用adam优化器
    20.         optimizer = optim.Adam(pg0, lr=hyp['lr0'], betas=(hyp['momentum'], 0.999))  # adjust beta1 to momentum
    21.     else:  # SGD优化器
    22.         optimizer = optim.SGD(pg0, lr=hyp['lr0'], momentum=hyp['momentum'], nesterov=True)
    23.  
    24.     optimizer.add_param_group({'params': pg1, 'weight_decay': hyp['weight_decay']})  # add pg1 with weight_decay
    25.     optimizer.add_param_group({'params': pg2})  # add pg2 (biases)
    26.     logger.info('Optimizer groups: %g .bias, %g conv.weight, %g other' % (len(pg2), len(pg1), len(pg0)))
    27.     del pg0, pg1, pg2
    28.     # Scheduler https://arxiv.org/pdf/1812.01187.pdf
    29.     # https://pytorch.org/docs/stable/_modules/torch/optim/lr_scheduler.html#OneCycleLR
    30.     lf = lambda x: ((1 + math.cos(x * math.pi / epochs)) / 2) * (1 - hyp['lrf']) + hyp['lrf']  # cosine
    31.     scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)
    32.     # plot_lr_scheduler(optimizer, scheduler, epochs)

    optimizer 预权重加载：

    1.     start_epoch, best_fitness = 0, 0.0
    2.     # 精确率p，召回率R，AP50，AP，F值，存储预权重模型参数
    3.     best_fitness_p, best_fitness_r, best_fitness_ap50, best_fitness_ap, best_fitness_f = 0.0, 0.0, 0.0, 0.0, 0.0
    4.     if pretrained:  # 如果加载预权重的相关配置
    5.         # Optimizer  加载预权重中的参数继续训练
    6.         # if opt.pt:
    7.         #     ckpt['optimizer'] = None
    8.         if ckpt['optimizer'] is not None:
    9.             optimizer_dict = optimizer.state_dict()
    10.             pred_optimizer = ckpt['optimizer']  # 获得优化器参数
    11.             pred_optimizer = {k: v for k, v in pred_optimizer.items() if np.shape(optimizer_dict[k]) == np.shape(pred_optimizer[k])}
    12.             #optimizer.load_state_dict(ckpt['optimizer'])
    13.             optimizer.load_state_dict(optimizer_dict)
    14.             best_fitness = ckpt['best_fitness']
    15.             best_fitness_p = ckpt['best_fitness_p']
    16.             best_fitness_r = ckpt['best_fitness_r']
    17.             best_fitness_ap50 = ckpt['best_fitness_ap50']
    18.             best_fitness_ap = ckpt['best_fitness_ap']
    19.             best_fitness_f = ckpt['best_fitness_f']
    20.  
    21.         # Results
    22.         if ckpt.get('training_results') is not None:
    23.             with open(results_file, 'w') as file:
    24.                 file.write(ckpt['training_results'])  # write results.txt
    25.  
    26.         # Epochs
    27.         start_epoch = ckpt['epoch'] + 1
    28.         if opt.resume:  # 继续中断后的训练
    29.             assert start_epoch > 0, '%s training to %g epochs is finished, nothing to resume.' % (weights, epochs)
    30.         if epochs < start_epoch:
    31.             logger.info('%s has been trained for %g epochs. Fine-tuning for %g additional epochs.' %
    32.                         (weights, ckpt['epoch'], epochs))
    33.             epochs += ckpt['epoch']  # finetune additional epochs
    34.         if opt.pt:
    35.             del ckpt
    36.         else:
    37.             del ckpt, state_dict

    训练集的加载 

    1.     # Trainloader 训练集的加载
    2.     dataloader, dataset = create_dataloader(train_path, imgsz, batch_size, gs, opt,
    3.                                             hyp=hyp, augment=True, cache=opt.cache_images, rect=opt.rect,
    4.                                             rank=rank, world_size=opt.world_size, workers=opt.workers)
    5.     """
    6.     dataset中的label形式：【类，box】
    7.     """
    8.     mlc = np.concatenate(dataset.labels, 0)[:, 0].max()  # max label class
    9.     nb = len(dataloader)  # number of batches
    10.     assert mlc < nc, 'Label class %g exceeds nc=%g in %s. Possible class labels are 0-%g' % (mlc, nc, opt.data, nc - 1)

    验证集的加载

    1.     # Process 0
    2.     if rank in [-1, 0]:
    3.         ema.updates = start_epoch * nb // accumulate  # set EMA updates
    4.         # 测试集的加载
    5.         testloader = create_dataloader(test_path, imgsz_test, batch_size*2, gs, opt,
    6.                                        hyp=hyp, cache=opt.cache_images and not opt.notest, rect=True,
    7.                                        rank=-1, world_size=opt.world_size, workers=opt.workers)[0]  # testloader

     绘制labels

    该功能是通过plot_labels函数进行实现的。该功能将会绘制3个子图：

    1.柱状图：横坐标为classes,纵坐标为数据集中各个类的数量；

    2.散点图：将所有boxes的center_x,center_y进行绘制，可以看target的中心点分布情况【这里是进行了归一化的】，横坐标是center_x，纵坐标是center_y；

    3.散点图：绘制数据集所有boxes的w,h，横坐标为width,纵坐标为height；

    1.             if plots:  # 是否开启绘制功能
    2.                 """
    3.                     该函数会绘制3个子图：
    4.                     1.柱状图：横坐标为classes,纵坐标为数据集中各个类的数量
    5.                     2.散点图：将所有boxes的center_x,center_y进行绘制，可以看target的中心点分布情况【这里是进行了归一化的】，横坐标是center_x，纵坐标是center_y
    6.                     3.散点图：绘制数据集所有boxes的w,h，横坐标为width,纵坐标为height
    8.                 """
    9.                 plot_labels(labels, save_dir=save_dir)  # 绘制标签

     效果图如下：

     第一幅图反应了当前训练数据集中总的目标数【我这里只有一个类】。第二个图是我所有图像中中心点坐标的分布情况，横坐标是center_x，纵坐标是center_y。第三幅图是图像的height和width分布。第四图就是一个空白图而已。

     

    同时还会再绘制一个label_correlogram.png的图像，这副图绘制的是label中各个属性的相关性。

    主对角线的柱状图是各自属性的相关性，其他的则是不同属性间的相关分布。有关代码如下：

    1.         import seaborn as sns
    2.         import pandas as pd
    3.         x = pd.DataFrame(b.transpose(), columns=['x', 'y', 'width', 'height'])
    4.         sns.pairplot(x, corner=True, diag_kind='hist', kind='scatter', markers='o',
    5.                      plot_kws=dict(s=3, edgecolor=None, color='b', linewidth=3, alpha=0.5),
    6.                      diag_kws=dict(bins=50))
    7.         plt.savefig(Path(save_dir) / 'labels_correlogram.png', dpi=200)
    8.         plt.close()
    """
此时b的形式为：box1:center_x, center_y, w, h
            box2:center_x, center_y, w, h.....
使用pd.DataFrame,生成数据表：
_________________________________________
|  x        |  y      |  width | height |
|  center_x |center_y |   w    |   h    |......

sns.pairplot主要展现的是变量两两之间的关系（线性或非线性，有无较为明显的相关关系)
    diag_kind：控制对角线上的图的类型，可选"hist"与"kde"
    kind：用于控制非对角线上的图的类型，可选"scatter"与"reg",参数设置为 "reg" 会为非对角线上的散点图拟合出一条回归直线，更直观地显示变量之间的关系。
    markers：控制散点的样式
    plot_kws：用于控制非对角线上的图的样式
    diag_kws：用于控制对角线上图的样式
    绘制出来的图祝对角线是各自属性的直方图，非对角线是不同属性间的相关性
"""

     

     

    训练

    从下面的代码开始就是正式的训练代码。包含了损失函数的计算，前向传播和反向传播，以及各个权重的保存，这里的权重会保存best.pt(权衡了mAP@.5，mAP@0.5~.95)，last.pt，以及精确率P最高的权重，召回率R最高的权重等，可能大家会注意到保存的权重很大，这是因为里面还保存了优化器中的权重所以会保存的大一些。

    1.     # 从下面开始就是正式训练
    2.     for epoch in range(start_epoch, epochs):  # epoch ------------------------------------------------------------------
    3.         model.train()
    4.  
    5.         # Update image weights (optional)
    6.         if opt.image_weights:
    7.             # Generate indices
    8.             if rank in [-1, 0]:
    9.                 cw = model.class_weights.cpu().numpy() * (1 - maps) ** 2  # class weights
    10.                 iw = labels_to_image_weights(dataset.labels, nc=nc, class_weights=cw)  # image weights
    11.                 dataset.indices = random.choices(range(dataset.n), weights=iw, k=dataset.n)  # rand weighted idx
    12.             # Broadcast if DDP
    13.             if rank != -1:
    14.                 indices = (torch.tensor(dataset.indices) if rank == 0 else torch.zeros(dataset.n)).int()
    15.                 dist.broadcast(indices, 0)
    16.                 if rank != 0:
    17.                     dataset.indices = indices.cpu().numpy()
    18.         # 存储mean loss
    19.         mloss = torch.zeros(4, device=device)  # mean losses
    20.         if rank != -1:
    21.             dataloader.sampler.set_epoch(epoch)
    22.         pbar = enumerate(dataloader)
    23.         logger.info(('\n' + '%10s' * 8) % ('Epoch', 'gpu_mem', 'box', 'obj', 'cls', 'total', 'targets', 'img_size'))
    24.         if rank in [-1, 0]:
    25.             pbar = tqdm(pbar, total=nb)  # progress bar
    26.         optimizer.zero_grad()
    27.         for i, (imgs, targets, paths, _) in pbar:  # batch -------------------------------------------------------------
    28.             ni = i + nb * epoch  # number integrated batches (since train start)
    29.             imgs = imgs.to(device, non_blocking=True).float() / 255.0  # uint8 to float32, 0-255 to 0.0-1.0
    30.  
    31.             # Warmup
    32.             if ni <= nw:
    33.                 xi = [0, nw]  # x interp
    34.                 # model.gr = np.interp(ni, xi, [0.0, 1.0])  # iou loss ratio (obj_loss = 1.0 or iou)
    35.                 accumulate = max(1, np.interp(ni, xi, [1, nbs / total_batch_size]).round())
    36.                 for j, x in enumerate(optimizer.param_groups):
    37.                     # bias lr falls from 0.1 to lr0, all other lrs rise from 0.0 to lr0
    38.                     x['lr'] = np.interp(ni, xi, [hyp['warmup_bias_lr'] if j == 2 else 0.0, x['initial_lr'] * lf(epoch)])
    39.                     if 'momentum' in x:
    40.                         x['momentum'] = np.interp(ni, xi, [hyp['warmup_momentum'], hyp['momentum']])
    41.  
    42.             # Multi-scale，多尺度缩放训练
    43.             if opt.multi_scale:
    44.                 sz = random.randrange(imgsz * 0.5, imgsz * 1.5 + gs) // gs * gs  # size
    45.                 sf = sz / max(imgs.shape[2:])  # scale factor
    46.                 if sf != 1:
    47.                     ns = [math.ceil(x * sf / gs) * gs for x in imgs.shape[2:]]  # new shape (stretched to gs-multiple)
    48.                     imgs = F.interpolate(imgs, size=ns, mode='bilinear', align_corners=False)
    49.  
    50.             # Forward 前向传播
    51.             with amp.autocast(enabled=cuda):
    52.                 pred = model(imgs)  # forward
    53.                 # 求loss
    54.                 loss, loss_items = compute_loss(pred, targets.to(device), model)  # loss scaled by batch_size
    55.                 if rank != -1:
    56.                     loss *= opt.world_size  # gradient averaged between devices in DDP mode
    57.  
    58.             # Backward 反向传播
    59.             scaler.scale(loss).backward()
    60.  
    61.             # Optimize 优化器的更新
    62.             if ni % accumulate == 0:
    63.                 scaler.step(optimizer)  # optimizer.step
    64.                 scaler.update()
    65.                 optimizer.zero_grad()
    66.                 if ema:
    67.                     ema.update(model)
    68.  
    69.             # Print功能
    70.             """
    71.             打印：epoch,总epochs, cuda占用情况mem,平均loss，当前batch中target数量，图像shape
    72.             Plot功能：保存3张训练过程数据集可视化图，就是在数据集上绘制box和类
    73.             """
    74.             if rank in [-1, 0]:
    75.                 mloss = (mloss * i + loss_items) / (i + 1)  # 更新 mean losses
    76.                 # 打印内存
    77.                 mem = '%.3gG' % (torch.cuda.memory_reserved() / 1E9 if torch.cuda.is_available() else 0)  # (GB)
    78.                 s = ('%10s' * 2 + '%10.4g' * 6) % (
    79.                     '%g/%g' % (epoch, epochs - 1), mem, *mloss, targets.shape[0], imgs.shape[-1])
    80.                 pbar.set_description(s)
    81.  
    82.                 # Plot
    83.                 if plots and ni < 3:
    84.                     f = save_dir / f'train_batch{ni}.jpg'  # filename
    85.                     plot_images(images=imgs, targets=targets, paths=paths, fname=f)
    86.                     # if tb_writer:
    87.                     #     tb_writer.add_image(f, result, dataformats='HWC', global_step=epoch)
    88.                     #     tb_writer.add_graph(model, imgs)  # add model to tensorboard
    89.                 elif plots and ni == 3 and wandb:
    90.                     wandb.log({"Mosaics": [wandb.Image(str(x), caption=x.name) for x in save_dir.glob('train*.jpg')]})
    91.  
    92.             # end batch ------------------------------------------------------------------------------------------------
    93.         # end epoch ----------------------------------------------------------------------------------------------------
    94.  
    95.         # Scheduler
    96.         lr = [x['lr'] for x in optimizer.param_groups]  # for tensorboard
    97.         scheduler.step()
    98.  
    99.         # DDP process 0 or single-GPU
    100.         if rank in [-1, 0]:
    101.             # mAP
    102.             if ema:
    103.                 ema.update_attr(model)
    104.             final_epoch = epoch + 1 == epochs
    105.             if not opt.notest or final_epoch:  # Calculate mAP
    106.                 if epoch >= opt.period:  # 大于3轮的时候才测试
    107.                     results, maps, times = test.test(opt.data,
    108.                                                  batch_size=batch_size*2,
    109.                                                  imgsz=imgsz_test,
    110.                                                  model=ema.ema.module if hasattr(ema.ema, 'module') else ema.ema,
    111.                                                  single_cls=opt.single_cls,
    112.                                                  dataloader=testloader,
    113.                                                  save_dir=save_dir,
    114.                                                  plots=plots and final_epoch,
    115.                                                  log_imgs=opt.log_imgs if wandb else 0)
    116.  
    117.             # Write
    118.             with open(results_file, 'a') as f:
    119.                 f.write(s + '%10.4g' * 7 % results + '\n')  # P, R, mAP@.5, mAP@.5-.95, val_loss(box, obj, cls)
    120.             if len(opt.name) and opt.bucket:
    121.                 os.system('gsutil cp %s gs://%s/results/results%s.txt' % (results_file, opt.bucket, opt.name))
    122.  
    123.             # Log
    124.             tags = ['train/box_loss', 'train/obj_loss', 'train/cls_loss',  # train loss
    125.                     'metrics/precision', 'metrics/recall', 'metrics/mAP_0.5', 'metrics/mAP_0.5:0.95',
    126.                     'val/box_loss', 'val/obj_loss', 'val/cls_loss',  # val loss
    127.                     'x/lr0', 'x/lr1', 'x/lr2']  # params
    128.             for x, tag in zip(list(mloss[:-1]) + list(results) + lr, tags):
    129.                 if tb_writer:
    130.                     tb_writer.add_scalar(tag, x, epoch)  # tensorboard
    131.                 if wandb:
    132.                     wandb.log({tag: x})  # W&B
    133.  
    134.             # Update best mAP
    135.             """
    136.             fitness:给mAP@0.5和mAP@0.5~.95分别分配0.1，0.9的权重后求和
    137.             fitness_p：给p分配100%的权重，其他权重为0
    138.             fitness_r：给R分配100%权重，其他为0
    139.             fitness_ap50：仅对mAP@0.5分配100%权重
    140.             fitness_ap：仅对mAP@0.5~.95分配100%权重
    141.             """
    142.             fi = fitness(np.array(results).reshape(1, -1))  # weighted combination of [P, R, mAP@.5, mAP@.5-.95]
    143.             fi_p = fitness_p(np.array(results).reshape(1, -1))  # weighted combination of [P, R, mAP@.5, mAP@.5-.95]
    144.             fi_r = fitness_r(np.array(results).reshape(1, -1))  # weighted combination of [P, R, mAP@.5, mAP@.5-.95]
    145.             fi_ap50 = fitness_ap50(np.array(results).reshape(1, -1))  # weighted combination of [P, R, mAP@.5, mAP@.5-.95]
    146.             fi_ap = fitness_ap(np.array(results).reshape(1, -1))  # weighted combination of [P, R, mAP@.5, mAP@.5-.95]
    147.             if (fi_p > 0.0) or (fi_r > 0.0):  # 当P 或 R大于0的时候，计算F值
    148.                 fi_f = fitness_f(np.array(results).reshape(1, -1))  # weighted combination of [P, R, mAP@.5, mAP@.5-.95]
    149.             else:
    150.                 fi_f = 0.0
    151.             if fi > best_fitness:  # 记录最好的best_fitness
    152.                 best_fitness = fi
    153.             if fi_p > best_fitness_p:  # 记录最好的P
    154.                 best_fitness_p = fi_p
    155.             if fi_r > best_fitness_r:  # 记录最好的R
    156.                 best_fitness_r = fi_r
    157.             if fi_ap50 > best_fitness_ap50:  # 记录最好的mAP@0.5
    158.                 best_fitness_ap50 = fi_ap50
    159.             if fi_ap > best_fitness_ap: # 记录最好的mAP@0.5~.95
    160.                 best_fitness_ap = fi_ap
    161.             if fi_f > best_fitness_f:  # 记录最好的F值
    162.                 best_fitness_f = fi_f
    163.  
    164.             # Save model
    165.             """
    166.             model的保存会将epoch，上述的各个best以及训练结果，模型权重(不含结构图)，优化器参数进行保存，因此推理中需要调用['model']keys
    167.             """
    168.             save = (not opt.nosave) or (final_epoch and not opt.evolve)
    169.             if save:
    170.                 with open(results_file, 'r') as f:  # create checkpoint
    171.                     if opt.pt:
    172.                         ckpt = {'epoch': epoch,
    173.                                 'best_fitness': best_fitness,
    174.                                 'best_fitness_p': best_fitness_p,
    175.                                 'best_fitness_r': best_fitness_r,
    176.                                 'best_fitness_ap50': best_fitness_ap50,
    177.                                 'best_fitness_ap': best_fitness_ap,
    178.                                 'best_fitness_f': best_fitness_f,
    179.                                 'training_results': f.read(),
    180.                                 'model': ema.ema.module if hasattr(ema,'module') else ema.ema,
    181.                                 'optimizer': None if final_epoch else optimizer.state_dict(),
    182.                                 'wandb_id': wandb_run.id if wandb else None}
    183.                     else:
    184.                         ckpt = {'epoch': epoch,
    185.                                 'best_fitness': best_fitness,
    186.                                 'best_fitness_p': best_fitness_p,
    187.                                 'best_fitness_r': best_fitness_r,
    188.                                 'best_fitness_ap50': best_fitness_ap50,
    189.                                 'best_fitness_ap': best_fitness_ap,
    190.                                 'best_fitness_f': best_fitness_f,
    191.                                 'training_results': f.read(),
    192.                                 'model': ema.ema.module.state_dict() if hasattr(ema, 'module') else ema.ema.state_dict(),
    193.                                 'optimizer': None if final_epoch else optimizer.state_dict(),
    194.                                 'wandb_id': wandb_run.id if wandb else None}
    195.  
    196.                 # Save last, best and delete
    197.                 torch.save(ckpt, last)   # 保存last.pt
    198.                 if best_fitness == fi:  # 保存best.pt (这个权重保存的是map0.5和map0.5~.95综合加权后的)
    199.                     torch.save(ckpt, best)
    200.                 if (best_fitness == fi) and (epoch >= 200): # 大于200epoch后每个epoch保存一次
    201.                     torch.save(ckpt, wdir / 'best_{:03d}.pt'.format(epoch))
    202.                 if best_fitness == fi:
    203.                     torch.save(ckpt, wdir / 'best_overall.pt')  # 保存best_overall.pt
    204.                 if best_fitness_p == fi_p:  # 保存p值最好的权重
    205.                     torch.save(ckpt, wdir / 'best_p.pt')
    206.                 if best_fitness_r == fi_r:  # 保存R最好的权重
    207.                     torch.save(ckpt, wdir / 'best_r.pt')
    208.                 if best_fitness_ap50 == fi_ap50:  # 保存mAP0.5最好的权重
    209.                     torch.save(ckpt, wdir / 'best_ap50.pt')
    210.                 if best_fitness_ap == fi_ap:  # 保存mAP@.5~.95最好的权重
    211.                     torch.save(ckpt, wdir / 'best_ap.pt')
    212.                 if best_fitness_f == fi_f:  # 保存F值最好的权重
    213.                     torch.save(ckpt, wdir / 'best_f.pt')
    214.                 if epoch == 0:  # 保存第一轮的权重
    215.                     torch.save(ckpt, wdir / 'epoch_{:03d}.pt'.format(epoch))
    216.                 if ((epoch+1) % 25) == 0:  # 每25个epoch保存一轮
    217.                     torch.save(ckpt, wdir / 'epoch_{:03d}.pt'.format(epoch))
    218.                 if epoch >= (epochs-5):
    219.                     torch.save(ckpt, wdir / 'last_{:03d}.pt'.format(epoch))
    220.                 elif epoch >= 420: 
    221.                     torch.save(ckpt, wdir / 'last_{:03d}.pt'.format(epoch))
    222.                 del ckpt
    223.         # end epoch ----------------------------------------------------------------------------------------------------
    224.     # end training
    225.  
    226.     if rank in [-1, 0]:
    227.         # Strip optimizers
    228.         n = opt.name if opt.name.isnumeric() else ''
    229.         fresults, flast, fbest = save_dir / f'results{n}.txt', wdir / f'last{n}.pt', wdir / f'best{n}.pt'
    230.         for f1, f2 in zip([wdir / 'last.pt', wdir / 'best.pt', results_file], [flast, fbest, fresults]):
    231.             if f1.exists():
    232.                 os.rename(f1, f2)  # rename
    233.                 if str(f2).endswith('.pt'):  # is *.pt
    234.                     strip_optimizer(f2)  # strip optimizer
    235.                     os.system('gsutil cp %s gs://%s/weights' % (f2, opt.bucket)) if opt.bucket else None  # upload
    236.         # Finish
    237.         if plots:
    238.             plot_results(save_dir=save_dir)  # save as results.png

     

    训练自己的数据集

    在工程下面中的dataset文件下放入自己的数据集。目录形式如下：

    dataset
    |-- Annotations
    |-- ImageSets
    |-- images
    |-- labels

    Annotations是存放xml标签文件的，images是存放图像的，ImageSets存放四个txt文件【后面运行代码的时候会自动生成】，labels是将xml转txt文件。

    1.运行makeTXT.py。这将会在ImageSets文件夹下生成  trainval.txt，test.txt，train.txt，val.txt四个文件【如果你打开这些txt文件，里面仅有图像的名字】。

    2.打开voc_label.py，并修改代码 classes=[""]填入自己的类名，比如你的是训练猫和狗，那么就是classes=["dog","cat"]，然后运行该程序。此时会在labels文件下生成对应每个图像的txt文件，形式如下：【最前面的0是类对应的索引，我这里只有一个类，后面的四个数为box的参数，均归一化以后的，分别表示box的左上和右下坐标，等训练的时候会处理成center_x，center_y，w, h】

    0 0.4723557692307693 0.5408653846153847 0.34375 0.8990384615384616
0 0.8834134615384616 0.5793269230769231 0.21875 0.8221153846153847 

     3.在data文件夹下新建一个mydata.yaml文件。内容如下【你也可以把coco.yaml复制过来】。

    你只需要修改nc以及names即可，nc是类的数量，names是类的名字。

    train: ./dataset/train.txt
val: ./dataset/val.txt
test: ./dataset/test.txt

# number of classes
nc: 1

# class names
names: ['target']

    4.以yolor_csp为例。打开cfg下的yolor_csp.cfg，搜索两个内容，搜索classes【有三个地方】，将classes修改为自己的类别数量 。再继续搜索255【6个地方】，这个255指的是coco数据集，为3 * (5 + 80)，如果是你自己的类，你需要自己计算一下,3*(5+你自己类的数量)。比如我这里是1个类，就是改成18.

    5.在data/下新建一个myclasses.names文件，写入自己的类【这个是为了后面检测的时候读取类名】

    6.终端输入参数，开始训练。

    python train.py --weights yolor_csp.pt --cfg cfg/yolor_csp.cfg --data data/mydata.yaml --batch-size 8 --device 0

    训练的权重会存储在当前工程下runs/train/exp的weights中。每次运行train.py的时候会在runs/train/下生成exp,exp1,exp2...为的防止权重的覆盖。 

    检测推理

    终端输入参数，开始检测。

    python detect.py --source 【你的图像路径】 --cfg cfg/yolor_csp.cfg --weights 【你训练好的权重路径】 --conf 0.2 --img-size 640 --device 0

    剪枝

    在利用剪枝功能前，需要安装一下剪枝的库。需要安装0.2.7版本，0.2.8有些人说有问题。

    pip install torch_pruning==0.2.7

    1.保存完整的权重文件

    剪枝之前先要保存一下网络的权重和网络结构【非剪枝训练的权重仅含有权值，也就是通过torch.save(model.state_dict())形式保存的】。

    修改tools/save_whole_model.py中的--weights，改为自己的训练后的权权重路径，修改--save_whole_model 为True，运行代码后会生成一个whole_model.pt，如果还想得到onnx模型，可以将--onnx设置为True。

    2.网络剪枝

    剪枝之前需要自己熟悉网络结构，也可以通过tools/printmodel.py 打印网络结构。

    这里的剪枝操作支持两种类型：

    1.单独卷积的剪枝

    在Conv_pruning这个函数中，修改三个地方：

    通过keys指筛选层：

    if k == 'module_list.22.Conv2d.weight':  # 筛选出该层 (根据自己需求)

    amount是剪枝率，可以按续修改。 

    pruning_idxs = strategy(v, amount=0.4)  # or manually selected pruning_idxs=[2, 6, 9, ...]

    修改DG.get_pruning_plan model参数，改为需要剪枝的层 

    # 放入要剪枝的层
pruning_plan = DG.get_pruning_plan(model.module_list[22].Conv2d, tp.prune_conv, idxs=pruning_idxs)

    我这里是仅对第22个卷积进行剪枝。看到如下参数有变化就是剪枝成功了，如果参数没变就说明你剪的不对【可以看到你参数变化不大，因为你仅仅对一个层剪枝了而已，当然变化不大】

    1. -------------
    2. [  prune_conv on module_list.22.Conv2d (Conv2d(128, 128, kernel_size=(1, 1), stride=(1, 1), bias=False))>, Index=[1, 5, 8, 10, 11, 12, 13, 14, 15, 18, 20, 21, 22, 23, 25, 28, 32, 38, 40, 44, 52, 54, 55, 57, 58, 59, 61, 63, 64, 66, 72, 73, 77, 82, 84, 86, 88, 96, 97, 98, 99, 101, 102, 103, 109, 113, 114, 120, 123, 126, 127], NumPruned=6528]
    3. [  prune_batchnorm on module_list.22.BatchNorm2d (BatchNorm2d(128, eps=0.0001, momentum=0.03, affine=True, track_running_stats=True))>, Index=[1, 5, 8, 10, 11, 12, 13, 14, 15, 18, 20, 21, 22, 23, 25, 28, 32, 38, 40, 44, 52, 54, 55, 57, 58, 59, 61, 63, 64, 66, 72, 73, 77, 82, 84, 86, 88, 96, 97, 98, 99, 101, 102, 103, 109, 113, 114, 120, 123, 126, 127], NumPruned=102]
    4. [  _prune_elementwise_op on _ElementWiseOp()>, Index=[1, 5, 8, 10, 11, 12, 13, 14, 15, 18, 20, 21, 22, 23, 25, 28, 32, 38, 40, 44, 52, 54, 55, 57, 58, 59, 61, 63, 64, 66, 72, 73, 77, 82, 84, 86, 88, 96, 97, 98, 99, 101, 102, 103, 109, 113, 114, 120, 123, 126, 127], NumPruned=0]
    5. [  _prune_elementwise_op on _ElementWiseOp()>, Index=[1, 5, 8, 10, 11, 12, 13, 14, 15, 18, 20, 21, 22, 23, 25, 28, 32, 38, 40, 44, 52, 54, 55, 57, 58, 59, 61, 63, 64, 66, 72, 73, 77, 82, 84, 86, 88, 96, 97, 98, 99, 101, 102, 103, 109, 113, 114, 120, 123, 126, 127], NumPruned=0]
    6. [  prune_related_conv on module_list.23.Conv2d (Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False))>, Index=[1, 5, 8, 10, 11, 12, 13, 14, 15, 18, 20, 21, 22, 23, 25, 28, 32, 38, 40, 44, 52, 54, 55, 57, 58, 59, 61, 63, 64, 66, 72, 73, 77, 82, 84, 86, 88, 96, 97, 98, 99, 101, 102, 103, 109, 113, 114, 120, 123, 126, 127], NumPruned=58752]
    7. 65382 parameters will be pruned
    8. -------------
    9.  
    10. 2022-09-19 11:01:55.563 | INFO     | __main__:Conv_pruning:42 -   Params: 52497868 => 52432486
    11.  
    12. 2022-09-19 11:01:56.361 | INFO     | __main__:Conv_pruning:55 - 剪枝完成

    剪枝完在model_data/下会保存一个 Conv_pruning.pt权重。这个就是剪枝后的权重。

    2.卷积层(某个模块)的剪枝

    通过运行layer_pruning()函数。修改两个地方：

    included_layers是需要剪枝的层，比如我这里是对前60层进行剪枝。

    1. included_layers = [layer.Conv2d for layer in model.module_list[:61] if
    2.                       type(layer) is torch.nn.Sequential and layer.Conv2d]

    修改amount剪枝率。 

    pruning_plan = DG.get_pruning_plan(m, tp.prune_conv, idxs=strategy(m.weight, amount=0.9))

    看到如下参数的变化说明剪枝成功了。 将会在model_data/下生成一个layer_pruning.pt。

    2022-09-19 11:12:40.519 | INFO     | __main__:layer_pruning:81 - 
    -------------
    [ prune_conv on module_list.60.Conv2d (Conv2d(26, 1, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False))>, Index=[], NumPruned=0]
    0 parameters will be pruned
    -------------

    2022-09-19 11:12:40.522 | INFO     | __main__:layer_pruning:87 -   Params: 52497868 => 43633847

    2022-09-19 11:12:41.709 | INFO     | __main__:layer_pruning:102 - 剪枝完成
     

     

     

    3.剪枝后的微调训练 

     与上面的训练一样，只不过weights需要改为自己的剪枝后的权重路径，同时再加一个--pt参数，如下：

    参数--pt：指剪枝后的训练

    这里默认的epochs还是300，自己可以修改。

    python train.py --weights model_data/layer_pruning.pt --cfg cfg/yolor_csp.cfg --data data/mydata.yaml --batch-size 8 --device 0 --pt

    4.剪枝后的推理检测

     --weights 为剪枝后的权重，在加一个--pd表示剪枝后的检测。

    python detect.py --source 【你的图像路径】 --cfg cfg/yolor_csp.cfg --weights 【剪枝的权重路径】 --conf 0.2 --img-size 640 --device 0 --pd

    github代码：

     https://github.com/YINYIPENG-EN/Pruning_for_YOLOR_pytorch

    权重百度云：

    链接：https://pan.baidu.com/s/1uQflOXCQtffkD5kHAAF94A 
    提取码：yypn

    所遇问题：

    1.剪枝训练期间在测mAP的时候报错：

    RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

    解决办法：在models/modes.py中的第416行和417行，强行修改 :

    io[..., :2] = (io[..., :2] * 2. - 0.5 + self.grid)
io[..., 2:4] = (io[..., 2:4] * 2) ** 2 * self.anchor_wh

    改为：

    io[..., :2] = (io[..., :2] * 2. - 0.5 + self.grid.cuda())
io[..., 2:4] = (io[..., 2:4] * 2) ** 2 * self.anchor_wh.cuda()

    在训练完需要推理测试的时候需要把上面的再改回去(不是不能检测，只是会特别的不准)

     

  • 相关阅读:
    计算机毕业设计Java珠宝首饰进销存管理系统（源码+系统+mysql数据库+Lw文档）
    基于Adaboost的数据分类算法matlab仿真
    玩转Vue3全家桶02丨上手：一个清单应用帮你入门Vue
    SQL多个字段拼接组合成新字段的常用方法
    智能家居安防监控系统设计： 使用MATLAB进行智能家居安防监控系统建模和仿真，包括视频监控、入侵检测和报警响应等
    浅谈放大器交调失真对系统的影响
    102.二叉树的层序遍历
    压力测试工具Jmeter的下载与使用
    建表时如何合理选择字段类型
    微服务项目：尚融宝（46）（核心业务流程：借款申请（3））
  • 原文地址：https://blog.csdn.net/z240626191s/article/details/126902178