Spaces:

YuePanEdward
/

RAP

Sleeping

RAP / rectified_point_flow /utils /logging.py

Squash history: release superseded example-data blobs

be88765 about 2 months ago

20.8 kB

from collections import defaultdict
from itertools import chain
import re
from typing import Dict, List, Any

import lightning as L
import torch
import torch.distributed as dist
from rich.console import Console
from rich.table import Table


def log_metrics_on_step(
    module: L.LightningModule,
    metrics: Dict[str, float],
    prefix: str = "train",
):
    """Log per‐step scalars only on rank 0."""
    for name, value in metrics.items():
        module.log(
            f"{prefix}/{name}",
            value,
            on_step=True,
            on_epoch=False,
            prog_bar=True,
            rank_zero_only=True,
        )


def log_metrics_on_epoch(
    module: L.LightningModule,
    metrics: Dict[str, torch.Tensor],
    prefix: str = "val",
):
    """Log per‐epoch scalars, automatically synced & averaged across ranks.
    
    For single dataset: only log overall metrics to avoid redundancy.
    For multiple datasets: log both per-dataset and overall metrics.
    """
    # Separate dataset-specific and overall metrics
    dataset_metrics = {}
    overall_metrics = {}
    
    for name, value in metrics.items():
        if name.startswith("overall/"):
            overall_metrics[name.replace("overall/", "")] = value
        else:
            dataset_metrics[name] = value
    
    # Determine if we have multiple datasets
    dataset_names = set()
    for name in dataset_metrics.keys():
        if "/" in name:
            dataset_name = name.split("/")[0]
            dataset_names.add(dataset_name)
    
    # Always log overall metrics with "overall/" prefix to maintain consistency
    # This ensures ModelCheckpoint can always find "val/overall/object_chamfer"
    for name, value in overall_metrics.items():
        module.log(
            f"{prefix}/overall/{name}",
            value,
            on_step=False,
            on_epoch=True,
            prog_bar=True,
            rank_zero_only=True,
        )
    
    # For multiple datasets, also log per-dataset metrics
    if len(dataset_names) > 1:
        for name, value in dataset_metrics.items():
            module.log(
                f"{prefix}/{name}",
                value,
                on_step=False,
                on_epoch=True,
                prog_bar=True,
                rank_zero_only=True,
            )

