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

前言

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

问题根源

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

if rank in[-1,0]:# Plotsif 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 COCOfor 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 modelfor f in last, best:if f.exists():
             strip_optimizer(f)# strip optimizersif opt.bucket:
         os.system(f'gsutil cp {final} gs://{opt.bucket}/weights')# uploadelse:
     dist.destroy_process_group()
 torch.cuda.empty_cache()return results

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

strip_optimizer(f)# strip optimizers

这个方法定义在

/utils/general.py

文件中：

defstrip_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 emafor k in'optimizer','training_results','wandb_id','ema','updates':# keys
        x[k]=None
    x['epoch']=-1
    x['model'].half()# to FP16for p in x['model'].parameters():
        p.requires_grad =False
    torch.save(x, s or f)
    mb = os.path.getsize(s or f)/1E6# filesizeprint(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 torch

if __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 torch

if __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

函数

defattempt_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 ifisinstance(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():iftype(m)in[nn.Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6, nn.SiLU]:
            m.inplace =True# pytorch 1.7.0 compatibilityeliftype(m)is Conv:
            m._non_persistent_buffers_set =set()# pytorch 1.6.0 compatibility# 如果一个模型就直接返回iflen(model)==1:return model[-1]# return modelelse: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转换成FP32
• fuse()：将conv和bn层合并，提速模型推理速度
• eval()：eval()是模型进行预测推理时关闭BN（预测数据均值方差计算）和Dropout,从而让结果稳定 训练过程中，BN会不断计算均值和方差，Dropout比例会使一部分的网络连接不进行计算 预测过程中，需要让均值和方差稳定不变化，同时会使所有网络连接参与计算