19 lines
496 B
Python
19 lines
496 B
Python
from ultralytics import YOLO
|
|
import torch
|
|
|
|
# 加载模型时禁用数据验证
|
|
model = YOLO("/home/admin-root/haotian/rk3588/pytorch模型转rknn/models/yolov8m.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,
|
|
"yolov8m.onnx",
|
|
input_names=["input"],
|
|
output_names=["output"],
|
|
opset_version=11,
|
|
) |