### 将YOLO PyTorch (.pt) 模型转换为 ONNX 格式 为了将YOLO模型从PyTorch的`.pt`格式转换为ONNX格式，可以利用官方提供的工具或自定义脚本完成此操作。对于YOLOv8及其之前的版本，超参数和API可能有所不同，但基本思路一致。 #### 使用Ultralytics库进行转换 如果使用的是由Ultralytics开发的新版YOLO模型，则可以直接调用内置的方法来实现这一目标： ```python from ultralytics import YOLO model = YOLO('best.pt') # 加载预训练好的YOLO模型 success = model.export(format='onnx') # 导出为ONNX格式 print(f"Export successful: {success}") # 打印导出状态 ``` 这段代码会尝试加载名为`best.pt`的YOLO权重文件，并将其保存为对应的ONNX格式[^2]。 #### 自定义转换过程 当不适用上述简便方式时，也可以手动编写转换逻辑。这通常涉及到设置输入张量形状、指定动态轴以及其他必要的选项。下面是一个更通用的例子： ```python import torch import onnx dummy_input = torch.randn(1, 3, 640, 640).cuda() # 创建虚拟输入数据，假设图像大小为640x640像素 model = torch.load('best.pt')['model'].float().fuse().eval().half().to('cuda') # 载入并处理模型 for layer in model.modules(): if isinstance(layer, torch.nn.BatchNorm2d): layer.eps = 1e-3 # 修改BatchNorm层的小常数项以提高精度 output_names = ['output'] # 定义输出节点名称 dynamic_axes = {'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}} # 设置动态维度 with torch.no_grad(): # 关闭梯度计算 torch.onnx.export( model, dummy_input, "yolov5s.onnx", opset_version=11, do_constant_folding=True, input_names=['input'], output_names=output_names, dynamic_axes=dynamic_axes) print("Model has been successfully converted to ONNX format.") ``` 这里的关键在于正确配置`torch.onnx.export()`函数中的各个参数，特别是关于输入/输出的名字以及是否允许某些尺寸随批次变化而改变等问题[^4]。