17 lines
531 B
Python
17 lines
531 B
Python
from ultralytics import YOLO
|
|
import torch
|
|
# 加载模型时禁用数据验证
|
|
model = YOLO("/home/admin-root/haotian/康达瑞贝斯机器狗/YoloV8Obj/dataset_20250819/train2/weights/best.pt", task="detect")
|
|
# 手动设置模型为推理模式
|
|
model.model.eval()
|
|
# 导出 ONNX
|
|
dummy_input = torch.randn(1, 3, 640, 640)
|
|
torch.onnx.export(
|
|
model.model, # 使用 model.model 访问底层 PyTorch 模型
|
|
dummy_input,
|
|
"yolov8_20250820.onnx",
|
|
input_names=["input"],
|
|
output_names=["output"],
|
|
opset_version=11,
|
|
)
|