Spaces:
Sleeping
Sleeping
| import logging | |
| import torch | |
| from pathlib import Path | |
| from sklearn.metrics import accuracy_score, recall_score, f1_score | |
| from tqdm import tqdm | |
| from config import Config | |
| from dataset import MelanomaDataLoaders | |
| from model import MetadataMelanomaModel | |
| from evaluate import Evaluator | |
| from file_io_manager import FileIOManager | |
| logger = logging.getLogger(__name__) | |
| class Trainer: | |
| def __init__(self): | |
| self._device = Config.get_training_config()['device'] | |
| loaders = MelanomaDataLoaders() | |
| self._train_loader = loaders.get_train_loader() | |
| self._val_loader = loaders.get_val_loader() | |
| num_metadata_features = loaders.num_metadata_features | |
| self._model = MetadataMelanomaModel.build(num_metadata_features=num_metadata_features) | |
| self._criterion = MetadataMelanomaModel.get_criterion() | |
| self._optimizer = MetadataMelanomaModel.get_optimizer(self._model) | |
| self._scheduler = MetadataMelanomaModel.get_scheduler(self._optimizer) | |
| self._run_name = self._build_run_name() | |
| self._io = FileIOManager.for_run(self._output_name()) | |
| self._io.save_preprocessor(loaders.preprocessor) | |
| self._preprocessor = loaders.preprocessor | |
| def _output_name() -> str: | |
| exp = Config.get_experiment() | |
| if exp is not None: | |
| return f"experiment{exp}" | |
| return Config.MODEL_NAME | |
| def _build_run_name(self) -> str: | |
| cfg = Config.get_training_config() | |
| base = ( | |
| f"{Config.MODEL_NAME}_Meta" | |
| f"_LR{cfg['learning_rate']}_BS{cfg['batch_size']}_Ep{cfg['num_epochs']}" | |
| ) | |
| aug = Config.get_augmentation_config() | |
| if 'affine' in aug: | |
| suffix = "_AffAug" | |
| elif aug.get('random_erasing_prob', 0) > 0: | |
| suffix = "_EraseAug" | |
| else: | |
| suffix = "_BasicAug" | |
| return base + suffix | |
| def _train_epoch(self) -> tuple: | |
| self._model.train() | |
| total_loss = 0.0 | |
| all_preds: list[int] = [] | |
| all_labels: list[int] = [] | |
| for images, metadata, labels in tqdm(self._train_loader, desc="Training"): | |
| images = images.to(self._device) | |
| metadata = metadata.to(self._device) | |
| labels = labels.to(self._device) | |
| self._optimizer.zero_grad() | |
| outputs = self._model(images, metadata) | |
| loss = self._criterion(outputs, labels.unsqueeze(1).float()) | |
| loss.backward() | |
| self._optimizer.step() | |
| total_loss += loss.detach().item() | |
| preds = (torch.sigmoid(outputs.detach()) > 0.5).int() | |
| all_preds.extend(preds.cpu().numpy().flatten()) | |
| all_labels.extend(labels.cpu().numpy().flatten()) | |
| return total_loss / len(self._train_loader), all_preds, all_labels | |
| def _save_checkpoint(self, epoch: int) -> str: | |
| path = self._io.save_checkpoint(self._model, self._run_name, epoch) | |
| return str(path) | |
| def _final_evaluation(self, best_model_path: str) -> None: | |
| logger.info("Loading best model for final evaluation plots: %s", best_model_path) | |
| final_model = MetadataMelanomaModel.build(num_metadata_features=self._model.num_metadata_features) | |
| self._io.load_checkpoint(final_model, best_model_path, map_location=self._device) | |
| final_model = final_model.to(self._device) | |
| evaluator = Evaluator(final_model, MetadataMelanomaModel.get_criterion(), io=self._io) | |
| _, preds, labels, probs = evaluator.evaluate( | |
| self._val_loader, use_tta=Config.get_evaluation_config()['tta_enabled'] | |
| ) | |
| evaluator.plot_roc_curve(labels, probs) | |
| evaluator.plot_confusion_matrix(labels, preds) | |
| evaluator.compute_ood_stats(self._val_loader) | |
| if not final_model._image_only: | |
| evaluator.plot_shap(self._val_loader, self._preprocessor._feature_cols) | |
| def train(self) -> None: | |
| num_epochs = Config.get_training_config()['num_epochs'] | |
| best_val_f1 = 0.0 | |
| best_path = None | |
| best_epoch = 0 | |
| for epoch in range(num_epochs): | |
| current_lr = self._optimizer.param_groups[0]['lr'] | |
| logger.info("Epoch %d/%d | LR: %s", epoch + 1, num_epochs, current_lr) | |
| train_loss, train_preds, train_labels = self._train_epoch() | |
| val_loss, val_preds, val_labels, _ = Evaluator(self._model, self._criterion).evaluate( | |
| self._val_loader, use_tta=Config.get_evaluation_config()['tta_enabled'] | |
| ) | |
| train_acc = accuracy_score(train_labels, train_preds) | |
| train_recall = recall_score(train_labels, train_preds) | |
| train_f1 = f1_score(train_labels, train_preds) | |
| val_acc = accuracy_score(val_labels, val_preds) | |
| val_recall = recall_score(val_labels, val_preds) | |
| val_f1 = f1_score(val_labels, val_preds) | |
| logger.info("Train: Loss=%.4f | Acc=%.4f | Recall=%.4f | F1=%.4f", train_loss, train_acc, train_recall, train_f1) | |
| logger.info("Val: Loss=%.4f | Acc=%.4f | Recall=%.4f | F1=%.4f", val_loss, val_acc, val_recall, val_f1) | |
| self._io.append_epoch_metrics({ | |
| "epoch": epoch + 1, | |
| "experiment": self._output_name(), | |
| "learning_rate": current_lr, | |
| "train_loss": train_loss, "train_acc": train_acc, | |
| "train_recall": train_recall, "train_f1": train_f1, | |
| "val_loss": val_loss, "val_acc": val_acc, | |
| "val_recall": val_recall, "val_f1": val_f1, | |
| }) | |
| if val_f1 > best_val_f1: | |
| best_val_f1 = val_f1 | |
| prev_path = best_path | |
| best_epoch = epoch + 1 | |
| best_path = self._save_checkpoint(best_epoch) | |
| if prev_path: | |
| Path(prev_path).unlink(missing_ok=True) | |
| logger.info("Removed previous checkpoint: %s", prev_path) | |
| logger.info("New best model saved: %s (Val F1: %.4f)", best_path, best_val_f1) | |
| if self._scheduler: | |
| self._scheduler.step() | |
| if best_path: | |
| self._io.save_gradcam_checkpoint(self._model) | |
| logger.info("Inference checkpoint saved: %s", self._io.gradcam_checkpoint_path()) | |
| self._final_evaluation(best_path) | |
| else: | |
| logger.warning("No best model saved; skipping final plots.") | |