趁热打铁。上节分析了 v5.0 的 yolov5s 模型架构,本节顺便把 v6.0的图也画下。官方代码中贴心的给提供了 onnx 文件,如下图:
但是,当我打开 onnx 的时候,我麻了...
所以,还是需要自己生成下 onnx 文件。
打开 export.py 惊喜发现,代码中提供了onnx-simplifier导出,这种导出的模型更加简洁。https://github.com/daquexian/onnx-simplifierhttps://github.com/daquexian/onnx-simplifier
作者知乎:onnx simplifier 和 optimizer - 知乎
v6.0 版本中依旧将 conv + BN 进行融合,融合的原因如下:
Conv + BN + Act算子融合 - 知乎
为了细节都可视化出来,我们和v5.0 一样将BN分离出来。
v6.0 代码中,export.py 函数中增加了一条:
model = attempt_load(weights, map_location=device, inplace=True, fuse=False) # load FP32 model
通过 fuse=False 参数屏蔽 fuse BN操作。
和前文一样,由于 onnx 的 training = torch.onnx.TrainingMode.EVAL 默认合并了BN,所以修改为:
training=torch.onnx.TrainingMode.TRAINING
那么,导出 onnx:
python export.py --weights yolov5s.pt --simplify
由于 onnx-simplifier 不支持 TrainingMode.TRAINING 模式,所以以上会导出失败。
错误:
ONNX: simplifier failure: ......onnxruntime::BatchNorm
::BatchNorm !is_train_ was false. Training mode is not supported in this build.
所以,为了得到 BN的onnx,采用和 v5.0 一样的办法。
from onnx import shape_inference
# Checks
model_onnx = onnx.load(f) # load onnx model
onnx.save(shape_inference.infer_shapes(model_onnx), f)
# # Checks
# model_onnx = onnx.load(f) # load onnx model
python export.py --weights yolov5s.pt
可得到带BN的onnx。当然,还有一些操作比较复杂需要简化。简化不做了。
和 v5.0 一样,除了最后的3个 conv带bias外,其他的conv的bias=0,这些后面都带BN层。所以还是使用简化版本,把conv和BN融合,然后显示。改回去:
training=torch.onnx.TrainingMode.TRAINING if train else torch.onnx.TrainingMode.EVAL,
# from onnx import shape_inference
# Checks
model_onnx = onnx.load(f) # load onnx model
# onnx.save(shape_inference.infer_shapes(model_onnx), f)
model = attempt_load(weights, map_location=device, inplace=True, fuse=True) # load FP32 model
python export.py --weights yolov5s.pt --simplify