model/src/optimization/utils.py

"""
유틸리티 함수
"""
import os
import io
import json
import pickle
import logging
import contextlib
import traceback
import numpy as np
import tensorflow as tf
from pathlib import Path
from tqdm import tqdm

class TqdmProgressCallback(tf.keras.callbacks.Callback):
    """
    TensorFlow 훈련을 위한 커스텀 콜백
    """
    def __init__(self, epochs, verbose=1):
        super(TqdmProgressCallback, self).__init__()
        self.epochs = epochs
        self.verbose = verbose
        self.tqdm_bar = None
    
    def on_train_begin(self, logs=None):
        if self.verbose:
            self.tqdm_bar = tqdm(total=self.epochs, desc="Training", unit="epoch")
    
    def on_epoch_end(self, epoch, logs=None):
        if self.verbose:
            logs = logs or {}
            log_items = []
            for k, v in logs.items():
                if 'val_' not in k:  # 훈련 지표만 출력
                    log_items.append(f"{k}: {v:.4f}")
            
            desc = ", ".join(log_items)
            self.tqdm_bar.set_description(desc)
            self.tqdm_bar.update(1)
    
    def on_train_end(self, logs=None):
        if self.verbose and self.tqdm_bar is not None:
            self.tqdm_bar.close()
            print("학습 완료!")

def get_project_root():
    """
    프로젝트 루트 디렉토리를 반환합니다
    """
    return Path(__file__).parent.parent.parent

def ensure_directory(directory_path):
    """
    디렉토리가 존재하지 않으면 생성합니다.
    """
    Path(directory_path).mkdir(parents=True, exist_ok=True)
    return Path(directory_path)

def normalize_path(path_str, base_dir=None):
    """
    상대 경로를 절대 경로로 변환합니다.
    """
    path = Path(path_str)
    
    if path.is_absolute():
        return path
    
    if base_dir is None:
        base_dir = get_project_root()
    
    return Path(base_dir) / path

def save_model(model, model_path, config=None, encoders=None):
    """
    모델을 TensorFlow Lite 형식으로 저장합니다.
    """
    model_path = Path(model_path)
    ensure_directory(model_path.parent)
    
    # TensorFlow 로그 레벨 임시 조정
    original_tf_log_level = os.environ.get('TF_CPP_MIN_LOG_LEVEL', '')
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
    
    tf_logger = logging.getLogger('tensorflow')
    original_tf_level = tf_logger.level
    tf_logger.setLevel(logging.ERROR)
    
    try:
        # .tflite 확장자로 변경
        if not str(model_path).endswith('.tflite'):
            model_path = Path(str(model_path).replace('.keras', '').replace('.h5', '') + '.tflite')
        
        # TensorFlow Lite 변환
        print("TensorFlow Lite 모델로 변환 중...")
        converter = tf.lite.TFLiteConverter.from_keras_model(model)
        
        # LSTM 호환성을 위한 설정
        converter.optimizations = [tf.lite.Optimize.DEFAULT]
        converter.target_spec.supported_ops = [
            tf.lite.OpsSet.TFLITE_BUILTINS,    # 기본 TFLite 연산
            tf.lite.OpsSet.SELECT_TF_OPS       # 추가 TensorFlow 연산 허용
        ]
        converter._experimental_lower_tensor_list_ops = False  # TensorList 변환 비활성화
        converter.allow_custom_ops = True  # 커스텀 연산 허용
        
        # 변환 실행
        with contextlib.redirect_stdout(io.StringIO()):
            with contextlib.redirect_stderr(io.StringIO()):
                tflite_model = converter.convert()
        
        # TFLite 모델 저장
        with open(model_path, 'wb') as f:
            f.write(tflite_model)
        
        print(f"TensorFlow Lite 모델이 {model_path}에 저장되었습니다.")
        
        # 모델 파일 기본명 추출
        model_stem = model_path.stem
        
        # 인코더 정보 저장
        if encoders is not None:
            encoder_path = model_path.parent / f"{model_stem}_encoders.json"
            with open(encoder_path, 'w') as f:
                json.dump(encoders, f, indent=2)
            print(f"인코더 정보가 {encoder_path}에 저장되었습니다.")
        
        # 모델 설정 저장
        if config is not None:
            config_path = model_path.parent / f"{model_stem}_config.json"
            with open(config_path, 'w') as f:
                # 직렬화 가능한 형태로 변환
                json_safe_config = {k: str(v) if not isinstance(v, (int, float, str, bool, list, dict)) else v 
                                  for k, v in config.items()}
                json.dump(json_safe_config, f, indent=2)
            print(f"모델 설정이 {config_path}에 저장되었습니다.")
        
        return True
    
    finally:
        # 원래 로그 설정 복원
        tf_logger.setLevel(original_tf_level)
        if original_tf_log_level:
            os.environ['TF_CPP_MIN_LOG_LEVEL'] = original_tf_log_level
        else:
            os.environ.pop('TF_CPP_MIN_LOG_LEVEL', None)

