Function valid_multi_gpu

valid.py:696–832

Validate a separator model across multiple GPUs with a unified API. Runs validation either in Distributed Data Parallel (DDP) mode—detected via `torch.distributed.is_initialized()`—or, if DDP is not active, via multi-processing / single-GPU execution using the provided `device

(
    model: torch.nn.Module,
    args,
    config: ConfigDict,
    device_ids: Optional[List[int]] = None,
    verbose: bool = False
)

def valid_multi_gpu(
    model: torch.nn.Module,
    args,
    config: ConfigDict,
    device_ids: Optional[List[int]] = None,
    verbose: bool = False
) -> Tuple[Dict[str, float], Dict]:
    """
    Validate a separator model across multiple GPUs with a unified API.

    Runs validation either in Distributed Data Parallel (DDP) mode—detected via
    `torch.distributed.is_initialized()`—or, if DDP is not active, via
    multi-processing / single-GPU execution using the provided `device_ids`.
    Collects per-track metrics, aggregates them into per-instrument/per-metric
    arrays, and computes per-metric averages.

    Behavior:
      * DDP mode: splits the dataset across ranks and gathers metrics; only rank 0
        returns results, while other ranks return `(None, None)`.
      * Non-DDP: launches parallel workers when `len(device_ids) > 1`, otherwise
        runs on a single device/CPU.

    Args:
        model (torch.nn.Module): Trained model to evaluate.
        args: Runtime arguments (e.g., metrics list, store dir).
        config (ConfigDict): Configuration with inference/training settings.
        device_ids (Optional[List[int]]): GPU device IDs for non-DDP parallelism.
            If None or length is 1, runs on a single device.
        verbose (bool, optional): If True, print progress/logs. Defaults to False.

    Returns:
        Tuple[Dict[str, float], Dict]: A pair `(metric_avg, all_metrics)` where
            - `metric_avg` maps metric name to its average score,
            - `all_metrics` is a nested dict `{metric: {instrument: List[float]}}`.
          In DDP mode, non-zero ranks return `(None, None)`.
    """

    start_time = time.time()

    inference = getattr(config, "inference", None)
    if inference is None and isinstance(config, dict):
        inference = config.get("inference", {})
    extension = getattr(inference, "extension", None)
    if extension is None:
        if isinstance(inference, dict):
            extension = inference.get("extension", getattr(args, "extension", "wav"))
        else:
            extension = getattr(args, "extension", "wav")

    all_mixtures_path = get_mixture_paths(args, verbose, config, extension)

    ddp_mode = dist.is_initialized()

    if ddp_mode:
        rank = dist.get_rank()
        world_size = dist.get_world_size()

        device = torch.device(f"cuda:{rank}")