【目标检测】YOLOv5模型从大变小，发生了什么？
前言
在某次使用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 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函数
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转换成FP32
- fuse()：将conv和bn层合并，提速模型推理速度
- eval()：eval()是模型进行预测推理时关闭BN（预测数据均值方差计算）和Dropout,从而让结果稳定
训练过程中，BN会不断计算均值和方差，Dropout比例会使一部分的网络连接不进行计算
预测过程中，需要让均值和方差稳定不变化，同时会使所有网络连接参与计算