diff --git a/models/onnx_export.py b/models/onnx_export.py index 9906693..a78aab9 100644 --- a/models/onnx_export.py +++ b/models/onnx_export.py @@ -1,3 +1,7 @@ +# Exports a pytorch *.pt model to *.onnx format. Example usage: +# $ export PYTHONPATH="$PWD" +# $ python models/onnx_export.py --weights ./weights/yolov5s.pt --img 640 --batch 1 + import argparse import onnx @@ -6,7 +10,7 @@ from models.common import * if __name__ == '__main__': parser = argparse.ArgumentParser() - parser.add_argument('--weights', default='../weights/yolov5s.pt', help='model path RELATIVE to ./models/') + parser.add_argument('--weights', default='./weights/yolov5s.pt', help='model path RELATIVE to ./models/') parser.add_argument('--img-size', default=640, help='inference size (pixels)') parser.add_argument('--batch-size', default=1, help='batch size') opt = parser.parse_args()