46 lines
1.8 KiB
Python
46 lines
1.8 KiB
Python
import onnxruntime as ort
|
|
import sys
|
|
|
|
def check_onnx_batch_support(model_path):
|
|
"""
|
|
ONNX 모델을 로드하여 배치 처리 지원 여부를 확인하고,
|
|
입력 및 출력 정보를 출력하는 함수
|
|
"""
|
|
try:
|
|
session = ort.InferenceSession(model_path)
|
|
print(f"✅ 모델 로드 성공: {model_path}")
|
|
|
|
print("\n[입력 정보]")
|
|
for i, input_meta in enumerate(session.get_inputs()):
|
|
shape = input_meta.shape
|
|
print(f" - 입력 #{i}")
|
|
print(f" - 이름: {input_meta.name}")
|
|
print(f" - 형태: {shape}")
|
|
# 배치 차원이 문자열(동적)인지 확인
|
|
if isinstance(shape[0], str) or shape[0] is None:
|
|
print(f" - ✅ 배치 처리 지원 (동적 배치 크기 '{shape[0]}')")
|
|
else:
|
|
print(f" - ⚠️ 배치 처리를 지원하지 않거나, 배치 크기가 '{shape[0]}'으로 고정됨")
|
|
|
|
print("\n[출력 정보]")
|
|
for i, output_meta in enumerate(session.get_outputs()):
|
|
shape = output_meta.shape
|
|
print(f" - 출력 #{i}")
|
|
print(f" - 이름: {output_meta.name}")
|
|
print(f" - 형태: {shape}")
|
|
if isinstance(shape[0], str) or shape[0] is None:
|
|
print(f" - ✅ 배치 처리 지원 (동적 배치 크기 '{shape[0]}')")
|
|
|
|
except Exception as e:
|
|
print(f"❌ 오류 발생: 모델을 로드하거나 분석하는 중 문제가 발생했습니다.")
|
|
print(f" {model_path}: {e}")
|
|
|
|
if __name__ == '__main__':
|
|
if len(sys.argv) < 2:
|
|
print("사용법: python check_onnx.py <ONNX_FILE_PATH_1> [ONNX_FILE_PATH_2] ...")
|
|
else:
|
|
for model_path in sys.argv[1:]:
|
|
print("-" * 50)
|
|
check_onnx_batch_support(model_path)
|
|
print("-" * 50)
|