39 lines
855 B
Python
39 lines
855 B
Python
import torch
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
def export(pt_path):
|
|
pt_path = Path(pt_path)
|
|
|
|
# 输出 ONNX
|
|
onnx_path = pt_path.with_suffix('.onnx')
|
|
print(f"===> 加载 PT: {pt_path}")
|
|
|
|
# 正确加载 YOLOv5 模型
|
|
device = "cpu"
|
|
ckpt = torch.load(pt_path, map_location=device)
|
|
model = ckpt['model'].float() # 包含完整结构
|
|
model.eval()
|
|
|
|
# dummy input
|
|
dummy = torch.zeros((1, 3, 640, 640))
|
|
|
|
print("===> 导出 ONNX...")
|
|
torch.onnx.export(
|
|
model,
|
|
dummy,
|
|
str(onnx_path),
|
|
opset_version=12,
|
|
input_names=['images'],
|
|
output_names=['output'],
|
|
)
|
|
|
|
print(f"ONNX 导出成功: {onnx_path}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
if len(sys.argv) != 2:
|
|
print("用法: python 1_export_pt_to_onnx.py best.pt")
|
|
exit(1)
|
|
export(sys.argv[1])
|