train/train.py

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()