PYTORCH 가중치 파일을 ONNX로 변환하기 (Feat. YOLO)

Pytorch의 모델 가중치 파일을 Onnx로 변환하는 방법을 알아보자

 

Pytorch에서 Onnx 타입으로 변환하는건 여러 이유가 있습니다.

 

• 양자화
• 다른 타입으로 변환

등이 있습니다.

 

오늘은 YOLO의 가중치 파일(.pt)을 ONNX(.onnx)로 바꾸는 방법을 알아보겠습니다.

 

Ultralytics YOLO는 이미 ONNX 로 변환하는 기능을 라이브러리에 포함하여 제공하고 있습니다.

Ultralytics의 onnx변환 그리고 pytorch에서 제공하는 onnx 변환 방식 두가지 모두 실행하고 결과를 비교해보겠습니다.

 

1. Pytorch 모델 가중치 파일 준비

from ultralytics import YOLO

origin_model = YOLO('yolo11n.pt')

※ 나중에 검증을 빠르게 끝내기 위해서 가장 작은 nano 모델로 선정

 

2. Ultralytics onnx 변환 (export)

origin_model.export(format='onnx')

 

3. Pytorch onnx 변환 (export)

# https://docs.pytorch.org/tutorials/beginner/onnx/export_simple_model_to_onnx_tutorial.html

import torch
import os

try:
    torch_origin_model = origin_model.model
    torch_origin_model.eval()
except Exception as e:
    print(f'Error : {e}')
example_inputs = (torch.randn(1, 3, 640, 640))
onnx_path = "pure_yolo11n.onnx"

# 4. torch.onnx.export 실행
try:
    # ONNX 내보내기 함수
    torch.onnx.export(
        torch_origin_model,                         # 실행할 모델
        example_inputs,                         # 모델 입력으로 사용할 더미 데이터
        onnx_path,                           # ONNX 모델 저장 경로
        export_params=True,                  # 모델 매개변수 저장
        opset_version=17,                    # ONNX Operator Set 버전 (호환성을 위해 17 사용)
        do_constant_folding=True,            # 상수 폴딩 최적화 수행
        input_names=['images'],              # 입력 노드 이름 정의
        output_names=['output'],             # 출력 노드 이름 정의 
        dynamic_axes={
            'images': {0: 'batch_size'},     # 배치 크기 동적 설정 (선택 사항)
            'output': {0: 'batch_size'}
        }
    )
    print(f"ONNX 모델 변환 성공: {onnx_path}")
    print(f"파일 크기: {os.path.getsize(onnx_path) / (1024*1024):.2f} MB")
except Exception as e:
    print(f"ONNX 모델 변환 중 오류 발생: {e}")

 

4. "COCO128" & "Ultralytics 라이브러리"로 성능 검증

# 1. 파일 경로 설정
onnx_model_1_path = 'yolo11n.onnx'
onnx_model_2_path = 'pure_yolo11n.onnx'

# 2. 데이터셋 YAML 파일 경로 설정
data_yaml_path = 'coco128.yaml' # coco128 데이터셋 사용

print("=== ONNX 모델 검증 시작 ===")
print(f"검증에 사용할 데이터셋: {data_yaml_path}")
print("-" * 30)


# --- 모델 1 검증: yolo11n.onnx ---
print(f"1. {onnx_model_1_path} 검증 중...")

try:
    # ONNX 모델 로드
    model_1 = YOLO(onnx_model_1_path)
    
    # 검증 실행
    # 'data' 인자에 데이터셋 YAML 파일 경로를 지정합니다.
    metrics_1 = model_1.val(data=data_yaml_path, device='cpu', batch=16, plots=False) 
    
    # 주요 지표 저장
    map50_1 = metrics_1.box.map50
    map_1 = metrics_1.box.map  # map@0.5:0.95
    print(f"   -> mAP@0.5: {map50_1:.4f}")
    print(f"   -> mAP@0.5:0.95: {map_1:.4f}")

except Exception as e:
    print(f"    {onnx_model_1_path} 검증 중 오류 발생: {e}")
    map50_1, map_1 = None, None
    
print("-" * 30)

# --- 모델 2 검증: pure_yolo11n.onnx ---
print(f"2. {onnx_model_2_path} 검증 중...")

try:
    # ONNX 모델 로드
    model_2 = YOLO(onnx_model_2_path)
    
    # 검증 실행
    metrics_2 = model_2.val(data=data_yaml_path, device='cpu', batch=16, plots=False)
    
    # 주요 지표 저장
    map50_2 = metrics_2.box.map50
    map_2 = metrics_2.box.map # map@0.5:0.95
    print(f"   -> mAP@0.5: {map50_2:.4f}")
    print(f"   -> mAP@0.5:0.95: {map_2:.4f}")

except Exception as e:
    print(f"    {onnx_model_2_path} 검증 중 오류 발생: {e}")
    map50_2, map_2 = None, None

print("-" * 30)

# 3. 결과 비교
print("=== 최종 비교 결과 ===")
if map_1 is not None and map_2 is not None:
    print(f"모델 1 ({onnx_model_1_path}) mAP@0.5:0.95: {map_1:.4f}")
    print(f"모델 2 ({onnx_model_2_path}) mAP@0.5:0.95: {map_2:.4f}")

    if map_1 > map_2:
        print(f"모델 1이 모델 2보다 성능이 우수. (차이: {(map_1 - map_2):.4f})")
    elif map_2 > map_1:
        print(f"모델 2가 모델 1보다 성능이 우수. (차이: {(map_2 - map_1):.4f})")
    else:
        print("두 모델의 성능이 동일.")
else:
    print("하나 이상의 모델 검증에 실패하여 비교할 수 없음.")

 

5. 검증 실행 결과

Origin Onnx Yolo11n Model
Speed: 0.6ms preprocess, 23.3ms inference, 0.0ms loss, 2.0ms postprocess per image
-> mAP@0.5: 0.6761
-> mAP@0.5:0.95: 0.5077

Pure Onnx Yolo11n Model
Speed: 0.6ms preprocess, 23.6ms inference, 0.0ms loss, 1.9ms postprocess per image
   -> mAP@0.5: 0.6761
   -> mAP@0.5:0.95: 0.5077
   
둘 다 동일 성능