yolov5导出onnx模型 fp16
时间: 2025-04-08 22:25:23 浏览: 36
### YOLOv5 模型导出为 FP16 精度 ONNX 格式的教程
要将训练好的 YOLOv5 模型导出为 FP16 精度的 ONNX 格式,可以按照以下方法操作:
#### 安装必要的依赖库
为了支持 FP16 的量化以及 ONNX 转换功能,需要先安装相关的 Python 库。可以通过运行以下命令完成环境准备:
```bash
pip install onnx onnxconverter-common
```
此步骤确保环境中具备了 ONNX 和其扩展工具的支持[^2]。
#### 修改 `export.py` 文件
YOLOv5 提供了一个脚本文件用于模型导出,即 `export.py`。如果希望将其转换为 FP16 精度,则需调整该脚本中的参数配置部分。具体来说,在调用 PyTorch 的 `.to(torch.float16)` 或者通过 ONNX 自带的功能启用半精度模式时,应设置对应的选项。
以下是修改后的关键代码片段示例:
```python
import torch.onnx as onnx_utils
from models.common import Detect, AutoShape
def main():
model = attempt_load(weights, map_location=device) # 加载权重到指定设备上
model.eval() # 切换至评估模式
# 启用FP16支持 (仅当硬件兼容时有效)
half_precision = True # 设置标志位开启半精度计算
if half_precision and device.type != 'cpu': # 如果GPU可用则切换至float16
model.half()
dummy_input = torch.randn(1, 3, imgsz, imgsz).to(device).type_as(next(model.parameters()))
output_file = opt.output or Path(opt.weights).with_suffix('.onnx')
dynamic_axes = {'input' : {0:'batch_size', 2:'height', 3:'width'},
'output' : {0:'batch_size'}}
with torch.no_grad(): # 阻断梯度传播提升效率
onnx_utils.export(
model,
dummy_input,
str(output_file),
opset_version=opt.opset,
do_constant_folding=True,
input_names=['images'],
output_names=['outputs'],
dynamic_axes=dynamic_axes)
if __name__ == "__main__":
parser.add_argument('--half', action='store_true') # 添加新的CLI参数控制是否采用FP16
...
args = parser.parse_args()
main(args)
```
上述代码展示了如何在原有基础上增加对 FP16 支持的能力,并允许用户通过命令行参数灵活决定是否应用浮点数压缩技术[^1]^。
#### 执行导出过程
准备好一切之后,就可以执行实际的模型转储工作啦!只需打开终端并输入类似下面这样的指令即可顺利完成整个流程:
```bash
python export.py --weights yolov5s.pt --img-size 640 --include onnx --half
```
这里指定了目标预训练权值路径(`--weights`)、图像尺寸大小(`--img-size`)还有期望获得的结果形式(`--include`)再加上刚才新增加的那个布尔开关变量(`--half`)用来表明我们确实想要得到经过优化处理过的版本^。
完成后会在当前目录下生成一个新的`.onnx`结尾命名规则遵循原版名称加上后缀的形式保存下来的轻量级推理引擎专用格式文档副本^。
---
###
阅读全文
相关推荐








