yolov5-pytorch导出模型问题

在官网提供的代码中,很方便的可以把pytorch的模型转为libtorch的模型,但是在转换前要明白自己转换后的模型是仅仅为了推理部署,还是说用转换后的模型作为libtorch的预训练模型继续使用呢?
明白目的就很好解决了。
如果仅仅是为了推理,那么export.py文件中下面这个地方不用改动,
目的:推理部署

model = attempt_load(opt.weights, map_location=device,fuse=True)  # load FP32 model

fuse参数的意思是融合,在attempt_load函数中,conv层和bn层的参数进行了参数融合,融合的目的是为了部署的时候加速推理。
目的:导出模型作为libtorch的预训练模型

model = attempt_load(opt.weights, map_location=device,fuse=False)  # load FP32 model

fuse这个参数设置为false,禁止参数融合,这样的话,conv层的参数和bn层的参数就不会进行融合,各自的参数是各自的,
具体见下面这张图:
yolov5-pytorch导出模型问题_第1张图片如果取消参数融合,图上会出现batch_norm这个参数,并且_convolution只显示一个参数weight,如果没有取消的话,batch_norm这个参数没有,并且_convolution会显示两个参数weigth和bias。

关于参数融合的具体讲解,可以参考这为博主的内容
yolov5 conv和bn层参数融合

你可能感兴趣的