소스 검색

onnx_export.py

5.0
Glenn Jocher 4 년 전
부모
커밋
7a00a5eff2
1개의 변경된 파일5개의 추가작업 그리고 1개의 파일을 삭제
  1. +5
    -1
      models/onnx_export.py

+ 5
- 1
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()

Loading…
취소
저장