# trainer.py - 모델 훈련 및 평가 로직 (모듈화 및 깔끔하게 재정리) import os, json import sys import numpy as np import shutil from pathlib import Path from typing import Dict, Optional, List import pandas as pd import torch from transformers.models.auto.tokenization_auto import AutoTokenizer import pytorch_lightning as pl from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint, LearningRateMonitor from pytorch_lightning.loggers import TensorBoardLogger, WandbLogger # 프로젝트 루트 경로 추가 _current_dir = os.path.dirname(os.path.abspath(__file__)) _project_root = os.path.dirname(os.path.dirname(_current_dir)) if _project_root not in sys.path: sys.path.insert(0, _project_root) from admet_ft._modules.model import ChemBERTaMultiTaskLightning from admet_ft._modules.dataset import ChemMultiTaskDataModule from admet_ft._modules.utils import get_task_list, DIDB_FILTER_COLS from admet_ft._modules.scaler import (SCALER_FILE_NAME, reverse_scaling, reverse_scaling_power, reverse_scaling_minmax, reverse_scaling_adaptive) class ChemBERTaTrainer: def __init__(self, model_name: str, output_dir: str = "./results", batch_size: int = 16, learning_rate: float = 2e-5, epochs: int = 10, weight_decay: float = 0.01, warmup_steps: int = 500, scaling: bool = True, task_type: str = 'cls', missing_label_strategy: str = 'any', hidden_dim: List[int] = [128, 256, 128, 64], task_weights: Optional[Dict[str, float]] = None, early_stopping_patience: int = 10, data_type: str = "didb", load_type: str = "default", precision: str = '32', num_workers: int = 8, devices: int = 1, accelerator: str = 'gpu', strategy: str = 'auto', mode:str = 'train', grad_accum: int = 1, num_nodes: int = 1, scaler_type: str = 'power', loss_type: str = 'mse'): self.model_name = model_name self.output_dir = output_dir self.batch_size = batch_size self.learning_rate = learning_rate self.epochs = epochs self.weight_decay = weight_decay self.warmup_steps = warmup_steps self.task_type = task_type self.missing_label_strategy = missing_label_strategy self.scaling = scaling self.hidden_dim = hidden_dim self.early_stopping_patience = early_stopping_patience self.data_type = data_type self.load_type = load_type self.precision = precision self.num_workers = max(0, int(num_workers)) self.devices = devices self.accelerator = accelerator self.strategy = strategy self.mode = mode self.grad_accum = max(1, int(grad_accum)) self.num_nodes = max(1, int(num_nodes)) self.scaler_type = scaler_type self.loss_type = loss_type self.task_list = get_task_list(task_type) # 태스크별 유형을 지정 (classification, regression, multi_reg 지원) if task_type == 'cls': self.task_types = {task: 'classification' for task in self.task_list} elif task_type == 'reg': self.task_types = {task: 'regression' for task in self.task_list} elif task_type == 'multi_reg': self.task_types = {task: 'multi_layer_regression' for task in self.task_list} else: raise ValueError(f"알 수 없는 task_type: {task_type}") self.task_weights = task_weights or {task: 1.0 for task in self.task_list} self.model = None self.data_module = None self.trainer = None self.scaler_save_path = None def _build_deepspeed_config(self) -> Dict: cfg_path = Path(__file__).resolve().parent.parent / "config" / "deepspeed_zero2.json" with cfg_path.open() as f: base_cfg = json.load(f) if isinstance(self.devices, int): device_count = max(1, self.devices) elif isinstance(self.devices, (list, tuple, set)): device_count = max(1, len(self.devices)) else: device_count = 1 train_micro_batch = max(1, int(self.batch_size)) grad_accum = max(1, int(self.grad_accum)) train_batch = max(1, train_micro_batch * grad_accum * device_count) variable_cfg = { "train_batch_size": train_batch, "train_micro_batch_size_per_gpu": train_micro_batch } merged = base_cfg.copy() for key, value in variable_cfg.items(): if isinstance(value, dict) and key in merged and isinstance(merged[key], dict): merged[key].update(value) else: merged[key] = value return merged def setup(self, use_wandb=False, wandb_project=None, wandb_entity=None): if self.mode == 'train' or self.mode == 'test': # Load dataset if self.scaling: self.scaler_save_path = os.path.join(self.output_dir, "scaler") os.makedirs(self.scaler_save_path, exist_ok=True) else: self.scaler_save_path = None self.data_module = ChemMultiTaskDataModule( batch_size=self.batch_size, scaling=self.scaling, scaler_path=self.scaler_save_path, task_type=self.task_type, model_name=self.model_name, missing_label_strategy=self.missing_label_strategy, data_type=self.data_type, load_type=self.load_type, num_workers=self.num_workers, output_dir=self.output_dir, scaler_type=self.scaler_type ) self.data_module.setup() scaler_file_path = getattr(self.data_module, "scaler_file", None) vocab_dir = os.path.join(self.output_dir, "vocab") os.makedirs(vocab_dir, exist_ok=True) for task, vocab in self.data_module.all_vocabs.items(): with open(os.path.join(vocab_dir, f"{task}.json"), "w") as f: json.dump(vocab, f) num_classes = None if self.task_type == 'cls': num_classes = { task: len(self.data_module.all_vocabs.get(task, [0, 1])) for task in self.task_list } self.filter_cols = DIDB_FILTER_COLS self.model = ChemBERTaMultiTaskLightning( model_name=self.model_name, filter_cols=self.filter_cols, task_list=self.task_list, task_types=self.task_types, num_classes=num_classes, hidden_dim=self.hidden_dim, learning_rate=self.learning_rate, weight_decay=self.weight_decay, warmup_steps=self.warmup_steps, task_weights=self.task_weights, scaling=self.scaling, scaler_path=scaler_file_path, scaler_type=self.scaler_type, loss_type=self.loss_type ) if os.path.exists(self.output_dir): os.makedirs(self.output_dir, exist_ok=True) # Logger 선택 (기본: TensorBoard, 요청 시 WandB) if use_wandb: logger = WandbLogger(project=wandb_project, entity=wandb_entity) else: logger = TensorBoardLogger(save_dir=self.output_dir, name="logs") checkpoint_dir = os.path.join(self.output_dir, "checkpoints") os.makedirs(checkpoint_dir, exist_ok=True) callbacks = [ EarlyStopping(monitor='val_loss', patience=self.early_stopping_patience, mode='min'), ModelCheckpoint( monitor='val_loss', dirpath=checkpoint_dir, filename='best-{epoch:02d}-{val_loss:.4f}', save_top_k=1, mode='min', save_last=True, every_n_epochs=1 ), ModelCheckpoint( dirpath=checkpoint_dir, filename='epoch-{epoch:02d}', save_top_k=-1, every_n_epochs=1, monitor=None, save_last=False ), LearningRateMonitor(logging_interval='step') ] # 멀티-GPU 자동 지원: GPU가 2개 이상이면 DDPStrategy(find_unused_parameters=True) use_gpu = torch.cuda.is_available() num_gpus = torch.cuda.device_count() if use_gpu else 0 strategy = self.strategy if self.strategy == 'deepspeed': from pytorch_lightning.strategies import DeepSpeedStrategy strategy = DeepSpeedStrategy(config=self._build_deepspeed_config()) self.trainer = pl.Trainer( max_epochs=self.epochs, logger=logger, callbacks=callbacks, default_root_dir=self.output_dir, devices=self.devices, accelerator=self.accelerator, num_nodes=self.num_nodes, strategy=strategy, log_every_n_steps=10, accumulate_grad_batches=self.grad_accum, gradient_clip_val=1.0, precision=self.precision ) def train(self): if self.data_module is None or self.model is None: self.setup() train_loader, val_loader, test_loader = self.data_module.get_dataloaders() if self.trainer is not None: total_batches = self.trainer.estimated_stepping_batches print(f"[Trainer] Estimated stepping batches: {total_batches}") if self.model is None: raise ValueError("모델이 초기화되지 않았습니다.") if train_loader is None or val_loader is None: raise ValueError("학습 또는 검증 데이터로더가 없습니다.") self.trainer.fit(self.model, train_dataloaders=train_loader, val_dataloaders=val_loader) results = [] # Test loader 평가 if test_loader is not None: print("\n[Trainer] Test 데이터셋 평가 중...") test_result = self.trainer.test(self.model, dataloaders=test_loader) # 캐시된 결과 가져오기 (SMILES와 predictions 포함) if hasattr(self.model, 'test_results_cache') and self.model.test_results_cache: cached_results = self.model.test_results_cache print(f"[Debug] 캐시된 결과 keys: {list(cached_results.keys())}") # DataFrame 변환 및 저장 self.save_results_to_dataframe(cached_results, save_name="test_results.csv") results.append(cached_results) # 캐시 초기화 self.model.test_results_cache = None else: print("[Warning] test_results_cache가 비어있습니다!") if test_result and len(test_result) > 0: results.append(test_result[0]) # Validation loader 평가 if val_loader is not None: print("\n[Trainer] Validation 데이터셋 평가 중...") valid_result = self.trainer.test(self.model, dataloaders=val_loader) # 캐시된 결과 가져오기 (SMILES와 predictions 포함) if hasattr(self.model, 'test_results_cache') and self.model.test_results_cache: cached_results = self.model.test_results_cache print(f"[Debug] 캐시된 결과 keys: {list(cached_results.keys())}") # DataFrame 변환 및 저장 self.save_results_to_dataframe(cached_results, save_name="valid_results.csv") results.append(cached_results) # 캐시 초기화 self.model.test_results_cache = None else: print("[Warning] test_results_cache가 비어있습니다!") if valid_result and len(valid_result) > 0: results.append(valid_result[0]) return results if results else None def predict(self, smiles_input, path=None): # model mode setting self.model.eval() tokenizer = AutoTokenizer.from_pretrained(self.model_name) device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") self.model.to(device) # Reverse scaling logic scaler_path = None scaler_filename = SCALER_FILE_NAME.get(self.scaler_type, 'scaler_zscore.csv') if path and os.path.exists(os.path.join(path, scaler_filename)): scaler_path = os.path.join(path, scaler_filename) if isinstance(smiles_input, str): pred_output = {} with torch.no_grad(): enc = tokenizer([smiles_input], padding='max_length', max_length=510, truncation=True, return_tensors='pt') input_ids = enc['input_ids'].to(device) attention_mask = enc['attention_mask'].to(device) outputs = self.model(input_ids, attention_mask) pred_output = {prop: float(outputs[prop][0].detach().cpu().numpy()) for prop in DIDB_FILTER_COLS} # Reverse scaling logic try: if scaler_path and os.path.exists(scaler_path): df = pd.DataFrame({k: [v] for k, v in pred_output.items()}) if self.scaler_type == 'power': df = reverse_scaling_power(df, scaler_path) elif self.scaler_type in 'adapt': df = reverse_scaling_adaptive(df, scaler_path) elif self.scaler_type == 'minmax': df = reverse_scaling_minmax(df, scaler_path) else: df = reverse_scaling(df, scaler_path) return {k: (float(df[k].iloc[0]) if k in df.columns and pd.notna(df[k].iloc[0]) else None) for k in DIDB_FILTER_COLS} except Exception as e: print(f"역스케일링 중 오류 발생: {e}. 원본 스케일링된 값을 반환합니다.") return pred_output elif isinstance(smiles_input, list): batch_size = min(len(smiles_input), 16) predictions = [] with torch.no_grad(): for i in range(0, len(smiles_input), batch_size): batch = smiles_input[i:i+batch_size] enc = tokenizer(batch, padding='max_length', max_length=510, truncation=True, return_tensors='pt') input_ids, attention_mask = enc['input_ids'].to(device), enc['attention_mask'].to(device) outputs = self.model(input_ids, attention_mask) batch_preds = {} for task in self.task_list: if self.task_types[task] == 'classification': probs = torch.nn.functional.softmax(outputs[task], dim=1) preds = torch.argmax(probs, dim=1) batch_preds[task] = { 'probabilities': probs.cpu().numpy(), 'predictions': preds.cpu().numpy() } else: batch_preds[task] = outputs[task].cpu().numpy() predictions.append(batch_preds) merged = {} # 1) 먼저 모든 태스크에 대해 배치 예측을 병합 for task in self.task_list: if self.task_types[task] == 'classification': merged[task] = { 'probabilities': np.concatenate([p[task]['probabilities'] for p in predictions], axis=0), 'predictions': np.concatenate([p[task]['predictions'] for p in predictions], axis=0) } else: merged[task] = np.concatenate([p[task] for p in predictions], axis=0) # 2) 회귀 태스크에 한해 역스케일링을 한 번만 수행 if self.scaling and scaler_path: reg_tasks = [t for t, ttype in self.task_types.items() if 'regression' in ttype] if reg_tasks: df_data = {t: np.array(merged[t]).reshape(-1) for t in reg_tasks} scaled_preds_df = pd.DataFrame(df_data) try: if self.scaler_type == 'power': unscaled_preds_df = reverse_scaling_power(scaled_preds_df, scaler_path) elif self.scaler_type in 'adapt': unscaled_preds_df = reverse_scaling_adaptive(scaled_preds_df, scaler_path) elif self.scaler_type == 'minmax': unscaled_preds_df = reverse_scaling_minmax(scaled_preds_df, scaler_path) else: unscaled_preds_df = reverse_scaling(scaled_preds_df, scaler_path) for t in reg_tasks: merged[t] = unscaled_preds_df[t].to_numpy() except Exception as e: print(f"역스케일링 중 오류 발생: {e}. 원본 스케일링된 값을 반환합니다.") return merged def save_model(self, path=None): path = path or os.path.join(self.output_dir, "final_model") os.makedirs(path, exist_ok=True) best_ckpt = None for cb in getattr(self.trainer, "callbacks", []): if isinstance(cb, ModelCheckpoint): if cb.best_model_path and os.path.exists(cb.best_model_path): best_ckpt = cb.best_model_path break if best_ckpt: shutil.copy(best_ckpt, os.path.join(path, "best_model.ckpt")) else: print("경고: 최적 모델 체크포인트를 찾을 수 없습니다.") torch.save(self.model.state_dict(), os.path.join(path, "model_weights.pt")) if self.scaling and self.scaler_save_path and os.path.exists(self.scaler_save_path): import shutil as _sh # scaler_type에 따라 해당하는 scaler 파일만 복사 if self.scaler_type == 'adapt': scaler_filename = 'scaler_adapt.csv' elif self.scaler_type == 'minmax': scaler_filename = 'scaler_minmax.csv' elif self.scaler_type == 'power': scaler_filename = 'scaler_power_config.csv' else: # zscore (default) scaler_filename = 'scaler_config.csv' scaler_src = os.path.join(self.scaler_save_path, scaler_filename) if os.path.exists(scaler_src): scaler_dst = os.path.join(path, scaler_filename) _sh.copyfile(scaler_src, scaler_dst) print(f"✅ Scaler 파일 저장 완료: {scaler_filename}") else: print(f"⚠️ 경고: Scaler 파일을 찾을 수 없습니다: {scaler_src}") if hasattr(self.model, 'hparams') and hasattr(self.model.hparams, 'to_dict'): import json with open(os.path.join(path, "model_config.json"), 'w') as f: json.dump(dict(self.model.hparams), f) # hparam.json 저장: 트레이너 초기화에 사용된 주요 하이퍼파라미터 기록 try: import json as _json hparams = { "model_name": self.model_name, "output_dir": self.output_dir, "batch_size": self.batch_size, "learning_rate": self.learning_rate, "epochs": self.epochs, "weight_decay": self.weight_decay, "warmup_steps": self.warmup_steps, "scaling": self.scaling, "task_type": self.task_type, "missing_label_strategy": self.missing_label_strategy, "hidden_dim": self.hidden_dim, "early_stopping_patience": self.early_stopping_patience, "data_type": self.data_type, "num_workers": self.num_workers, "precision": self.precision, "devices": self.devices, "accelerator": self.accelerator, } with open(os.path.join(path, 'hparam.json'), 'w') as fp: _json.dump(hparams, fp, indent=2) except Exception as e: print(f"하이퍼파라미터 저장 중 경고: {e}") def save_results_to_dataframe(self, results: Dict, save_name: str) -> Dict: """ 테스트 결과를 DataFrame 형식으로 변환하고 CSV로 저장 Args: results: on_test_epoch_end에서 반환된 결과 딕셔너리 save_name: 저장할 CSV 파일명 Returns: DataFrame이 추가된 결과 딕셔너리 """ print(f"\n[DataFrame 변환] 시작...") print(f" - Results keys: {list(results.keys())}") all_smiles = results.get('smiles', []) all_predictions = results.get('predictions', {}) print(f" - SMILES 개수: {len(all_smiles)}") print(f" - Predictions tasks: {list(all_predictions.keys())}") # DataFrame 변환을 위한 구조화된 데이터 생성 df_data = {'SMILES': all_smiles} # 각 task별 predictions를 columns로 추가 for task in self.task_list: if task in all_predictions and all_predictions[task].get('preds'): preds_list = all_predictions[task]['preds'] print(f" - Task '{task}': {len(preds_list)}개 배치 (텐서)") # Tensor 리스트를 numpy array로 변환 if len(preds_list) > 0: if isinstance(preds_list[0], torch.Tensor): # 여러 배치의 텐서를 하나로 합치기 preds_concat = torch.cat([p.unsqueeze(0) if p.dim() == 0 else p for p in preds_list], dim=0) preds_np = preds_concat.numpy() else: preds_np = np.concatenate([np.atleast_1d(p) for p in preds_list], axis=0) df_data[task] = preds_np print(f" - 예측값 shape: {preds_np.shape}, 길이: {len(preds_np)}") # Labels도 포함 (비교용) if all_predictions[task].get('labels'): labels_list = all_predictions[task]['labels'] if len(labels_list) > 0: if isinstance(labels_list[0], torch.Tensor): # 여러 배치의 텐서를 하나로 합치기 labels_concat = torch.cat([l.unsqueeze(0) if l.dim() == 0 else l for l in labels_list], dim=0) labels_np = labels_concat.numpy() else: labels_np = np.concatenate([np.atleast_1d(l) for l in labels_list], axis=0) df_data[f'{task}_label'] = labels_np print(f" - 라벨 shape: {labels_np.shape}, 길이: {len(labels_np)}") # DataFrame 생성 print(f"\n[DataFrame 생성] Columns: {list(df_data.keys())}") df_results = pd.DataFrame(df_data) print(f" - Shape: {df_results.shape}") print(f" - 첫 3행:\n{df_results.head(3)}") # 저장 data_save_path = os.path.join(self.output_dir, "valid_test_split/") os.makedirs(data_save_path, exist_ok=True) output_path = os.path.join(data_save_path, save_name) df_results.to_csv(output_path, index=False) print(f"\n✅ 예측 결과 저장 완료: {output_path}") return results def load_model(self, path): if self.model is None: self.setup() import json vocab_dir = os.path.join(path, "vocab") if os.path.exists(vocab_dir): for file in os.listdir(vocab_dir): if file.endswith(".json"): task = file.replace(".json", "") with open(os.path.join(vocab_dir, file), "r") as f: self.data_module.all_vocabs[task] = json.load(f) ckpt_path = os.path.join(path, "best_model.ckpt") weights_path = os.path.join(path, "model_weights.pt") if os.path.exists(ckpt_path): self.model = ChemBERTaMultiTaskLightning.load_from_checkpoint(ckpt_path) elif os.path.exists(weights_path): self.model.load_state_dict(torch.load(weights_path, map_location='cpu')) return self.model