[ONNX 변환 코드]
import sys
import os
import torch
import argparse
from pathlib import Path
# 현재 파일 기준으로 상위 경로 추가
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
# 프로젝트 루트 경로를 PYTHONPATH에 추가
sys.path.append(BASE_DIR)
# 프로젝트 내 모듈 불러오기
try:
from mods.ocr.detection.network import CustomModel, CPthIO
except ImportError as e:
print(f"모듈 임포트 실패: {e}")
print(f"현재 sys.path: {sys.path}")
sys.exit(1)
def convert_to_onnx(model_path: str, output_path: str):
"""PyTorch 모델을 ONNX로 변환"""
device = torch.device("cpu") # GPU 대신 CPU 사용
# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"[INFO] Using device: {device}")
# 모델 로드
model = CustomModel(effinet="efficientnet-b3", use_amp=False).to(device)
try:
state_dict = torch.load(CPthIO(model_path).load(), map_location=device)
model.load_state_dict(state_dict)
except Exception as e:
print(f"[ERROR] 모델 로딩 실패: {e}")
sys.exit(1)
model.eval()
print("[INFO] 모델 로딩 완료")
# 더미 입력 생성
dummy_input = torch.randn(1, 3, 928, 1408, device=device)
# ONNX 변환
try:
torch.onnx.export(
model,
dummy_input,
output_path,
input_names=["input"],
output_names=["output"],
opset_version=13,
)
print(f"✅ ONNX 변환 완료: {output_path}")
except Exception as e:
print(f"[ERROR] ONNX 변환 실패: {e}")
sys.exit(1)
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="PyTorch 모델을 ONNX로 변환하는 스크립트")
# 기본 모델 경로
default_model_path = "/home/clapi/ocr_npu/innocr/data/model/freeze_detector.pt"
parser.add_argument("--model_path", type=str, help="변환할 PyTorch 모델 파일 경로 (.pt)")
parser.add_argument("--output_path", type=str, help="변환된 ONNX 모델 저장 경로 (.onnx)")
args = parser.parse_args()
model_path = Path(args.model_path)
if not model_path.exists():
print(f"[ERROR] 모델 파일을 찾을 수 없음: {model_path}")
sys.exit(1)
convert_to_onnx(str(model_path), args.output_path)
[양자화 실행 코드]
import onnx
import numpy as np
from furiosa.quantizer import quantize
from PIL import Image
import os
import sys
import argparse
# 전처리 함수 - PyTorch 모델 기준 정규화 방식 반영
def normalize_mean_variance(img_np):
"""
PyTorch 모델에서 사용하는 normalizeMeanVariance와 동일한 정규화
"""
mean = np.array([0.485, 0.456, 0.406], dtype=np.float32) # 평균
std = np.array([0.229, 0.224, 0.225], dtype=np.float32) # 표준편차
img_np = img_np / 255.0 # 원본이미지 img_np를 0~1 사이의 실수 값으로 스케일링
img_np = (img_np - mean[None, None, :]) / std[None, None, :] # 평균과 표준편차를 사용하여 정규화
return img_np
def preprocess_image(image_path: str, target_size=(928, 1408)):
"""
NPU 추론 시 사용되는 928*1408 정규화 입력과 동일한 전처리
"""
img = Image.open(image_path).convert('RGB')
# 1. 224x224로 리사이즈
img = img.resize(target_size, Image.Resampling.LANCZOS)
# 2. numpy 배열로 변환 및 float32 캐스팅
img_np = np.array(img).astype(np.float32)
# 3. 정규화
img_np = normalize_mean_variance(img_np)
# 4. 채널 순서 변경: HWC → CHW
img_np = img_np.transpose((2, 0, 1))
# 5. 배치 차원 추가: CHW → NCHW
img_np = np.expand_dims(img_np, axis=0)
return img_np
def find_images_in_directory(directory_path: str):
"""지정된 디렉토리에서 이미지 파일 목록을 찾습니다."""
image_paths = []
supported_extensions = ('.png', '.jpg', '.jpeg', '.gif', '.bmp')
if not os.path.isdir(directory_path):
return []
for filename in os.listdir(directory_path):
if filename.lower().endswith(supported_extensions):
full_path = os.path.join(directory_path, filename)
image_paths.append(full_path)
return image_paths
def quantize_with_ranges(model_path: str, calibration_image_dir: str, output_path: str):
"""보정 데이터셋을 사용하여 양자화를 수행합니다."""
# 1. ONNX 모델 로드
model = onnx.load(model_path)
# 2. 보정 데이터셋 생성
calibration_dataset = []
calibration_image_paths = find_images_in_directory(calibration_image_dir)
if not calibration_image_paths:
print(f"[ERROR] '{calibration_image_dir}' 디렉토리에서 보정 이미지를 찾을 수 없습니다.")
sys.exit(1)
for img_path in calibration_image_paths:
try:
preprocessed_img = preprocess_image(img_path)
calibration_dataset.append(preprocessed_img)
print(f"보정 데이터 로드: {img_path}, shape: {preprocessed_img.shape}")
except Exception as e:
print(f"[경고] 이미지 로딩/전처리 실패: {img_path}, 오류: {e}")
if not calibration_dataset:
print("[ERROR] 보정 데이터셋이 비어 있습니다. 양자화를 진행할 수 없습니다.")
sys.exit(1)
# 3. 보정 데이터셋으로 입력 텐서의 min/max 범위 계산
print("[INFO] 보정 데이터셋으로 입력 텐서의 범위를 계산 중...")
all_images = np.concatenate(calibration_dataset, axis=0)
input_min_val = np.min(all_images)
input_max_val = np.max(all_images)
# 4. 모든 텐서 이름 추출
all_tensor_names = [t.name for t in model.graph.input] \
+ [t.name for t in model.graph.output] \
+ [t.name for t in model.graph.value_info]
# 5. 모든 텐서에 대한 양자화 범위 딕셔너리 생성
tensor_name_to_range = {}
input_tensor_name = model.graph.input[0].name
for name in all_tensor_names:
if name == input_tensor_name:
# 입력 텐서는 보정 데이터로 계산한 정확한 범위를 사용
tensor_name_to_range[name] = (input_min_val, input_max_val)
else:
# 나머지 텐서들은 넉넉한 범위로 설정
tensor_name_to_range[name] = (-127.0, 127.0)
print(f"[INFO] 생성된 텐서별 양자화 범위: {tensor_name_to_range}")
# 6. 양자화 수행
print("[INFO] 계산된 범위를 기반으로 양자화 진행 중...")
quantized_model = quantize(model, tensor_name_to_range)
# 7. 결과 저장
onnx.save(quantized_model, output_path)
print(f"✅ 양자화 완료: {output_path}")
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="ONNX 모델을 보정 이미지로 양자화하는 스크립트")
parser.add_argument("--model_path", type=str, required=True, help="ONNX 모델 파일 경로")
parser.add_argument("--calibration_image_dir", type=str, required=True, help="보정 이미지가 들어있는 디렉토리 경로")
parser.add_argument("--output_path", type=str, required=True, help="양자화된 ONNX 모델 저장 경로")
args = parser.parse_args()
quantize_with_ranges(args.model_path, args.calibration_image_dir, args.output_path)