add py to onnx tool

This commit is contained in:
liangyuxuan
2025-11-27 16:48:33 +08:00
parent 694272080a
commit 664652d144

23
tools/pt2onnx.py Normal file
View File

@@ -0,0 +1,23 @@
import argparse
from ultralytics import YOLO
def main(checkpoint_path):
model = YOLO(checkpoint_path)
model.export(
format="onnx",
imgsz=(640, 480),
dynamic=True,
simplify=True,
half=True,
workspace=0.8,
batch=1,
device="cpu"
)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("checkpoint_path", type=str)
args = parser.parse_args()
main(args.checkpoint_path)