import logging import os import random import time from datetime import datetime from pathlib import Path import numpy as np import torch import torch.distributed as dist import torch.nn as nn import torch.nn.functional as F import torch.optim as optim from torch.nn.parallel import DistributedDataParallel as DDP from torch.utils.data import DataLoader from torch.utils.data.distributed import DistributedSampler from tqdm import tqdm from dataset.DeepchoiceDataset import DeepChoiceDataset from dataset.RandomSubsetBatchDataset import RandomSubsetBatchDataset from model.deepchoice_transformer import DeepChoiceTransformer from model.deepchoice_mlp import DeepChoiceMLP from utils.compute_metrics import compute_metrics from utils.dataset_contract import select_visibility_indices from utils.utilities import compute_proba_batch logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s") try: import wandb except ImportError: # pragma: no cover - optional dependency wandb = None try: from transformers import get_cosine_schedule_with_warmup except ImportError: # pragma: no cover - optional dependency get_cosine_schedule_with_warmup = None def collate_prebatched_samples(samples): if not samples: raise ValueError("Cannot collate an empty sample list") per_point_tensors = { "visibility": [], "logits": [], "mask": [], "target": [], "coords_int": [], "coords_scale": [], "coords_offset": [], "coords_tile_offset": [], } tile_names = [] source_paths = [] for sample in samples: num_points = int(sample["target"].shape[0]) per_point_tensors["visibility"].append(sample["visibility"]) per_point_tensors["logits"].append(sample["logits"]) per_point_tensors["mask"].append(sample["mask"]) per_point_tensors["target"].append(sample["target"]) per_point_tensors["coords_int"].append(sample["coords_int"]) scale = torch.as_tensor(sample["coords_scale"], dtype=torch.float64) if scale.ndim == 0: scale = scale.reshape(1).repeat(num_points) elif scale.ndim == 1 and scale.shape[0] == 1: scale = scale.repeat(num_points) elif scale.ndim != 1 or scale.shape[0] != num_points: raise ValueError(f"Unsupported coords_scale shape {tuple(scale.shape)} for {num_points} points") per_point_tensors["coords_scale"].append(scale) offset = torch.as_tensor(sample["coords_offset"], dtype=torch.float64) if offset.ndim == 1 and offset.shape[0] == 3: offset = offset.reshape(1, 3).repeat(num_points, 1) elif offset.ndim != 2 or offset.shape != (num_points, 3): raise ValueError(f"Unsupported coords_offset shape {tuple(offset.shape)} for {num_points} points") per_point_tensors["coords_offset"].append(offset) tile_offset = torch.as_tensor(sample.get("coords_tile_offset", torch.zeros(3, dtype=torch.float64)), dtype=torch.float64) if tile_offset.ndim == 1 and tile_offset.shape[0] == 3: tile_offset = tile_offset.reshape(1, 3).repeat(num_points, 1) elif tile_offset.ndim != 2 or tile_offset.shape != (num_points, 3): raise ValueError(f"Unsupported coords_tile_offset shape {tuple(tile_offset.shape)} for {num_points} points") per_point_tensors["coords_tile_offset"].append(tile_offset) if isinstance(sample["tile_name"], list): tile_names.extend(sample["tile_name"]) else: tile_names.extend([sample["tile_name"]] * num_points) if isinstance(sample["source_path"], list): source_paths.extend(sample["source_path"]) else: source_paths.extend([sample["source_path"]] * num_points) batch = {key: torch.cat(values, dim=0) for key, values in per_point_tensors.items()} batch["tile_name"] = tile_names batch["source_path"] = source_paths return batch def is_distributed(): return dist.is_available() and dist.is_initialized() def get_rank(): return dist.get_rank() if is_distributed() else 0 def get_world_size(): return dist.get_world_size() if is_distributed() else 1 def is_main_process(): return get_rank() == 0 def setup_distributed_training(config): world_size = int(os.environ.get("WORLD_SIZE", "1")) use_ddp = world_size > 1 training_cfg = config["training"] requested_device = str(training_cfg["device"]).lower() if not use_ddp: return False backend = training_cfg.get("ddp_backend") if backend is None: backend = "nccl" if requested_device.startswith("cuda") else "gloo" dist.init_process_group(backend=backend) local_rank = int(os.environ.get("LOCAL_RANK", "0")) if requested_device.startswith("cuda"): torch.cuda.set_device(local_rank) training_cfg["device"] = f"cuda:{local_rank}" else: training_cfg["device"] = "cpu" return True def cleanup_distributed_training(): if is_distributed(): dist.barrier() dist.destroy_process_group() def create_run_dir(base_dir): run_dir = None if is_main_process(): run_dir = str(Path(base_dir + datetime.now().strftime("_%m_%d_%H_%M_%S"))) Path(run_dir).mkdir(parents=True, exist_ok=True) if is_distributed(): payload = [run_dir] dist.broadcast_object_list(payload, src=0) run_dir = payload[0] Path(run_dir).mkdir(parents=True, exist_ok=True) return Path(run_dir) def set_global_seed(seed, deterministic=False): seed = int(seed) random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) if deterministic: torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False def unwrap_model(model): return model.module if isinstance(model, DDP) else model def _split_folder(config, split_name): split_roots = config.get("data", {}).get("split_roots", {}) or {} override = split_roots.get(split_name) if override: return Path(override) return Path(config["data"]["batches_root"]) / split_name def list_split_batches(config, split_name, limit=None): folder = _split_folder(config, split_name) paths = sorted(str(path) for path in folder.glob("*.pt")) if limit is not None: paths = paths[: int(limit)] return paths def build_split_loader(config, split_name, shuffle=False, limit=None, distributed=None, file_batch_size=None): paths = list_split_batches(config, split_name, limit=limit) if not paths: raise FileNotFoundError(f"No .pt batches found in {str(_split_folder(config, split_name))}") return build_loader_from_paths(config, paths, shuffle=shuffle, distributed=distributed, file_batch_size=file_batch_size) def build_loader_from_paths(config, paths, shuffle=False, distributed=None, file_batch_size=None): dataset = DeepChoiceDataset(paths, shuffle=False) device = str(config["training"]["device"]).lower() num_workers = int(config["training"].get("num_workers", 0)) file_batch_size = int(file_batch_size or config["training"].get("file_batch_size", 1)) if distributed is None: distributed = is_distributed() sampler = None if distributed: sampler = DistributedSampler(dataset, shuffle=shuffle, drop_last=False) return DataLoader( dataset, batch_size=file_batch_size, shuffle=False if sampler is not None else shuffle, sampler=sampler, num_workers=num_workers, pin_memory=device.startswith("cuda"), collate_fn=collate_prebatched_samples, persistent_workers=num_workers > 0, ) def build_train_loader_from_paths(config, paths, distributed=None): random_limit = config["training"].get("train_limit_files") randomize_each_epoch = bool(config["training"].get("randomize_train_limit_each_epoch", False)) if random_limit is not None and randomize_each_epoch: dataset = RandomSubsetBatchDataset( paths, subset_size=int(random_limit), seed=int(config["training"].get("seed", 42)), ) device = str(config["training"]["device"]).lower() num_workers = int(config["training"].get("num_workers", 0)) file_batch_size = int(config["training"].get("file_batch_size", 1)) if distributed is None: distributed = is_distributed() sampler = None if distributed: sampler = DistributedSampler(dataset, shuffle=True, drop_last=False) return DataLoader( dataset, batch_size=file_batch_size, shuffle=False if sampler is not None else True, sampler=sampler, num_workers=num_workers, pin_memory=device.startswith("cuda"), collate_fn=collate_prebatched_samples, persistent_workers=num_workers > 0, ) return build_loader_from_paths( config, paths, shuffle=True, distributed=distributed, file_batch_size=config["training"].get("file_batch_size", 1), ) def _split_train_val_paths(config): train_limit = config["training"].get("train_limit_files") val_limit = config["training"].get("val_limit_files") train_limit_for_listing = None if bool(config["training"].get("randomize_train_limit_each_epoch", False)) else train_limit train_paths = list_split_batches(config, config["training"]["train_split"], limit=train_limit_for_listing) if not train_paths: raise FileNotFoundError(f"No .pt batches found in {str(_split_folder(config, config['training']['train_split']))}") val_paths = list_split_batches(config, config["training"]["val_split"], limit=val_limit) if val_paths: return train_paths, val_paths if len(train_paths) < 2: raise FileNotFoundError("Validation split is empty and there are not enough train batches to create a fallback split") seed = int(config["training"].get("split_seed", 42)) fraction = float(config["training"].get("val_from_train_fraction", 0.05)) num_val = max(1, int(round(len(train_paths) * fraction))) num_val = min(num_val, len(train_paths) - 1) rng = np.random.default_rng(seed) perm = rng.permutation(len(train_paths)) val_indices = set(int(idx) for idx in perm[:num_val]) fallback_train = [path for idx, path in enumerate(train_paths) if idx not in val_indices] fallback_val = [path for idx, path in enumerate(train_paths) if idx in val_indices] if is_main_process(): logging.warning( "Validation split '%s' is empty; using %s train batch files as fallback validation set", config["training"]["val_split"], len(fallback_val), ) return fallback_train, fallback_val def _selected_visibility_indices(config): feature_names = config["dataset"]["visibility_feature_names"] selected = config["model"].get("viscrit") return select_visibility_indices(feature_names, selected) def _select_visibility(sample, config): indices = _selected_visibility_indices(config) max_views = int(config["model"]["max_views"]) return sample["visibility"][:, :max_views, indices] def infer_model_dims(config): num_features = len(_selected_visibility_indices(config)) max_views = int(config["model"]["max_views"]) include_logits = bool(config["model"].get("use_logit_features", False)) include_argmax = bool(config["model"].get("use_argmax_feature", False)) extra_features = 0 if include_logits: extra_features += int(config["model"]["num_classes"]) if include_argmax: extra_features += 1 return { "num_visibility_features": num_features, "mlp_input_dim": max_views * (num_features + extra_features), "transformer_token_dim": num_features + extra_features, } def synchronize_model_dims(config): dims = infer_model_dims(config) model_cfg = config.setdefault("model", {}) if str(model_cfg.get("type", "")).upper() == "MLP": model_cfg["input_dim"] = int(dims["mlp_input_dim"]) else: model_cfg["token_dim"] = int(dims["transformer_token_dim"]) return dims def validate_sample_contract(sample, config): if sample["visibility"].ndim != 3: raise ValueError(f"Expected visibility shape [B, V, F], got {tuple(sample['visibility'].shape)}") if sample["logits"].ndim != 3: raise ValueError(f"Expected logits shape [B, V, C], got {tuple(sample['logits'].shape)}") if sample["mask"].ndim != 2: raise ValueError(f"Expected mask shape [B, V], got {tuple(sample['mask'].shape)}") if sample["target"].ndim != 1: raise ValueError(f"Expected target shape [B], got {tuple(sample['target'].shape)}") batch_size, num_views, num_features = sample["visibility"].shape logits_batch, logits_views, num_classes = sample["logits"].shape mask_batch, mask_views = sample["mask"].shape target_batch = sample["target"].shape[0] if logits_batch != batch_size or mask_batch != batch_size or target_batch != batch_size: raise ValueError("Visibility/logits/mask/target batch dimensions are inconsistent") if logits_views != num_views or mask_views != num_views: raise ValueError("Visibility/logits/mask view dimensions are inconsistent") if num_views < int(config["model"]["max_views"]): raise ValueError(f"Expected at least {config['model']['max_views']} views, got {num_views}") if num_classes != int(config["model"]["num_classes"]): raise ValueError(f"Expected {config['model']['num_classes']} logit classes, got {num_classes}") dims = infer_model_dims(config) selected_features = len(_selected_visibility_indices(config)) if selected_features != dims["num_visibility_features"]: raise ValueError("Selected visibility feature count is inconsistent") if config["model"]["type"] == "MLP" and int(config["model"]["input_dim"]) != dims["mlp_input_dim"]: raise ValueError( f"MLP input_dim={config['model']['input_dim']} does not match max_views*features={dims['mlp_input_dim']}" ) if config["model"]["type"] != "MLP" and int(config["model"]["token_dim"]) != dims["transformer_token_dim"]: raise ValueError( f"Transformer token_dim={config['model']['token_dim']} does not match selected features={dims['transformer_token_dim']}" ) if num_features < selected_features: raise ValueError(f"Batch visibility has only {num_features} features but {selected_features} were requested") def _build_model_tokens(visibility, logits, config): parts = [visibility] if bool(config["model"].get("use_logit_features", False)): parts.append(logits) if bool(config["model"].get("use_argmax_feature", False)): argmax_feature = torch.argmax(logits, dim=2, keepdim=True).to(visibility.dtype) argmax_feature = argmax_feature / max(int(config["model"]["num_classes"]) - 1, 1) parts.append(argmax_feature) if len(parts) == 1: return visibility return torch.cat(parts, dim=2) def _prepare_batch(sample, config): validate_sample_contract(sample, config) device = config["training"]["device"] max_views = int(config["model"]["max_views"]) visibility = _select_visibility(sample, config).to(device, non_blocking=True) logits = sample["logits"][:, :max_views, :].to(device, non_blocking=True) mask = sample["mask"][:, :max_views].to(device, non_blocking=True) target = sample["target"].to(device, non_blocking=True) ignore_labels = config["model"].get("ignore_labels") if ignore_labels: ignore_index = int(config["model"].get("ignore_index", 255)) remapped = target.clone() for label in ignore_labels: remapped[target == int(label)] = ignore_index target = remapped model_inputs = _build_model_tokens(visibility, logits, config) return model_inputs, logits, mask, target def _forward_weights(model, visibility, mask, model_type): if model_type == "MLP": return model(visibility) return model(visibility, mask=mask) def fused_nll_loss(fused_probs, target, config, eps=1e-8): fused_probs = fused_probs.clamp_min(eps) ignore_index = int(config["model"].get("ignore_index", 255)) return F.nll_loss(torch.log(fused_probs), target, ignore_index=ignore_index) def build_model(config): dims = infer_model_dims(config) if config["model"]["type"] == "MLP": return DeepChoiceMLP( input_dim=dims["mlp_input_dim"], hidden_dim1=config["model"]["hidden_dim1"], hidden_dim2=config["model"]["hidden_dim2"], hidden_dim3=config["model"]["hidden_dim3"], hidden_dim4=config["model"]["hidden_dim4"], hidden_dim5=config["model"]["hidden_dim5"], output_dim=config["model"]["max_views"], ).to(config["training"]["device"]) return DeepChoiceTransformer( token_dim=dims["transformer_token_dim"], num_tokens=config["model"]["max_views"], model_dim=config["model"]["model_dim"], num_heads=config["model"]["num_heads"], ff_dim=config["model"]["ff_dim"], dropout=config["model"]["dropout"], num_layers=config["model"]["num_layers"], ).to(config["training"]["device"]) def wrap_model_for_training(model, config): if not is_distributed(): return model device = str(config["training"]["device"]).lower() if device.startswith("cuda"): local_rank = int(os.environ.get("LOCAL_RANK", "0")) return DDP(model, device_ids=[local_rank], output_device=local_rank, find_unused_parameters=False) return DDP(model, find_unused_parameters=False) def build_scheduler(optimizer, total_steps, warmup_steps): total_steps = max(int(total_steps), 1) warmup_steps = min(max(int(warmup_steps), 0), total_steps) if get_cosine_schedule_with_warmup is not None: return get_cosine_schedule_with_warmup( optimizer, num_warmup_steps=warmup_steps, num_training_steps=total_steps, ) def lr_lambda(step): if total_steps <= 1: return 1.0 if warmup_steps > 0 and step < warmup_steps: return float(step + 1) / float(warmup_steps) progress = (step - warmup_steps) / max(total_steps - warmup_steps, 1) progress = min(max(progress, 0.0), 1.0) return 0.5 * (1.0 + np.cos(np.pi * progress)) return optim.lr_scheduler.LambdaLR(optimizer, lr_lambda) def _reduce_scalar(value): tensor = torch.tensor(float(value), dtype=torch.float64, device=configured_device_for_reduce()) if is_distributed(): dist.all_reduce(tensor, op=dist.ReduceOp.SUM) return float(tensor.item()) def configured_device_for_reduce(): if torch.cuda.is_available() and str(os.environ.get("LOCAL_RANK", "0")).isdigit() and str(os.environ.get("WORLD_SIZE", "1")) != "1": return torch.device(f"cuda:{int(os.environ.get('LOCAL_RANK', '0'))}") return torch.device("cpu") def _gather_numpy(array): if not is_distributed(): return [array] gathered = [None for _ in range(get_world_size())] dist.all_gather_object(gathered, array) return gathered def _masked_mean_logits(logits, mask): weights = mask.to(logits.dtype).unsqueeze(-1) summed = (logits * weights).sum(dim=1) counts = weights.sum(dim=1).clamp_min(1.0) return summed / counts def _baseline_anyview(logits, mask, target): pred_views = torch.argmax(logits, dim=2) matches = (pred_views == target[:, None]) & mask has_match = matches.any(dim=1) majority_pred = torch.argmax(_masked_mean_logits(logits, mask), dim=1) return torch.where(has_match, target, majority_pred) def _baseline_hard_vote(logits, mask): pred_views = torch.argmax(logits, dim=2) num_classes = logits.shape[2] one_hot = torch.nn.functional.one_hot(pred_views, num_classes=num_classes).to(logits.dtype) weights = mask.to(logits.dtype).unsqueeze(-1) class_counts = (one_hot * weights).sum(dim=1) return torch.argmax(class_counts, dim=1) def _prepare_targets_cpu(target, config): target = target.clone() ignore_labels = config["model"].get("ignore_labels") if ignore_labels: ignore_index = int(config["model"].get("ignore_index", 255)) for label in ignore_labels: target[target == int(label)] = ignore_index return target def compute_split_baselines(config, split_name=None, paths=None, file_batch_size=None, desc="Computing Baselines"): if not is_main_process(): return None if paths is None: if split_name is None: raise ValueError("Either split_name or paths must be provided") paths = list_split_batches(config, split_name) if not paths: raise FileNotFoundError("No batch files found to compute baselines") loader = build_loader_from_paths( config, paths, shuffle=False, distributed=False, file_batch_size=file_batch_size or config["training"].get("eval_file_batch_size", config["training"].get("file_batch_size", 1)), ) max_views = int(config["model"]["max_views"]) majority_preds = [] hard_vote_preds = [] anyview_preds = [] labels = [] for sample in tqdm(loader, desc=desc): validate_sample_contract(sample, config) logits = sample["logits"][:, :max_views, :] mask = sample["mask"][:, :max_views] target = _prepare_targets_cpu(sample["target"], config) majority_pred = torch.argmax(_masked_mean_logits(logits, mask), dim=1) hard_vote_pred = _baseline_hard_vote(logits, mask) anyview_pred = _baseline_anyview(logits, mask, target) majority_preds.append(majority_pred.cpu().numpy()) hard_vote_preds.append(hard_vote_pred.cpu().numpy()) anyview_preds.append(anyview_pred.cpu().numpy()) labels.append(target.cpu().numpy()) y_true = np.concatenate(labels) majority_pred = np.concatenate(majority_preds) hard_vote_pred = np.concatenate(hard_vote_preds) anyview_pred = np.concatenate(anyview_preds) ignore_index = int(config["model"].get("ignore_index", 255)) majority_miou, majority_mf1, majority_ious = compute_metrics( y_true, majority_pred, int(config["model"]["num_classes"]), ignore_index=ignore_index ) hard_vote_miou, hard_vote_mf1, hard_vote_ious = compute_metrics( y_true, hard_vote_pred, int(config["model"]["num_classes"]), ignore_index=ignore_index ) anyview_miou, anyview_mf1, anyview_ious = compute_metrics( y_true, anyview_pred, int(config["model"]["num_classes"]), ignore_index=ignore_index ) return { "majority": { "miou": float(majority_miou), "mf1": float(majority_mf1), "ious": majority_ious, }, "hard_vote": { "miou": float(hard_vote_miou), "mf1": float(hard_vote_mf1), "ious": hard_vote_ious, }, "anyview": { "miou": float(anyview_miou), "mf1": float(anyview_mf1), "ious": anyview_ious, }, } def compute_validation_baselines(config, val_paths): baselines = compute_split_baselines( config, paths=val_paths, file_batch_size=config["training"].get("eval_file_batch_size", config["training"].get("file_batch_size", 1)), desc="Computing Validation Baselines", ) logging.info( "Validation baselines | majority mIoU=%.4f mF1=%.4f | hard_vote mIoU=%.4f mF1=%.4f | anyview mIoU=%.4f mF1=%.4f", baselines["majority"]["miou"], baselines["majority"]["mf1"], baselines["hard_vote"]["miou"], baselines["hard_vote"]["mf1"], baselines["anyview"]["miou"], baselines["anyview"]["mf1"], ) return baselines def resolve_validation_baselines(config, val_paths): precomputed = config.get("training", {}).get("precomputed_validation_baselines") if precomputed is not None: if "hard_vote" not in precomputed: logging.warning("Precomputed validation baselines missing 'hard_vote'; recomputing validation baselines") return compute_validation_baselines(config, val_paths) logging.info( "Using precomputed validation baselines | majority mIoU=%.4f mF1=%.4f | hard_vote mIoU=%.4f mF1=%.4f | anyview mIoU=%.4f mF1=%.4f", precomputed["majority"]["miou"], precomputed["majority"]["mf1"], precomputed["hard_vote"]["miou"], precomputed["hard_vote"]["mf1"], precomputed["anyview"]["miou"], precomputed["anyview"]["mf1"], ) return precomputed compute_val_baselines = bool(config["training"].get("compute_validation_baselines", True)) if not compute_val_baselines: return None return compute_validation_baselines(config, val_paths) def evaluate(model, loader, config, n_classes=11, desc="Evaluating"): model.eval() all_preds = [] all_labels = [] total_loss = 0.0 total_samples = 0 total_data_time = 0.0 total_compute_time = 0.0 show_progress = is_main_process() and not is_distributed() iterator = tqdm(loader, desc=desc) if show_progress else loader loop_end = time.perf_counter() with torch.no_grad(): for sample in iterator: total_data_time += time.perf_counter() - loop_end compute_start = time.perf_counter() visibility, logits, mask, target = _prepare_batch(sample, config) weights = _forward_weights(model, visibility, mask, config["model"]["type"]) fused_logits = compute_proba_batch(weights, logits, mask=mask) loss = fused_nll_loss(fused_logits, target, config) total_compute_time += time.perf_counter() - compute_start batch_size = int(target.shape[0]) total_loss += float(loss.item()) * batch_size total_samples += batch_size pred = torch.argmax(fused_logits, dim=1) all_preds.append(pred.cpu().numpy()) all_labels.append(target.cpu().numpy()) loop_end = time.perf_counter() if total_samples == 0: raise RuntimeError(f"No samples evaluated in loader {desc}") y_true_local = np.concatenate(all_labels) y_pred_local = np.concatenate(all_preds) gathered_true = _gather_numpy(y_true_local) gathered_pred = _gather_numpy(y_pred_local) gathered_counts = _gather_numpy(np.asarray([total_samples], dtype=np.int64)) gathered_losses = _gather_numpy(np.asarray([total_loss], dtype=np.float64)) gathered_data_times = _gather_numpy(np.asarray([total_data_time], dtype=np.float64)) gathered_compute_times = _gather_numpy(np.asarray([total_compute_time], dtype=np.float64)) if not is_main_process(): return None y_true = np.concatenate(gathered_true) y_pred = np.concatenate(gathered_pred) total_samples_global = int(np.concatenate(gathered_counts).sum()) total_loss_global = float(np.concatenate(gathered_losses).sum()) total_data_time_global = float(np.concatenate(gathered_data_times).sum()) total_compute_time_global = float(np.concatenate(gathered_compute_times).sum()) miou, mf1, ious = compute_metrics( y_true, y_pred, n_classes, ignore_index=int(config["model"].get("ignore_index", 255)), ) return { "miou": float(miou), "mf1": float(mf1), "ious": ious, "loss": total_loss_global / total_samples_global, "num_samples": total_samples_global, "data_time_s": total_data_time_global, "compute_time_s": total_compute_time_global, } def train_deepchoice(config, save_dir=None): use_ddp = setup_distributed_training(config) try: seed = int(config["training"].get("seed", 42)) deterministic = bool(config["training"].get("deterministic", False)) set_global_seed(seed, deterministic=deterministic) save_dir = create_run_dir(str(save_dir or config["training"]["output_dir"])) train_paths, val_paths = _split_train_val_paths(config) train_loader = build_train_loader_from_paths( config, train_paths, distributed=use_ddp, ) val_loader = build_loader_from_paths( config, val_paths, shuffle=False, distributed=use_ddp, file_batch_size=config["training"].get("eval_file_batch_size", config["training"].get("file_batch_size", 1)), ) baseline_metrics = resolve_validation_baselines(config, val_paths) if is_distributed(): dist.barrier() model = build_model(config) model = wrap_model_for_training(model, config) if config["model"]["type"] == "MLP": lr = config["training"]["lr"] else: lr = config["training"]["lr_transformer"] optimizer = optim.AdamW(model.parameters(), lr=lr, weight_decay=config["training"]["weight_decay"]) total_steps = int(config["training"]["epochs"]) * len(train_loader) warmup_steps = int(config["training"].get("sched_step", 0)) scheduler = build_scheduler(optimizer, total_steps, warmup_steps) wandb_enabled = config.get("wandb", {}).get("enabled", False) if wandb_enabled and wandb is None: raise ImportError("wandb is enabled in config but the package is not installed in the current environment") if wandb_enabled and is_main_process(): wandb.init(project=config["wandb"]["project"], entity=config["wandb"].get("entity"), config=config) wandb.watch(unwrap_model(model), log="all") best_miou = float("-inf") best_path = Path(save_dir) / "best_model.pt" last_path = Path(save_dir) / "last_model.pt" eval_every = int(config["training"].get("eval_every", 1)) for epoch in range(1, int(config["training"]["epochs"]) + 1): sampler = getattr(train_loader, "sampler", None) if isinstance(sampler, DistributedSampler): sampler.set_epoch(epoch) train_dataset = getattr(train_loader, "dataset", None) if hasattr(train_dataset, "set_epoch"): train_dataset.set_epoch(epoch) model.train() total_loss = 0.0 total_samples = 0 total_data_time = 0.0 total_compute_time = 0.0 show_progress = is_main_process() progress = tqdm(train_loader, desc=f"Epoch {epoch}/{config['training']['epochs']} Training", disable=not show_progress) loop_end = time.perf_counter() for sample in progress: total_data_time += time.perf_counter() - loop_end compute_start = time.perf_counter() visibility, logits, mask, target = _prepare_batch(sample, config) optimizer.zero_grad(set_to_none=True) weights = _forward_weights(model, visibility, mask, config["model"]["type"]) fused_logits = compute_proba_batch(weights, logits, mask=mask) loss = fused_nll_loss(fused_logits, target, config) loss.backward() optimizer.step() scheduler.step() total_compute_time += time.perf_counter() - compute_start batch_size = int(target.shape[0]) total_loss += float(loss.item()) * batch_size total_samples += batch_size if show_progress: progress.set_postfix(loss=f"{(total_loss / max(total_samples, 1)):.4f}") loop_end = time.perf_counter() train_loss_sum = _reduce_scalar(total_loss) train_data_time = _reduce_scalar(total_data_time) train_compute_time = _reduce_scalar(total_compute_time) train_samples = int(sum(int(x[0]) for x in _gather_numpy(np.asarray([total_samples], dtype=np.int64)))) train_loss = train_loss_sum / max(train_samples, 1) if is_main_process(): num_steps = max(len(train_loader), 1) logging.info( "Epoch %s - train loss %.4f over %s samples | avg_data=%.4fs avg_compute=%.4fs avg_step=%.4fs", epoch, train_loss, train_samples, train_data_time / num_steps, train_compute_time / num_steps, (train_data_time + train_compute_time) / num_steps, ) metrics = None if epoch % eval_every == 0: metrics = evaluate( model, val_loader, config, n_classes=int(config["model"]["num_classes"]), desc=f"Epoch {epoch} Validation", ) if is_main_process(): logging.info( "Epoch %s - val loss %.4f | val mIoU %.4f | val mF1 %.4f | majority %.4f/%.4f | anyview %.4f/%.4f | data=%.4fs compute=%.4fs", epoch, metrics["loss"], metrics["miou"], metrics["mf1"], baseline_metrics["majority"]["miou"] if baseline_metrics is not None else float("nan"), baseline_metrics["majority"]["mf1"] if baseline_metrics is not None else float("nan"), baseline_metrics["anyview"]["miou"] if baseline_metrics is not None else float("nan"), baseline_metrics["anyview"]["mf1"] if baseline_metrics is not None else float("nan"), metrics["data_time_s"], metrics["compute_time_s"], ) if metrics["miou"] > best_miou: best_miou = metrics["miou"] torch.save(unwrap_model(model).state_dict(), best_path) logging.info("Saved new best model to %s", best_path) if is_main_process(): torch.save(unwrap_model(model).state_dict(), last_path) if wandb_enabled: log_dict = { "epoch": epoch, "train_loss": train_loss, "current_lr": scheduler.get_last_lr()[0], } if metrics is not None: log_dict.update( { "val_loss": metrics["loss"], "val_mIoU": metrics["miou"], "val_mF1": metrics["mf1"], } ) log_dict.update({f"val_iou_class_{idx}": float(value) for idx, value in enumerate(metrics["ious"])}) if baseline_metrics is not None: log_dict.update( { "baseline_majority_mIoU": baseline_metrics["majority"]["miou"], "baseline_majority_mF1": baseline_metrics["majority"]["mf1"], "baseline_anyview_mIoU": baseline_metrics["anyview"]["miou"], "baseline_anyview_mF1": baseline_metrics["anyview"]["mf1"], } ) wandb.log(log_dict) if is_distributed(): dist.barrier() if is_main_process(): if best_miou == float("-inf"): torch.save(unwrap_model(model).state_dict(), best_path) if wandb_enabled: wandb.finish() logging.info("Training complete | best mIoU %.4f", best_miou if best_miou != float("-inf") else float("nan")) return unwrap_model(model) finally: cleanup_distributed_training()