pt 파일을 onnx 파일로 변경하기 위해 위와 같이 작업을 진행한다.
참고로 위와 같이 경로를 역슬래시를 쓸경우 오류가 나는데 역슬래시를 슬래시로 바꾸면 된다.
---------------------------------------------------------------------------------------
from ultralytics import YOLO
model = YOLO("C:/Users/L/runs/detect/train/weights/best.pt")
model.export(format="onnx", opset=12)
---------------------------------------------------------------------------------------
처리가 완료되면 pt 파일이 있던위치와 동일한 경로에 onnx 파일이 생성되게 된다.
pt 파일과 onnx 파일의 차이점은 위와 같다.
내가 생성한 onnx 파일이 정상적으로 작동하는지 확인하기 위해서는 위와 같이 입력하고 결과가 나오는지 확인하면 된다.
---------------------------------------------------------------------------------------
import onnxruntime as ort
import numpy as np
ort_session = ort.InferenceSession("C:/Users/L/runs/detect/train/weights/best.onnx")
dummy_input = np.random.randn(1, 3, 640, 640).astype(np.float32)
outputs = ort_session.run(None, {"images": dummy_input})
print(outputs)
---------------------------------------------------------------------------------------
매번 cmd 창에서 명령을 입력하기 귀찮은경우 train_yolo.py 파일을 만들어서 실행한것처럼 파일을 만들고 python 으로 실행하면 된다.
댓글