【目标检测】YOLOv5模型从大变小，发生了什么？(目标检测yolo)

【目标检测】YOLOv5模型从大变小，发生了什么？ 前言

推荐整理分享【目标检测】YOLOv5模型从大变小，发生了什么？(目标检测yolo)，希望有所帮助，仅作参考，欢迎阅读内容。

文章相关热门搜索词:目标检测算法yolo,目标检测如何入门,目标检测 知乎,目标检测tricks,目标检测 知乎,目标检测教程,目标检测yolov5,目标检测yolo,内容如对您有帮助，希望把文章链接给更多的朋友！

在某次使用YOLOv5进行实验时，看到模型已经收敛得差不多，于是想提前停止训练，就果断直接终止程序。然而在查看文件大小时，突然发现，正常训练的yolov5m模型大小为40M左右，而此时生成的yolov5m模型大小达到了160M，于是产生如题疑问：模型从大变小，发生了什么？

问题根源

回到train.py这个文件，发现在模型训练完成之后，还存在这样一段代码：

if rank in [-1, 0]: # Plots if plots: plot_results(save_dir=save_dir) # save as results.png # Test best.pt logger.info('%g epochs completed in %.3f hours.\n' % (epoch - start_epoch + 1, (time.time() - t0) / 3600)) if opt.data.endswith('coco.yaml') and nc == 80: # if COCO for m in (last, best) if best.exists() else (last): # speed, mAP tests results, _, _ = test.test(opt.data, batch_size=batch_size * 2, imgsz=imgsz_test, conf_thres=0.001, iou_thres=0.7, model=attempt_load(m, device).half(), single_cls=opt.single_cls, dataloader=testloader, save_dir=save_dir, save_json=True, plots=False, is_coco=is_coco) # Strip optimizers final = best if best.exists() else last # final model for f in last, best: if f.exists(): strip_optimizer(f) # strip optimizers if opt.bucket: os.system(f'gsutil cp {final} gs://{opt.bucket}/weights') # upload else: dist.destroy_process_group() torch.cuda.empty_cache() return results

和模型大小直接挂钩的是这一句：

strip_optimizer(f) # strip optimizers

这个方法定义在/utils/general.py文件中：

def strip_optimizer(f='best.pt', s=''): # from utils.general import *; strip_optimizer() # Strip optimizer from 'f' to finalize training, optionally save as 's' x = torch.load(f, map_location=torch.device('cpu')) if x.get('ema'): x['model'] = x['ema'] # replace model with ema for k in 'optimizer', 'training_results', 'wandb_id', 'ema', 'updates': # keys x[k] = None x['epoch'] = -1 x['model'].half() # to FP16 for p in x['model'].parameters(): p.requires_grad = False torch.save(x, s or f) mb = os.path.getsize(s or f) / 1E6 # filesize print(f"Optimizer stripped from {f},{(' saved as %s,' % s) if s else ''}{mb:.1f}MB")

阅读代码，不难发现，这一步，程序将模型文件中的'optimizer', 'training_results', 'wandb_id', 'ema', 'updates'这几个设为None，也就是去除这几个值，同时将模型从FP32转成FP16。

因此，早停的模型没有经过这个步骤，导致模型精度是FP32，同时包含了大量优化器信息，导致模型过于庞大。

实验验证

为了验证答案的正确性，重新来加载模型看看。

首先加载官方提供的yolov5m.pt模型

import torchif __name__ == '__main__': ckpt = torch.load('yolov5m.pt') print(ckpt)

输出：

{'epoch': -1,'best_fitness': array([0.45065]), 'training_results': None, 'model': Model(...)'optimizer': None,'wandb_id': None

可以看到，这个模型文件中，只有best_fitness以及model的结构和参数为有效信息，不包含优化器信息。

再加载160M的模型：

import torchif __name__ == '__main__': ckpt = torch.load(r'runs\train\exp\weights\last.pt') print(ckpt)

发现这里输出了大量内容，主要内容是training_results和optimizer，由此可见结论正确。

{'epoch': 0,'best_fitness': 0.0, 'training_results':'....''model': Model(...)'updates': 4'optimizer': {'state':...}...

同时，也可以发现，模型文件实际上是一个字典，例如，可以用下面的方式获取某层结构或参数信息：

print(ckpt.model[0].conv.conv) # 打印某层print(ckpt.model[0].conv.conv.state_dict()) # 打印该层参数信息模型加载解读

阅读代码，发现官方在加载模型时，并没有直接torch.load，而是单独写了一个attempt_load函数

def attempt_load(weights, map_location=None): # Loads an ensemble of models weights=[a,b,c] or a single model weights=[a] or weights=a model = Ensemble() for w in weights if isinstance(weights, list) else [weights]: attempt_download(w) ckpt = torch.load(w, map_location=map_location) # load model.append(ckpt['ema' if ckpt.get('ema') else 'model'].float().fuse().eval()) # FP32 model # 适配pytorch不同版本 for m in model.modules(): if type(m) in [nn.Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6, nn.SiLU]: m.inplace = True # pytorch 1.7.0 compatibility elif type(m) is Conv: m._non_persistent_buffers_set = set() # pytorch 1.6.0 compatibility# 如果一个模型就直接返回 if len(model) == 1: return model[-1] # return model else: print('Ensemble created with %s\n' % weights) for k in ['names', 'stride']: setattr(model, k, getattr(model[-1], k)) return model # return ensemble

注意到模型加载完成之后，还有.float().fuse().eval()这样一个操作。

这三个函数功能如下：

float()：FP16转换成FP32fuse()：将conv和bn层合并，提速模型推理速度eval()：eval()是模型进行预测推理时关闭BN（预测数据均值方差计算）和Dropout,从而让结果稳定 训练过程中，BN会不断计算均值和方差，Dropout比例会使一部分的网络连接不进行计算 预测过程中，需要让均值和方差稳定不变化，同时会使所有网络连接参与计算