|
|
import torch
|
|
|
import argparse
|
|
|
import os
|
|
|
from model import get_face_model
|
|
|
|
|
|
def export_to_onnx(pt_path, output_path):
|
|
|
print(f"🔄 Loading model from: {pt_path}")
|
|
|
|
|
|
# 1. 모델 구조 불러오기 및 가중치 로드
|
|
|
# 보드용 모델이므로 학습용 ArcFace 헤더는 버리고, Backbone만 가져옵니다.
|
|
|
model = get_face_model()
|
|
|
|
|
|
# CPU로 로드 (변환은 굳이 GPU 불필요)
|
|
|
checkpoint = torch.load(pt_path, map_location='cpu')
|
|
|
|
|
|
# state_dict 로드 (혹시 모를 키 불일치 방지를 위해 strict=False는 선택사항이나, 여기선 구조가 같으므로 True 권장)
|
|
|
try:
|
|
|
model.load_state_dict(checkpoint)
|
|
|
except RuntimeError as e:
|
|
|
print(f"⚠️ Key mismatch detected. Trying to load with strict=False...")
|
|
|
model.load_state_dict(checkpoint, strict=False)
|
|
|
|
|
|
# 2. Eval 모드 전환 (매우 중요)
|
|
|
# 이걸 안 하면 BatchNorm, Dropout 등이pip 학습 모드로 동작하여 결과가 이상해집니다.
|
|
|
model.eval()
|
|
|
|
|
|
# 3. Dummy Input 생성 (Static Shape: 128x128)
|
|
|
# 보드 사양에 맞춰 배치 사이즈는 1로 고정합니다. [1, 3, 128, 128]
|
|
|
dummy_input = torch.randn(1, 3, 128, 128)
|
|
|
|
|
|
print(f"Target ONNX Path: {output_path}")
|
|
|
|
|
|
# 4. ONNX Export
|
|
|
# external_data=False는 PyTorch export에서 기본적으로 2GB 미만 모델에 대해 적용되어 단일 파일로 나옵니다.
|
|
|
# dynamic_axes 옵션을 뺌으로써 Static Shape을 강제합니다.
|
|
|
torch.onnx.export(
|
|
|
model, # 실행될 모델
|
|
|
dummy_input, # 모델 입력값 (차원 체크용)
|
|
|
output_path, # 저장될 경로
|
|
|
# export_params=True, # 모델 파일 안에 웨이트 저장 (external_data=False 효과)
|
|
|
# opset_version=11, # 임베디드 보드에서 가장 호환성 좋은 버전 (11 추천)
|
|
|
# do_constant_folding=True, # 상수 폴딩 최적화
|
|
|
input_names=['input'], # 입력 노드 이름
|
|
|
output_names=['output'], # 출력 노드 이름
|
|
|
external_data=False
|
|
|
# dynamic_axes={...} <-- 이 옵션을 사용하지 않음으로써 Static Shape으로 고정됨!
|
|
|
)
|
|
|
|
|
|
print(f"✅ Conversion Completed! Model saved at: {output_path}")
|
|
|
print(f"ℹ️ Input Shape: {dummy_input.shape} (Static)")
|
|
|
print(f"ℹ️ Please check if '{output_path}' is a single file.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
parser = argparse.ArgumentParser(description='Convert PyTorch model to ONNX')
|
|
|
|
|
|
# 입력받을 .pt 파일 경로
|
|
|
parser.add_argument('--input', type=str, required=True, help='Input .pt file path')
|
|
|
|
|
|
# 출력할 .onnx 파일 경로 (옵션)
|
|
|
parser.add_argument('--output', type=str, default=None, help='Output .onnx file path')
|
|
|
|
|
|
args = parser.parse_args()
|
|
|
|
|
|
# Output 경로가 없으면 Input 경로에서 확장자만 바꿔서 자동 지정
|
|
|
if args.output is None:
|
|
|
args.output = args.input.replace('.pt', '.onnx')
|
|
|
|
|
|
if not os.path.exists(args.input):
|
|
|
print(f"❌ Error: Input file not found: {args.input}")
|
|
|
else:
|
|
|
export_to_onnx(args.input, args.output) |
