ONNX 모델 진단 기능 추가 및 TensorRT 세션 생성 로직 개선. 이미지 경로 변경 및 관련 설정 업데이트. 전체적인 코드 정리 및 주석 보강.
This commit is contained in:
7600M
parent
e3bd4bb50b
commit
baf63e2a02
Binary file not shown.
|
Before Width: | Height: | Size: 951 KiB After Width: | Height: | Size: 857 KiB |
|
|
@ -16,7 +16,7 @@ import requests
|
|||
|
||||
|
||||
API_ROOT = "http://localhost:7890" # 메인 서버 주소
|
||||
IMAGE_PATH = pathlib.Path("6.jpg")
|
||||
IMAGE_PATH = pathlib.Path("8.jpg")
|
||||
TIMEOUT = 120 # 초
|
||||
|
||||
unwanted_texts = {
|
||||
|
|
|
|||
|
|
@ -19,6 +19,12 @@ import cv2
|
|||
import numpy as np
|
||||
import onnxruntime as ort
|
||||
|
||||
try:
|
||||
import onnx
|
||||
ONNX_AVAILABLE = True
|
||||
except ImportError:
|
||||
ONNX_AVAILABLE = False
|
||||
|
||||
# OpenCV 내부 최적화 off (호환성)
|
||||
cv2.setUseOptimized(False)
|
||||
|
||||
|
|
@ -146,6 +152,52 @@ class MIGANInpaintingModule:
|
|||
self._TRT_ENGINE_CACHE_DIR = cache_dir
|
||||
self.logger.log(f"TensorRT 엔진 캐시 디렉토리: {cache_dir}", level=logging.INFO)
|
||||
|
||||
def _diagnose_onnx_model(self, onnx_path: str) -> Dict[str, Any]:
|
||||
"""ONNX 모델 TensorRT 호환성 진단"""
|
||||
diagnosis = {
|
||||
'compatible': True,
|
||||
'issues': [],
|
||||
'dynamic_axes': [],
|
||||
'has_shape_info': False
|
||||
}
|
||||
|
||||
if not ONNX_AVAILABLE:
|
||||
diagnosis['issues'].append("ONNX 패키지 미설치 - 진단 불가")
|
||||
return diagnosis
|
||||
|
||||
try:
|
||||
model = onnx.load(onnx_path)
|
||||
|
||||
# 동적 축 확인
|
||||
for input_info in model.graph.input:
|
||||
for dim in input_info.type.tensor_type.shape.dim:
|
||||
if dim.HasField('dim_param'):
|
||||
diagnosis['dynamic_axes'].append(f"{input_info.name}:{dim.dim_param}")
|
||||
|
||||
# Shape inference 상태 확인
|
||||
try:
|
||||
onnx.checker.check_model(model)
|
||||
inferred_model = onnx.shape_inference.infer_shapes(model)
|
||||
if inferred_model.graph.value_info:
|
||||
diagnosis['has_shape_info'] = True
|
||||
else:
|
||||
diagnosis['issues'].append("Shape inference 정보 부족")
|
||||
diagnosis['compatible'] = False
|
||||
except Exception as e:
|
||||
diagnosis['issues'].append(f"Shape inference 실패: {str(e)}")
|
||||
diagnosis['compatible'] = False
|
||||
|
||||
# Pad 연산 확인 (TensorRT 문제 원인)
|
||||
for node in model.graph.node:
|
||||
if node.op_type == 'Pad':
|
||||
diagnosis['issues'].append(f"Pad 연산 발견: {node.name} - TensorRT 문제 가능성")
|
||||
|
||||
except Exception as e:
|
||||
diagnosis['issues'].append(f"ONNX 모델 로드 실패: {str(e)}")
|
||||
diagnosis['compatible'] = False
|
||||
|
||||
return diagnosis
|
||||
|
||||
def _get_session_key(self, config: Dict[str, Any]) -> tuple:
|
||||
"""세션 캐시 키 생성"""
|
||||
return (
|
||||
|
|
@ -158,7 +210,7 @@ class MIGANInpaintingModule:
|
|||
)
|
||||
|
||||
def _create_session(self, config: Dict[str, Any]) -> ort.InferenceSession:
|
||||
"""ONNX Runtime 세션 생성"""
|
||||
"""ONNX Runtime 세션 생성 - TensorRT → CUDA → CPU 순서로 폴백"""
|
||||
onnx_path = config['onnx_path']
|
||||
|
||||
if not os.path.exists(onnx_path):
|
||||
|
|
@ -174,72 +226,137 @@ class MIGANInpaintingModule:
|
|||
# 최적화 레벨 설정
|
||||
session_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
|
||||
|
||||
# Provider 설정
|
||||
providers = []
|
||||
provider_options = []
|
||||
|
||||
if config['use_cuda']:
|
||||
if config['use_tensorrt']:
|
||||
# TensorRT Provider 설정
|
||||
trt_options = {
|
||||
# 1차 시도: TensorRT (단계적 접근)
|
||||
if config.get('use_cuda', True) and config.get('use_tensorrt', True):
|
||||
# ONNX 모델 TensorRT 호환성 사전 진단
|
||||
diagnosis = self._diagnose_onnx_model(onnx_path)
|
||||
if diagnosis['issues']:
|
||||
self.logger.log(f"ONNX 모델 진단 결과 - 문제점 {len(diagnosis['issues'])}개 발견:", level=logging.WARNING)
|
||||
for issue in diagnosis['issues']:
|
||||
self.logger.log(f" - {issue}", level=logging.WARNING)
|
||||
|
||||
# 호환성 문제가 심각한 경우 TensorRT 건너뛰기
|
||||
if not diagnosis.get('compatible', True):
|
||||
self.logger.log("TensorRT 호환성 문제로 CUDA로 직접 폴백", level=logging.WARNING)
|
||||
else:
|
||||
if not diagnosis['issues']:
|
||||
self.logger.log("ONNX 모델 TensorRT 호환성 진단 통과", level=logging.INFO)
|
||||
else:
|
||||
self.logger.log("경고가 있지만 TensorRT 시도 계속", level=logging.INFO)
|
||||
|
||||
# 1-1. 기본 TensorRT 시도 (최소 옵션)
|
||||
try:
|
||||
trt_basic_options = {
|
||||
'device_id': 0,
|
||||
'trt_fp16_enable': config['trt_fp16_enable'],
|
||||
'trt_engine_cache_enable': config['trt_engine_cache_enable'],
|
||||
}
|
||||
|
||||
if config['trt_engine_cache_enable'] and self._TRT_ENGINE_CACHE_DIR:
|
||||
trt_basic_options['trt_engine_cache_path'] = self._TRT_ENGINE_CACHE_DIR
|
||||
|
||||
self.logger.log(f"TensorRT 기본 세션 생성 시도: FP16={config['trt_fp16_enable']}", level=logging.INFO)
|
||||
|
||||
session = ort.InferenceSession(
|
||||
onnx_path,
|
||||
sess_options=session_options,
|
||||
providers=['TensorrtExecutionProvider'],
|
||||
provider_options=[trt_basic_options]
|
||||
)
|
||||
|
||||
actual_providers = session.get_providers()
|
||||
self.logger.log(f"TensorRT 기본 세션 생성 성공: {actual_providers}", level=logging.INFO)
|
||||
return session
|
||||
|
||||
except Exception as e:
|
||||
self.logger.log(f"TensorRT 기본 세션 실패: {e}", level=logging.WARNING)
|
||||
|
||||
# 1-2. 고급 TensorRT 시도 (상세 옵션)
|
||||
try:
|
||||
trt_advanced_options = {
|
||||
'device_id': 0,
|
||||
'trt_fp16_enable': config['trt_fp16_enable'],
|
||||
'trt_engine_cache_enable': config['trt_engine_cache_enable'],
|
||||
'trt_max_workspace_size': 2 * 1024 * 1024 * 1024, # 2GB
|
||||
'trt_max_partition_iterations': 1000,
|
||||
'trt_min_subgraph_size': 5,
|
||||
'trt_dla_enable': False,
|
||||
'trt_int8_enable': False,
|
||||
'trt_builder_optimization_level': 3,
|
||||
# MIGAN 실제 사용 패턴에 맞춤
|
||||
'trt_profile_min_shapes': 'image:1x3x256x256,mask:1x1x256x256',
|
||||
'trt_profile_max_shapes': 'image:1x3x1600x1600,mask:1x1x1600x1600',
|
||||
'trt_profile_opt_shapes': 'image:1x3x816x1200,mask:1x1x816x1200',
|
||||
}
|
||||
|
||||
if config['trt_engine_cache_enable'] and self._TRT_ENGINE_CACHE_DIR:
|
||||
trt_advanced_options['trt_engine_cache_path'] = self._TRT_ENGINE_CACHE_DIR
|
||||
|
||||
self.logger.log("TensorRT 고급 세션 생성 시도 (동적 shape 프로파일 포함)", level=logging.INFO)
|
||||
|
||||
session = ort.InferenceSession(
|
||||
onnx_path,
|
||||
sess_options=session_options,
|
||||
providers=['TensorrtExecutionProvider'],
|
||||
provider_options=[trt_advanced_options]
|
||||
)
|
||||
|
||||
actual_providers = session.get_providers()
|
||||
self.logger.log(f"TensorRT 고급 세션 생성 성공: {actual_providers}", level=logging.INFO)
|
||||
return session
|
||||
|
||||
except Exception as e2:
|
||||
self.logger.log(f"TensorRT 고급 세션도 실패: {e2}", level=logging.WARNING)
|
||||
# TensorRT 실패 시 자세한 오류 정보 로깅
|
||||
if "shape inference" in str(e2).lower():
|
||||
self.logger.log("TensorRT shape inference 오류 - ONNX 모델 재생성 필요할 수 있음", level=logging.WARNING)
|
||||
elif "pad_output" in str(e2).lower():
|
||||
self.logger.log("Pad_output 오류 - 동적 축 문제, CUDA로 폴백", level=logging.WARNING)
|
||||
|
||||
# 2차 시도: CUDA
|
||||
if config.get('use_cuda', True):
|
||||
try:
|
||||
cuda_options = {
|
||||
'device_id': 0,
|
||||
'trt_fp16_enable': config['trt_fp16_enable'],
|
||||
'trt_engine_cache_enable': config['trt_engine_cache_enable'],
|
||||
'arena_extend_strategy': 'kNextPowerOfTwo',
|
||||
'gpu_mem_limit': 2 * 1024 * 1024 * 1024, # 2GB
|
||||
'cudnn_conv_algo_search': 'EXHAUSTIVE',
|
||||
'do_copy_in_default_stream': True,
|
||||
}
|
||||
|
||||
if config['trt_engine_cache_enable'] and self._TRT_ENGINE_CACHE_DIR:
|
||||
trt_options['trt_engine_cache_path'] = self._TRT_ENGINE_CACHE_DIR
|
||||
self.logger.log("CUDA 세션 생성 시도", level=logging.INFO)
|
||||
|
||||
providers.append('TensorrtExecutionProvider')
|
||||
provider_options.append(trt_options)
|
||||
session = ort.InferenceSession(
|
||||
onnx_path,
|
||||
sess_options=session_options,
|
||||
providers=['CUDAExecutionProvider'],
|
||||
provider_options=[cuda_options]
|
||||
)
|
||||
|
||||
self.logger.log(f"TensorRT 설정: FP16={config['trt_fp16_enable']}, 캐시={config['trt_engine_cache_enable']}", level=logging.INFO)
|
||||
|
||||
# CUDA Provider
|
||||
cuda_options = {
|
||||
'device_id': 0,
|
||||
'arena_extend_strategy': 'kNextPowerOfTwo',
|
||||
'gpu_mem_limit': 2 * 1024 * 1024 * 1024, # 2GB
|
||||
'cudnn_conv_algo_search': 'EXHAUSTIVE',
|
||||
'do_copy_in_default_stream': True,
|
||||
}
|
||||
providers.append('CUDAExecutionProvider')
|
||||
provider_options.append(cuda_options)
|
||||
|
||||
# CPU Provider (폴백)
|
||||
providers.append('CPUExecutionProvider')
|
||||
provider_options.append({})
|
||||
actual_providers = session.get_providers()
|
||||
self.logger.log(f"CUDA 세션 생성 성공: {actual_providers}", level=logging.INFO)
|
||||
return session
|
||||
|
||||
except Exception as e:
|
||||
self.logger.log(f"CUDA 세션 생성 실패: {e}", level=logging.WARNING)
|
||||
|
||||
# 3차 시도: CPU (최종 폴백)
|
||||
try:
|
||||
if provider_options:
|
||||
session = ort.InferenceSession(
|
||||
onnx_path,
|
||||
sess_options=session_options,
|
||||
providers=providers,
|
||||
provider_options=provider_options
|
||||
)
|
||||
else:
|
||||
session = ort.InferenceSession(
|
||||
onnx_path,
|
||||
sess_options=session_options,
|
||||
providers=providers
|
||||
)
|
||||
|
||||
actual_providers = session.get_providers()
|
||||
self.logger.log(f"ONNX Runtime 세션 생성 완료: {actual_providers}", level=logging.INFO)
|
||||
self.logger.log("CPU 세션 생성 시도 (최종 폴백)", level=logging.INFO)
|
||||
|
||||
return session
|
||||
|
||||
except Exception as e:
|
||||
self.logger.log(f"GPU 세션 생성 실패, CPU로 폴백: {e}", level=logging.WARNING)
|
||||
# CPU 전용 폴백
|
||||
session = ort.InferenceSession(
|
||||
onnx_path,
|
||||
sess_options=session_options,
|
||||
providers=['CPUExecutionProvider']
|
||||
)
|
||||
self.logger.log("CPU 세션으로 폴백 완료", level=logging.INFO)
|
||||
|
||||
actual_providers = session.get_providers()
|
||||
self.logger.log(f"CPU 세션 생성 성공: {actual_providers}", level=logging.INFO)
|
||||
return session
|
||||
|
||||
except Exception as e:
|
||||
self.logger.log(f"모든 provider에서 세션 생성 실패: {e}", level=logging.ERROR)
|
||||
raise RuntimeError(f"ONNX 세션을 생성할 수 없습니다: {e}")
|
||||
|
||||
def _get_or_create_session(self, config: Dict[str, Any]) -> ort.InferenceSession:
|
||||
"""세션 가져오기 또는 생성 (캐시 활용)"""
|
||||
|
|
@ -563,7 +680,7 @@ def build_migan_from_toggle(toggle_states: Dict[str, Any], logger=None) -> MIGAN
|
|||
config = {
|
||||
'onnx_path': toggle_states.get('migan_onnx_path', '/app/worker/models/migan_pipeline_v2.onnx'),
|
||||
'use_cuda': toggle_states.get('migan_use_cuda', True),
|
||||
'use_tensorrt': toggle_states.get('migan_use_tensorrt', True),
|
||||
'use_tensorrt': toggle_states.get('migan_use_tensorrt', False),
|
||||
'trt_fp16_enable': toggle_states.get('migan_trt_fp16_enable', True),
|
||||
'trt_engine_cache_enable': toggle_states.get('migan_trt_engine_cache_enable', True),
|
||||
'intra_threads': int(toggle_states.get('migan_intra_threads', 0) or 0),
|
||||
|
|
|
|||
Loading…
Reference in New Issue