diff --git a/tda_ml/losses.py b/tda_ml/losses.py new file mode 100644 index 0000000..8a07423 --- /dev/null +++ b/tda_ml/losses.py @@ -0,0 +1,198 @@ +import logging + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch_topological.nn import VietorisRipsComplex, WassersteinDistance + +from tda_ml.distance_backend import compute_distance_matrix_batch +from tda_ml.topology import compute_anisotropic_distance_matrix + +# Distance-mode aliases for topology loss (kept for config/test compatibility). +logger = logging.getLogger(__name__) + +DISTANCE_MODE_MAHALANOBIS = "mahalanobis" +DISTANCE_MODE_ELLPHI = "ellphi" + + +def normalize_topo_distance_mode(mode: str) -> str: + m = str(mode).strip().lower() + if m == "mahalanobis": + return DISTANCE_MODE_MAHALANOBIS + if m == "ellphi": + return DISTANCE_MODE_ELLPHI + raise ValueError(f"Unknown topo distance mode: {mode!r}") + + +def compute_topo_distance_matrix( + points: torch.Tensor, + params: torch.Tensor, + *, + distance_mode: str = "mahalanobis", + ellphi_backend: str = "auto", +) -> torch.Tensor: + """ + Shared topology-loss entrypoint: map batched points/ellipse params to ``(B,N,N)`` distances. + + ``distance_mode`` is ``mahalanobis`` or ``ellphi``. + When ``ellphi_backend='auto'``, a differentiable ellphi path is preferred and + falls back to NumPy when unavailable. + """ + backend = normalize_topo_distance_mode(distance_mode) + eb = str(ellphi_backend).strip().lower() + ellphi_diff = eb in ("auto", "torch", "grad", "differentiable", "1", "true", "yes") + return compute_distance_matrix_batch( + points, + params, + probs=None, + symmetrize="max", + backend=backend, + ellphi_differentiable=ellphi_diff, + ) + + +def mahalanobis_distance_matrix_batched(points: torch.Tensor, params: torch.Tensor) -> torch.Tensor: + """Mahalanobis-style batched distance matrix without probability weighting.""" + return compute_anisotropic_distance_matrix( + points, params, probs=None, symmetrize="max" + ) + + +class ClassificationLoss(nn.Module): + """ + Standard Binary Cross-Entropy with Logits for inlier/outlier classification. + """ + def __init__(self, pos_weight=None): + super().__init__() + self.loss_fn = nn.BCEWithLogitsLoss(pos_weight=pos_weight) + + def forward(self, logits, labels): + return self.loss_fn(logits.squeeze(-1), labels.float()) + +class SizeRegularizationLoss(nn.Module): + """ + Penalizes the size of estimated ellipses to prevent over-expansion. + + Formula: L = lambda_major * a^2 + lambda_minor * b^2 + """ + def __init__(self, w_major=0.1, w_minor=0.1): + super().__init__() + self.w_major = w_major + self.w_minor = w_minor + + def forward(self, params): + axes = params[..., 0:2] + major_axis = axes.max(dim=-1)[0] + minor_axis = axes.min(dim=-1)[0] + loss = (self.w_major * (major_axis**2) + self.w_minor * (minor_axis**2)).mean() + return loss + +class AnisotropyPenaltyLoss(nn.Module): + """ + Prevents ellipses from becoming too elongated by penalizing high aspect ratios. + + Modes: + - linear: Penalizes aspect ratio (major/minor) directly. + - barrier: Penalizes aspect ratio squared only above a certain threshold. + """ + def __init__(self, weight=0.01, mode='linear', barrier_threshold=6.0): + super().__init__() + self.weight = weight + self.mode = mode + self.barrier_threshold = barrier_threshold + + def forward(self, params): + if abs(self.weight) < 1e-9: + return torch.tensor(0.0, device=params.device) + + axes = params[..., 0:2] + major_axis = axes.max(dim=-1)[0] + minor_axis = axes.min(dim=-1)[0] + + aspect_ratios = major_axis / (minor_axis + 1e-6) + + if self.mode == 'barrier': + barrier_term = F.relu(aspect_ratios - self.barrier_threshold).pow(2).mean() + loss = 10.0 * barrier_term + else: + loss = aspect_ratios.mean() + + return self.weight * loss + +class TopologicalLoss(nn.Module): + """ + Computes the Topological Loss between the predicted anisotropic filtration + and the clean ground truth persistence diagram using Wasserstein distance. + + distance_backend: + - ``mahalanobis``: differentiable anisotropic distance + - ``ellphi``: tangency distance. With ``ellphi_differentiable=True``, + gradients can flow to ellipse parameters via ``ellphi.grad`` (numerical + singularities may still produce NaN/zero gradients). + + ellphi_differentiable: + If ``False``, use NumPy ``EllipseCloud.pdist_tangency`` only (no gradients). + """ + def __init__( + self, + weight=0.1, + distance_backend: str = "mahalanobis", + ellphi_differentiable: bool = True, + ): + super().__init__() + self.weight = weight + self.distance_backend = distance_backend.lower().strip() + self.ellphi_differentiable = ellphi_differentiable + self.vr_complex = VietorisRipsComplex(dim=1) + self.wasserstein = WassersteinDistance(q=2) + + def forward(self, points, params, logits, clean_pd_info): + if self.weight <= 0: + return torch.tensor(0.0, device=points.device) + + batch_size = points.shape[0] + probs_outlier = torch.sigmoid(logits).squeeze(-1) + + D_prime = compute_distance_matrix_batch( + points, + params, + probs=probs_outlier, + symmetrize="max", + backend=self.distance_backend, + ellphi_differentiable=self.ellphi_differentiable, + ) + + total_loss = 0.0 + valid_samples = 0 + topo_failures: list[tuple[int, str]] = [] + + for i in range(batch_size): + d_mat = D_prime[i] + try: + pd_pred_info = self.vr_complex(d_mat, treat_as_distances=True) + loss_sample = self.wasserstein(pd_pred_info, clean_pd_info[i]) ** 2 + + if not torch.isnan(loss_sample): + total_loss += loss_sample + valid_samples += 1 + else: + topo_failures.append((i, "nan or inf Wasserstein loss")) + except Exception as exc: + topo_failures.append((i, f"{type(exc).__name__}: {exc}")) + continue + + if topo_failures: + batch_skipped = batch_size - valid_samples + first_i, first_msg = topo_failures[0] + logger.warning( + "TopologicalLoss: skipped %d/%d batch items (first batch_index=%s: %s)", + batch_skipped, + batch_size, + first_i, + first_msg, + ) + + if valid_samples == 0: + return torch.tensor(0.0, device=points.device, requires_grad=True) + + return self.weight * (total_loss / valid_samples) diff --git a/tda_ml/main.py b/tda_ml/main.py new file mode 100644 index 0000000..f79de25 --- /dev/null +++ b/tda_ml/main.py @@ -0,0 +1,294 @@ +import argparse +import csv +import json +import logging +import os + +import torch + +from tda_ml.checkpoint_io import extract_model_state_dict, load_torch_checkpoint +from tda_ml.config import deep_update, load_config, model_kwargs_from_config +from tda_ml.data_loader import NoisyMNISTDataset, create_data_loader +from tda_ml.models import AnisotropicOutlierClassifier +from tda_ml.seed_utils import set_global_seed +from tda_ml.runtime_profile import build_runtime_profile +from tda_ml.trainer import Trainer + +logger = logging.getLogger(__name__) + + +def _resolve_dataloader_settings(config, device): + data_cfg = config["data"] + cpu_threads = os.cpu_count() or 8 + default_workers = min(16, max(4, cpu_threads // 2)) + + num_workers = int(data_cfg.get("num_workers", default_workers)) + num_workers = max(0, min(num_workers, cpu_threads)) + + pin_memory = bool(data_cfg.get("pin_memory", device.type == "cuda")) + persistent_workers = bool(data_cfg.get("persistent_workers", num_workers > 0)) + prefetch_factor = data_cfg.get("prefetch_factor", 4 if num_workers > 0 else None) + + if num_workers == 0: + persistent_workers = False + prefetch_factor = None + + return num_workers, pin_memory, persistent_workers, prefetch_factor + + +def _configure_torch_runtime(config, device): + perf_cfg = config.get("performance", {}) + if device.type != "cuda": + return + + if perf_cfg.get("enable_tf32", True): + torch.backends.cuda.matmul.allow_tf32 = True + torch.backends.cudnn.allow_tf32 = True + torch.set_float32_matmul_precision(perf_cfg.get("matmul_precision", "high")) + torch.backends.cudnn.benchmark = bool(perf_cfg.get("cudnn_benchmark", True)) + +def main(config_name=None, config=None, trial=None, config_overrides=None): + logging.basicConfig( + level=logging.INFO, + format="%(levelname)s %(name)s: %(message)s", + ) + if config is None: + if config_name is None: + raise ValueError("Either config_name or config must be provided") + config = load_config(config_name) + + if config_overrides: + deep_update(config, config_overrides) + + logger.info("Loaded config: %s", config["meta"].get("config_id", "unknown")) + + if config.get('device') and config['device'] != 'auto': + device = torch.device(config['device']) + elif torch.cuda.is_available(): + device = torch.device('cuda') + elif torch.backends.mps.is_available(): + device = torch.device('mps') + else: + device = torch.device('cpu') + + import datetime + timestamp = datetime.datetime.now().strftime('%Y%m%d_%H%M%S') + config_id = config['meta'].get('config_id', 'unknown') + + # --- Directory Setup --- + if 'outputs' in config: + base_dir = config['outputs'].get('base_dir', 'outputs') + run_dir_name = f"{config_id}_{timestamp}" + run_dir = os.path.join(base_dir, run_dir_name) + + # Consistent directory names across the project + config['outputs']['log_dir'] = os.path.join(run_dir, 'logs') + config['outputs']['image_dir'] = os.path.join(run_dir, 'images') + + # Note: Trainer also calls makedirs to ensure robustness + os.makedirs(config['outputs']['log_dir'], exist_ok=True) + os.makedirs(config['outputs']['image_dir'], exist_ok=True) + + logger.info("Project structure created at: %s", run_dir) + + data_cfg = config['data'] + seed = data_cfg.get('seed', 42) + deterministic_algorithms = bool(config.get("reproducibility", {}).get("deterministic_algorithms", False)) + set_global_seed(seed, deterministic_algorithms=deterministic_algorithms) + logger.info( + "Global seed initialized: seed=%s, deterministic_algorithms=%s", + seed, + deterministic_algorithms, + ) + train_size = data_cfg.get('train_size', 4500) + val_size = data_cfg.get('val_size', 500) + test_size = data_cfg.get('test_size', 1000) + + generator = torch.Generator().manual_seed(seed) + full_train_indices = torch.randperm(60000, generator=generator)[:train_size + val_size] + + train_indices = full_train_indices[:train_size] + val_indices = full_train_indices[train_size:] + + num_workers, use_pin_memory, persistent_workers, prefetch_factor = _resolve_dataloader_settings(config, device) + _configure_torch_runtime(config, device) + logger.info( + "DataLoader settings: workers=%s, pin_memory=%s, persistent_workers=%s, prefetch_factor=%s", + num_workers, + use_pin_memory, + persistent_workers, + prefetch_factor, + ) + + train_dataset = NoisyMNISTDataset( + root='./data', train=True, max_points=data_cfg['max_points'], + num_outliers=data_cfg['num_outliers'], indices=train_indices, + deterministic=True, noise_seed=seed + ) + + data_loader = create_data_loader( + train_dataset, + batch_size=data_cfg['batch_size'], + shuffle=True, + num_workers=num_workers, + pin_memory=use_pin_memory, + persistent_workers=persistent_workers, + prefetch_factor=prefetch_factor, + ) + + val_dataset = NoisyMNISTDataset( + root='./data', train=True, max_points=data_cfg['max_points'], + num_outliers=data_cfg['num_outliers'], indices=val_indices, + deterministic=True, noise_seed=seed + ) + + val_loader = create_data_loader( + val_dataset, + batch_size=data_cfg['batch_size'], + shuffle=False, + num_workers=num_workers, + pin_memory=use_pin_memory, + persistent_workers=persistent_workers, + prefetch_factor=prefetch_factor, + ) + + test_indices = torch.randperm(10000, generator=generator)[:test_size] + test_dataset = NoisyMNISTDataset( + root='./data', train=False, max_points=data_cfg['max_points'], + num_outliers=data_cfg['num_outliers'], indices=test_indices, + deterministic=True, noise_seed=seed + ) + + test_loader = create_data_loader( + test_dataset, + batch_size=data_cfg["batch_size"], + shuffle=False, + num_workers=num_workers, + pin_memory=use_pin_memory, + persistent_workers=persistent_workers, + prefetch_factor=prefetch_factor, + ) + + model = AnisotropicOutlierClassifier(**model_kwargs_from_config(config)) + model.to(device) + + trainer = Trainer(model, config, device=device) + + init_checkpoint = config.get('init_checkpoint') + if init_checkpoint and os.path.exists(init_checkpoint): + logger.info("Loading initial weights from %s", init_checkpoint) + checkpoint = load_torch_checkpoint(init_checkpoint, map_location=device) + model.load_state_dict(extract_model_state_dict(checkpoint), strict=False) + elif init_checkpoint: + logger.warning("Initial checkpoint not found at %s", init_checkpoint) + + log_dir = config['outputs']['log_dir'] + metrics_path = os.path.join(log_dir, 'metrics.csv') + runtime_profile = build_runtime_profile( + config=config, + device=device, + num_workers=num_workers, + pin_memory=use_pin_memory, + persistent_workers=persistent_workers, + prefetch_factor=prefetch_factor, + use_amp_effective=trainer.use_amp, + amp_dtype_effective=str(trainer.amp_dtype).replace("torch.", ""), + ) + runtime_profile_path = os.path.join(log_dir, "runtime_profile.json") + with open(runtime_profile_path, "w", encoding="utf-8") as f: + json.dump(runtime_profile, f, ensure_ascii=True, indent=2) + logger.info("Runtime profile saved: %s", runtime_profile_path) + with open(metrics_path, 'w', newline='') as f: + writer = csv.writer(f) + writer.writerow(['epoch', 'train_loss', 'train_class_loss', 'train_topo_loss', 'train_aniso_loss', 'train_size_loss', 'val_loss', 'val_recall', 'val_mcc', 'val_aniso', 'val_size']) + + # --- Training Loop --- + epochs = config['training']['epochs'] + save_every = config['outputs'].get('save_every', 10) + best_val_mcc = -1.0 + + for epoch in range(1, epochs + 1): + # res returns (avg_loss, class_loss, topo_loss, aniso_loss, size_loss, ...) + res = trainer.train_epoch(data_loader, epoch) + val_res = trainer.validate(val_loader) + + val_mcc = val_res[4] # MCC is at index 4 + print(f"Epoch {epoch}: Val MCC={val_mcc:.4f}, Aniso={val_res[5]:.4f}") + + # Save best model logic + if val_mcc > best_val_mcc: + best_val_mcc = val_mcc + best_model_path = os.path.join(run_dir, 'best_model.pth') + torch.save({ + 'epoch': epoch, + 'model_state_dict': model.state_dict(), + 'val_mcc': val_mcc, + }, best_model_path) + print(f"Saved best model (MCC: {val_mcc:.4f}) to {best_model_path}") + + if epoch % save_every == 0: + checkpoint_path = os.path.join(run_dir, f'checkpoint_epoch_{epoch}.pth') + torch.save({ + 'epoch': epoch, + 'model_state_dict': model.state_dict(), + 'val_mcc': val_mcc, + }, checkpoint_path) + print(f"Saved checkpoint to {checkpoint_path}") + + with open(metrics_path, 'a', newline='') as f: + writer = csv.writer(f) + writer.writerow([epoch, res[0], res[1], res[2], res[3], res[4], val_res[0], val_res[1], val_res[4], val_res[5], val_res[6]]) + + final_model_path = os.path.join(run_dir, 'final_model.pth') + torch.save(model.state_dict(), final_model_path) + print(f"Saved final model to {final_model_path}") + + test_res = trainer.validate(test_loader) + ( + test_loss, + test_recall, + test_specificity, + test_gmean, + test_mcc, + test_aniso, + test_size, + ) = test_res + test_metrics = { + "test_loss": test_loss, + "test_recall": test_recall, + "test_specificity": test_specificity, + "test_gmean": test_gmean, + "test_mcc": test_mcc, + "test_aniso": test_aniso, + "test_size": test_size, + } + test_metrics_path = os.path.join(log_dir, "test_metrics.json") + with open(test_metrics_path, "w", encoding="utf-8") as f: + json.dump(test_metrics, f, ensure_ascii=True, indent=2) + logger.info( + "Test (held-out): MCC=%.4f, recall=%.4f, specificity=%.4f, G-mean=%.4f, loss=%.4f", + test_mcc, + test_recall, + test_specificity, + test_gmean, + test_loss, + ) + logger.info("Test metrics saved: %s", test_metrics_path) + + return { + "val_gmean": val_res[3], + "val_mcc": val_res[4], + "run_dir": run_dir, + } + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument('--config', type=str, default='dev', help='Config path') + parser.add_argument('--init_checkpoint', type=str, default=None, help='Path to initial checkpoint') + args = parser.parse_args() + + config_overrides = {} + if args.init_checkpoint: + config_overrides['init_checkpoint'] = args.init_checkpoint + + main(config_name=args.config, config_overrides=config_overrides) diff --git a/tda_ml/trainer.py b/tda_ml/trainer.py new file mode 100644 index 0000000..72b5775 --- /dev/null +++ b/tda_ml/trainer.py @@ -0,0 +1,319 @@ +import logging +import os + +import torch +from torch.optim import Adam +from torch_topological.nn import VietorisRipsComplex +from tda_ml.visualization import visualize +from tda_ml.losses import ( + ClassificationLoss, + TopologicalLoss, + SizeRegularizationLoss, + AnisotropyPenaltyLoss +) +from tda_ml.metrics import compute_recall_specificity_gmean_mcc +import tqdm +from sklearn.metrics import f1_score, precision_score, recall_score + +logger = logging.getLogger(__name__) + + +class Trainer: + def __init__(self, model, config, device=None, trial=None): + self.model = model + self.config = config + self.trial = trial + if device: + self.device = device + elif torch.cuda.is_available(): + self.device = torch.device('cuda') + elif torch.backends.mps.is_available(): + self.device = torch.device('mps') + else: + self.device = torch.device('cpu') + self.model.to(self.device) + + self.optimizer = Adam(model.parameters(), lr=config['training']['lr']) + training_cfg = config.get('training', {}) + perf_cfg = config.get('performance', {}) + self.use_amp = ( + self.device.type == "cuda" + and bool(training_cfg.get("use_amp", perf_cfg.get("use_amp", True))) + ) + self.autocast_device_type = "cuda" if self.device.type == "cuda" else "cpu" + amp_dtype_name = str(training_cfg.get("amp_dtype", perf_cfg.get("amp_dtype", "float16"))).lower() + self.amp_dtype = torch.float16 if amp_dtype_name == "float16" else torch.bfloat16 + self.scaler = torch.amp.GradScaler("cuda", enabled=self.use_amp) + + loss_cfg = config.get('loss', {}) + + self.lambda_class = loss_cfg.get('w_class', training_cfg.get('lambda_class', 1.0)) + self.lambda_topo = loss_cfg.get('w_topo', training_cfg.get('lambda_topo', 0.1)) + self.lambda_aniso = loss_cfg.get('w_aniso', training_cfg.get('lambda_aniso', 0.01)) + + size_default = loss_cfg.get( + "w_size", training_cfg.get("lambda_size", 0.1) + ) + self.lambda_major = training_cfg.get("lambda_major", size_default) + self.lambda_minor = training_cfg.get("lambda_minor", size_default) + + self.aniso_mode = loss_cfg.get("aniso_mode", training_cfg.get("aniso_mode", "linear")) + logger.info("Anisotropy penalty mode: %s", self.aniso_mode) + + pos_weight_val = config.get('loss', {}).get('pos_weight', 1.0) + pos_weight = torch.tensor([pos_weight_val], device=self.device) if pos_weight_val != 1.0 else None + + # Initialize Losses + self.class_loss_fn = ClassificationLoss(pos_weight=pos_weight) + _topo = config.get("model", {}).get("topology_loss", {}) + self.distance_backend = _topo.get("distance_backend", "mahalanobis") + self.ellphi_differentiable = _topo.get("ellphi_differentiable", True) + logger.info( + "Topological distance backend: %s%s", + self.distance_backend, + ( + f" (ellphi_differentiable={self.ellphi_differentiable})" + if self.distance_backend == "ellphi" + else "" + ), + ) + self.topo_loss_fn = TopologicalLoss( + weight=self.lambda_topo, + distance_backend=self.distance_backend, + ellphi_differentiable=self.ellphi_differentiable, + ) + self.size_loss_fn = SizeRegularizationLoss(w_major=self.lambda_major, w_minor=self.lambda_minor) + self.aniso_loss_fn = AnisotropyPenaltyLoss( + weight=self.lambda_aniso, + mode=self.aniso_mode, + barrier_threshold=config['training'].get('barrier_threshold', 6.0) + ) + + self.visualize_every = config['training'].get('visualize_every', 5) + self.output_dir = config['outputs']['image_dir'] + self.log_dir = config['outputs']['log_dir'] + + os.makedirs(self.output_dir, exist_ok=True) + os.makedirs(self.log_dir, exist_ok=True) + + self.fixed_indices = None + self.threshold = config['model'].get('threshold', 0.5) + + self.vr_complex = VietorisRipsComplex(dim=1) + + self.warmup_epochs = training_cfg.get('warmup_epochs', 0) + + self._val_aniso_accum = 0.0 + self._val_size_accum = 0.0 + + # _compute_regularization_loss is now handled by classes in losses.py + + def _compute_clean_pd_info(self, clean_pc: torch.Tensor) -> list: + """Compute clean persistence diagrams without gradient tracking.""" + clean_pd_info = [] + with torch.no_grad(): + for j in range(clean_pc.shape[0]): + c = clean_pc[j] + valid_mask = torch.abs(c).sum(dim=1) > 1e-6 + clean_pd_info.append(self.vr_complex(c[valid_mask])) + return clean_pd_info + + def train_epoch(self, data_loader, epoch): + self.model.train() + total_loss = 0 + total_class_loss = 0 + total_topo_loss = 0 + total_aniso_loss = 0 + total_size_loss = 0 + steps_completed = 0 + + all_train_preds = [] + all_train_labels = [] + + pbar = tqdm.tqdm(data_loader, desc=f"Epoch {epoch}") + + if self.fixed_indices is None: + import random + self.fixed_indices = random.sample(range(len(data_loader.dataset)), 3) + + for i, (data, labels, clean_pc) in enumerate(pbar): + data = data.to(self.device, non_blocking=True) + labels = labels.to(self.device, non_blocking=True) + clean_pc = clean_pc.to(self.device, non_blocking=True) + + if self.config.get('training', {}).get('rotation_augmentation', False): + theta = torch.rand(1, device=self.device) * 2 * 3.141592653589793 + cos_t, sin_t = torch.cos(theta), torch.sin(theta) + rotation_matrix = torch.stack([ + torch.stack([cos_t, -sin_t], dim=-1), + torch.stack([sin_t, cos_t], dim=-1) + ], dim=-2).squeeze(0) + + data = torch.matmul(data, rotation_matrix.T) + clean_pc = torch.matmul(clean_pc, rotation_matrix.T) + + self.optimizer.zero_grad(set_to_none=True) + + clean_pd_info = None + if self.lambda_topo > 0 and epoch > self.warmup_epochs: + # Topological target PD does not require autograd; keep it out of AMP/grad graph. + clean_pd_info = self._compute_clean_pd_info(clean_pc) + + with torch.amp.autocast( + device_type=self.autocast_device_type, + dtype=self.amp_dtype, + enabled=self.use_amp, + ): + logits, params = self.model(data) + class_loss = self.class_loss_fn(logits, labels) + + topo_loss = torch.tensor(0.0, device=self.device) + if clean_pd_info is not None: + topo_loss = self.topo_loss_fn(data, params, logits, clean_pd_info) + + # Regularization Losses (Size and Anisotropy only, as per slides) + if epoch > self.warmup_epochs: + size_loss = self.size_loss_fn(params) + aniso_loss = self.aniso_loss_fn(params) + else: + size_loss = torch.tensor(0.0, device=self.device) + aniso_loss = torch.tensor(0.0, device=self.device) + + loss = class_loss + topo_loss + aniso_loss + size_loss + + if torch.isnan(loss): + logger.warning("NaN loss at epoch=%s batch_index=%s; skipping step", epoch, i) + continue + + steps_completed += 1 + if self.use_amp: + self.scaler.scale(loss).backward() + self.scaler.unscale_(self.optimizer) + torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.config['training']['grad_clip_value']) + self.scaler.step(self.optimizer) + self.scaler.update() + else: + loss.backward() + torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.config['training']['grad_clip_value']) + self.optimizer.step() + + total_loss += loss.item() + total_class_loss += class_loss.item() + total_topo_loss += topo_loss.item() + total_aniso_loss += aniso_loss.item() + total_size_loss += size_loss.item() + + probs = torch.sigmoid(logits).squeeze(-1) + preds = (probs > self.threshold).long() + all_train_preds.extend(preds.cpu().numpy().flatten()) + all_train_labels.extend(labels.cpu().numpy().flatten()) + + pbar.set_postfix(loss=f"{loss.item():.4f}", cls=f"{class_loss.item():.4f}", topo=f"{topo_loss.item():.4f}", aniso=f"{aniso_loss.item():.4f}", size=f"{size_loss.item():.4f}") + if i % 10 == 0: + logger.debug( + "Step %s: loss=%.4f class=%.4f topo=%.4f aniso=%.4f size=%.4f", + i, + loss.item(), + class_loss.item(), + topo_loss.item(), + aniso_loss.item(), + size_loss.item(), + ) + + if steps_completed == 0 or not all_train_labels: + raise RuntimeError( + f"All training batches were skipped at epoch={epoch}; loss was NaN for every batch." + ) + + denom = steps_completed if steps_completed > 0 else 1 + avg_loss = total_loss / denom + avg_class_loss = total_class_loss / denom + avg_topo_loss = total_topo_loss / denom + avg_aniso_loss = total_aniso_loss / denom + avg_size_loss = total_size_loss / denom + + train_f1 = f1_score(all_train_labels, all_train_preds, zero_division=0) + train_precision = precision_score(all_train_labels, all_train_preds, zero_division=0) + train_recall = recall_score(all_train_labels, all_train_preds, zero_division=0) + + _, train_specificity, train_gmean, train_mcc = compute_recall_specificity_gmean_mcc( + all_train_labels, all_train_preds + ) + + logger.info( + "Epoch %s avg loss=%.4f (class=%.4f topo=%.4f) train F1=%.4f spec=%.4f " + "G-mean=%.4f MCC=%.4f", + epoch, + avg_loss, + avg_class_loss, + avg_topo_loss, + train_f1, + train_specificity, + train_gmean, + train_mcc, + ) + + if epoch % self.visualize_every == 0: + visualize(self.model, self.device, data_loader.dataset, epoch, output_dir=self.output_dir, title_prefix=self.config['meta'].get('config_id', 'train'), sample_indices=self.fixed_indices, threshold=self.threshold) + + return ( + avg_loss, + avg_class_loss, + avg_topo_loss, + avg_aniso_loss, + avg_size_loss, + train_f1, + train_precision, + train_recall, + train_specificity, + train_gmean, + train_mcc, + ) + + def validate(self, data_loader): + self.model.eval() + total_loss = 0 + all_labels = [] + all_preds = [] + self._val_aniso_accum = 0.0 + self._val_size_accum = 0.0 + + with torch.no_grad(): + for data, labels, _ in data_loader: + data = data.to(self.device, non_blocking=True) + labels = labels.to(self.device, non_blocking=True) + with torch.amp.autocast( + device_type=self.autocast_device_type, + dtype=self.amp_dtype, + enabled=self.use_amp, + ): + logits, params = self.model(data) + class_loss = self.class_loss_fn(logits, labels) + total_loss += class_loss.item() + + aniso_loss = self.aniso_loss_fn(params) + size_loss = self.size_loss_fn(params) + + probs = torch.sigmoid(logits).squeeze(-1) + preds = (probs > self.threshold).long() + + all_labels.extend(labels.cpu().numpy().flatten()) + all_preds.extend(preds.cpu().numpy().flatten()) + + self._val_aniso_accum += aniso_loss.item() + self._val_size_accum += size_loss.item() + + num_batches = len(data_loader) if len(data_loader) > 0 else 1 + avg_loss = total_loss / num_batches + + avg_aniso = self._val_aniso_accum / num_batches + avg_size = self._val_size_accum / num_batches + + recall = recall_score(all_labels, all_preds, zero_division=0) + + _, specificity, gmean, mcc = compute_recall_specificity_gmean_mcc( + all_labels, all_preds + ) + + return avg_loss, recall, specificity, gmean, mcc, avg_aniso, avg_size + diff --git a/tda_ml/visualization.py b/tda_ml/visualization.py new file mode 100644 index 0000000..a346ac0 --- /dev/null +++ b/tda_ml/visualization.py @@ -0,0 +1,123 @@ +"""Training-time visualization helpers.""" + +import os + +import matplotlib.pyplot as plt +import numpy as np +import torch + + +def visualize( + model, + device, + dataset, + epoch, + output_dir=".", + title_prefix="", + sample_indices=None, + threshold=0.5, +): + """ + Visualizes model predictions on 3 random samples from the dataset. + """ + model.eval() + fig, axes = plt.subplots(3, 3, figsize=(15, 15)) + fig.suptitle(f"Epoch {epoch} - {title_prefix} Results", fontsize=16) + + os.makedirs(output_dir, exist_ok=True) + + with torch.no_grad(): + for i in range(3): + if sample_indices is not None and i < len(sample_indices): + idx = sample_indices[i] + else: + idx = torch.randint(0, len(dataset), (1,)).item() + data, labels, clean_pc = dataset[idx] + + data_np = data.numpy() + labels_np = labels.numpy() + + data_batch = data.to(device).unsqueeze(0) + logits, params = model(data_batch) + + probs = torch.sigmoid(logits).squeeze(0).cpu().numpy() + pred_labels = (probs > threshold).astype(int).flatten() + + params_np = params.squeeze(0).cpu().numpy() + + axes[i, 0].scatter( + data_np[labels_np == 1, 0], + data_np[labels_np == 1, 1], + c="red", + s=10, + label="Outlier (GT)", + ) + axes[i, 0].scatter( + data_np[labels_np == 0, 0], + data_np[labels_np == 0, 1], + c="blue", + s=10, + label="Inlier (GT)", + ) + axes[i, 0].set_title(f"Sample {i + 1}: GT Labels") + axes[i, 0].legend() + + axes[i, 1].scatter( + data_np[pred_labels == 1, 0], + data_np[pred_labels == 1, 1], + c="red", + s=10, + marker="x", + label="Pred Outlier", + ) + axes[i, 1].scatter( + data_np[pred_labels == 0, 0], + data_np[pred_labels == 0, 1], + c="blue", + s=10, + label="Pred Inlier", + ) + axes[i, 1].set_title(f"Sample {i + 1}: Prediction") + axes[i, 1].legend() + + axes[i, 2].scatter( + data_np[pred_labels == 1, 0], + data_np[pred_labels == 1, 1], + c="red", + s=10, + marker="x", + label="Pred Outlier", + ) + axes[i, 2].scatter( + data_np[pred_labels == 0, 0], + data_np[pred_labels == 0, 1], + c="blue", + s=10, + label="Pred Inlier", + ) + + t = np.linspace(0, 2 * np.pi, 50) + if len(data_np) > 0: + for k in range(len(data_np)): + a, b, theta = params_np[k] + cx = data_np[k, 0] + cy = data_np[k, 1] + + x_e = a * np.cos(t) + y_e = b * np.sin(t) + x_r = x_e * np.cos(theta) - y_e * np.sin(theta) + cx + y_r = x_e * np.sin(theta) + y_e * np.cos(theta) + cy + + line_color = "blue" if pred_labels[k] == 0 else "red" + axes[i, 2].plot(x_r, y_r, color=line_color, alpha=0.3, linewidth=1) + + axes[i, 2].set_title(f"Sample {i + 1}: Predicted Inliers & Ellipses") + for ax in axes[i]: + ax.set_aspect("equal") + ax.set_xlim(-1.2, 1.2) + ax.set_ylim(-1.2, 1.2) + + plt.tight_layout(rect=[0, 0, 1, 0.96]) + filename = os.path.join(output_dir, f"{title_prefix}_result_epoch_{epoch}.png") + plt.savefig(filename) + plt.close() diff --git a/tests/test_losses_topo_distance.py b/tests/test_losses_topo_distance.py new file mode 100644 index 0000000..82838c8 --- /dev/null +++ b/tests/test_losses_topo_distance.py @@ -0,0 +1,72 @@ +"""Regression tests for training.topo_distance_mode (mahalanobis vs ellphi).""" + +import unittest +import torch + +from tda_ml.losses import ( + DISTANCE_MODE_ELLPHI, + DISTANCE_MODE_MAHALANOBIS, + compute_topo_distance_matrix, + mahalanobis_distance_matrix_batched, + normalize_topo_distance_mode, +) + + +class TestTopoDistanceMode(unittest.TestCase): + def test_normalize_aliases(self): + self.assertEqual(normalize_topo_distance_mode("Mahalanobis"), DISTANCE_MODE_MAHALANOBIS) + self.assertEqual(normalize_topo_distance_mode("ellphi"), DISTANCE_MODE_ELLPHI) + + def test_mahalanobis_shape(self): + torch.manual_seed(0) + b, n = 2, 9 + pt = torch.randn(b, n, 2) + par = torch.randn(b, n, 3) + par[:, :, 0:2] = par[:, :, 0:2].abs() + 0.1 + d = compute_topo_distance_matrix(pt, par, distance_mode="mahalanobis") + self.assertEqual(d.shape, (b, n, n)) + + def test_ellphi_finite_and_grad(self): + torch.manual_seed(0) + b, n = 1, 7 + pt = torch.randn(b, n, 2, requires_grad=True) + par = torch.randn(b, n, 3) + par[:, :, 0:2] = par[:, :, 0:2].abs() + 0.15 + par.requires_grad_(True) + d_e = compute_topo_distance_matrix(pt, par, distance_mode="ellphi", ellphi_backend="auto") + d_m = mahalanobis_distance_matrix_batched(pt, par) + self.assertEqual(d_e.shape, (b, n, n)) + self.assertTrue(torch.isfinite(d_e).all()) + loss = d_e.sum() + loss.backward() + self.assertIsNotNone(pt.grad) + self.assertIsNotNone(par.grad) + # Forward definitions differ; should not be identical in general + self.assertGreater((d_e - d_m).abs().mean().item(), 1e-6) + + def test_normalize_invalid_mode_raises_value_error(self): + with self.assertRaises(ValueError): + normalize_topo_distance_mode("unknown-mode") + + def test_mahalanobis_extreme_params_remain_finite(self): + """極小/極大軸長でも距離行列が非有限値にならないことを確認。""" + b, n = 1, 5 + pt = torch.tensor( + [[[1e-9, -1e-9], [1.0, 2.0], [3.0, -1.0], [0.5, 0.1], [-2.0, 1.5]]], + dtype=torch.float64, + ) + par = torch.tensor( + [[[1e-8, 1e8, 0.0], + [2e-8, 5e7, 0.3], + [1e7, 2e-7, -0.4], + [5e-8, 8e7, 0.7], + [2e7, 3e-8, -1.0]]], + dtype=torch.float64, + ) + d = compute_topo_distance_matrix(pt, par, distance_mode="mahalanobis") + self.assertEqual(d.shape, (b, n, n)) + self.assertTrue(torch.isfinite(d).all()) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_trainer.py b/tests/test_trainer.py new file mode 100644 index 0000000..a8c8cd8 --- /dev/null +++ b/tests/test_trainer.py @@ -0,0 +1,82 @@ +import torch +from torch.utils.data import DataLoader +import unittest +import os +import shutil + +from tda_ml.models import AnisotropicOutlierClassifier +from tda_ml.trainer import Trainer +from torch.utils.data import Dataset + +class MockOutlierDataset(Dataset): + def __init__(self, num_samples=10, max_points=30, num_outliers=5): + self.num_samples = num_samples + self.max_points = max_points + self.num_outliers = num_outliers + + def __len__(self): + return self.num_samples + + def __getitem__(self, idx): + # Create random data + total_points = self.max_points + self.num_outliers + data = torch.randn(total_points, 2) + labels = torch.zeros(total_points, dtype=torch.long) + labels[self.max_points:] = 1 # Outliers + clean_pc = torch.randn(self.max_points, 2) + return data, labels, clean_pc + +class TestTrainer(unittest.TestCase): + def setUp(self): + self.device = torch.device("cpu") + self.model = AnisotropicOutlierClassifier().to(self.device) + self.config = { + 'training': { + 'lr': 0.001, + 'lambda_class': 1.0, + 'lambda_topo': 0.1, + 'lambda_aniso': 0.01, + 'grad_clip_value': 1.0, + 'visualize_every': 10 + }, + 'model': { + 'topology_loss': { + 'distance_backend': 'mahalanobis', + } + }, + 'outputs': { + 'image_dir': 'test_outputs/images', + 'log_dir': 'test_outputs/logs' + }, + 'meta': { + 'config_id': 'test' + } + } + + dataset = MockOutlierDataset(num_samples=5, max_points=30, num_outliers=5) + self.train_loader = DataLoader(dataset, batch_size=1) + self.trainer = Trainer(self.model, self.config, self.device) + + def tearDown(self): + if os.path.exists('test_outputs'): + shutil.rmtree('test_outputs') + + def test_train_epoch(self): + """ + Tests if the train_epoch function runs for one epoch and updates weights. + """ + initial_params = [p.clone() for p in self.model.parameters()] + + avg_loss = self.trainer.train_epoch(self.train_loader, epoch=1)[0] + + params_updated = False + for initial_p, final_p in zip(initial_params, self.model.parameters()): + if not torch.equal(initial_p, final_p): + params_updated = True + break + + self.assertTrue(params_updated, "Model weights were not updated after a training step.") + self.assertGreater(avg_loss, 0.0) + +if __name__ == '__main__': + unittest.main() \ No newline at end of file