def print_eval_table(
    results: list[dict[str, float]],
    dataset_names: list[str],
    digits: int = 4,
    sample_counts: list[int] = None,
    part_count_ranges: list[tuple[int, int]] = None,
):
    """
    Pretty-print a list of evaluation-result dicts with Rich, split into Avg and BoN sections.
    For single dataset, use a more concise format.
    Args:
        results: List of dicts from trainer.test(), each dict corresponds to one dataloader.
        dataset_names: List of dataset names corresponding to each dataloader.
        digits: Number of decimal places for floats.
        sample_counts: Optional list of sample counts for each dataset.
        part_count_ranges: Optional list of (min, max) part count tuples for each dataset.
    """
    if not results:
        print("No results to display")
        return
    
    # PyTorch Lightning test() returns a list of dicts, one per dataloader
    # Each dict contains metrics like "avg/object_chamfer", "best_of_n/rotation_error", etc.
    
    # Collect all unique metric names across all dataloaders
    all_metrics = set()
    for result_dict in results:
        all_metrics.update(result_dict.keys())
    
    # Split into avg and best-of-n metrics
    avg_metrics = sorted(m for m in all_metrics if m.startswith("avg/"))
    bon_metrics = sorted(m for m in all_metrics if re.match(r"best_of_\w+/", m))
    rigidity_selected_metrics = sorted(m for m in all_metrics if m.startswith("rigidity_selected/"))
    
    fmt = f"{{:.{digits}f}}"
    
    # For single dataset, use a more concise format
    if len(results) == 1 and len(dataset_names) >= 1:
        dataset_name = dataset_names[0]
        result_dict = results[0]
        
        table = Table()
        table.add_column("Metric", style="bold magenta", justify="left", no_wrap=True)
        table.add_column(dataset_name, style="cyan", justify="right")
        
        # Add sample count row if provided
        if sample_counts and len(sample_counts) >= 1:
            table.add_row("Sample Count", str(sample_counts[0]))
        
        # Add part count range row if provided
        if part_count_ranges and len(part_count_ranges) >= 1:
            min_parts, max_parts = part_count_ranges[0]
            if min_parts == max_parts:
                table.add_row("Part Count", str(min_parts))
            else:
                table.add_row("Part Count Range", f"[{min_parts}, {max_parts}]")
        
        # Add section separator if we added any info rows
        if (sample_counts and len(sample_counts) >= 1) or (part_count_ranges and len(part_count_ranges) >= 1):
            table.add_section()
        
        # Add average metrics
        for metric in avg_metrics:
            val = result_dict.get(metric)
            if val is not None:
                display_name = metric.replace("avg/", "")
                if isinstance(val, (int, float)):
                    table.add_row(display_name, fmt.format(float(val)))
                else:
                    table.add_row(display_name, str(val))
        
        # Add section separator if both avg and bon metrics exist
        if avg_metrics and bon_metrics:
            table.add_section()
        
        # Add best-of-N metrics
        for metric in bon_metrics:
            val = result_dict.get(metric)
            if val is not None:
                # Extract the actual n value from best_of_n/ or best_of_<number>/
                match = re.match(r"best_of_(\d+)/", metric)
                if match:
                    n_value = match.group(1)
                    display_name = metric.replace(f"best_of_{n_value}/", f"best_{n_value}_")
                else:
                    # Fallback to original behavior if no number found
                    display_name = metric.replace("best_of_n/", "best_")
                if isinstance(val, (int, float)):
                    table.add_row(display_name, fmt.format(float(val)))
                else:
                    table.add_row(display_name, str(val))
        
        # Add section separator if rigidity-selected metrics exist
        if rigidity_selected_metrics:
            table.add_section()
        
        # Add rigidity-selected metrics
        for metric in rigidity_selected_metrics:
            val = result_dict.get(metric)
            if val is not None:
                display_name = metric.replace("rigidity_selected/", "selected_")
                if isinstance(val, (int, float)):
                    table.add_row(display_name, fmt.format(float(val)))
                else:
                    table.add_row(display_name, str(val))
    else:
        # Multiple datasets:
        table = Table()
        table.add_column("Metrics", style="bold magenta", justify="left", no_wrap=True)
        
        # Add columns for each dataset
        for dataset_name in dataset_names:
            table.add_column(dataset_name, style="cyan")

        # Add sample count row if provided
        if sample_counts and len(sample_counts) == len(dataset_names):
            count_row = ["Sample Count"]
            for count in sample_counts:
                count_row.append(str(count))
            table.add_row(*count_row)
        
        # Add part count range row if provided
        if part_count_ranges and len(part_count_ranges) == len(dataset_names):
            parts_row = ["Part Count Range"]
            for min_parts, max_parts in part_count_ranges:
                if min_parts == max_parts:
                    parts_row.append(str(min_parts))
                else:
                    parts_row.append(f"[{min_parts}, {max_parts}]")
            table.add_row(*parts_row)
        
        # Add section separator if we added any info rows
        if ((sample_counts and len(sample_counts) == len(dataset_names)) or 
            (part_count_ranges and len(part_count_ranges) == len(dataset_names))):
            table.add_section()

        # Collect all unique base metric names (e.g., "avg/object_chamfer")
        # and their original full names with dataloader index suffix
        base_to_full_metrics = {}
        for i, result_dict in enumerate(results):
            dataloader_idx_suffix = f"/dataloader_idx_{i}" if len(results) > 1 else ""
            for full_metric_key in result_dict.keys():
                if full_metric_key.endswith(dataloader_idx_suffix):
                    base_metric_key = full_metric_key[:-len(dataloader_idx_suffix)]
                else:
                    base_metric_key = full_metric_key
                
                if base_metric_key not in base_to_full_metrics:
                    base_to_full_metrics[base_metric_key] = {}
                base_to_full_metrics[base_metric_key][i] = full_metric_key
        
        # Split into avg and best-of-n base metrics
        base_avg_metrics = sorted([m for m in base_to_full_metrics if m.startswith("avg/")])
        base_bon_metrics = sorted([m for m in base_to_full_metrics if re.match(r"best_of_\w+/", m)])
        base_rigidity_selected_metrics = sorted([m for m in base_to_full_metrics if m.startswith("rigidity_selected/")])
        
        # Add average metrics
        for base_metric in base_avg_metrics:
            row = [base_metric.replace("avg/", "")]
            for i, dataset_name in enumerate(dataset_names):
                full_metric_key = base_to_full_metrics[base_metric].get(i)
                if full_metric_key and i < len(results):
                    val = results[i].get(full_metric_key)
                    if val is None:
                        row.append("-")
                    elif isinstance(val, (int, float)):
                        row.append(fmt.format(float(val)))
                    else:
                        row.append(str(val))
                else:
                    row.append("-")
            table.add_row(*row)

        if base_avg_metrics and base_bon_metrics:
            table.add_section()

        # Add best-of-N metrics
        for base_metric in base_bon_metrics:
            # Extract the actual n value from best_of_n/ or best_of_<number>/
            match = re.match(r"best_of_(\d+)/", base_metric)
            if match:
                n_value = match.group(1)
                display_name = base_metric.replace(f"best_of_{n_value}/", f"best_{n_value}_")
            else:
                # Fallback to original behavior if no number found
                display_name = base_metric.replace("best_of_n/", "best_")
            row = [display_name]
            for i, dataset_name in enumerate(dataset_names):
                full_metric_key = base_to_full_metrics[base_metric].get(i)
                if full_metric_key and i < len(results):
                    val = results[i].get(full_metric_key)
                    if val is None:
                        row.append("-")
                    elif isinstance(val, (int, float)):
                        row.append(fmt.format(float(val)))
                    else:
                        row.append(str(val))
                else:
                    row.append("-")
            table.add_row(*row)

        if base_bon_metrics and base_rigidity_selected_metrics:
            table.add_section()

        # Add rigidity-selected metrics
        for base_metric in base_rigidity_selected_metrics:
            row = [base_metric.replace("rigidity_selected/", "selected_")]
            for i, dataset_name in enumerate(dataset_names):
                full_metric_key = base_to_full_metrics[base_metric].get(i)
                if full_metric_key and i < len(results):
                    val = results[i].get(full_metric_key)
                    if val is None:
                        row.append("-")
                    elif isinstance(val, (int, float)):
                        row.append(fmt.format(float(val)))
                    else:
                        row.append(str(val))
                else:
                    row.append("-")
            table.add_row(*row)
    console = Console()
    console.print(table)