def load_tflite_model(model_path):
    """
    TensorFlow Lite 모델을 로드합니다.
    """
    try:
        # TFLite 인터프리터 생성
        interpreter = tf.lite.Interpreter(model_path=str(model_path))
        interpreter.allocate_tensors()
        
        print(f"TensorFlow Lite 모델이 {model_path}에서 성공적으로 로드되었습니다.")
        return interpreter
    
    except Exception as e:
        print(f"TensorFlow Lite 모델 로드 실패: {e}")
        print(traceback.format_exc())
        return None

def predict_with_tflite(interpreter, inputs, verbose=False):
    """
    TensorFlow Lite 모델로 예측 수행
    """
    try:
        # 입력 텐서 정보 가져오기
        input_details = interpreter.get_input_details()
        output_details = interpreter.get_output_details()
        
        # 각 입력 설정
        for i, input_tensor in enumerate(inputs):
            if i < len(input_details):
                interpreter.set_tensor(input_details[i]['index'], input_tensor)
        
        # 실행
        interpreter.invoke()
        
        # 출력 가져오기
        outputs = []
        for output_detail in output_details:
            output = interpreter.get_tensor(output_detail['index'])
            outputs.append(output)
        
        return outputs if len(outputs) > 1 else outputs[0]
        
    except Exception as e:
        print("예측 실패")
        return None

def save_results(results, output_path, include_model=False):
    """
    최적화 결과를 저장합니다.
    """
    output_path = Path(output_path)
    ensure_directory(output_path.parent)
    
    # 결과 복사본 생성
    pickle_safe_results = {
        'grid_results': [],
        'best_config': results.get('best_config', {})
    }
    
    # 모델 객체 제거한 결과 복사
    results_list = results.get('results', [])
    if not results_list and 'best_result' in results:
        results_list = [results['best_result']]
    
    for result in results_list:
        result_copy = result.copy()
        if not include_model and 'model' in result_copy:
            del result_copy['model']
        pickle_safe_results['grid_results'].append(result_copy)
    
    # 테스트 결과 추가
    if 'test_backtest' in results:
        pickle_safe_results['test_backtest'] = results['test_backtest']
    
    # 결과 저장
    with open(output_path, 'wb') as f:
        pickle.dump(pickle_safe_results, f)
    print(f"최적화 결과가 {output_path}에 저장되었습니다.")
    
    return True

def save_metadata(metadata, output_path):
    """
    메타데이터를 JSON 형식으로 저장합니다.
    """
    output_path = Path(output_path)
    ensure_directory(output_path.parent)
    
    # 직렬화 가능한 형태로 변환
    json_safe_metadata = {k: str(v) if not isinstance(v, (int, float, str, bool, list, dict)) else v 
                       for k, v in metadata.items()}
    
    with open(output_path, 'w') as f:
        json.dump(json_safe_metadata, f, indent=2)
    print(f"메타데이터가 {output_path}에 저장되었습니다.")
    
    return True