本文中的剪枝采用的是通道剪枝，为离线剪枝的一种，也就是可以直接对已经训练好的模型进行剪枝后再微调训练，不用稀疏训练【也就是不用边训练边剪枝】。
剪枝参考的paper：Pruning Filters for Efficient ConvNets。
本文主要实现：
1.可训练自己的数据集
2.对单独某一个卷积的剪枝，
3.对某些层的剪枝。
PS：最终的效果本文并不保证，需要根据自己剪枝方案进行评估，剪的地方不同以及剪枝率的不同都会有影响，本文只是把功能进行了实现。这里需要说一下的是，我在实际测试的时候，发现YOLOR其实并没多好，而且参数量也挺大的，所以大家选这个算法要慎重。
本文暂时不讲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
parser = argparse.ArgumentParser() parser.add_argument('--weights', type=str, default='yolor_csp.pt', help='initial weights path') parser.add_argument('--cfg', type=str, default='cfg/yolor_csp.cfg', help='model.yaml path') parser.add_argument('--data', type=str, default='data/coco.yaml', help='data.yaml path') parser.add_argument('--hyp', type=str, default='data/hyp.scratch.640.yaml', help='hyperparameters path') parser.add_argument('--epochs', type=int, default=300) parser.add_argument('--batch-size', type=int, default=4, 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') parser.add_argument('--resume', nargs='?', const=True, default=False, help='resume most recent training') parser.add_argument('--nosave', action='store_true', help='only save final checkpoint') parser.add_argument('--notest', action='store_true', help='only test final epoch') parser.add_argument('--noautoanchor', action='store_true', help='disable autoanchor check') parser.add_argument('--evolve', action='store_true', help='evolve hyperparameters') parser.add_argument('--bucket', type=str, default='', help='gsutil bucket') parser.add_argument('--cache-images', action='store_true', help='cache images for faster training') parser.add_argument('--image-weights', action='store_true', help='use weighted image selection for training') parser.add_argument('--device', default='0', 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%%') parser.add_argument('--single-cls', action='store_true', help='train as single-class dataset') parser.add_argument('--adam', action='store_true', help='use torch.optim.Adam() optimizer') parser.add_argument('--sync-bn', action='store_true', help='use SyncBatchNorm, only available in DDP mode') parser.add_argument('--local_rank', type=int, default=-1, help='DDP parameter, do not modify') parser.add_argument('--log-imgs', type=int, default=16, help='number of images for W&B logging, max 100') parser.add_argument('--workers', type=int, default=4, help='maximum number of dataloader workers') parser.add_argument('--project', default='runs/train', help='save to project/name') 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('--pt', action='store_true', default=False, help='pruned model train') parser.add_argument('--period', type=int, default=2, help='test period') opt = parser.parse_args()
训练阶段：
训练阶段的代码会调用train()函数，主要是hyp【超参数】，opt【传入参数】，device【设备】。
# Train logger.info(opt) if not opt.evolve: tb_writer = None # init loggers if opt.global_rank in [-1, 0]: # tensorboard的写入 logger.info(f'Start Tensorboard with "tensorboard --logdir {opt.project}", view at http://localhost:6006/') tb_writer = SummaryWriter(opt.save_dir) # Tensorboard train(hyp, opt, device, tb_writer, wandb) # 训练过程
相关路径代码：
# 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路径(记录训练过程)
保存训练过程参数的yaml文件
# 保存训练过程中的参数文件，方便中断训练，yaml格式 with open(save_dir / 'hyp.yaml', 'w') as f: yaml.dump(hyp, f, sort_keys=False) with open(save_dir / 'opt.yaml', 'w') as f: yaml.dump(vars(opt), f, sort_keys=False)
相关配置
相关配置包括了是否开启绘图功能、yaml文件的读取，获取训练集、验证集，类的数量。
# 相关配置 plots = not opt.evolve # create plots cuda = device.type != 'cpu' init_seeds(2 + rank) # 初始化随机种子 with open(opt.data) as f: # 读取数据集的yaml文件 data_dict = yaml.load(f, Loader=yaml.FullLoader) # data dict with torch_distributed_zero_first(rank): # 分布式的初始化 check_dataset(data_dict) # check 检查数据集 train_path = data_dict['train'] # 读取训练集 test_path = data_dict['val'] # 读取验证集 # nc:类的数量，names:类名 nc, names = (1, ['item']) if opt.single_cls else (int(data_dict['nc']), data_dict['names']) # number classes, names assert len(names) == nc, '%g names found for nc=%g dataset in %s' % (len(names), nc, opt.data) # check
Model的定义
包含了预权重的加载以及model的定义，其中opt.pt的含义是是否采用剪枝微调训练。
# Model的相关定义 pretrained = weights.endswith('.pt') # 读取预训练权重 if pretrained: if opt.pt: ckpt = torch.load(weights) model = ckpt['model'] model.to(device) else: with torch_distributed_zero_first(rank): attempt_download(weights) # download if not found locally ckpt = torch.load(weights, map_location=device) # load checkpoint model = Darknet(opt.cfg).to(device) # 创建网络 state_dict = {k: v for k, v in ckpt['model'].items() if model.state_dict()[k].numel() == v.numel()} model.load_state_dict(state_dict, strict=False) # 加载预权重 print('Transferred %g/%g items from %s' % (len(state_dict), len(model.state_dict()), weights)) # report else: model = Darknet(opt.cfg).to(device) # create
optimizer相关配置代码
这里包括了常用的优化器以及超参数配置代码，代码默认采用的SGD【Adam发现有时候在剪枝训练的有问题，还没搞清楚什么问题】。
# 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 pg0, pg1, pg2 = [], [], [] # optimizer parameter groups for k, v in dict(model.named_parameters()).items(): if '.bias' in k: pg2.append(v) # biases elif 'Conv2d.weight' in k: pg1.append(v) # apply weight_decay elif 'm.weight' in k: pg1.append(v) # apply weight_decay elif 'w.weight' in k: pg1.append(v) # apply weight_decay else: pg0.append(v) # all else if opt.adam: # 采用adam优化器 optimizer = optim.Adam(pg0, lr=hyp['lr0'], betas=(hyp['momentum'], 0.999)) # adjust beta1 to momentum else: # SGD优化器 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 # Scheduler https://arxiv.org/pdf/1812.01187.pdf # https://pytorch.org/docs/stable/_modules/torch/optim/lr_scheduler.html#OneCycleLR lf = lambda x: ((1 + math.cos(x * math.pi / epochs)) / 2) * (1 - hyp['lrf']) + hyp['lrf'] # cosine scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf) # plot_lr_scheduler(optimizer, scheduler, epochs)
optimizer 预权重加载：
start_epoch, best_fitness = 0, 0.0 # 精确率p，召回率R，AP50，AP，F值，存储预权重模型参数 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 if pretrained: # 如果加载预权重的相关配置 # Optimizer 加载预权重中的参数继续训练 # if opt.pt: # ckpt['optimizer'] = None if ckpt['optimizer'] is not None: optimizer_dict = optimizer.state_dict() pred_optimizer = ckpt['optimizer'] # 获得优化器参数 pred_optimizer = {k: v for k, v in pred_optimizer.items() if np.shape(optimizer_dict[k]) == np.shape(pred_optimizer[k])} #optimizer.load_state_dict(ckpt['optimizer']) optimizer.load_state_dict(optimizer_dict) best_fitness = ckpt['best_fitness'] best_fitness_p = ckpt['best_fitness_p'] best_fitness_r = ckpt['best_fitness_r'] best_fitness_ap50 = ckpt['best_fitness_ap50'] best_fitness_ap = ckpt['best_fitness_ap'] best_fitness_f = ckpt['best_fitness_f'] # Results if ckpt.get('training_results') is not None: with open(results_file, 'w') as file: file.write(ckpt['training_results']) # write results.txt # Epochs start_epoch = ckpt['epoch'] + 1 if opt.resume: # 继续中断后的训练 assert start_epoch > 0, '%s training to %g epochs is finished, nothing to resume.' % (weights, epochs) if epochs < start_epoch: logger.info('%s has been trained for %g epochs. Fine-tuning for %g additional epochs.' % (weights, ckpt['epoch'], epochs)) epochs += ckpt['epoch'] # finetune additional epochs if opt.pt: del ckpt else: del ckpt, state_dict
训练集的加载
# 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) """ dataset中的label形式：【类，box】 """ 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)
验证集的加载
# Process 0 if rank in [-1, 0]: ema.updates = start_epoch * nb // accumulate # set EMA updates # 测试集的加载 testloader = create_dataloader(test_path, imgsz_test, batch_size*2, gs, opt, hyp=hyp, cache=opt.cache_images and not opt.notest, rect=True, 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；
if plots: # 是否开启绘制功能 """ 该函数会绘制3个子图： 1.柱状图：横坐标为classes,纵坐标为数据集中各个类的数量 2.散点图：将所有boxes的center_x,center_y进行绘制，可以看target的中心点分布情况【这里是进行了归一化的】，横坐标是center_x，纵坐标是center_y 3.散点图：绘制数据集所有boxes的w,h，横坐标为width,纵坐标为height """ plot_labels(labels, save_dir=save_dir) # 绘制标签
效果图如下：
第一幅图反应了当前训练数据集中总的目标数【我这里只有一个类】。第二个图是我所有图像中中心点坐标的分布情况，横坐标是center_x，纵坐标是center_y。第三幅图是图像的height和width分布。第四图就是一个空白图而已。
同时还会再绘制一个label_correlogram.png的图像，这副图绘制的是label中各个属性的相关性。
主对角线的柱状图是各自属性的相关性，其他的则是不同属性间的相关分布。有关代码如下：
import seaborn as sns import pandas as pd x = pd.DataFrame(b.transpose(), columns=['x', 'y', 'width', 'height']) sns.pairplot(x, corner=True, diag_kind='hist', kind='scatter', markers='o', plot_kws=dict(s=3, edgecolor=None, color='b', linewidth=3, alpha=0.5), diag_kws=dict(bins=50)) plt.savefig(Path(save_dir) / 'labels_correlogram.png', dpi=200) 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：用于控制对角线上图的样式
绘制出来的图祝对角线是各自属性的直方图，非对角线是不同属性间的相关性
"""