class MetricsMeter:
    """Helper class for accumulating metrics for each dataset.
    Example:
        >>> metrics_meter = MetricsMeter(module)
        >>> metrics_meter.add_metrics(
                dataset_names=["A", "B", "A"], 
                loss=torch.tensor([0.1, 0.2, 0.3]),
                acc=torch.tensor([0.9, 0.8, 0.7]),
            )
        >>> metrics_meter.add_metrics(
                dataset_names=["A", "B", "C"], 
                loss=torch.tensor([0.4, 0.5, 0.6]),
                acc=torch.tensor([0.6, 0.5, 0.4]),
            )
        >>> results = metrics_meter.log_on_epoch_end()
        >>> print(results)
        {
            "A/loss": 0.2667,
            "A/acc": 0.7333,
            "B/loss": 0.35,
            "B/acc": 0.65,
            "C/loss": 0.6,
            "C/acc": 0.4,
            "overall/loss": 0.35,
            "overall/acc": 0.65,
        }
    """

    def __init__(self, module: L.LightningModule):
        self.module = module
        self.reset()

    def reset(self):
        self._sums = defaultdict(lambda: defaultdict(float))
        self._counts = defaultdict(lambda: defaultdict(int))
        self._metrics_seen = set()
        self._part_counts = defaultdict(lambda: defaultdict(list))  # Track part counts per dataset

    def add_metrics(self, dataset_names: List[str], num_parts: torch.Tensor = None, **metrics: torch.Tensor):
        """Accumulate a batch of per-sample metrics."""
        if not metrics:
            return
        
        if any(ds == "overall" for ds in dataset_names):
            raise ValueError("'overall' is a reserved dataset name and cannot be used.")

        B = next(iter(metrics.values())).shape[0]
        if len(dataset_names) != B:
            raise ValueError(f"len(dataset_names)={len(dataset_names)} != batch size {B}")
        for k, t in metrics.items():
            if t.shape[0] != B:
                raise ValueError(f"metric '{k}' has shape {t.shape} != ({B},)")
            self._metrics_seen.add(k)

        for i, ds in enumerate(dataset_names):
            for k, t in metrics.items():
                v = t[i].item()
                self._sums[k][ds] += v
                self._counts[k][ds] += 1
                self._sums[k]["_overall"] += v
                self._counts[k]["_overall"] += 1
            
            # Track part counts if provided
            if num_parts is not None:
                try:
                    if isinstance(num_parts, torch.Tensor):
                        part_count = int(num_parts[i].item())
                    else:
                        # Handle case where num_parts is a list or single value
                        if isinstance(num_parts, (list, tuple)):
                            part_count = int(num_parts[i])
                        else:
                            part_count = int(num_parts)
                    
                    self._part_counts[ds]["parts"].append(part_count)
                    self._part_counts["_overall"]["parts"].append(part_count)
                except Exception as e:
                    print(f"Error processing num_parts: {e}")
                    print(f"num_parts type: {type(num_parts)}, value: {num_parts}")
                    print(f"dataset: {ds}, sample index: {i}")

    def compute_average(self) -> Dict[str, torch.Tensor]:
        """Gather per-dataset sums/counts, and compute global averages."""
        # local dataset list
        local_ds = sorted(
            set(chain.from_iterable(self._counts[k].keys() for k in self._metrics_seen))
            - {"_overall"}
        )

        # gather global dataset list
        world_size = dist.get_world_size() if dist.is_initialized() else 1
        if world_size > 1:
            gathered = [None] * world_size
            dist.all_gather_object(gathered, local_ds)
            global_ds = sorted(set(chain.from_iterable(gathered)))
        else:
            global_ds = local_ds

        # flatten sums and counts with fixed order
        metrics = sorted(self._metrics_seen)
        N = len(global_ds) + 1
        flat_sums = []
        flat_counts = []
        for k in metrics:
            for ds in global_ds:
                flat_sums.append(self._sums[k].get(ds, 0.0))
                flat_counts.append(self._counts[k].get(ds, 0))
            flat_sums.append(self._sums[k].get("_overall", 0.0))
            flat_counts.append(self._counts[k].get("_overall", 0))

        device = getattr(self.module, "device", torch.device("cpu")) or torch.device("cpu")
        sums_t = torch.tensor(flat_sums, dtype=torch.float64, device=device)
        counts_t = torch.tensor(flat_counts, dtype=torch.float64, device=device)

        # all-reduce
        if world_size > 1:
            dist.all_reduce(sums_t, op=dist.ReduceOp.SUM)
            dist.all_reduce(counts_t, op=dist.ReduceOp.SUM)

        # compute average metrics
        results: Dict[str, torch.Tensor] = {}
        for idx, k in enumerate(metrics):
            for j, ds in enumerate(global_ds + ["overall"]):
                pos = idx * N + j
                total_sum = sums_t[pos]
                total_count = counts_t[pos]
                avg = (
                    total_sum / total_count if total_count > 0 
                    else torch.tensor(float("nan"), device=device)
                )
                results[f"{ds}/{k}"] = avg

        # reset for next epoch
        self.reset()
        return results

    def get_sample_counts(self) -> Dict[str, int]:
        """Get sample counts per dataset."""
        # local dataset list
        local_ds = sorted(
            set(chain.from_iterable(self._counts[k].keys() for k in self._metrics_seen))
            - {"_overall"}
        )

        # gather global dataset list
        world_size = dist.get_world_size() if dist.is_initialized() else 1
        if world_size > 1:
            gathered = [None] * world_size
            dist.all_gather_object(gathered, local_ds)
            global_ds = sorted(set(chain.from_iterable(gathered)))
        else:
            global_ds = local_ds

        if not self._metrics_seen:
            return {ds: 0 for ds in global_ds}

        # Use the first metric to get counts (all metrics should have same counts per dataset)
        first_metric = sorted(self._metrics_seen)[0]
        
        # flatten counts with fixed order
        flat_counts = []
        for ds in global_ds:
            flat_counts.append(self._counts[first_metric].get(ds, 0))

        device = getattr(self.module, "device", torch.device("cpu")) or torch.device("cpu")
        counts_t = torch.tensor(flat_counts, dtype=torch.int64, device=device)

        # all-reduce
        if world_size > 1:
            dist.all_reduce(counts_t, op=dist.ReduceOp.SUM)

        # return sample counts per dataset
        sample_counts = {}
        for j, ds in enumerate(global_ds):
            sample_counts[ds] = int(counts_t[j].item())
        
        return sample_counts

    def get_part_count_ranges(self) -> Dict[str, tuple[int, int]]:
        """Get part count ranges (min, max) per dataset."""
        # local dataset list
        local_ds = sorted(
            set(self._part_counts.keys()) - {"_overall"}
        )

        # gather global dataset list
        world_size = dist.get_world_size() if dist.is_initialized() else 1
        if world_size > 1:
            gathered = [None] * world_size
            dist.all_gather_object(gathered, local_ds)
            global_ds = sorted(set(chain.from_iterable(gathered)))
        else:
            global_ds = local_ds

        if not self._part_counts:
            return {ds: (0, 0) for ds in global_ds}

        # Gather part counts from all processes
        part_count_ranges = {}
        for ds in global_ds:
            local_parts = self._part_counts[ds].get("parts", [])
            
            if world_size > 1:
                # Gather part counts from all processes
                gathered_parts = [None] * world_size
                dist.all_gather_object(gathered_parts, local_parts)
                all_parts = []
                for parts_list in gathered_parts:
                    all_parts.extend(parts_list)
            else:
                all_parts = local_parts
            
            if all_parts:
                part_count_ranges[ds] = (min(all_parts), max(all_parts))
            else:
                part_count_ranges[ds] = (0, 0)
        
        return part_count_ranges

	from collections import defaultdict
	from itertools import chain
	import re
	from typing import Dict, List, Any

	import lightning as L
	import torch
	import torch.distributed as dist
	from rich.console import Console
	from rich.table import Table


	def log_metrics_on_step(
	module: L.LightningModule,
	metrics: Dict[str, float],
	prefix: str = "train",
	):
	"""Log per‐step scalars only on rank 0."""
	for name, value in metrics.items():
	module.log(
	f"{prefix}/{name}",
	value,
	on_step=True,
	on_epoch=False,
	prog_bar=True,
	rank_zero_only=True,
	)


	def log_metrics_on_epoch(
	module: L.LightningModule,
	metrics: Dict[str, torch.Tensor],
	prefix: str = "val",
	):
	"""Log per‐epoch scalars, automatically synced & averaged across ranks.

	For single dataset: only log overall metrics to avoid redundancy.
	For multiple datasets: log both per-dataset and overall metrics.
	"""
	# Separate dataset-specific and overall metrics
	dataset_metrics = {}
	overall_metrics = {}

	for name, value in metrics.items():
	if name.startswith("overall/"):
	overall_metrics[name.replace("overall/", "")] = value
	else:
	dataset_metrics[name] = value

	# Determine if we have multiple datasets
	dataset_names = set()
	for name in dataset_metrics.keys():
	if "/" in name:
	dataset_name = name.split("/")[0]
	dataset_names.add(dataset_name)

	# Always log overall metrics with "overall/" prefix to maintain consistency
	# This ensures ModelCheckpoint can always find "val/overall/object_chamfer"
	for name, value in overall_metrics.items():
	module.log(
	f"{prefix}/overall/{name}",
	value,
	on_step=False,
	on_epoch=True,
	prog_bar=True,
	rank_zero_only=True,
	)

	# For multiple datasets, also log per-dataset metrics
	if len(dataset_names) > 1:
	for name, value in dataset_metrics.items():
	module.log(
	f"{prefix}/{name}",
	value,
	on_step=False,
	on_epoch=True,
	prog_bar=True,
	rank_zero_only=True,
	)

	def print_eval_table(
	results: list[dict[str, float]],
	dataset_names: list[str],
	digits: int = 4,
	sample_counts: list[int] = None,
	part_count_ranges: list[tuple[int, int]] = None,
	):
	"""
	Pretty-print a list of evaluation-result dicts with Rich, split into Avg and BoN sections.
	For single dataset, use a more concise format.

	Args:
	results: List of dicts from trainer.test(), each dict corresponds to one dataloader.
	dataset_names: List of dataset names corresponding to each dataloader.
	digits: Number of decimal places for floats.
	sample_counts: Optional list of sample counts for each dataset.
	part_count_ranges: Optional list of (min, max) part count tuples for each dataset.
	"""
	if not results:
	print("No results to display")
	return

	# PyTorch Lightning test() returns a list of dicts, one per dataloader
	# Each dict contains metrics like "avg/object_chamfer", "best_of_n/rotation_error", etc.

	# Collect all unique metric names across all dataloaders
	all_metrics = set()
	for result_dict in results:
	all_metrics.update(result_dict.keys())

	# Split into avg and best-of-n metrics
	avg_metrics = sorted(m for m in all_metrics if m.startswith("avg/"))
	bon_metrics = sorted(m for m in all_metrics if re.match(r"best_of_\w+/", m))
	rigidity_selected_metrics = sorted(m for m in all_metrics if m.startswith("rigidity_selected/"))

	fmt = f"{{:.{digits}f}}"

	# For single dataset, use a more concise format
	if len(results) == 1 and len(dataset_names) >= 1:
	dataset_name = dataset_names[0]
	result_dict = results[0]

	table = Table()
	table.add_column("Metric", style="bold magenta", justify="left", no_wrap=True)
	table.add_column(dataset_name, style="cyan", justify="right")

	# Add sample count row if provided
	if sample_counts and len(sample_counts) >= 1:
	table.add_row("Sample Count", str(sample_counts[0]))

	# Add part count range row if provided
	if part_count_ranges and len(part_count_ranges) >= 1:
	min_parts, max_parts = part_count_ranges[0]
	if min_parts == max_parts:
	table.add_row("Part Count", str(min_parts))
	else:
	table.add_row("Part Count Range", f"[{min_parts}, {max_parts}]")

	# Add section separator if we added any info rows
	if (sample_counts and len(sample_counts) >= 1) or (part_count_ranges and len(part_count_ranges) >= 1):
	table.add_section()

	# Add average metrics
	for metric in avg_metrics:
	val = result_dict.get(metric)
	if val is not None:
	display_name = metric.replace("avg/", "")
	if isinstance(val, (int, float)):
	table.add_row(display_name, fmt.format(float(val)))
	else:
	table.add_row(display_name, str(val))

	# Add section separator if both avg and bon metrics exist
	if avg_metrics and bon_metrics:
	table.add_section()

	# Add best-of-N metrics
	for metric in bon_metrics:
	val = result_dict.get(metric)
	if val is not None:
	# Extract the actual n value from best_of_n/ or best_of_<number>/
	match = re.match(r"best_of_(\d+)/", metric)
	if match:
	n_value = match.group(1)
	display_name = metric.replace(f"best_of_{n_value}/", f"best_{n_value}_")
	else:
	# Fallback to original behavior if no number found
	display_name = metric.replace("best_of_n/", "best_")
	if isinstance(val, (int, float)):
	table.add_row(display_name, fmt.format(float(val)))
	else:
	table.add_row(display_name, str(val))

	# Add section separator if rigidity-selected metrics exist
	if rigidity_selected_metrics:
	table.add_section()

	# Add rigidity-selected metrics
	for metric in rigidity_selected_metrics:
	val = result_dict.get(metric)
	if val is not None:
	display_name = metric.replace("rigidity_selected/", "selected_")
	if isinstance(val, (int, float)):
	table.add_row(display_name, fmt.format(float(val)))
	else:
	table.add_row(display_name, str(val))
	else:
	# Multiple datasets:
	table = Table()
	table.add_column("Metrics", style="bold magenta", justify="left", no_wrap=True)

	# Add columns for each dataset
	for dataset_name in dataset_names:
	table.add_column(dataset_name, style="cyan")

	# Add sample count row if provided
	if sample_counts and len(sample_counts) == len(dataset_names):
	count_row = ["Sample Count"]
	for count in sample_counts:
	count_row.append(str(count))
	table.add_row(*count_row)

	# Add part count range row if provided
	if part_count_ranges and len(part_count_ranges) == len(dataset_names):
	parts_row = ["Part Count Range"]
	for min_parts, max_parts in part_count_ranges:
	if min_parts == max_parts:
	parts_row.append(str(min_parts))
	else:
	parts_row.append(f"[{min_parts}, {max_parts}]")
	table.add_row(*parts_row)

	# Add section separator if we added any info rows
	if ((sample_counts and len(sample_counts) == len(dataset_names)) or
	(part_count_ranges and len(part_count_ranges) == len(dataset_names))):
	table.add_section()

	# Collect all unique base metric names (e.g., "avg/object_chamfer")
	# and their original full names with dataloader index suffix
	base_to_full_metrics = {}
	for i, result_dict in enumerate(results):
	dataloader_idx_suffix = f"/dataloader_idx_{i}" if len(results) > 1 else ""
	for full_metric_key in result_dict.keys():
	if full_metric_key.endswith(dataloader_idx_suffix):
	base_metric_key = full_metric_key[:-len(dataloader_idx_suffix)]
	else:
	base_metric_key = full_metric_key

	if base_metric_key not in base_to_full_metrics:
	base_to_full_metrics[base_metric_key] = {}
	base_to_full_metrics[base_metric_key][i] = full_metric_key

	# Split into avg and best-of-n base metrics
	base_avg_metrics = sorted([m for m in base_to_full_metrics if m.startswith("avg/")])
	base_bon_metrics = sorted([m for m in base_to_full_metrics if re.match(r"best_of_\w+/", m)])
	base_rigidity_selected_metrics = sorted([m for m in base_to_full_metrics if m.startswith("rigidity_selected/")])

	# Add average metrics
	for base_metric in base_avg_metrics:
	row = [base_metric.replace("avg/", "")]
	for i, dataset_name in enumerate(dataset_names):
	full_metric_key = base_to_full_metrics[base_metric].get(i)
	if full_metric_key and i < len(results):
	val = results[i].get(full_metric_key)
	if val is None:
	row.append("-")
	elif isinstance(val, (int, float)):
	row.append(fmt.format(float(val)))
	else:
	row.append(str(val))
	else:
	row.append("-")
	table.add_row(*row)

	if base_avg_metrics and base_bon_metrics:
	table.add_section()

	# Add best-of-N metrics
	for base_metric in base_bon_metrics:
	# Extract the actual n value from best_of_n/ or best_of_<number>/
	match = re.match(r"best_of_(\d+)/", base_metric)
	if match:
	n_value = match.group(1)
	display_name = base_metric.replace(f"best_of_{n_value}/", f"best_{n_value}_")
	else:
	# Fallback to original behavior if no number found
	display_name = base_metric.replace("best_of_n/", "best_")
	row = [display_name]
	for i, dataset_name in enumerate(dataset_names):
	full_metric_key = base_to_full_metrics[base_metric].get(i)
	if full_metric_key and i < len(results):
	val = results[i].get(full_metric_key)
	if val is None:
	row.append("-")
	elif isinstance(val, (int, float)):
	row.append(fmt.format(float(val)))
	else:
	row.append(str(val))
	else:
	row.append("-")
	table.add_row(*row)

	if base_bon_metrics and base_rigidity_selected_metrics:
	table.add_section()

	# Add rigidity-selected metrics
	for base_metric in base_rigidity_selected_metrics:
	row = [base_metric.replace("rigidity_selected/", "selected_")]
	for i, dataset_name in enumerate(dataset_names):
	full_metric_key = base_to_full_metrics[base_metric].get(i)
	if full_metric_key and i < len(results):
	val = results[i].get(full_metric_key)
	if val is None:
	row.append("-")
	elif isinstance(val, (int, float)):
	row.append(fmt.format(float(val)))
	else:
	row.append(str(val))
	else:
	row.append("-")
	table.add_row(*row)
	console = Console()
	console.print(table)


	class MetricsMeter:
	"""Helper class for accumulating metrics for each dataset.

	Example:
	>>> metrics_meter = MetricsMeter(module)
	>>> metrics_meter.add_metrics(
	dataset_names=["A", "B", "A"],
	loss=torch.tensor([0.1, 0.2, 0.3]),
	acc=torch.tensor([0.9, 0.8, 0.7]),
	)
	>>> metrics_meter.add_metrics(
	dataset_names=["A", "B", "C"],
	loss=torch.tensor([0.4, 0.5, 0.6]),
	acc=torch.tensor([0.6, 0.5, 0.4]),
	)
	>>> results = metrics_meter.log_on_epoch_end()
	>>> print(results)
	{
	"A/loss": 0.2667,
	"A/acc": 0.7333,
	"B/loss": 0.35,
	"B/acc": 0.65,
	"C/loss": 0.6,
	"C/acc": 0.4,
	"overall/loss": 0.35,
	"overall/acc": 0.65,
	}
	"""

	def __init__(self, module: L.LightningModule):
	self.module = module
	self.reset()

	def reset(self):
	self._sums = defaultdict(lambda: defaultdict(float))
	self._counts = defaultdict(lambda: defaultdict(int))
	self._metrics_seen = set()
	self._part_counts = defaultdict(lambda: defaultdict(list)) # Track part counts per dataset

	def add_metrics(self, dataset_names: List[str], num_parts: torch.Tensor = None, **metrics: torch.Tensor):
	"""Accumulate a batch of per-sample metrics."""
	if not metrics:
	return

	if any(ds == "overall" for ds in dataset_names):
	raise ValueError("'overall' is a reserved dataset name and cannot be used.")

	B = next(iter(metrics.values())).shape[0]
	if len(dataset_names) != B:
	raise ValueError(f"len(dataset_names)={len(dataset_names)} != batch size {B}")
	for k, t in metrics.items():
	if t.shape[0] != B:
	raise ValueError(f"metric '{k}' has shape {t.shape} != ({B},)")
	self._metrics_seen.add(k)

	for i, ds in enumerate(dataset_names):
	for k, t in metrics.items():
	v = t[i].item()
	self._sums[k][ds] += v
	self._counts[k][ds] += 1
	self._sums[k]["_overall"] += v
	self._counts[k]["_overall"] += 1

	# Track part counts if provided
	if num_parts is not None:
	try:
	if isinstance(num_parts, torch.Tensor):
	part_count = int(num_parts[i].item())
	else:
	# Handle case where num_parts is a list or single value
	if isinstance(num_parts, (list, tuple)):
	part_count = int(num_parts[i])
	else:
	part_count = int(num_parts)

	self._part_counts[ds]["parts"].append(part_count)
	self._part_counts["_overall"]["parts"].append(part_count)
	except Exception as e:
	print(f"Error processing num_parts: {e}")
	print(f"num_parts type: {type(num_parts)}, value: {num_parts}")
	print(f"dataset: {ds}, sample index: {i}")

	def compute_average(self) -> Dict[str, torch.Tensor]:
	"""Gather per-dataset sums/counts, and compute global averages."""
	# local dataset list
	local_ds = sorted(
	set(chain.from_iterable(self._counts[k].keys() for k in self._metrics_seen))
	- {"_overall"}
	)

	# gather global dataset list
	world_size = dist.get_world_size() if dist.is_initialized() else 1
	if world_size > 1:
	gathered = [None] * world_size
	dist.all_gather_object(gathered, local_ds)
	global_ds = sorted(set(chain.from_iterable(gathered)))
	else:
	global_ds = local_ds

	# flatten sums and counts with fixed order
	metrics = sorted(self._metrics_seen)
	N = len(global_ds) + 1
	flat_sums = []
	flat_counts = []
	for k in metrics:
	for ds in global_ds:
	flat_sums.append(self._sums[k].get(ds, 0.0))
	flat_counts.append(self._counts[k].get(ds, 0))
	flat_sums.append(self._sums[k].get("_overall", 0.0))
	flat_counts.append(self._counts[k].get("_overall", 0))

	device = getattr(self.module, "device", torch.device("cpu")) or torch.device("cpu")
	sums_t = torch.tensor(flat_sums, dtype=torch.float64, device=device)
	counts_t = torch.tensor(flat_counts, dtype=torch.float64, device=device)

	# all-reduce
	if world_size > 1:
	dist.all_reduce(sums_t, op=dist.ReduceOp.SUM)
	dist.all_reduce(counts_t, op=dist.ReduceOp.SUM)

	# compute average metrics
	results: Dict[str, torch.Tensor] = {}
	for idx, k in enumerate(metrics):
	for j, ds in enumerate(global_ds + ["overall"]):
	pos = idx * N + j
	total_sum = sums_t[pos]
	total_count = counts_t[pos]
	avg = (
	total_sum / total_count if total_count > 0
	else torch.tensor(float("nan"), device=device)
	)
	results[f"{ds}/{k}"] = avg

	# reset for next epoch
	self.reset()
	return results

	def get_sample_counts(self) -> Dict[str, int]:
	"""Get sample counts per dataset."""
	# local dataset list
	local_ds = sorted(
	set(chain.from_iterable(self._counts[k].keys() for k in self._metrics_seen))
	- {"_overall"}
	)

	# gather global dataset list
	world_size = dist.get_world_size() if dist.is_initialized() else 1
	if world_size > 1:
	gathered = [None] * world_size
	dist.all_gather_object(gathered, local_ds)
	global_ds = sorted(set(chain.from_iterable(gathered)))
	else:
	global_ds = local_ds

	if not self._metrics_seen:
	return {ds: 0 for ds in global_ds}

	# Use the first metric to get counts (all metrics should have same counts per dataset)
	first_metric = sorted(self._metrics_seen)[0]

	# flatten counts with fixed order
	flat_counts = []
	for ds in global_ds:
	flat_counts.append(self._counts[first_metric].get(ds, 0))

	device = getattr(self.module, "device", torch.device("cpu")) or torch.device("cpu")
	counts_t = torch.tensor(flat_counts, dtype=torch.int64, device=device)

	# all-reduce
	if world_size > 1:
	dist.all_reduce(counts_t, op=dist.ReduceOp.SUM)

	# return sample counts per dataset
	sample_counts = {}
	for j, ds in enumerate(global_ds):
	sample_counts[ds] = int(counts_t[j].item())

	return sample_counts

	def get_part_count_ranges(self) -> Dict[str, tuple[int, int]]:
	"""Get part count ranges (min, max) per dataset."""
	# local dataset list
	local_ds = sorted(
	set(self._part_counts.keys()) - {"_overall"}
	)

	# gather global dataset list
	world_size = dist.get_world_size() if dist.is_initialized() else 1
	if world_size > 1:
	gathered = [None] * world_size
	dist.all_gather_object(gathered, local_ds)
	global_ds = sorted(set(chain.from_iterable(gathered)))
	else:
	global_ds = local_ds

	if not self._part_counts:
	return {ds: (0, 0) for ds in global_ds}

	# Gather part counts from all processes
	part_count_ranges = {}
	for ds in global_ds:
	local_parts = self._part_counts[ds].get("parts", [])

	if world_size > 1:
	# Gather part counts from all processes
	gathered_parts = [None] * world_size
	dist.all_gather_object(gathered_parts, local_parts)
	all_parts = []
	for parts_list in gathered_parts:
	all_parts.extend(parts_list)
	else:
	all_parts = local_parts

	if all_parts:
	part_count_ranges[ds] = (min(all_parts), max(all_parts))
	else:
	part_count_ranges[ds] = (0, 0)

	return part_count_ranges