add py to onnx tool
This commit is contained in:
23
tools/pt2onnx.py
Normal file
23
tools/pt2onnx.py
Normal 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)
|
||||
Reference in New Issue
Block a user