DanLing

danling

AverageMeter

A lightweight utility to compute and store running averages of values.

AverageMeter provides an efficient way to track running statistics (current value, sum, count, average) with minimal memory overhead and support for distributed environments.

Attributes:

Name	Type	Description
val	float	Most recent value added to the meter
bat	float	Most recent value, synchronized across distributed processes
avg	float	Running average of all values, weighted by counts
sum	float	Sum of all values added to the meter
count	int	Total count of values added (considering weights)

Examples:

>>> meter = AverageMeter()
>>> meter.update(0.7)
>>> meter.val
0.7
>>> meter.bat  # Same as val in non-distributed settings
0.7
>>> meter.avg
0.7
>>> meter.update(0.9)
>>> meter.val
0.9
>>> meter.avg
0.8
>>> meter.sum
1.6
>>> meter.count
2
>>> # Weighted update
>>> meter.update(value=0.5, n=3)
>>> meter.avg
0.62
>>> meter.reset()
AverageMeter(val=nan, avg=nan)

See Also

• MetricMeter: Memory-efficient metric tracker that averages metrics batch-by-batch.

Source code in danling/metric/average_meter.py

class AverageMeter:
    r"""
    A lightweight utility to compute and store running averages of values.

    AverageMeter provides an efficient way to track running statistics (current value, sum, count, average)
    with minimal memory overhead and support for distributed environments.

    Attributes:
        val: Most recent value added to the meter
        bat: Most recent value, synchronized across distributed processes
        avg: Running average of all values, weighted by counts
        sum: Sum of all values added to the meter
        count: Total count of values added (considering weights)

    Examples:
        >>> meter = AverageMeter()
        >>> meter.update(0.7)
        >>> meter.val
        0.7
        >>> meter.bat  # Same as val in non-distributed settings
        0.7
        >>> meter.avg
        0.7
        >>> meter.update(0.9)
        >>> meter.val
        0.9
        >>> meter.avg
        0.8
        >>> meter.sum
        1.6
        >>> meter.count
        2
        >>> # Weighted update
        >>> meter.update(value=0.5, n=3)
        >>> meter.avg
        0.62
        >>> meter.reset()
        AverageMeter(val=nan, avg=nan)

    See Also:
        - [`MetricMeter`][danling.metric.metric_meter.MetricMeter]:
            Memory-efficient metric tracker that averages metrics batch-by-batch.
    """

    v: float = 0.0
    n: int = 1
    sum: float = 0.0
    count: int = 0

    def __init__(self) -> None:
        self.reset()

    def reset(self) -> Self:
        r"""
        Resets the meter.
        """

        self.v = 0.0
        self.n = 1
        self.sum = 0.0
        self.count = 0
        return self

    def update(self, value: float | int, n: int = 1) -> None:
        r"""
        Updates the average and current value in the meter.

        Args:
            value: Value to be added to the average.
            n: Number of values to be added.
        """

        self.v = value
        self.n = n
        self.sum += value * n
        self.count += n

    def value(self) -> float:
        if self.count == 0:
            return nan
        return self.v

    def batch(self) -> float:
        world_size = get_world_size()
        if world_size <= 1:
            return self.value()
        synced_tensor = torch.tensor([self.v, self.n], dtype=torch.float64).cuda()
        dist.all_reduce(synced_tensor)
        val, count = synced_tensor
        if count == 0:
            return nan
        return (val / count).item()

    def average(self) -> float:
        world_size = get_world_size()
        if world_size <= 1:
            return self.sum / self.count if self.count != 0 else nan
        synced_tensor = torch.tensor([self.sum, self.count], dtype=torch.float64).cuda()
        dist.all_reduce(synced_tensor)
        val, count = synced_tensor
        if count == 0:
            return nan
        return (val / count).item()

    @property
    def val(self) -> float:
        return self.value()

    @property
    def bat(self) -> float:
        return self.batch()

    @property
    def avg(self) -> float:
        return self.average()

    def __format__(self, format_spec: str) -> str:
        return f"{self.val.__format__(format_spec)} ({self.avg.__format__(format_spec)})"

    def __repr__(self) -> str:
        return f"{self.__class__.__name__}(val={self.val}, avg={self.avg})"

reset

Python

reset() -> Self

Resets the meter.

Source code in danling/metric/average_meter.py

def reset(self) -> Self:
    r"""
    Resets the meter.
    """

    self.v = 0.0
    self.n = 1
    self.sum = 0.0
    self.count = 0
    return self

update

Python

update(value: float | int, n: int = 1) -> None

Updates the average and current value in the meter.

Parameters:

Name	Type	Description	Default
value	float | int	Value to be added to the average.	required
n	int	Number of values to be added.	1

Source code in danling/metric/average_meter.py

def update(self, value: float | int, n: int = 1) -> None:
    r"""
    Updates the average and current value in the meter.

    Args:
        value: Value to be added to the average.
        n: Number of values to be added.
    """

    self.v = value
    self.n = n
    self.sum += value * n
    self.count += n

AverageMeters

Bases: MetricsDict

Manages multiple average meters in one object.

Examples:

>>> meters = AverageMeters()
>>> meters.update({"loss": 0.6, "auroc": 0.7, "r2": 0.8})
>>> f"{meters:.4f}"
'loss: 0.6000 (0.6000)\tauroc: 0.7000 (0.7000)\tr2: 0.8000 (0.8000)'
>>> meters['loss'].update(value=0.9, n=1)
>>> f"{meters:.4f}"
'loss: 0.9000 (0.7500)\tauroc: 0.7000 (0.7000)\tr2: 0.8000 (0.8000)'
>>> meters.sum.dict()
{'loss': 1.5, 'auroc': 0.7, 'r2': 0.8}
>>> meters.count.dict()
{'loss': 2, 'auroc': 1, 'r2': 1}
>>> meters.reset()
AverageMeters(...)
>>> f"{meters:.4f}"
'loss: nan (nan)\tauroc: nan (nan)\tr2: nan (nan)'

See Also

• MetricMeters: Memory-efficient metric tracker that averages multiple metrics batch-by-batch.

Source code in danling/metric/average_meter.py

class AverageMeters(MetricsDict):
    r"""
    Manages multiple average meters in one object.

    Examples:
        >>> meters = AverageMeters()
        >>> meters.update({"loss": 0.6, "auroc": 0.7, "r2": 0.8})
        >>> f"{meters:.4f}"
        'loss: 0.6000 (0.6000)\tauroc: 0.7000 (0.7000)\tr2: 0.8000 (0.8000)'
        >>> meters['loss'].update(value=0.9, n=1)
        >>> f"{meters:.4f}"
        'loss: 0.9000 (0.7500)\tauroc: 0.7000 (0.7000)\tr2: 0.8000 (0.8000)'
        >>> meters.sum.dict()
        {'loss': 1.5, 'auroc': 0.7, 'r2': 0.8}
        >>> meters.count.dict()
        {'loss': 2, 'auroc': 1, 'r2': 1}
        >>> meters.reset()
        AverageMeters(...)
        >>> f"{meters:.4f}"
        'loss: nan (nan)\tauroc: nan (nan)\tr2: nan (nan)'

    See Also:
        - [`MetricMeters`][danling.metric.metric_meter.MetricMeters]:
            Memory-efficient metric tracker that averages multiple metrics batch-by-batch.
    """

    def __init__(self, default_factory: Type[AverageMeter] = AverageMeter, **meters) -> None:
        super().__init__(default_factory=default_factory, **meters)
        for name, meter in self.items():
            if not isinstance(meter, AverageMeter):
                raise ValueError(f"Expected {name} to be an instance of AverageMeter, but got {type(meter)}")

    @property
    def sum(self) -> RoundDict[str, float]:
        return RoundDict({key: meter.sum for key, meter in self.all_items()})

    @property
    def count(self) -> RoundDict[str, int]:
        return RoundDict({key: meter.count for key, meter in self.all_items()})

    def update(self, *args: Dict, **values: int | float) -> None:  # pylint: disable=W0237
        r"""
        Updates the average and current value in all meters.

        Args:
            values: Dict of values to be added to the average.
            n: Number of values to be added.

        Raises:
            ValueError: If the value is not an instance of (int, float).
        """  # noqa: E501

        if args:
            if len(args) > 1:
                raise ValueError("Expected only one positional argument, but got multiple.")
            values = args[0].update(values) or args[0] if values else args[0]

        for meter, value in values.items():
            if isinstance(value, Tensor):
                value = value.item()
            if not isinstance(value, (int, float)):
                raise ValueError(f"Expected values to be int or float, but got {type(value)}")
            self[meter].update(value)

    def set(self, name: str, meter: AverageMeter) -> None:  # pylint: disable=W0237
        if not isinstance(meter, AverageMeter):
            raise ValueError(f"Expected meter to be an instance of AverageMeter, but got {type(meter)}")
        super().set(name, meter)

update

Python

update(*args: Dict, **values: int | float) -> None

Updates the average and current value in all meters.

Parameters:

Name	Type	Description	Default
values	int | float	Dict of values to be added to the average.	{}
n		Number of values to be added.	required

Raises:

Type	Description
ValueError	If the value is not an instance of (int, float).

Source code in danling/metric/average_meter.py

def update(self, *args: Dict, **values: int | float) -> None:  # pylint: disable=W0237
    r"""
    Updates the average and current value in all meters.

    Args:
        values: Dict of values to be added to the average.
        n: Number of values to be added.

    Raises:
        ValueError: If the value is not an instance of (int, float).
    """  # noqa: E501

    if args:
        if len(args) > 1:
            raise ValueError("Expected only one positional argument, but got multiple.")
        values = args[0].update(values) or args[0] if values else args[0]

    for meter, value in values.items():
        if isinstance(value, Tensor):
            value = value.item()
        if not isinstance(value, (int, float)):
            raise ValueError(f"Expected values to be int or float, but got {type(value)}")
        self[meter].update(value)

MetricMeter

Bases: AverageMeter

A memory-efficient metric tracker that computes and averages metrics across batches.

MetricMeter applies a metric function to each batch and maintains running averages without storing the complete history of predictions and labels. This makes it ideal for metrics that can be meaningfully averaged across batches (like accuracy or loss).

Attributes:

Name	Type	Description
metric	Callable	The metric function to compute on each batch
preprocess	Callable	Optional preprocessing function to apply to inputs and targets
val	float	Result from the most recent batch
bat	float	Result from the most recent batch, synchronized across devices
avg	float	Weighted average of all results so far
sum	float	Running sum of (metric × batch_size) values
count	int	Running sum of batch sizes

Parameters:

Name	Type	Description	Default
metric	Callable	Function that computes a metric given input and target tensors	required
preprocess		Function to preprocess inputs before computing the metric	required

Examples:

>>> import torch
>>> from danling.metric.functional import accuracy
>>> meter = MetricMeter(accuracy)
>>> meter.update(torch.tensor([0.1, 0.8, 0.6, 0.2]), torch.tensor([0, 1, 0, 0]))
>>> meter.val
0.75
>>> meter.avg
0.75
>>> meter.update(torch.tensor([0.1, 0.7, 0.3, 0.2, 0.8, 0.4]), torch.tensor([0, 1, 1, 0, 0, 1]))
>>> meter.val
0.5
>>> meter.avg
0.6
>>> meter.sum
6.0
>>> meter.count
10
>>> meter.reset()
MetricMeter(accuracy)
>>> meter.val
nan
>>> meter.avg
nan

Notes

• MetricMeter is more memory-efficient than Metrics because it only stores running statistics
• Only suitable for metrics that can be meaningfully averaged batch-by-batch
• Not suitable for metrics like AUROC that need the entire dataset
• The metric function should accept input and target tensors and return a scalar value
• For multiple metrics, use MetricMeters

See Also

• AverageMeter: A lightweight utility to compute and store running averages of values.

Source code in danling/metric/metric_meter.py

class MetricMeter(AverageMeter):
    r"""
    A memory-efficient metric tracker that computes and averages metrics across batches.

    MetricMeter applies a metric function to each batch and maintains running averages
    without storing the complete history of predictions and labels. This makes it ideal for
    metrics that can be meaningfully averaged across batches (like accuracy or loss).

    Attributes:
        metric: The metric function to compute on each batch
        preprocess: Optional preprocessing function to apply to inputs and targets
        val: Result from the most recent batch
        bat: Result from the most recent batch, synchronized across devices
        avg: Weighted average of all results so far
        sum: Running sum of (metric × batch_size) values
        count: Running sum of batch sizes

    Args:
        metric: Function that computes a metric given input and target tensors
        preprocess: Function to preprocess inputs before computing the metric

    Examples:
        >>> import torch
        >>> from danling.metric.functional import accuracy
        >>> meter = MetricMeter(accuracy)
        >>> meter.update(torch.tensor([0.1, 0.8, 0.6, 0.2]), torch.tensor([0, 1, 0, 0]))
        >>> meter.val
        0.75
        >>> meter.avg
        0.75
        >>> meter.update(torch.tensor([0.1, 0.7, 0.3, 0.2, 0.8, 0.4]), torch.tensor([0, 1, 1, 0, 0, 1]))
        >>> meter.val
        0.5
        >>> meter.avg
        0.6
        >>> meter.sum
        6.0
        >>> meter.count
        10
        >>> meter.reset()
        MetricMeter(accuracy)
        >>> meter.val
        nan
        >>> meter.avg
        nan

    Notes:
        - MetricMeter is more memory-efficient than [`Metrics`][danling.metric.metrics.Metrics]
          because it only stores running statistics
        - Only suitable for metrics that can be meaningfully averaged batch-by-batch
        - Not suitable for metrics like AUROC that need the entire dataset
        - The metric function should accept input and target tensors and return a scalar value
        - For multiple metrics, use [`MetricMeters`][danling.metric.metric_meter.MetricMeters]

    See Also:
        - [`AverageMeter`][danling.metric.average_meter.AverageMeter]:
            A lightweight utility to compute and store running averages of values.
    """

    metric: Callable

    def __init__(
        self,
        metric: Callable,
    ) -> None:
        super().__init__()
        if not callable(metric):
            raise ValueError(f"Expected metric to be callable, but got {type(metric)}")
        self.metric = metric

    def update(  # type: ignore[override] # pylint: disable=W0237
        self,
        input: Tensor | NestedTensor,  # pylint: disable=W0622
        target: Tensor | NestedTensor,
    ) -> None:
        r"""
        Updates the average and current value in the meter.

        Args:
            value: Value to be added to the average.
            n: Number of values to be added.
        """
        n = len(input)
        value = self.metric(input, target)
        if isinstance(value, Tensor):
            value = value.item()
        super().update(value=value, n=n)

    def __repr__(self):
        metric = self.metric
        if isinstance(metric, partial):
            metric = metric.func
        return f"{self.__class__.__name__}({metric.__name__})"

update

Python

update(input: Tensor | NestedTensor, target: Tensor | NestedTensor) -> None

Updates the average and current value in the meter.

Parameters:

Name	Type	Description	Default
value		Value to be added to the average.	required
n		Number of values to be added.	required

Source code in danling/metric/metric_meter.py

def update(  # type: ignore[override] # pylint: disable=W0237
    self,
    input: Tensor | NestedTensor,  # pylint: disable=W0622
    target: Tensor | NestedTensor,
) -> None:
    r"""
    Updates the average and current value in the meter.

    Args:
        value: Value to be added to the average.
        n: Number of values to be added.
    """
    n = len(input)
    value = self.metric(input, target)
    if isinstance(value, Tensor):
        value = value.item()
    super().update(value=value, n=n)

MetricMeters

Bases: AverageMeters

A container for managing multiple MetricMeter instances with shared preprocessing.

MetricMeters allows you to organize and track multiple metrics in a unified interface, with consistent preprocessing applied to all inputs before computing each metric. This is particularly useful when you want to track several metrics that can be meaningfully averaged across batches.

Attributes:

Name	Type	Description
preprocess		Shared preprocessing function for all meters
val	RoundDict[str, float]	Dictionary of current values from all meters
avg	RoundDict[str, float]	Dictionary of running averages from all meters
sum	RoundDict[str, float]	Dictionary of sums from all meters
count	RoundDict[str, int]	Dictionary of counts from all meters

Parameters:

Name	Type	Description	Default
*args		Either metric functions or a Metrics instance to extract metrics from	()
preprocess	Callable	Preprocessing function to apply to inputs before computing metrics	base_preprocess
**meters		Named MetricMeter instances or metric functions	{}

Examples:

>>> import torch
>>> from danling.metric.functional import accuracy, auroc, auprc
>>> meters = MetricMeters(acc=accuracy, auroc=auroc, auprc=auprc)
>>> meters.update([0.1, 0.8, 0.6, 0.2], [0, 1, 0, 0])
>>> meters.sum.dict()
{'acc': 3.0, 'auroc': 4.0, 'auprc': 4.0}
>>> meters.count.dict()
{'acc': 4, 'auroc': 4, 'auprc': 4}
>>> meters['auroc'].update(torch.tensor([0.2, 0.8]), torch.tensor([0, 1]))
>>> meters.sum.dict()
{'acc': 3.0, 'auroc': 6.0, 'auprc': 4.0}
>>> meters.count.dict()
{'acc': 4, 'auroc': 6, 'auprc': 4}
>>> meters.update([[0.1, 0.7, 0.3, 0.2], [0.8, 0.4]], [[0, 0, 1, 0], [0, 0]])
>>> meters.sum.round(2).dict()
{'acc': 6.0, 'auroc': 8.4, 'auprc': 5.5}
>>> meters.count.dict()
{'acc': 10, 'auroc': 12, 'auprc': 10}
>>> meters['auroc'].update(torch.tensor([0.4, 0.8, 0.6, 0.2]), torch.tensor([0, 1, 1, 0]))
>>> meters.avg.round(4).dict()
{'acc': 0.6, 'auroc': 0.775, 'auprc': 0.55}
>>> meters.update(dict(loss=""))
Traceback (most recent call last):
TypeError: ...update() missing 1 required positional argument: 'target'

Notes

• MetricMeters manages multiple MetricMeter instances with shared preprocessing
• Each metric is computed independently but uses the same inputs
• All meters are updated simultaneously when you call update()
• Individual meters can be accessed like dictionary items or attributes

See Also

• AverageMeters: A container for managing multiple average meters in one object.
• Metrics: Metric tracker that stores the complete prediction and target history.

Source code in danling/metric/metric_meter.py

class MetricMeters(AverageMeters):
    r"""
    A container for managing multiple MetricMeter instances with shared preprocessing.

    MetricMeters allows you to organize and track multiple metrics in a unified interface,
    with consistent preprocessing applied to all inputs before computing each metric.
    This is particularly useful when you want to track several metrics that can be
    meaningfully averaged across batches.

    Attributes:
        preprocess: Shared preprocessing function for all meters
        val: Dictionary of current values from all meters
        avg: Dictionary of running averages from all meters
        sum: Dictionary of sums from all meters
        count: Dictionary of counts from all meters

    Args:
        *args: Either metric functions or a Metrics instance to extract metrics from
        preprocess: Preprocessing function to apply to inputs before computing metrics
        **meters: Named MetricMeter instances or metric functions

    Examples:
        >>> import torch
        >>> from danling.metric.functional import accuracy, auroc, auprc
        >>> meters = MetricMeters(acc=accuracy, auroc=auroc, auprc=auprc)
        >>> meters.update([0.1, 0.8, 0.6, 0.2], [0, 1, 0, 0])
        >>> meters.sum.dict()
        {'acc': 3.0, 'auroc': 4.0, 'auprc': 4.0}
        >>> meters.count.dict()
        {'acc': 4, 'auroc': 4, 'auprc': 4}
        >>> meters['auroc'].update(torch.tensor([0.2, 0.8]), torch.tensor([0, 1]))
        >>> meters.sum.dict()
        {'acc': 3.0, 'auroc': 6.0, 'auprc': 4.0}
        >>> meters.count.dict()
        {'acc': 4, 'auroc': 6, 'auprc': 4}
        >>> meters.update([[0.1, 0.7, 0.3, 0.2], [0.8, 0.4]], [[0, 0, 1, 0], [0, 0]])
        >>> meters.sum.round(2).dict()
        {'acc': 6.0, 'auroc': 8.4, 'auprc': 5.5}
        >>> meters.count.dict()
        {'acc': 10, 'auroc': 12, 'auprc': 10}
        >>> meters['auroc'].update(torch.tensor([0.4, 0.8, 0.6, 0.2]), torch.tensor([0, 1, 1, 0]))
        >>> meters.avg.round(4).dict()
        {'acc': 0.6, 'auroc': 0.775, 'auprc': 0.55}
        >>> meters.update(dict(loss=""))  # doctest: +ELLIPSIS
        Traceback (most recent call last):
        TypeError: ...update() missing 1 required positional argument: 'target'

    Notes:
        - `MetricMeters` manages multiple `MetricMeter` instances with shared preprocessing
        - Each metric is computed independently but uses the same inputs
        - All meters are updated simultaneously when you call `update()`
        - Individual meters can be accessed like dictionary items or attributes

    See Also:
        - [`AverageMeters`][danling.metric.average_meter.AverageMeters]:
            A container for managing multiple average meters in one object.
        - [`Metrics`][danling.metric.metrics.Metrics]:
            Metric tracker that stores the complete prediction and target history.
    """

    preprocess = base_preprocess
    default_cls = MetricMeter

    def __init__(
        self,
        *args,
        preprocess: Callable = base_preprocess,
        **meters,
    ) -> None:
        if args:
            from .metrics import Metrics

            if len(args) == 1 and isinstance(args[0], Metrics):
                metrics = args[0]
                for name, metric in metrics.metrics.items():
                    meters.setdefault(name, metric)
                if preprocess is base_preprocess:
                    preprocess = metrics.preprocess
            else:
                for metric in args:
                    if not callable(metric):
                        raise ValueError(f"Expected metric to be callable, but got {type(metric)}")
                    meters.setdefault(metric.__name__, metric)
        self.setattr("preprocess", preprocess)
        super().__init__(default_factory=None, **meters)  # type: ignore[arg-type]

    def update(  # type: ignore[override] # pylint: disable=W0221
        self,
        input: Tensor | NestedTensor | Sequence,  # pylint: disable=W0622
        target: Tensor | NestedTensor | Sequence,
    ) -> None:
        r"""
        Updates the average and current value in all meters.

        Args:
            input: Input values to compute the metrics.
            target: Target values to compute the metrics.
        """

        input, target = self.preprocess(input, target)  # type: ignore[arg-type]
        if (
            isinstance(input, (Tensor, NestedTensor))
            and isinstance(target, (Tensor, NestedTensor))
            and input.ndim == target.ndim + 1
        ):
            input = input.squeeze(-1)
        if isinstance(input, (Tensor, NestedTensor)):
            input = input.detach().to("cpu")
        if isinstance(target, (Tensor, NestedTensor)):
            target = target.detach().to("cpu")
        for meter in self.values():
            meter.update(input, target)

    def set(self, name: str, meter: MetricMeter | Callable) -> None:  # type: ignore[override] # pylint: disable=W0237
        if callable(meter):
            meter = self.getattr("default_cls", MetricMeter)(meter)
        if not isinstance(meter, MetricMeter):
            raise ValueError(f"Expected meter to be an instance of MetricMeter, but got {type(meter)}")
        super().set(name, meter)

    def __repr__(self):
        keys = tuple(i for i in self.keys())
        return f"{self.__class__.__name__}{keys}"

update

Python

update(input: Tensor | NestedTensor | Sequence, target: Tensor | NestedTensor | Sequence) -> None

Updates the average and current value in all meters.

Parameters:

Name	Type	Description	Default
input	Tensor | NestedTensor | Sequence	Input values to compute the metrics.	required
target	Tensor | NestedTensor | Sequence	Target values to compute the metrics.	required

Source code in danling/metric/metric_meter.py

def update(  # type: ignore[override] # pylint: disable=W0221
    self,
    input: Tensor | NestedTensor | Sequence,  # pylint: disable=W0622
    target: Tensor | NestedTensor | Sequence,
) -> None:
    r"""
    Updates the average and current value in all meters.

    Args:
        input: Input values to compute the metrics.
        target: Target values to compute the metrics.
    """

    input, target = self.preprocess(input, target)  # type: ignore[arg-type]
    if (
        isinstance(input, (Tensor, NestedTensor))
        and isinstance(target, (Tensor, NestedTensor))
        and input.ndim == target.ndim + 1
    ):
        input = input.squeeze(-1)
    if isinstance(input, (Tensor, NestedTensor)):
        input = input.detach().to("cpu")
    if isinstance(target, (Tensor, NestedTensor)):
        target = target.detach().to("cpu")
    for meter in self.values():
        meter.update(input, target)

LRScheduler

Bases: _LRScheduler

General learning rate scheduler.

PyTorch LRScheduler is hard to extend. This class is a wrapper of PyTorch LRScheduler, which provides a more general interface. You only needs to add a new scaling which calculates a learning rate ratio (range from 0 to 1) with total progress (range from 0 to 1), and everything else will be done automatically.

Moreover, this class has warmup and cooldown built-in. By default, the first 5% and last 20% of training steps will be warmup and cooldown respectively. You can alternate by passing warmup_steps and cooldown_steps, or disable them by setting them to 0.

Parameters:

Name	Type	Description	Default
optimizer	Optimizer	Wrapped optimizer.	required
total_steps	int	Total number of trainable steps.	required
final_lr_ratio	Optional[float]	Final learning rate ratio to initial learning rate. Defaults to 1e-3.	None
final_lr	Optional[float]	Final learning rate.	None
min_lr	float	Minimal learning rate. Defaults to 1e-9.	1e-09
method	str	Scaling method. Defaults to “cosine”.	'cosine'
warmup_steps	Optional[int]	Number of warmup steps. Defaults to steps // 20.	None
cooldown_steps	Optional[int]	Number of cooldown steps. Defaults to steps // 5.	None
last_epoch	int	The index of last epoch. Defaults to -1.	-1
scaling	Optional[str]	Method to calculate learning rate given ratio, should be one of “percentile” or “numerical”. Defaults to “percentile” if final_lr_ratio is set, otherwise “numerical”.	None

Examples:

>>> from danling.optim import LRScheduler
>>> import torch
>>> from torch import optim
>>> optimizer = optim.SGD([{'params': torch.tensor([0])}], lr=1, momentum=0.9)
>>> scheduler = LRScheduler(optimizer, total_steps=5, final_lr_ratio=1e-5, method='linear')
>>> lrs = []
>>> for epoch in range(5):
...     lrs.append(scheduler.get_lr()[0])
...     scheduler.step()
>>> [round(lr, 10) for lr in lrs]
[0.1, 0.01, 0.001, 0.0001, 1e-09]
>>> scheduler = LRScheduler(optimizer, total_steps=5, final_lr_ratio=1e-5, method='cosine')
>>> lrs = []
>>> for epoch in range(5):
...     lrs.append(scheduler.get_lr()[0])
...     scheduler.step()
>>> [round(lr, 10) for lr in lrs]
[0.3330753446, 0.0187302031, 0.000533897, 3.00232e-05, 1e-09]
>>> scheduler = LRScheduler(optimizer, total_steps=5, final_lr_ratio=1e-5, method='linear', scaling='numerical')
>>> lrs = []
>>> for epoch in range(5):
...     lrs.append(scheduler.get_lr()[0])
...     scheduler.step()
>>> [round(lr, 2) for lr in lrs]
[0.8, 0.6, 0.4, 0.2, 0.0]

Source code in danling/optim/lr_scheduler/lr_scheduler.py

class LRScheduler(lr_scheduler._LRScheduler):  # pylint: disable=protected-access
    r"""
    General learning rate scheduler.

    PyTorch LRScheduler is hard to extend.
    This class is a wrapper of PyTorch LRScheduler, which provides a more general interface.
    You only needs to add a new scaling which calculates a learning rate ratio (range from 0 to 1)
    with total progress (range from 0 to 1), and everything else will be done automatically.

    Moreover, this class has warmup and cooldown built-in.
    By default, the first 5% and last 20% of training steps will be warmup and cooldown respectively.
    You can alternate by passing `warmup_steps` and `cooldown_steps`, or disable them by setting them to 0.

    Args:
        optimizer: Wrapped optimizer.
        total_steps: Total number of trainable steps.
        final_lr_ratio: Final learning rate ratio to initial learning rate.
            Defaults to 1e-3.
        final_lr: Final learning rate.
        min_lr: Minimal learning rate.
            Defaults to 1e-9.
        method: Scaling method.
            Defaults to "cosine".
        warmup_steps: Number of warmup steps.
            Defaults to `steps // 20`.
        cooldown_steps: Number of cooldown steps.
            Defaults to `steps // 5`.
        last_epoch: The index of last epoch.
            Defaults to -1.
        scaling: Method to calculate learning rate given ratio, should be one of "percentile" or "numerical".
            Defaults to "percentile" if `final_lr_ratio` is set, otherwise "numerical".

    Examples:
        >>> from danling.optim import LRScheduler
        >>> import torch
        >>> from torch import optim
        >>> optimizer = optim.SGD([{'params': torch.tensor([0])}], lr=1, momentum=0.9)
        >>> scheduler = LRScheduler(optimizer, total_steps=5, final_lr_ratio=1e-5, method='linear')
        >>> lrs = []
        >>> for epoch in range(5):
        ...     lrs.append(scheduler.get_lr()[0])
        ...     scheduler.step()
        >>> [round(lr, 10) for lr in lrs]
        [0.1, 0.01, 0.001, 0.0001, 1e-09]
        >>> scheduler = LRScheduler(optimizer, total_steps=5, final_lr_ratio=1e-5, method='cosine')
        >>> lrs = []
        >>> for epoch in range(5):
        ...     lrs.append(scheduler.get_lr()[0])
        ...     scheduler.step()
        >>> [round(lr, 10) for lr in lrs]
        [0.3330753446, 0.0187302031, 0.000533897, 3.00232e-05, 1e-09]
        >>> scheduler = LRScheduler(optimizer, total_steps=5, final_lr_ratio=1e-5, method='linear', scaling='numerical')
        >>> lrs = []
        >>> for epoch in range(5):
        ...     lrs.append(scheduler.get_lr()[0])
        ...     scheduler.step()
        >>> [round(lr, 2) for lr in lrs]
        [0.8, 0.6, 0.4, 0.2, 0.0]
    """  # noqa: E501

    def __init__(
        self,
        optimizer: Optimizer,
        total_steps: int,
        final_lr_ratio: Optional[float] = None,
        final_lr: Optional[float] = None,
        min_lr: float = 1e-9,
        method: str = "cosine",
        warmup_steps: Optional[int] = None,
        cooldown_steps: Optional[int] = None,
        last_epoch: int = -1,
        scaling: Optional[str] = None,
        step_with_optimizer: bool = True,
    ):
        if total_steps <= 0:
            raise ValueError(f"Total steps must be positive, but got {total_steps}")
        if warmup_steps is None:
            warmup_steps = total_steps // 20
        elif warmup_steps > total_steps:
            raise ValueError(f"Warmup steps must be less than total steps, but got {warmup_steps} > {total_steps}")
        elif warmup_steps < 0:
            raise ValueError(f"Warmup steps must be positive, but got {warmup_steps}")
        if cooldown_steps is None:
            cooldown_steps = total_steps // 5
        elif cooldown_steps > total_steps:
            raise ValueError(f"Cooldown steps must be less than total steps, but got {cooldown_steps} > {total_steps}")
        elif cooldown_steps < 0:
            raise ValueError(f"Cooldown steps must be positive, but got {cooldown_steps}")
        if warmup_steps + cooldown_steps > total_steps:
            raise ValueError(
                "Warmup steps + cooldown steps must be less than total steps, "
                f"but got {warmup_steps} + {cooldown_steps} > {total_steps}"
            )
        if final_lr_ratio is not None:
            if final_lr is not None:
                raise ValueError("Only one of `final_lr_ratio` and `final_lr` should be set, but not both")
            if final_lr_ratio < 0:
                raise ValueError(f"`final_lr_ratio` must be positive, but got {final_lr_ratio}")
            if scaling is None:
                scaling = "percentile"
        if final_lr is not None and final_lr < 0:
            raise ValueError(f"`final_lr` must be positive, but got {final_lr}")
        if min_lr < 0:
            raise ValueError(f"`min_lr` must be positive, but got {min_lr}")
        self.strategies = {
            k: v for k, v in self.__class__.__dict__.items() if callable(v) and (not k.startswith("_") or k in "get_lr")
        }
        if method not in self.strategies:
            raise ValueError(f"Scaling method must be one of {self.strategies.keys()}, but got {method}")

        if final_lr_ratio is None and final_lr is None:
            final_lr_ratio = 1e-3
            if scaling is None:
                scaling = "percentile"
        if final_lr is not None and min_lr > final_lr:
            min_lr = final_lr
        if scaling is None:
            scaling = "numerical"

        self.final_lr_ratio = final_lr_ratio
        self.final_lr = final_lr
        self.total_steps = total_steps
        self.min_lr = min_lr
        self.method = method
        self.scaling = scaling
        self.warmup_steps = warmup_steps
        self.cooldown_steps = cooldown_steps
        self.cooldown_steps_begin = self.total_steps - self.cooldown_steps
        super().__init__(optimizer, last_epoch)
        if step_with_optimizer:
            self.optimizer.register_step_post_hook(self.step_post_hook)

    def step_post_hook(self, optimizer: Optimizer, *args, **kwargs):
        self.step()

    def get_lr(self) -> List[float]:
        step_count = self._step_count
        if step_count > self.total_steps + 1 or step_count < 1:
            warn(
                f"Step count {step_count} is out of range [1, {self.total_steps + 1}]",
                category=RuntimeWarning,
                stacklevel=2,
            )
        return [self._get_lr(lr, step_count) for lr in self.base_lrs]

    def _get_lr(
        self,
        lr: float,
        step_count: Optional[int] = None,
        progress: Optional[float] = None,
        warmup_ratio: Optional[float] = None,
        cooldown_ratio: Optional[float] = None,
        scaling: Optional[str] = None,
    ) -> float:
        scaling = scaling or self.scaling
        step_count = step_count or self._step_count
        progress = progress or min(max(step_count / self.total_steps, 0.0), 1.0)
        final_lr = self.final_lr if self.final_lr is not None else lr * self.final_lr_ratio  # type: ignore[operator]
        ratio = getattr(self, self.method)(progress)
        if scaling == "percentile":
            lr *= pow(final_lr / lr, ratio)
        elif scaling == "numerical":
            lr = (1 - ratio) * (lr - final_lr) + final_lr
        else:
            raise ValueError(f"Method must be one of ['percentile', 'numerical'], but got {scaling}")
        if self.warmup_steps > step_count > 0:
            warmup_ratio = warmup_ratio or step_count / self.warmup_steps
            lr = warmup_ratio * (lr - self.min_lr) + self.min_lr
        elif self.cooldown_steps > 0 and step_count > self.cooldown_steps_begin:
            cooldown_ratio = cooldown_ratio or 1 - (step_count - self.cooldown_steps_begin) / self.cooldown_steps
            lr = cooldown_ratio * (lr - self.min_lr) + self.min_lr
        return max(self.min_lr, lr)

    def linear(self, progress: float) -> float:
        return progress

    def cosine(self, progress: float) -> float:
        return 1 - ((1 + cos(pi * progress)) / 2)

    def constant(self, progress: float) -> float:  # pylint: disable=unused-argument
        return 0.0

    def __repr__(self) -> str:
        return (
            f"{self.__class__.__name__}({self.method}, scaling={self.scaling}, "
            f"final_lr_ratio={self.final_lr_ratio}, total_steps={self.total_steps}, "
            f"warmup_steps={self.warmup_steps}, cooldown_steps={self.cooldown_steps})"
        )

AccelerateRunner

Bases: TorchRunner, Accelerator

HuggingFace Accelerate-powered unified model workflow runner.

AccelerateRunner integrates with 🤗 Accelerate to provide a simplified interface for model operations that works across various platforms and hardware configurations with minimal code changes.

Key features:

• Seamless multi-platform execution (CPU, GPU, TPU, multiple GPUs)
• Mixed precision support across hardware types
• DeepSpeed integration without complex configuration
• Simplified model, optimizer, and dataloader preparation
• Hardware abstraction for consistent workflow code

AccelerateRunner is particularly useful for:

• Researchers who need to work across different computing environments
• Teams with diverse hardware setups who need consistent workflow code
• Projects that need flexible deployment without platform-specific optimizations
• Quick prototyping that can scale from local testing to distributed environments

This runner inherits from both TorchRunner and Accelerator, providing the comprehensive experiment management of DanLing with the hardware abstraction of Accelerate.

Note

AccelerateRunner prioritizes ease of use over maximum performance. For highly optimized large-scale operations, consider DeepSpeedRunner instead.

See Also

• TorchRunner: PyTorch DDP runner.
• DeepSpeedRunner: DeepSpeed-optimized runner.
• Accelerate Documentation: Official Accelerate docs.

Source code in danling/runner/accelerate_runner.py

class AccelerateRunner(TorchRunner, Accelerator):  # pylint: disable=too-many-public-methods
    r"""
    HuggingFace Accelerate-powered unified model workflow runner.

    AccelerateRunner integrates with 🤗 Accelerate to provide a simplified interface for model operations
    that works across various platforms and hardware configurations with minimal code changes.

    Key features:

    * Seamless multi-platform execution (CPU, GPU, TPU, multiple GPUs)
    * Mixed precision support across hardware types
    * DeepSpeed integration without complex configuration
    * Simplified model, optimizer, and dataloader preparation
    * Hardware abstraction for consistent workflow code

    AccelerateRunner is particularly useful for:

    * Researchers who need to work across different computing environments
    * Teams with diverse hardware setups who need consistent workflow code
    * Projects that need flexible deployment without platform-specific optimizations
    * Quick prototyping that can scale from local testing to distributed environments

    This runner inherits from both TorchRunner and Accelerator, providing the comprehensive
    experiment management of DanLing with the hardware abstraction of Accelerate.

    Note:
        AccelerateRunner prioritizes ease of use over maximum performance. For highly
        optimized large-scale operations, consider DeepSpeedRunner instead.

    See Also:
        - [`TorchRunner`][danling.runner.TorchRunner]: PyTorch DDP runner.
        - [`DeepSpeedRunner`][danling.runner.DeepSpeedRunner]: DeepSpeed-optimized runner.
        - [Accelerate Documentation](https://huggingface.co/docs/accelerate/): Official Accelerate docs.
    """

    _accelerate: FlatDict | None = None

    def __init__(self, config: Config) -> None:
        ac.check()
        TorchRunner.__init__(self, config)
        Accelerator.__init__(self, **self.accelerate)
        if self.distributed:
            object_list = [self.id, self.timestamp]
            dist.broadcast_object_list(object_list)
            self.id, self.timestamp = object_list

    def __post_init__(self) -> None:
        self.project_configuration.set_directories(self.dir)
        if self.datasets:
            self.build_dataloaders()
        self.model, self.ema, self.criterion, self.optimizer, self.scheduler = self.prepare(
            self.model, self.ema, self.criterion, self.optimizer, self.scheduler
        )

    def train_step(self, data) -> torch.Tensor:
        with self.autocast(), self.accumulate():
            input = data["input"] if isinstance(data, Mapping) else data[0]
            target = data["target"] if isinstance(data, Mapping) else data[1]
            pred = self.model(**input) if isinstance(input, Mapping) else self.model(input)
            loss = self.criterion(pred, target)
            if self.metrics is not None:
                self.metrics.update(pred.squeeze(-1), target)
            self.advance(loss)
        return loss

    def advance(self, loss) -> None:
        r"""
        Backward loss and step optimizer & scheduler.

        Args:
            loss: The loss tensor from which to backpropagate.
        """

        self.backward(loss)
        if self.sync_gradients:
            if self.config.get("max_grad_value") is not None:
                self.clip_grad_value_(self.model.parameters(), self.config["max_grad_value"])
            if self.config.get("max_grad_norm") is not None:
                self.clip_grad_norm_(self.model.parameters(), self.config["max_grad_norm"])
        self.optimizer.step()
        self.optimizer.zero_grad()
        self.steps += 1

    def unwrap(self, model: nn.Module) -> nn.Module:
        return self.unwrap_model(model)

    @property
    def accelerate(self) -> FlatDict:
        if self._accelerate is None:
            self._accelerate = self.get_accelerate_config(self.config)
        return self._accelerate

    @accelerate.setter
    def accelerate(self, config: FlatDict) -> None:
        self._accelerate = config

    @property
    def deepspeed(self) -> dict | None:
        if self.state.deepspeed_plugin is not None:
            return self.state.deepspeed_plugin.deepspeed_config
        return None

    @contextmanager
    def accumulate(self, *models: nn.Module):
        if not models:
            models = (self.model,)
        yield Accelerator.accumulate(self, *models)

    @property
    def device(self) -> torch.device:
        return self.state.device

    @property
    def world_size(self) -> int:
        if "state" in self.__dict__:
            return self.state.num_processes
        return 1

    @property
    def rank(self) -> int:
        if "state" in self.__dict__:
            return self.state.process_index
        return 0

    @property
    def local_rank(self) -> int:
        if "state" in self.__dict__:
            return self.state.local_process_index
        return 0

    @cached_property
    def accum_steps(self) -> int:
        return self.gradient_accumulation_steps

    def get_accelerate_config(self, config) -> FlatDict:
        accelerate = FlatDict(step_scheduler_with_optimizer=False)
        if "accelerate" in config:
            accelerate.update(config.accelerate)
        if "precision" in config:
            accelerate.mixed_precision = config.precision
        if "dynamo" in config:
            accelerate.dynamo_backend = config.dynamo.upper()
        if "accum_steps" in config:
            accelerate.gradient_accumulation_steps = config.accum_steps
        if "kwargs_handlers" not in accelerate:
            accelerate.kwargs_handlers = []
        # Must NOT set project_dir here as timestamp is not synced yet
        # config.project_dir = self.dir
        if os.getenv("ACCELERATE_USE_DEEPSPEED", "false").lower() == "true":
            deepspeed_config = config.get("deepspeed", os.getenv("ACCELERATE_DEEPSPEED_CONFIG_FILE"))
            accelerate.deepspeed_plugin = DeepSpeedPlugin(hf_ds_config=self.get_deepspeed_config(deepspeed_config))
        return accelerate

    def build_dataloaders(self):
        datasets = {k: d for k, d in self.datasets.items() if k not in self.dataloaders}
        default_kwargs = self.config.setdefault("dataloader", NestedDict())
        dataloader_kwargs = NestedDict({k: default_kwargs.pop(k) for k in self.datasets if k in default_kwargs})
        for k, d in datasets.items():
            dataloader_kwargs.setdefault(k, NestedDict())
            dataloader_kwargs[k].merge(default_kwargs, overwrite=False)
            dataloader_kwargs[k].setdefault("shuffle", getattr(d, "train", True))
            dataloader_kwargs[k].setdefault("drop_last", not getattr(d, "train", True))
            self.dataloaders[k] = utils.data.DataLoader(d, collate_fn=self.collate_fn, **dataloader_kwargs[k])
        default_kwargs.update(dataloader_kwargs)
        for k, d in self.dataloaders.items():
            self.dataloaders[k] = self.prepare(d)

advance

Python

advance(loss) -> None

Backward loss and step optimizer & scheduler.

Parameters:

Name	Type	Description	Default
loss		The loss tensor from which to backpropagate.	required

Source code in danling/runner/accelerate_runner.py

def advance(self, loss) -> None:
    r"""
    Backward loss and step optimizer & scheduler.

    Args:
        loss: The loss tensor from which to backpropagate.
    """

    self.backward(loss)
    if self.sync_gradients:
        if self.config.get("max_grad_value") is not None:
            self.clip_grad_value_(self.model.parameters(), self.config["max_grad_value"])
        if self.config.get("max_grad_norm") is not None:
            self.clip_grad_norm_(self.model.parameters(), self.config["max_grad_norm"])
    self.optimizer.step()
    self.optimizer.zero_grad()
    self.steps += 1

BaseRunner

Base class for DanLing runners.

The BaseRunner class provides a complete infrastructure for managing:

1. Experiment Environment: Sets up reproducible training environment with:

   • Random seed management
   • Deterministic execution options
   • Configurable logging systems

2. Experiment Lifecycle: Handles the complete experiment workflow:

   • Model, optimizer, and scheduler setup
   • Dataset and dataloader management
   • Training, evaluation, and inference loops
   • Checkpointing and restoration

3. Experiment Tracking: Provides comprehensive experiment monitoring:

   • Progress tracking
   • Metrics computation and logging
   • Results collection and comparison
   • Best model selection

4. Distributed Training: Supports various parallel training configurations:

   • Multi-process distributed training
   • Gradient accumulation
   • Mixed precision training

This class defines standard attributes and methods that can be extended by more specialized runners. It is designed to be modular and configurable to support various experiment requirements.

ID:

Name	Type	Description
timestamp	str	Time string of when the runner was created. Format: WDHMSUUU- (Week, Weekday, Hour, Minute, Second, Microsecond)
name	str	Human-readable name combining experiment and run names ({experiment_name}-{run_name}).
id	str	Unique identifier generated from experiment and run IDs.
uuid	UUID	UUID5 generated from the run ID and combined ID.

See Also

Config: Configuration for runners.

RunnerMode: Enumeration of runner modes.

Model:

Name	Type	Description
model	Callable | None	The model being trained/evaluated.
criterion	Callable | None	Loss function used for optimization.
optimizer	Any | None	Optimizer used for parameter updates.
scheduler	Any | None	Learning rate scheduler for adaptive learning rates.

Data:

Name	Type	Description
datasets	FlatDict	Dictionary mapping split names to dataset objects.
datasamplers	FlatDict	Dictionary mapping split names to data samplers.
dataloaders	FlatDict	Dictionary mapping split names to dataloader objects.
split	str | None	Current active data split.
batch_size	int	Number of samples per batch for current split.
batch_size_equivalent	int	Total effective batch size accounting for distribution.

Progress:

Name	Type	Description
progress	float	Running progress as a value between 0 and 1.
steps	int	Current step count.
epochs	int	Current epoch count.
total_steps	int	Total steps to be performed.
total_epochs	int	Total epochs to be performed.

Results:

Name	Type	Description
results	NestedDict	Hierarchical results organized by epoch, split, and metric.
latest_result	NestedDict | None	Most recent evaluation results.
best_result	NestedDict | None	Best evaluation results achieved.
scores	FlatDict | None	Main metric values used for model comparison.
latest_score	float | None	Most recent score.
best_score	float | None	Best score achieved.
is_best	bool	Whether the latest result is the best so far.

I/O:

Name	Type	Description
dir	str	Root directory for all experiment outputs.
checkpoint_dir	str	Directory for saving model checkpoints.
log_file	str	Path to the log file.

Distributed Training:

Name	Type	Description
world_size	int	Number of processes in distributed training.
rank	int	Global process rank.
local_rank	int	Local process rank on current node.
distributed	bool	Whether running in distributed mode.
is_main_process	bool	Whether current process is the main process.

Logging and Visualization:

Name	Type	Description
meters	AverageMeters	Meters for tracking running averages of metrics.
metrics	Metrics | MetricMeters | None	Metric computation objects.
logger	Logger | None	Logger for capturing runtime information.
writer	Any | None	Writer for visualizing metrics (e.g., TensorBoard).

Source code in danling/runner/base_runner.py

class BaseRunner(metaclass=RunnerMeta):  # pylint: disable=too-many-public-methods
    r"""
    Base class for DanLing runners.

    The `BaseRunner` class provides a complete infrastructure for managing:

    1. **Experiment Environment**: Sets up reproducible training environment with:
        - Random seed management
        - Deterministic execution options
        - Configurable logging systems

    2. **Experiment Lifecycle**: Handles the complete experiment workflow:
        - Model, optimizer, and scheduler setup
        - Dataset and dataloader management
        - Training, evaluation, and inference loops
        - Checkpointing and restoration

    3. **Experiment Tracking**: Provides comprehensive experiment monitoring:
        - Progress tracking
        - Metrics computation and logging
        - Results collection and comparison
        - Best model selection

    4. **Distributed Training**: Supports various parallel training configurations:
        - Multi-process distributed training
        - Gradient accumulation
        - Mixed precision training

    This class defines standard attributes and methods that can be extended by more
    specialized runners. It is designed to be modular and configurable to support
    various experiment requirements.

    Attributes: ID:
        timestamp: Time string of when the runner was created.
            Format: `WDHMSUUU-` (Week, Weekday, Hour, Minute, Second, Microsecond)
        name: Human-readable name combining experiment and run names (`{experiment_name}-{run_name}`).
        id: Unique identifier generated from experiment and run IDs.
        uuid: UUID5 generated from the run ID and combined ID.

    See Also:
        [`Config`][danling.runner.Config]: Configuration for runners.

        [`RunnerMode`][danling.runner.utils.RunnerMode]: Enumeration of runner modes.

    Attributes: Model:
        model: The model being trained/evaluated.
        criterion: Loss function used for optimization.
        optimizer: Optimizer used for parameter updates.
        scheduler: Learning rate scheduler for adaptive learning rates.

    Attributes: Data:
        datasets: Dictionary mapping split names to dataset objects.
        datasamplers: Dictionary mapping split names to data samplers.
        dataloaders: Dictionary mapping split names to dataloader objects.
        split: Current active data split.
        batch_size: Number of samples per batch for current split.
        batch_size_equivalent: Total effective batch size accounting for distribution.

    Attributes: Progress:
        progress: Running progress as a value between 0 and 1.
        steps: Current step count.
        epochs: Current epoch count.
        total_steps: Total steps to be performed.
        total_epochs: Total epochs to be performed.

    Attributes: Results:
        results: Hierarchical results organized by epoch, split, and metric.
        latest_result: Most recent evaluation results.
        best_result: Best evaluation results achieved.
        scores: Main metric values used for model comparison.
        latest_score: Most recent score.
        best_score: Best score achieved.
        is_best: Whether the latest result is the best so far.

    Attributes: I/O:
        dir: Root directory for all experiment outputs.
        checkpoint_dir: Directory for saving model checkpoints.
        log_file: Path to the log file.

    Attributes: Distributed Training:
        world_size: Number of processes in distributed training.
        rank: Global process rank.
        local_rank: Local process rank on current node.
        distributed: Whether running in distributed mode.
        is_main_process: Whether current process is the main process.

    Attributes: Logging and Visualization:
        meters: Meters for tracking running averages of metrics.
        metrics: Metric computation objects.
        logger: Logger for capturing runtime information.
        writer: Writer for visualizing metrics (e.g., TensorBoard).
    """

    # DO NOT set default value in class, as they won't be stored in `__dict__`.

    timestamp: str

    _mode: RunnerMode
    _config: Config
    inited: bool = False

    model: Callable | None = None
    ema: Callable | None = None
    criterion: Callable | None = None
    optimizer: Any | None = None
    scheduler: Any | None = None

    datasets: FlatDict
    datasamplers: FlatDict
    dataloaders: FlatDict
    split: str | None = None
    _train_splits: list[str] | None = None
    _evaluate_splits: list[str] | None = None

    results: NestedDict
    meters: AverageMeters
    metrics: Metrics | MetricMeters | None = None
    train_metrics: Metrics | MetricMeters | None = None
    evaluate_metrics: Metrics | MetricMeters | None = None
    logger: logging.Logger | None = None
    writer: Any | None = None

    def __init__(self, config: Config) -> None:
        self.timestamp = get_time_str()
        if "datasets" not in self.__dict__:
            self.datasets = FlatDict()
        if "datasamplers" not in self.__dict__:
            self.datasamplers = FlatDict()
        if "dataloaders" not in self.__dict__:
            self.dataloaders = FlatDict()
        if "results" not in self.__dict__:
            self.results = NestedDict()
        self.meters = AverageMeters()
        self._mode = RunnerMode.train
        # must init config at last to avoid name conflicts
        if not isinstance(config, Config):
            config = Config(config)
        self._config = config
        self.init_distributed()
        if self.config.log:
            self.init_logging()
        self.init_print()
        if self.config.seed is not None:
            self.set_seed()
        if self.config.deterministic:
            self.set_deterministic()
        if self.config.tensorboard:
            self.init_tensorboard()
        self.inited = True
        if "checkpoint" in config:
            self.load_config(config["checkpoint"])

    def __post_init__(self):
        pass

    def init_distributed(self) -> None:
        r"""
        Initialise distributed running environment.
        """

    @on_main_process
    def init_logging(self) -> None:
        r"""
        Set up logging.
        """

        # Why is setting up proper logging so !@?#! ugly?
        logging.config.dictConfig(
            {
                "version": 1,
                "disable_existing_loggers": False,
                "formatters": {
                    "standard": {"format": "%(asctime)s [%(levelname)s] %(name)s: %(message)s"},
                },
                "handlers": {
                    "stdout": {
                        "level": "INFO",
                        "formatter": "standard",
                        "class": "logging.StreamHandler",
                        "stream": "ext://sys.stdout",
                    },
                    "logfile": {
                        "level": "DEBUG",
                        "formatter": "standard",
                        "class": "logging.FileHandler",
                        "filename": self.log_file,
                        "mode": "a",
                    },
                },
                "loggers": {
                    "": {
                        "handlers": ["stdout", "logfile"],
                        "level": "DEBUG",
                        "propagate": True,
                    },
                },
            }
        )
        logging.captureWarnings(True)
        self.logger = logging.getLogger("runner")
        self.logger.flush = lambda: [h.flush() for h in self.logger.handlers]  # type: ignore

    def init_print(self, process: int = 0) -> None:
        r"""
        Set up `print`.

        Only print on a specific `process` or when `force = True`.

        Args:
            process: The process to `print` on.

        Note:
            If `self.config.log = True`, the default `print` function will be override by `logging.info`.
        """

        logger = logging.getLogger("print")
        logger.flush = lambda: [h.flush for h in logger.handlers]  # type: ignore
        import builtins as __builtin__  # pylint: disable=C0415

        builtin_print = __builtin__.print

        @catch
        def print(*args, force=False, end="\n", file=None, flush=False, **kwargs):  # pylint: disable=redefined-builtin
            if self.rank == process or force:
                if self.config.log:
                    if not args:
                        args = [""]
                    logger.info(*args, **kwargs)
                else:
                    builtin_print(*args, end=end, file=file, flush=flush, **kwargs)

        __builtin__.print = print

    @on_main_process
    def init_tensorboard(self, *args, **kwargs) -> None:
        r"""
        Set up Tensoraoard SummaryWriter.
        """
        raise NotImplementedError

    def set_seed(self, seed: int = None, bias: int = None) -> int:  # type: ignore[assignment]
        r"""
        Set random seeds for reproducibility across the application.

        This method configures the random number generators for various libraries
        to ensure reproducible behavior in experiments. In the base implementation,
        it sets seeds for Python's built-in `random` module and NumPy's random
        generators if available.

        Derived runners typically extend this method to set seeds for additional
        libraries (e.g., PyTorch, TensorFlow).

        Args:
            seed: The base random seed to use.
                If None, uses the value from `self.config.seed`.

            bias: An offset to add to the seed for process-specific randomness.
                - If specified, adds this value to the seed
                - If None, uses `self.rank` (recommended for distributed training)
                - If False, disables bias (all processes use identical seed)

                Using a per-process bias is important in distributed training to avoid
                all processes applying identical data augmentation.

        Returns:
            int: The actual seed that was set (after applying bias)

        Examples:
            # Set seed from config
            runner.set_seed()

            # Set specific seed
            runner.set_seed(42)

            # Set seed with explicit bias
            runner.set_seed(42, bias=10)

            # Set same seed for all processes (not recommended for data augmentation)
            runner.set_seed(42, bias=False)

        Note:
            - Setting seeds is essential for reproducible experiments
            - Different biases across processes ensure diverse data augmentation
            - Subclasses should call `super().set_seed()` if overriding
        """

        seed = seed or self.config.seed  # type: ignore[assignment]
        bias = bias or self.rank
        if bias:
            seed += bias
        if np_random is not None:
            np_random.seed(seed)
        random.seed(seed)
        return seed

    def set_deterministic(self) -> None:
        r"""
        Set up deterministic.
        """

        raise NotImplementedError

    def scale_lr(
        self,
        lr: float,
        lr_scale_factor: float | None = None,
        batch_size_base: int | None = None,
    ) -> float:
        r"""
        Scale learning rate according to [linear scaling rule](https://arxiv.org/abs/1706.02677).
        """

        if lr_scale_factor in self.config:
            lr_scale_factor = self.config.lr_scale_factor

        if lr_scale_factor is None:
            if batch_size_base is None:
                batch_size_base = getattr(self, "batch_size_base", None)
                if batch_size_base is None:
                    raise ValueError("batch_size_base must be specified to auto scale lr")
            lr_scale_factor = self.batch_size_equivalent / batch_size_base
        elif batch_size_base is not None:
            warn(
                "batch_size_base will be ignored if lr_scale_factor is specified", category=RuntimeWarning, stacklevel=2
            )
        lr = lr * lr_scale_factor
        self.config.lr_scale_factor = lr_scale_factor
        return lr

    def advance(self, loss, *args, **kwargs) -> None:
        r"""
        Backward loss and step optimizer & scheduler.

        Args:
            loss: loss.
        """

        raise NotImplementedError

    def state_dict(self, cls: Callable = dict) -> Mapping:
        r"""
        Return dict of all attributes for checkpoint.
        """

        return cls(self.config)

    def dict(self, cls: Callable = dict) -> Mapping:
        r"""
        Convert config to Mapping.

        Args:
            cls: Target `clc to convert to.
        """

        return self.config.dict(cls)

    @catch
    def save(self, obj: Any, file: PathStr, main_process_only: bool = True, *args, **kwargs) -> File:
        r"""
        Save any file with supported extensions.

        `Runner.save` internally calls `dl.save`,
        but with additional arguments to allow it save only on the main process.
        Moreover, any error raised by `Runner.save` will be caught and logged.
        """

        if main_process_only and self.is_main_process or not main_process_only:
            return save(obj, file, *args, **kwargs)
        return file

    @staticmethod
    def load(file: PathStr, *args, **kwargs) -> Any:
        r"""
        Load any file with supported extensions.

        `Runner.load` is identical to `dl.load`.
        """

        return load(file, *args, **kwargs)

    @catch
    def json(self, file: File, main_process_only: bool = True, *args, **kwargs) -> None:  # pylint: disable=R1710
        r"""
        Dump Runner config to json file.
        """

        if main_process_only and self.is_main_process or not main_process_only:
            return self.config.json(file, *args, **kwargs)

    @classmethod
    def from_json(cls, file: File, *args, **kwargs) -> BaseRunner:
        r"""
        Construct Runner from json file.

        This function calls `self.from_jsons()` to construct object from json string.
        You may overwrite `from_jsons` in case something is not json serializable.
        """

        with FlatDict.open(file) as fp:
            return cls.from_jsons(fp.read(), *args, **kwargs)

    def jsons(self, *args, **kwargs) -> str:
        r"""
        Dump Runner config to json string.
        """

        return self.config.jsons(*args, **kwargs)

    @classmethod
    def from_jsons(cls, string: str, *args, **kwargs) -> BaseRunner:
        r"""
        Construct Runner from json string.
        """

        return cls(Config.from_jsons(string, *args, **kwargs))

    @catch
    def yaml(self, file: File, main_process_only: bool = True, *args, **kwargs) -> None:  # pylint: disable=R1710
        r"""
        Dump Runner config to yaml file.
        """

        if main_process_only and self.is_main_process or not main_process_only:
            return self.config.yaml(file, *args, **kwargs)

    @classmethod
    def from_yaml(cls, file: File, *args, **kwargs) -> BaseRunner:
        r"""
        Construct Runner from yaml file.

        This function calls `self.from_yamls()` to construct object from yaml string.
        You may overwrite `from_yamls` in case something is not yaml serializable.
        """

        with FlatDict.open(file) as fp:
            return cls.from_yamls(fp.read(), *args, **kwargs)

    def yamls(self, *args, **kwargs) -> str:
        r"""
        Dump Runner config to yaml string.
        """

        return self.config.yamls(*args, **kwargs)

    @classmethod
    def from_yamls(cls, string: str, *args, **kwargs) -> BaseRunner:
        r"""
        Construct Runner from yaml string.
        """

        return cls(Config.from_yamls(string, *args, **kwargs))

    def check_dir(self, action: str = "warn") -> bool:
        r"""
        Check if `self.dir` is not empty.

        Args:
            action (str): The action to perform if `self.dir` is not empty.
            Can be one of ("warn", "raise", "ignore"), default is "warn".
        """

        if action and action not in ("warn", "raise", "ignore"):
            raise ValueError(f"action should be one of warn, raise or ignore, but got {action}")
        if os.listdir(self.dir):
            if action == "warn":
                warn(
                    f"Directory `{self.dir}` is not empty",
                    category=RuntimeWarning,
                    stacklevel=2,
                )
            if action == "raise":
                raise RuntimeError(f"Directory `{self.dir}` is not empty")
            return False
        return True

    @catch
    @on_main_process
    def save_checkpoint(self, name: str = "latest", epochs: int | None = None, save_best: bool = True) -> None:
        r"""
        Save the current model state and runner configuration to a checkpoint file.

        This method saves the complete state of the runner, including model weights,
        optimizer state, scheduler state, and configuration to the checkpoint directory.
        The method handles various checkpoint naming strategies based on parameters.

        Args:
            name: Base name of the checkpoint file (without extension).
                Defaults to `"latest"`, which creates a checkpoint at `{checkpoint_dir}/latest.pth`.
            epochs: Epoch number to record in the checkpoint.
                If None, uses `self.epochs` (current epoch counter).
            save_best: Whether to also save a copy of the checkpoint as `best.pth`
                when the current model performance is the best so far (`self.is_best` is True).
                Defaults to True.

        Note:
            - Checkpoints are saved only on the main process in distributed training.
            - The format is determined by the extension, defaulting to `.pth`.
            - The checkpoint includes the complete runner state from `self.state_dict()`.
            - Periodic checkpoints are saved based on `save_interval` configuration:
              If `self.config.save_interval > 0` and `(epochs + 1) % save_interval == 0`,
              the checkpoint will be copied to `{checkpoint_dir}/epoch-{epochs}.pth`.

        See Also:
            [`load_checkpoint`][danling.BaseRunner.load_checkpoint]: Load a saved checkpoint.
            [`state_dict`][danling.BaseRunner.state_dict]: Get the state dictionary to save.
        """

        epochs = epochs or self.epochs
        save_interval = self.config.get("save_interval", -1)
        latest_path = os.path.join(self.checkpoint_dir, f"{name}.pth")
        self.save(self.state_dict(), latest_path)
        if save_interval > 0 and (epochs + 1) % save_interval == 0:
            save_path = os.path.join(self.checkpoint_dir, f"epoch-{epochs}.pth")
            shutil.copy(latest_path, save_path)
        if save_best and self.is_best:
            best_path = os.path.join(self.checkpoint_dir, "best.pth")
            shutil.copy(latest_path, best_path)

    def load_config(
        self, checkpoint: Mapping | bytes | str | os.PathLike, overwrite: bool = False, *args, **kwargs
    ) -> None:
        r"""
        Load config from checkpoint.

        Args:
            checkpoint: Checkpoint (or its path) to load.
            overwrite: If `True`, overwrite the current config with the loaded config.
                Defaults to `False`.
            *args: Additional arguments to pass to `self.load`.
            **kwargs: Additional keyword arguments to pass to `self.load`.

        Raises:
            FileNotFoundError: If `checkpoint` does not exists.
        """

        if isinstance(checkpoint, (bytes, str, os.PathLike)):
            if not os.path.exists(checkpoint):
                raise FileNotFoundError(f"checkpoint is set to {checkpoint!r} but does not exist")
            config = self.load(checkpoint, *args, **kwargs)
        elif isinstance(checkpoint, Mapping):
            config = checkpoint
        else:
            raise ValueError(f"checkpoint is set to {checkpoint!r} but is not a valid checkpoint")

        config = config.get("runner", config)
        self.config.merge(config, overwrite=overwrite)
        self.step_begin = config["steps"] + 1
        self.epoch_begin = config["epochs"] + 1

    def load_checkpoint(self, checkpoint: Mapping | bytes | str | os.PathLike, *args, **kwargs) -> None:
        """
        Restore model, optimizer, and scheduler states from a saved checkpoint.

        This method loads the complete training state from a checkpoint, including:
        - Model weights and parameters
        - Optimizer state (learning rates, momentum buffers, etc.)
        - Learning rate scheduler state

        The method supports loading from either:
        - A file path (str, bytes, PathLike)
        - An already loaded checkpoint dictionary (Mapping)

        Args:
            checkpoint: The checkpoint to load from, which can be:
                - A path to the checkpoint file
                - A pre-loaded checkpoint dictionary
            *args: Additional arguments passed to the underlying `self.load` method.
            **kwargs: Additional keyword arguments passed to the underlying `self.load` method.
                      Common options include `map_location` for controlling device mapping.

        Raises:
            ValueError: If `self.model` is not defined (must initialize model before loading).
            ValueError: If `checkpoint` is not a recognized checkpoint format.
            FileNotFoundError: If the checkpoint file specified does not exist.

        Note:
            - This method attempts to unwrap the model before loading (useful for distributed training wrappers)
            - Unlike `load_config`, this method loads both model parameters and optimizer/scheduler states
            - Missing optimizer or scheduler states in the checkpoint will produce warnings but not errors
            - The checkpoint is stored in `self.config.checkpoint` for reference

        Example:
            ```python
            # Load from a file
            runner.load_checkpoint('checkpoints/best.pth')

            # Load with device mapping
            runner.load_checkpoint('checkpoints/latest.pth', map_location='cuda:0')
            ```

        See Also:
            [`from_checkpoint`][danling.BaseRunner.from_checkpoint]: Build a new runner instance from checkpoint.
            [`load_pretrained`][danling.BaseRunner.load_pretrained]: Load only model parameters from a checkpoint.
            [`save_checkpoint`][danling.BaseRunner.save_checkpoint]: Save current state to a checkpoint.
        """

        if self.model is None:
            raise ValueError("model is not defined")
        if isinstance(checkpoint, (bytes, str, os.PathLike)):
            if not os.path.exists(checkpoint):
                raise FileNotFoundError(f"checkpoint is set to {checkpoint!r} but does not exist")
            ckpt = self.load(checkpoint, *args, **kwargs)
        elif isinstance(checkpoint, Mapping):
            ckpt = checkpoint
        else:
            raise ValueError(f"checkpoint is set to {checkpoint!r} but is not a valid checkpoint")

        state_dict = ckpt
        while "model" in state_dict or "module" in state_dict:
            state_dict = state_dict.get("model", state_dict)
            state_dict = state_dict.get("module", state_dict)
        self.unwrap(self.model).load_state_dict(state_dict)
        if self.optimizer is not None:
            if "optimizer" in ckpt:
                self.optimizer.load_state_dict(ckpt["optimizer"])
            else:
                warn("optimizer is not in checkpoint", category=RuntimeWarning, stacklevel=2)
        if self.scheduler is not None:
            if "scheduler" in ckpt:
                self.scheduler.load_state_dict(ckpt["scheduler"])
            else:
                warn("scheduler is not in checkpoint", category=RuntimeWarning, stacklevel=2)
        self.config.checkpoint = checkpoint

    @classmethod
    def from_checkpoint(cls, checkpoint: Mapping | bytes | str | os.PathLike, *args, **kwargs) -> BaseRunner:
        r"""
        Build BaseRunner from checkpoint.

        Args:
            checkpoint: Checkpoint (or its path) to load.
            *args: Additional arguments to pass to `cls.load`.
            **kwargs: Additional keyword arguments to pass to `cls.load`.

        Returns:
            (BaseRunner):
        """

        if isinstance(checkpoint, (bytes, str, os.PathLike)):
            ckpt = cls.load(checkpoint, *args, **kwargs)
        elif isinstance(checkpoint, Mapping):
            ckpt = checkpoint
        else:
            raise ValueError(f"checkpoint is set to {checkpoint} but is not a valid checkpoint")
        runner = cls(ckpt["runner"])
        runner.load_checkpoint(ckpt, override_config=False)
        return runner

    def load_pretrained(self, checkpoint: Mapping | bytes | str | os.PathLike, *args, **kwargs) -> None:
        """
        Load model from pretrained checkpoint.

        This method only loads the model weights.

        Args:
            checkpoint: Pretrained checkpoint (or its path) to load.
            *args: Additional arguments to pass to `self.load`.
            **kwargs: Additional keyword arguments to pass to `self.load`.

        Raises:
            ValueError: If `model` is not defined.
            ValueError: If `checkpoint` is not a valid checkpoint.
            FileNotFoundError: If `checkpoint` does not exists.

        See Also:
            [`load_checkpoint`][danling.BaseRunner.load_checkpoint]: Load model, optimizer, and scheduler from
                checkpoint.
        """

        if self.model is None:
            raise ValueError("model is not defined")
        if isinstance(checkpoint, (bytes, str, os.PathLike)):
            if not os.path.exists(checkpoint):
                raise FileNotFoundError(f"pretrained is set to {checkpoint!r} but does not exist")
            ckpt = self.load(checkpoint, *args, **kwargs)
        elif isinstance(checkpoint, Mapping):
            ckpt = checkpoint
        else:
            raise ValueError(f"pretrained is set to {checkpoint!r} but is not a valid checkpoint")

        state_dict = ckpt
        while "model" in state_dict or "module" in state_dict:
            state_dict = state_dict.get("model", state_dict)
            state_dict = state_dict.get("module", state_dict)
        self.unwrap(self.model).load_state_dict(state_dict)

    def get_step_result(self) -> NestedDict:
        result = self.meters.value()
        if self.metrics is not None:
            return self._merge_result(result, self.metrics.value())
        return result

    def get_epoch_result(self) -> NestedDict:
        result = self.meters.average()
        if self.metrics is not None:
            return self._merge_result(result, self.metrics.average())
        return result

    def _merge_result(self, meter_result, metric_result) -> NestedDict:
        for key, value in metric_result.items():
            if isinstance(value, (Mapping)) and len(value) == 1:
                value = next(iter(value.values()))
            metric_result[key] = value
        meter_result.update(metric_result)
        return meter_result

    def append_result(self, result: NestedDict, index: int | None = None) -> None:
        r"""
        Append result to `self.results`.

        Warnings:
            `self.results` is heavily relied upon for computing metrics.

            Failed to use this method may lead to unexpected behavior.
        """

        if index is None:
            index = self.epochs
            global __APPEND_RESULT_COUNTER__  # pylint: disable=global-statement
            __APPEND_RESULT_COUNTER__ += 1
            if index == 0 and __APPEND_RESULT_COUNTER__ > 1:
                warn(
                    """
                    Automatically set index to `self.epochs`.
                    Please ensure `self.epochs` updates before calling `append_result`
                    """,
                    category=RuntimeWarning,
                    stacklevel=2,
                )
        if index in self.results:
            self.results[index].merge(result)
        else:
            self.results[index] = result

    def print_result(self) -> None:
        r"""
        Print latest and best result.
        """

        print(f"latest result: {self.latest_result}")
        print(f"best result: {self.best_result}")

    def step_log(self, split: str, iteration: int, length: int | None = None):
        if length is None:
            length = len(self.dataloaders[split]) - 1
        result = self.get_step_result()
        print(self.format_step_result(result, split, iteration, length))
        if self.mode == "train":
            self.write_result(result, split)
        return result

    def format_step_result(
        self, result: NestedDict, split: str, steps: int, length: int, format_spec: str = ".4f"
    ) -> str:
        repr_str = ""
        if split is not None:
            if self.mode == "train":
                repr_str = f"training on {split} "
            elif self.mode == "evaluate":
                repr_str = f"evaluating on {split} "
            elif self.mode == "infer":
                repr_str = f"inferring on {split} "
            else:
                repr_str = f"running in {self.mode} mode on {split} "
        repr_str += f"[{steps}/{length}]\t"
        return repr_str + self.format_result(result, format_spec=format_spec)

    def format_epoch_result(
        self, result: NestedDict, epochs: int | None = None, epoch_end: int | None = None, format_spec: str = ".4f"
    ) -> str:
        epochs = epochs or self.epochs
        epoch_end = epoch_end or self.epoch_end
        repr_str = f"epoch [{epochs}/{epoch_end - 1}]" if epochs is not None and epoch_end else ""
        return repr_str + self.format_result(result, format_spec=format_spec)

    def format_result(self, result: Mapping, format_spec: str = ".4f") -> str:
        return format_result(result, format_spec=format_spec)

    def write_result(self, result: NestedDict, split: str, steps: int | None = None):
        if steps is None:
            steps = self.steps
        for name, score in result.all_items():
            name = name.replace(".", "/")
            if name == "loss" and isinstance(score, AverageMeter):
                score = score.avg
            if isinstance(score, Sequence):
                for i, s in enumerate(score):
                    self.write_score(f"{name}/{i}", s, split, steps)
            elif isinstance(score, Mapping):
                for k, s in score.items():
                    self.write_score(f"{name}/{k}", s, split, steps)
            else:
                self.write_score(name, score, split, steps)

    def write_score(self, name: str, score: float, split: str, steps: int):
        if self.writer:
            self.writer.add_scalar(f"{split}/{name}", score, steps)

    @catch
    @on_main_process
    def save_result(self) -> None:
        r"""
        Save result to `self.dir`.

        This method will save latest and best result to
        `self.dir/latest.json` and `self.dir/best.json` respectively.
        """

        results_path = os.path.join(self.dir, "results.json")
        self.save(
            {
                "name": self.name,
                "id": self.id,
                "timestamp": self.timestamp,
                "results": self.results,
            },
            results_path,
            indent=4,
        )
        ret = {"name": self.name, "id": self.id, "timestamp": self.timestamp}
        result = self.latest_result
        if isinstance(result, FlatDict):
            result = result.dict()
        # This is slower but ensure id is the first key
        if result is not None:
            ret.update(result)
        latest_path = os.path.join(self.dir, "latest.json")
        self.save(ret, latest_path, indent=4)
        if self.is_best:
            best_path = os.path.join(self.dir, "best.json")
            shutil.copy(latest_path, best_path)

    def unwrap(self, model: Any) -> Any:
        return model

    @cached_property
    def name(self) -> str:
        if "name" in self.config:
            return self.config["name"]
        return f"{self.config.experiment_name}-{self.config.run_name}"

    @cached_property
    def id(self) -> str:
        return f"{self.config.experiment_id:.8}{self.config.run_id:.8}"

    @cached_property
    def uuid(self) -> UUID:
        r"""
        UUID of the config.
        """

        return uuid5(self.run_uuid, self.id)

    @property
    def mode(self) -> RunnerMode:
        return self._mode

    @mode.setter
    def mode(self, mode: str | RunnerMode) -> None:
        if isinstance(mode, str):
            mode = RunnerMode(mode)
        self._mode = mode

    @property
    def config(self) -> Config:
        return self._config

    @property
    def batch_size(self) -> int:
        r"""
        Batch size.

        Note:
            If `self.dataloaders` is specified:
                - If `self.split` is specified, `batch_size` is the batch size of `self.split`.
                - If `self.train_splits` is specified, `batch_size` is the batch size of the first split.
                - If `self.evaluate_splits` is specified, `batch_size` is the batch size of the first split.
                - Otherwise, `batch_size` is the batch size of the first dataloader.
            Otherwise, `batch_size` is the batch size specified in `self.config.dataloader.batch_size`.
        """

        batch_size = None
        if self.dataloaders:
            if self.split:
                batch_size = self.dataloaders[self.split].batch_size
                if batch_size:
                    return batch_size
            train_splits = self.train_splits or []
            evaluate_splits = self.evaluate_splits or []
            splits = train_splits + evaluate_splits
            for split in splits:
                if split in self.dataloaders:
                    batch_size = self.dataloaders[split].batch_size
                    if batch_size:
                        return batch_size
            for dataloader in self.dataloaders.values():
                batch_size = dataloader.batch_size
                if batch_size:
                    return batch_size
        if not batch_size:
            batch_size = self.config.get("dataloader.batch_size")
        if batch_size:
            return batch_size
        raise AttributeError("batch_size could not be inferred and is not in config")

    @property
    def batch_size_equivalent(self) -> int:
        r"""
        The effective total batch size across all processes and gradient accumulation steps.

        This property calculates the total number of examples processed in a single
        optimization step by accounting for:
        - The base batch size per device
        - The number of processes (world_size for distributed training)
        - The number of gradient accumulation steps

        This is particularly useful for:
        - Learning rate scaling using the linear scaling rule
        - Calculating the total steps in an epoch
        - Estimating memory requirements across all devices

        Returns:
            (int): Effective batch size calculated as `batch_size * world_size * accum_steps`
        """

        return self.batch_size * self.world_size * self.accum_steps

    @property
    def epochs(self) -> int:
        r"""
        Current epoch.
        """
        return self.config.epochs

    @epochs.setter
    def epochs(self, epochs: int) -> None:
        self.config.epochs = epochs

    @property
    def epoch_begin(self) -> int:
        return self.config.epoch_begin

    @epoch_begin.setter
    def epoch_begin(self, epoch_begin: int) -> None:
        self.config.epoch_begin = epoch_begin

    @property
    def epoch_end(self) -> int | None:
        return self.config.epoch_end

    @epoch_end.setter
    def epoch_end(self, epoch_end: int) -> None:
        self.config.epoch_end = epoch_end

    @property
    def total_epochs(self) -> int:
        r"""
        Number of training epochs.
        """
        if self.epoch_end:
            return self.epoch_end - self.epoch_begin
        raise ValueError("epoch_end is not specified")

    @property
    def steps(self) -> int:
        r"""
        Current step.
        """
        return self.config.steps

    @steps.setter
    def steps(self, steps: int) -> None:
        self.config.steps = steps

    @property
    def step_begin(self) -> int:
        return self.config.step_begin

    @step_begin.setter
    def step_begin(self, step_begin: int) -> None:
        self.config.step_begin = step_begin

    @property
    def step_end(self) -> int | None:
        return self.config.step_end

    @step_end.setter
    def step_end(self, step_end: int) -> None:
        self.config.step_end = step_end

    @property
    def total_steps(self) -> int:
        r"""
        Number of training steps.
        """
        if self.step_end:
            return self.step_end - self.step_begin
        train_splits = self.train_splits
        if self.dataloaders and train_splits:
            return self.total_epochs * sum(len(self.dataloaders[split]) for split in train_splits)
        if self.datasets and train_splits:
            datasets_length = sum(len(self.datasets[split]) for split in self.train_splits)
            return self.total_epochs * ceil(datasets_length / self.batch_size / self.world_size)
        raise ValueError("total_steps could not be inferred and is not in config")

    @property
    def accum_steps(self) -> int:
        r"""
        Number of steps to accumulate gradients.
        """

        return self.config.get("accum_steps", 1)

    @property
    def progress(self) -> float:
        r"""
        Training Progress.

        Raises:
            RuntimeError: If no terminal is defined.
        """

        return self.steps / self.total_steps

    @property
    def train_splits(self) -> list[str]:
        if self._train_splits is not None:
            return self._train_splits
        if "train_splits" in self.config:
            self._train_splits = sorted(set(self.config["train_splits"]))
            return self._train_splits
        if self.datasets:
            dataset_splits = {k: getattr(d, "split", None) for k, d in self.datasets.items()}
            train_splits_1 = [k for k, d in self.datasets.items() if getattr(d, "train", False)]
            train_splits_2 = [k for k, s in dataset_splits.items() if s == "train"]
            self._train_splits = sorted(set(train_splits_1 + train_splits_2))
            return self._train_splits
        raise ValueError("train_splits could not be inferred and is not in config")

    @train_splits.setter
    def train_splits(self, train_splits: list[str]) -> None:
        self._train_splits = sorted(set(train_splits))

    @property
    def evaluate_splits(self) -> list[str]:
        if self._evaluate_splits is not None:
            return self._evaluate_splits
        if "evaluate_splits" in self.config:
            self._evaluate_splits = sorted(set(self.config["evaluate_splits"]))
            return self._evaluate_splits
        if self.datasets:
            evaluate_splits = [k for k in self.datasets if k not in self.train_splits]
            self._evaluate_splits = sorted(set(evaluate_splits))
            return self._evaluate_splits
        raise ValueError("evaluate_splits could not be inferred and is not in config")

    @evaluate_splits.setter
    def evaluate_splits(self, evaluate_splits: list[str]) -> None:
        self._evaluate_splits = sorted(set(evaluate_splits))

    @property
    def world_size(self) -> int:
        r"""
        Number of processes.
        """

        return 1

    @property
    def rank(self) -> int:
        r"""
        Process index of all processes.
        """

        return 0

    @property
    def local_rank(self) -> int:
        r"""
        Process index of local processes.
        """

        return 0

    @property
    def distributed(self) -> bool:
        r"""
        If runner is running in distributed mode.
        """

        return self.world_size > 1

    @property
    def is_main_process(self) -> bool:
        r"""
        If current process is the main process of all processes.
        """

        return self.rank == 0

    @property
    def is_local_main_process(self) -> bool:
        r"""
        If current process is the main process of local processes.
        """

        return self.local_rank == 0

    @property
    def best_fn(self) -> Callable:
        r"""
        Function to determine the best score from a list of scores.

        By default, the `best_fn` returns `min` if `self.config.score_name` is `loss`,
        otherwise, returns `max`.

        Subclass can override this method to accommodate needs, such as `min`.
        """

        return max if self.config.score_name != "loss" else min

    @property
    def best_index(self) -> int:
        r"""
        Find the best index from all scores.
        """

        if not self.scores:
            return 0
        return self.best_fn(reversed(self.scores), key=self.scores.get)

    @property
    def latest_result(self) -> NestedDict | None:
        r"""
        Latest result.
        """

        if not self.results:
            return None
        latest_index = next(reversed(self.results if PY38_PLUS else list(self.results)))  # type: ignore
        ret = self.results[latest_index].clone()
        ret["index"] = latest_index
        return ret

    @property
    def best_result(self) -> NestedDict | None:
        r"""
        Best result.
        """

        if not self.results:
            return None
        best_index = self.best_index
        ret = self.results[best_index].clone()
        ret["index"] = best_index
        return ret

    @property
    def scores(self) -> FlatDict | None:
        r"""
        All scores.

        Scores are extracted from results by `score_split` and `runner.config.score_name`,
        following `[r[score_split][self.config.score_name] for r in self.results]`.

        Scores are considered as the index of the performance of the model.
        It is useful to determine the best model and the best hyper-parameters.

        `score_split` is defined in `self.config.score_split`.
        If it is not set, `DanLing` will use `val` or `validate` if they appear in the `latest_result`.
        If `DanLing` still could not find, it will fall back to the second key in the `latest_result`
        if it contains more that one element, or the first key.

        Note that certain keys are ignored when falling back, they are defined in {defaults.IGNORED_NAMES_IN_METRICS}.
        """

        if not self.results:
            return None
        subsets = [i for i in self.latest_result.keys() if i not in defaults.IGNORED_NAMES_IN_METRICS]  # type: ignore
        score_split = self.config.get("score_split")
        if score_split is None and "val" in subsets:
            score_split = "val"
        if score_split is None and "validate" in subsets:
            score_split = "validate"
        if score_split is None:
            score_split = subsets[1] if len(subsets) > 1 else subsets[0]
        return FlatDict({k: v[score_split][self.config.score_name] for k, v in self.results.items()})

    @property
    def latest_score(self) -> float | None:
        r"""
        Latest score.
        """

        if not self.results:
            return None
        if not PY38_PLUS:
            return next(reversed(list(self.scores.values())))  # type: ignore
        return next(reversed(self.scores.values()))  # type: ignore

    @property
    def best_score(self) -> float | None:
        r"""
        Best score.
        """

        if not self.results:
            return None
        return self.scores[self.best_index]  # type: ignore

    @property
    def is_best(self) -> bool:
        r"""
        Determines whether the latest model checkpoint is the best performing one.

        This property compares the `latest_score` with the `best_score` to determine
        if the most recent model evaluation represents the best performance achieved
        so far. The comparison uses a small epsilon (1e-7) to handle floating point
        comparison issues.

        The property is used by:
        - `save_checkpoint` method to determine whether to save a "best.pth" checkpoint
        - Result logging and reporting mechanisms
        - Early stopping implementations

        Returns:
            (bool): True if the latest evaluation results represent the best performance
                 or if no results exist yet, False otherwise.

        Note:
            - Returns True if no results exist (first evaluation)
            - Uses numerical comparison with tolerance to avoid floating point issues
            - The definition of "best" depends on the metric (maximization for accuracy,
              minimization for loss, etc.) and is handled by the `best_fn` property

        See Also:
            [`best_score`][danling.BaseRunner.best_score]: The best score achieved so far.
            [`latest_score`][danling.BaseRunner.latest_score]: The most recent score.
            [`best_fn`][danling.BaseRunner.best_fn]: Function defining "best" (min/max).
        """

        if not self.results:
            return True
        try:
            return abs(self.latest_score - self.best_score) < 1e-7  # type: ignore
        except TypeError:
            return True

    @cached_property
    def log_interval(self) -> int:
        r"""
        Log interval.
        """

        if self.config.get("log_interval") is not None:
            return self.config.log_interval  # type: ignore[return-value]
        if self.dataloaders:
            return max(ceil(max(len(d) for d in self.dataloaders.values()) / 10), 1)
        return 1

    @cached_ensure_dir
    def project_root(self) -> str:
        r"""
        Directory of the project.
        """

        return self.config.get("project_root", "experiments")

    @cached_ensure_dir
    def dir(self) -> str:
        r"""
        Directory of the run.
        """

        if "dir" in self.config:
            return self.config.dir
        return os.path.join(self.project_root, f"{self.name}-{self.id}", self.timestamp)

    @cached_ensure_parent_dir
    def log_file(self) -> str:
        r"""
        Path of log file.
        """

        if "log_file" in self.config:
            return self.config.log_file
        return os.path.join(self.dir, "run.log")

    @cached_ensure_dir
    def checkpoint_dir(self) -> str:
        r"""
        Directory of checkpoints.
        """

        if "checkpoint_dir" in self.config:
            return self.config.checkpoint_dir
        return os.path.join(self.dir, self.config.checkpoint_dir_name)

    # def __getattribute__(self, name) -> Any:
    #     if name in ("__class__", "__dict__"):
    #         return super().__getattribute__(name)
    #     if name in self.__dict__:
    #         return self.__dict__[name]
    #     if name in dir(self):
    #         return super().__getattribute__(name)
    #     if "config" in self and name in self.config:
    #         return self.config[name]
    #     return super().__getattribute__(name)

    def __getattr__(self, name) -> Any:
        if self.inited:
            if name in self.config:
                return self.config[name]
            if name in dir(self.config):
                return getattr(self.config, name)
        return super().__getattribute__(name)

    def __setattr__(self, name, value) -> None:
        if name in self.__dict__:
            if isinstance(self.__dict__[name], Variable):
                self.__dict__[name].set(value)
            else:
                self.__dict__[name] = value
            return
        if name in dir(self):
            if isinstance(super().__getattribute__(name), Variable):
                super().__getattribute__(name).set(value)
            else:
                object.__setattr__(self, name, value)
            return
        if self.inited:
            if name in self.config:
                if isinstance(self.config[name], Variable):
                    self.config[name].set(value)
                else:
                    self.config[name] = value
                return
            if name in dir(self.config):
                setattr(self.config, name, value)
                return
        object.__setattr__(self, name, value)

    def __contains__(self, name) -> bool:
        return name in dir(self) or ("config" in self.__dict__ and name in dir(self.config))

    def __repr__(self):
        lines = []
        for key, value in self.__dict__.items():
            if key.startswith("_") or key.endswith("_") or key == "inited":
                continue
            value_str = repr(value)
            value_str = self._add_indent(value_str)
            lines.append("(" + key + "): " + value_str)

        main_str = self.__class__.__name__ + "("
        if lines:
            main_str += "\n  " + "\n  ".join(lines) + "\n"

        main_str += ")"
        return main_str

    def _add_indent(self, text):
        lines = text.split("\n")
        # don't do anything for single-line stuff
        if len(lines) == 1:
            return text
        first = lines.pop(0)
        # add 2 spaces to each line but the first
        lines = [(2 * " ") + line for line in lines]
        lines = "\n".join(lines)
        lines = first + "\n" + lines
        return lines

batch_size property

Python

batch_size: int

Batch size.

Note

If self.dataloaders is specified: - If self.split is specified, batch_size is the batch size of self.split. - If self.train_splits is specified, batch_size is the batch size of the first split. - If self.evaluate_splits is specified, batch_size is the batch size of the first split. - Otherwise, batch_size is the batch size of the first dataloader. Otherwise, batch_size is the batch size specified in self.config.dataloader.batch_size.

batch_size_equivalent property

Python

batch_size_equivalent: int

The effective total batch size across all processes and gradient accumulation steps.

This property calculates the total number of examples processed in a single optimization step by accounting for: - The base batch size per device - The number of processes (world_size for distributed training) - The number of gradient accumulation steps

This is particularly useful for: - Learning rate scaling using the linear scaling rule - Calculating the total steps in an epoch - Estimating memory requirements across all devices

Returns:

Type	Description
int	Effective batch size calculated as batch_size * world_size * accum_steps

epochs property writable

Python

epochs: int

Current epoch.

total_epochs property

Python

total_epochs: int

Number of training epochs.

steps property writable

Python

steps: int

Current step.

total_steps property

Python

total_steps: int

Number of training steps.

accum_steps property

Python

accum_steps: int

Number of steps to accumulate gradients.

progress property

Python

progress: float

Training Progress.

Raises:

Type	Description
RuntimeError	If no terminal is defined.

world_size property

Python

world_size: int

Number of processes.

rank property

Python

rank: int

Process index of all processes.

local_rank property

Python

local_rank: int

Process index of local processes.

distributed property

Python

distributed: bool

If runner is running in distributed mode.

is_main_process property

Python

is_main_process: bool

If current process is the main process of all processes.

is_local_main_process property

Python

is_local_main_process: bool

If current process is the main process of local processes.

best_fn property

Python

best_fn: Callable

Function to determine the best score from a list of scores.

By default, the best_fn returns min if self.config.score_name is loss, otherwise, returns max.

Subclass can override this method to accommodate needs, such as min.

best_index property

Python

best_index: int

Find the best index from all scores.

latest_result property

Python

latest_result: NestedDict | None

Latest result.

best_result property

Python

best_result: NestedDict | None

Best result.

scores property

Python

scores: FlatDict | None

All scores.

Scores are extracted from results by score_split and runner.config.score_name, following [r[score_split][self.config.score_name] for r in self.results].

Scores are considered as the index of the performance of the model. It is useful to determine the best model and the best hyper-parameters.

score_split is defined in self.config.score_split. If it is not set, DanLing will use val or validate if they appear in the latest_result. If DanLing still could not find, it will fall back to the second key in the latest_result if it contains more that one element, or the first key.

Note that certain keys are ignored when falling back, they are defined in {defaults.IGNORED_NAMES_IN_METRICS}.

latest_score property

Python

latest_score: float | None

Latest score.

best_score property

Python

best_score: float | None

Best score.

is_best property

Python

is_best: bool

Determines whether the latest model checkpoint is the best performing one.

This property compares the latest_score with the best_score to determine if the most recent model evaluation represents the best performance achieved so far. The comparison uses a small epsilon (1e-7) to handle floating point comparison issues.

The property is used by: - save_checkpoint method to determine whether to save a “best.pth” checkpoint - Result logging and reporting mechanisms - Early stopping implementations

Returns:

Type	Description
bool	True if the latest evaluation results represent the best performance or if no results exist yet, False otherwise.

Note

• Returns True if no results exist (first evaluation)
• Uses numerical comparison with tolerance to avoid floating point issues
• The definition of “best” depends on the metric (maximization for accuracy, minimization for loss, etc.) and is handled by the best_fn property

See Also

best_score: The best score achieved so far. latest_score: The most recent score. best_fn: Function defining “best” (min/max).

init_distributed

Python

init_distributed() -> None

Initialise distributed running environment.

Source code in danling/runner/base_runner.py

def init_distributed(self) -> None:
    r"""
    Initialise distributed running environment.
    """

init_logging

Python

init_logging() -> None

Set up logging.

Source code in danling/runner/base_runner.py

@on_main_process
def init_logging(self) -> None:
    r"""
    Set up logging.
    """

    # Why is setting up proper logging so !@?#! ugly?
    logging.config.dictConfig(
        {
            "version": 1,
            "disable_existing_loggers": False,
            "formatters": {
                "standard": {"format": "%(asctime)s [%(levelname)s] %(name)s: %(message)s"},
            },
            "handlers": {
                "stdout": {
                    "level": "INFO",
                    "formatter": "standard",
                    "class": "logging.StreamHandler",
                    "stream": "ext://sys.stdout",
                },
                "logfile": {
                    "level": "DEBUG",
                    "formatter": "standard",
                    "class": "logging.FileHandler",
                    "filename": self.log_file,
                    "mode": "a",
                },
            },
            "loggers": {
                "": {
                    "handlers": ["stdout", "logfile"],
                    "level": "DEBUG",
                    "propagate": True,
                },
            },
        }
    )
    logging.captureWarnings(True)
    self.logger = logging.getLogger("runner")
    self.logger.flush = lambda: [h.flush() for h in self.logger.handlers]  # type: ignore

init_print

Python

init_print(process: int = 0) -> None

Set up print.

Only print on a specific process or when force = True.

Parameters:

Name	Type	Description	Default
process	int	The process to print on.	0

Note

If self.config.log = True, the default print function will be override by logging.info.

Source code in danling/runner/base_runner.py

def init_print(self, process: int = 0) -> None:
    r"""
    Set up `print`.

    Only print on a specific `process` or when `force = True`.

    Args:
        process: The process to `print` on.

    Note:
        If `self.config.log = True`, the default `print` function will be override by `logging.info`.
    """

    logger = logging.getLogger("print")
    logger.flush = lambda: [h.flush for h in logger.handlers]  # type: ignore
    import builtins as __builtin__  # pylint: disable=C0415

    builtin_print = __builtin__.print

    @catch
    def print(*args, force=False, end="\n", file=None, flush=False, **kwargs):  # pylint: disable=redefined-builtin
        if self.rank == process or force:
            if self.config.log:
                if not args:
                    args = [""]
                logger.info(*args, **kwargs)
            else:
                builtin_print(*args, end=end, file=file, flush=flush, **kwargs)

    __builtin__.print = print

init_tensorboard

Python

init_tensorboard(*args, **kwargs) -> None

Set up Tensoraoard SummaryWriter.

Source code in danling/runner/base_runner.py

@on_main_process
def init_tensorboard(self, *args, **kwargs) -> None:
    r"""
    Set up Tensoraoard SummaryWriter.
    """
    raise NotImplementedError

set_seed

Python

set_seed(seed: int = None, bias: int = None) -> int

Set random seeds for reproducibility across the application.

This method configures the random number generators for various libraries to ensure reproducible behavior in experiments. In the base implementation, it sets seeds for Python’s built-in random module and NumPy’s random generators if available.

Derived runners typically extend this method to set seeds for additional libraries (e.g., PyTorch, TensorFlow).

Parameters:

Name	Type	Description	Default
seed	int	The base random seed to use. If None, uses the value from self.config.seed.	None
bias	int	An offset to add to the seed for process-specific randomness. - If specified, adds this value to the seed - If None, uses self.rank (recommended for distributed training) - If False, disables bias (all processes use identical seed) Using a per-process bias is important in distributed training to avoid all processes applying identical data augmentation.	None

Returns:

Name	Type	Description
int	int	The actual seed that was set (after applying bias)

Examples:

Set seed from config

runner.set_seed()

Set specific seed

runner.set_seed(42)

Set seed with explicit bias

runner.set_seed(42, bias=10)

Set same seed for all processes (not recommended for data augmentation)

runner.set_seed(42, bias=False)

Note

• Setting seeds is essential for reproducible experiments
• Different biases across processes ensure diverse data augmentation
• Subclasses should call super().set_seed() if overriding

Source code in danling/runner/base_runner.py

def set_seed(self, seed: int = None, bias: int = None) -> int:  # type: ignore[assignment]
    r"""
    Set random seeds for reproducibility across the application.

    This method configures the random number generators for various libraries
    to ensure reproducible behavior in experiments. In the base implementation,
    it sets seeds for Python's built-in `random` module and NumPy's random
    generators if available.

    Derived runners typically extend this method to set seeds for additional
    libraries (e.g., PyTorch, TensorFlow).

    Args:
        seed: The base random seed to use.
            If None, uses the value from `self.config.seed`.

        bias: An offset to add to the seed for process-specific randomness.
            - If specified, adds this value to the seed
            - If None, uses `self.rank` (recommended for distributed training)
            - If False, disables bias (all processes use identical seed)

            Using a per-process bias is important in distributed training to avoid
            all processes applying identical data augmentation.

    Returns:
        int: The actual seed that was set (after applying bias)

    Examples:
        # Set seed from config
        runner.set_seed()

        # Set specific seed
        runner.set_seed(42)

        # Set seed with explicit bias
        runner.set_seed(42, bias=10)

        # Set same seed for all processes (not recommended for data augmentation)
        runner.set_seed(42, bias=False)

    Note:
        - Setting seeds is essential for reproducible experiments
        - Different biases across processes ensure diverse data augmentation
        - Subclasses should call `super().set_seed()` if overriding
    """

    seed = seed or self.config.seed  # type: ignore[assignment]
    bias = bias or self.rank
    if bias:
        seed += bias
    if np_random is not None:
        np_random.seed(seed)
    random.seed(seed)
    return seed

set_deterministic

Python

set_deterministic() -> None

Set up deterministic.

Source code in danling/runner/base_runner.py

def set_deterministic(self) -> None:
    r"""
    Set up deterministic.
    """

    raise NotImplementedError

scale_lr

Python

scale_lr(lr: float, lr_scale_factor: float | None = None, batch_size_base: int | None = None) -> float

Scale learning rate according to linear scaling rule.

Source code in danling/runner/base_runner.py

def scale_lr(
    self,
    lr: float,
    lr_scale_factor: float | None = None,
    batch_size_base: int | None = None,
) -> float:
    r"""
    Scale learning rate according to [linear scaling rule](https://arxiv.org/abs/1706.02677).
    """

    if lr_scale_factor in self.config:
        lr_scale_factor = self.config.lr_scale_factor

    if lr_scale_factor is None:
        if batch_size_base is None:
            batch_size_base = getattr(self, "batch_size_base", None)
            if batch_size_base is None:
                raise ValueError("batch_size_base must be specified to auto scale lr")
        lr_scale_factor = self.batch_size_equivalent / batch_size_base
    elif batch_size_base is not None:
        warn(
            "batch_size_base will be ignored if lr_scale_factor is specified", category=RuntimeWarning, stacklevel=2
        )
    lr = lr * lr_scale_factor
    self.config.lr_scale_factor = lr_scale_factor
    return lr

advance

Python

advance(loss, *args, **kwargs) -> None

Backward loss and step optimizer & scheduler.

Parameters:

Name	Type	Description	Default
loss		loss.	required

Source code in danling/runner/base_runner.py

def advance(self, loss, *args, **kwargs) -> None:
    r"""
    Backward loss and step optimizer & scheduler.

    Args:
        loss: loss.
    """

    raise NotImplementedError

state_dict

Python

state_dict(cls: Callable = dict) -> Mapping

Return dict of all attributes for checkpoint.

Source code in danling/runner/base_runner.py

def state_dict(self, cls: Callable = dict) -> Mapping:
    r"""
    Return dict of all attributes for checkpoint.
    """

    return cls(self.config)

dict

Python

dict(cls: Callable = dict) -> Mapping

Convert config to Mapping.

Parameters:

Name	Type	Description	Default
cls	Callable	Target `clc to convert to.	dict

Source code in danling/runner/base_runner.py

def dict(self, cls: Callable = dict) -> Mapping:
    r"""
    Convert config to Mapping.

    Args:
        cls: Target `clc to convert to.
    """

    return self.config.dict(cls)

save

Python

save(obj: Any, file: PathStr, main_process_only: bool = True, *args, **kwargs) -> File

Save any file with supported extensions.

Runner.save internally calls dl.save, but with additional arguments to allow it save only on the main process. Moreover, any error raised by Runner.save will be caught and logged.

Source code in danling/runner/base_runner.py

@catch
def save(self, obj: Any, file: PathStr, main_process_only: bool = True, *args, **kwargs) -> File:
    r"""
    Save any file with supported extensions.

    `Runner.save` internally calls `dl.save`,
    but with additional arguments to allow it save only on the main process.
    Moreover, any error raised by `Runner.save` will be caught and logged.
    """

    if main_process_only and self.is_main_process or not main_process_only:
        return save(obj, file, *args, **kwargs)
    return file

load staticmethod

Python

load(file: PathStr, *args, **kwargs) -> Any

Load any file with supported extensions.

Runner.load is identical to dl.load.

Source code in danling/runner/base_runner.py

@staticmethod
def load(file: PathStr, *args, **kwargs) -> Any:
    r"""
    Load any file with supported extensions.

    `Runner.load` is identical to `dl.load`.
    """

    return load(file, *args, **kwargs)

json

Python

json(file: File, main_process_only: bool = True, *args, **kwargs) -> None

Dump Runner config to json file.

Source code in danling/runner/base_runner.py

@catch
def json(self, file: File, main_process_only: bool = True, *args, **kwargs) -> None:  # pylint: disable=R1710
    r"""
    Dump Runner config to json file.
    """

    if main_process_only and self.is_main_process or not main_process_only:
        return self.config.json(file, *args, **kwargs)

from_json classmethod

Python

from_json(file: File, *args, **kwargs) -> BaseRunner

Construct Runner from json file.

This function calls self.from_jsons() to construct object from json string. You may overwrite from_jsons in case something is not json serializable.

Source code in danling/runner/base_runner.py

@classmethod
def from_json(cls, file: File, *args, **kwargs) -> BaseRunner:
    r"""
    Construct Runner from json file.

    This function calls `self.from_jsons()` to construct object from json string.
    You may overwrite `from_jsons` in case something is not json serializable.
    """

    with FlatDict.open(file) as fp:
        return cls.from_jsons(fp.read(), *args, **kwargs)

jsons

Python

jsons(*args, **kwargs) -> str

Dump Runner config to json string.

Source code in danling/runner/base_runner.py

def jsons(self, *args, **kwargs) -> str:
    r"""
    Dump Runner config to json string.
    """

    return self.config.jsons(*args, **kwargs)

from_jsons classmethod

Python

from_jsons(string: str, *args, **kwargs) -> BaseRunner

Construct Runner from json string.

Source code in danling/runner/base_runner.py

@classmethod
def from_jsons(cls, string: str, *args, **kwargs) -> BaseRunner:
    r"""
    Construct Runner from json string.
    """

    return cls(Config.from_jsons(string, *args, **kwargs))

yaml

Python

yaml(file: File, main_process_only: bool = True, *args, **kwargs) -> None

Dump Runner config to yaml file.

Source code in danling/runner/base_runner.py

@catch
def yaml(self, file: File, main_process_only: bool = True, *args, **kwargs) -> None:  # pylint: disable=R1710
    r"""
    Dump Runner config to yaml file.
    """

    if main_process_only and self.is_main_process or not main_process_only:
        return self.config.yaml(file, *args, **kwargs)

from_yaml classmethod

Python

from_yaml(file: File, *args, **kwargs) -> BaseRunner

Construct Runner from yaml file.

This function calls self.from_yamls() to construct object from yaml string. You may overwrite from_yamls in case something is not yaml serializable.

Source code in danling/runner/base_runner.py

@classmethod
def from_yaml(cls, file: File, *args, **kwargs) -> BaseRunner:
    r"""
    Construct Runner from yaml file.

    This function calls `self.from_yamls()` to construct object from yaml string.
    You may overwrite `from_yamls` in case something is not yaml serializable.
    """

    with FlatDict.open(file) as fp:
        return cls.from_yamls(fp.read(), *args, **kwargs)

yamls

Python

yamls(*args, **kwargs) -> str

Dump Runner config to yaml string.

Source code in danling/runner/base_runner.py

def yamls(self, *args, **kwargs) -> str:
    r"""
    Dump Runner config to yaml string.
    """

    return self.config.yamls(*args, **kwargs)

from_yamls classmethod

Python

from_yamls(string: str, *args, **kwargs) -> BaseRunner

Construct Runner from yaml string.

Source code in danling/runner/base_runner.py

@classmethod
def from_yamls(cls, string: str, *args, **kwargs) -> BaseRunner:
    r"""
    Construct Runner from yaml string.
    """

    return cls(Config.from_yamls(string, *args, **kwargs))

check_dir

Python

check_dir(action: str = 'warn') -> bool

Check if self.dir is not empty.

Parameters:

Name	Type	Description	Default
action	str	The action to perform if self.dir is not empty.	'warn'

Source code in danling/runner/base_runner.py

def check_dir(self, action: str = "warn") -> bool:
    r"""
    Check if `self.dir` is not empty.

    Args:
        action (str): The action to perform if `self.dir` is not empty.
        Can be one of ("warn", "raise", "ignore"), default is "warn".
    """

    if action and action not in ("warn", "raise", "ignore"):
        raise ValueError(f"action should be one of warn, raise or ignore, but got {action}")
    if os.listdir(self.dir):
        if action == "warn":
            warn(
                f"Directory `{self.dir}` is not empty",
                category=RuntimeWarning,
                stacklevel=2,
            )
        if action == "raise":
            raise RuntimeError(f"Directory `{self.dir}` is not empty")
        return False
    return True

save_checkpoint

Python

save_checkpoint(name: str = 'latest', epochs: int | None = None, save_best: bool = True) -> None

Save the current model state and runner configuration to a checkpoint file.

This method saves the complete state of the runner, including model weights, optimizer state, scheduler state, and configuration to the checkpoint directory. The method handles various checkpoint naming strategies based on parameters.

Parameters:

Name	Type	Description	Default
name	str	Base name of the checkpoint file (without extension). Defaults to "latest", which creates a checkpoint at {checkpoint_dir}/latest.pth.	'latest'
epochs	int | None	Epoch number to record in the checkpoint. If None, uses self.epochs (current epoch counter).	None
save_best	bool	Whether to also save a copy of the checkpoint as best.pth when the current model performance is the best so far (self.is_best is True). Defaults to True.	True

Note

• Checkpoints are saved only on the main process in distributed training.
• The format is determined by the extension, defaulting to .pth.
• The checkpoint includes the complete runner state from self.state_dict().
• Periodic checkpoints are saved based on save_interval configuration: If self.config.save_interval > 0 and (epochs + 1) % save_interval == 0, the checkpoint will be copied to {checkpoint_dir}/epoch-{epochs}.pth.

See Also

load_checkpoint: Load a saved checkpoint. state_dict: Get the state dictionary to save.

Source code in danling/runner/base_runner.py

@catch
@on_main_process
def save_checkpoint(self, name: str = "latest", epochs: int | None = None, save_best: bool = True) -> None:
    r"""
    Save the current model state and runner configuration to a checkpoint file.

    This method saves the complete state of the runner, including model weights,
    optimizer state, scheduler state, and configuration to the checkpoint directory.
    The method handles various checkpoint naming strategies based on parameters.

    Args:
        name: Base name of the checkpoint file (without extension).
            Defaults to `"latest"`, which creates a checkpoint at `{checkpoint_dir}/latest.pth`.
        epochs: Epoch number to record in the checkpoint.
            If None, uses `self.epochs` (current epoch counter).
        save_best: Whether to also save a copy of the checkpoint as `best.pth`
            when the current model performance is the best so far (`self.is_best` is True).
            Defaults to True.

    Note:
        - Checkpoints are saved only on the main process in distributed training.
        - The format is determined by the extension, defaulting to `.pth`.
        - The checkpoint includes the complete runner state from `self.state_dict()`.
        - Periodic checkpoints are saved based on `save_interval` configuration:
          If `self.config.save_interval > 0` and `(epochs + 1) % save_interval == 0`,
          the checkpoint will be copied to `{checkpoint_dir}/epoch-{epochs}.pth`.

    See Also:
        [`load_checkpoint`][danling.BaseRunner.load_checkpoint]: Load a saved checkpoint.
        [`state_dict`][danling.BaseRunner.state_dict]: Get the state dictionary to save.
    """

    epochs = epochs or self.epochs
    save_interval = self.config.get("save_interval", -1)
    latest_path = os.path.join(self.checkpoint_dir, f"{name}.pth")
    self.save(self.state_dict(), latest_path)
    if save_interval > 0 and (epochs + 1) % save_interval == 0:
        save_path = os.path.join(self.checkpoint_dir, f"epoch-{epochs}.pth")
        shutil.copy(latest_path, save_path)
    if save_best and self.is_best:
        best_path = os.path.join(self.checkpoint_dir, "best.pth")
        shutil.copy(latest_path, best_path)

load_config

Python

load_config(checkpoint: Mapping | bytes | str | PathLike, overwrite: bool = False, *args, **kwargs) -> None

Load config from checkpoint.

Parameters:

Name	Type	Description	Default
checkpoint	Mapping | bytes | str | PathLike	Checkpoint (or its path) to load.	required
overwrite	bool	If True, overwrite the current config with the loaded config. Defaults to False.	False
*args		Additional arguments to pass to self.load.	()
**kwargs		Additional keyword arguments to pass to self.load.	{}

Raises:

Type	Description
FileNotFoundError	If checkpoint does not exists.

Source code in danling/runner/base_runner.py

def load_config(
    self, checkpoint: Mapping | bytes | str | os.PathLike, overwrite: bool = False, *args, **kwargs
) -> None:
    r"""
    Load config from checkpoint.

    Args:
        checkpoint: Checkpoint (or its path) to load.
        overwrite: If `True`, overwrite the current config with the loaded config.
            Defaults to `False`.
        *args: Additional arguments to pass to `self.load`.
        **kwargs: Additional keyword arguments to pass to `self.load`.

    Raises:
        FileNotFoundError: If `checkpoint` does not exists.
    """

    if isinstance(checkpoint, (bytes, str, os.PathLike)):
        if not os.path.exists(checkpoint):
            raise FileNotFoundError(f"checkpoint is set to {checkpoint!r} but does not exist")
        config = self.load(checkpoint, *args, **kwargs)
    elif isinstance(checkpoint, Mapping):
        config = checkpoint
    else:
        raise ValueError(f"checkpoint is set to {checkpoint!r} but is not a valid checkpoint")

    config = config.get("runner", config)
    self.config.merge(config, overwrite=overwrite)
    self.step_begin = config["steps"] + 1
    self.epoch_begin = config["epochs"] + 1

load_checkpoint

Python

load_checkpoint(checkpoint: Mapping | bytes | str | PathLike, *args, **kwargs) -> None

Restore model, optimizer, and scheduler states from a saved checkpoint.

This method loads the complete training state from a checkpoint, including: - Model weights and parameters - Optimizer state (learning rates, momentum buffers, etc.) - Learning rate scheduler state

The method supports loading from either: - A file path (str, bytes, PathLike) - An already loaded checkpoint dictionary (Mapping)

Parameters:

Name	Type	Description	Default
checkpoint	Mapping | bytes | str | PathLike	The checkpoint to load from, which can be: - A path to the checkpoint file - A pre-loaded checkpoint dictionary	required
*args		Additional arguments passed to the underlying self.load method.	()
**kwargs		Additional keyword arguments passed to the underlying self.load method. Common options include map_location for controlling device mapping.	{}

Raises:

Type	Description
ValueError	If self.model is not defined (must initialize model before loading).
ValueError	If checkpoint is not a recognized checkpoint format.
FileNotFoundError	If the checkpoint file specified does not exist.

Note

• This method attempts to unwrap the model before loading (useful for distributed training wrappers)
• Unlike load_config, this method loads both model parameters and optimizer/scheduler states
• Missing optimizer or scheduler states in the checkpoint will produce warnings but not errors
• The checkpoint is stored in self.config.checkpoint for reference

Example

# Load from a file
runner.load_checkpoint('checkpoints/best.pth')

# Load with device mapping
runner.load_checkpoint('checkpoints/latest.pth', map_location='cuda:0')

See Also

from_checkpoint: Build a new runner instance from checkpoint. load_pretrained: Load only model parameters from a checkpoint. save_checkpoint: Save current state to a checkpoint.

Source code in danling/runner/base_runner.py

def load_checkpoint(self, checkpoint: Mapping | bytes | str | os.PathLike, *args, **kwargs) -> None:
    """
    Restore model, optimizer, and scheduler states from a saved checkpoint.

    This method loads the complete training state from a checkpoint, including:
    - Model weights and parameters
    - Optimizer state (learning rates, momentum buffers, etc.)
    - Learning rate scheduler state

    The method supports loading from either:
    - A file path (str, bytes, PathLike)
    - An already loaded checkpoint dictionary (Mapping)

    Args:
        checkpoint: The checkpoint to load from, which can be:
            - A path to the checkpoint file
            - A pre-loaded checkpoint dictionary
        *args: Additional arguments passed to the underlying `self.load` method.
        **kwargs: Additional keyword arguments passed to the underlying `self.load` method.
                  Common options include `map_location` for controlling device mapping.

    Raises:
        ValueError: If `self.model` is not defined (must initialize model before loading).
        ValueError: If `checkpoint` is not a recognized checkpoint format.
        FileNotFoundError: If the checkpoint file specified does not exist.

    Note:
        - This method attempts to unwrap the model before loading (useful for distributed training wrappers)
        - Unlike `load_config`, this method loads both model parameters and optimizer/scheduler states
        - Missing optimizer or scheduler states in the checkpoint will produce warnings but not errors
        - The checkpoint is stored in `self.config.checkpoint` for reference

    Example:
        ```python
        # Load from a file
        runner.load_checkpoint('checkpoints/best.pth')

        # Load with device mapping
        runner.load_checkpoint('checkpoints/latest.pth', map_location='cuda:0')
        ```

    See Also:
        [`from_checkpoint`][danling.BaseRunner.from_checkpoint]: Build a new runner instance from checkpoint.
        [`load_pretrained`][danling.BaseRunner.load_pretrained]: Load only model parameters from a checkpoint.
        [`save_checkpoint`][danling.BaseRunner.save_checkpoint]: Save current state to a checkpoint.
    """

    if self.model is None:
        raise ValueError("model is not defined")
    if isinstance(checkpoint, (bytes, str, os.PathLike)):
        if not os.path.exists(checkpoint):
            raise FileNotFoundError(f"checkpoint is set to {checkpoint!r} but does not exist")
        ckpt = self.load(checkpoint, *args, **kwargs)
    elif isinstance(checkpoint, Mapping):
        ckpt = checkpoint
    else:
        raise ValueError(f"checkpoint is set to {checkpoint!r} but is not a valid checkpoint")

    state_dict = ckpt
    while "model" in state_dict or "module" in state_dict:
        state_dict = state_dict.get("model", state_dict)
        state_dict = state_dict.get("module", state_dict)
    self.unwrap(self.model).load_state_dict(state_dict)
    if self.optimizer is not None:
        if "optimizer" in ckpt:
            self.optimizer.load_state_dict(ckpt["optimizer"])
        else:
            warn("optimizer is not in checkpoint", category=RuntimeWarning, stacklevel=2)
    if self.scheduler is not None:
        if "scheduler" in ckpt:
            self.scheduler.load_state_dict(ckpt["scheduler"])
        else:
            warn("scheduler is not in checkpoint", category=RuntimeWarning, stacklevel=2)
    self.config.checkpoint = checkpoint

from_checkpoint classmethod

Python

from_checkpoint(checkpoint: Mapping | bytes | str | PathLike, *args, **kwargs) -> BaseRunner

Build BaseRunner from checkpoint.

Parameters:

Name	Type	Description	Default
checkpoint	Mapping | bytes | str | PathLike	Checkpoint (or its path) to load.	required
*args		Additional arguments to pass to cls.load.	()
**kwargs		Additional keyword arguments to pass to cls.load.	{}

Returns:

Type	Description
BaseRunner

Source code in danling/runner/base_runner.py

@classmethod
def from_checkpoint(cls, checkpoint: Mapping | bytes | str | os.PathLike, *args, **kwargs) -> BaseRunner:
    r"""
    Build BaseRunner from checkpoint.

    Args:
        checkpoint: Checkpoint (or its path) to load.
        *args: Additional arguments to pass to `cls.load`.
        **kwargs: Additional keyword arguments to pass to `cls.load`.

    Returns:
        (BaseRunner):
    """

    if isinstance(checkpoint, (bytes, str, os.PathLike)):
        ckpt = cls.load(checkpoint, *args, **kwargs)
    elif isinstance(checkpoint, Mapping):
        ckpt = checkpoint
    else:
        raise ValueError(f"checkpoint is set to {checkpoint} but is not a valid checkpoint")
    runner = cls(ckpt["runner"])
    runner.load_checkpoint(ckpt, override_config=False)
    return runner

load_pretrained

Python

load_pretrained(checkpoint: Mapping | bytes | str | PathLike, *args, **kwargs) -> None

Load model from pretrained checkpoint.

This method only loads the model weights.

Parameters:

Name	Type	Description	Default
checkpoint	Mapping | bytes | str | PathLike	Pretrained checkpoint (or its path) to load.	required
*args		Additional arguments to pass to self.load.	()
**kwargs		Additional keyword arguments to pass to self.load.	{}

Raises:

Type	Description
ValueError	If model is not defined.
ValueError	If checkpoint is not a valid checkpoint.
FileNotFoundError	If checkpoint does not exists.

See Also

load_checkpoint: Load model, optimizer, and scheduler from checkpoint.

Source code in danling/runner/base_runner.py

def load_pretrained(self, checkpoint: Mapping | bytes | str | os.PathLike, *args, **kwargs) -> None:
    """
    Load model from pretrained checkpoint.

    This method only loads the model weights.

    Args:
        checkpoint: Pretrained checkpoint (or its path) to load.
        *args: Additional arguments to pass to `self.load`.
        **kwargs: Additional keyword arguments to pass to `self.load`.

    Raises:
        ValueError: If `model` is not defined.
        ValueError: If `checkpoint` is not a valid checkpoint.
        FileNotFoundError: If `checkpoint` does not exists.

    See Also:
        [`load_checkpoint`][danling.BaseRunner.load_checkpoint]: Load model, optimizer, and scheduler from
            checkpoint.
    """

    if self.model is None:
        raise ValueError("model is not defined")
    if isinstance(checkpoint, (bytes, str, os.PathLike)):
        if not os.path.exists(checkpoint):
            raise FileNotFoundError(f"pretrained is set to {checkpoint!r} but does not exist")
        ckpt = self.load(checkpoint, *args, **kwargs)
    elif isinstance(checkpoint, Mapping):
        ckpt = checkpoint
    else:
        raise ValueError(f"pretrained is set to {checkpoint!r} but is not a valid checkpoint")

    state_dict = ckpt
    while "model" in state_dict or "module" in state_dict:
        state_dict = state_dict.get("model", state_dict)
        state_dict = state_dict.get("module", state_dict)
    self.unwrap(self.model).load_state_dict(state_dict)

append_result

Python

append_result(result: NestedDict, index: int | None = None) -> None

Append result to self.results.

Source code in danling/runner/base_runner.py

def append_result(self, result: NestedDict, index: int | None = None) -> None:
    r"""
    Append result to `self.results`.

    Warnings:
        `self.results` is heavily relied upon for computing metrics.

        Failed to use this method may lead to unexpected behavior.
    """

    if index is None:
        index = self.epochs
        global __APPEND_RESULT_COUNTER__  # pylint: disable=global-statement
        __APPEND_RESULT_COUNTER__ += 1
        if index == 0 and __APPEND_RESULT_COUNTER__ > 1:
            warn(
                """
                Automatically set index to `self.epochs`.
                Please ensure `self.epochs` updates before calling `append_result`
                """,
                category=RuntimeWarning,
                stacklevel=2,
            )
    if index in self.results:
        self.results[index].merge(result)
    else:
        self.results[index] = result

print_result

Python

print_result() -> None

Print latest and best result.

Source code in danling/runner/base_runner.py

def print_result(self) -> None:
    r"""
    Print latest and best result.
    """

    print(f"latest result: {self.latest_result}")
    print(f"best result: {self.best_result}")

save_result

Python

save_result() -> None

Save result to self.dir.

This method will save latest and best result to self.dir/latest.json and self.dir/best.json respectively.

Source code in danling/runner/base_runner.py

@catch
@on_main_process
def save_result(self) -> None:
    r"""
    Save result to `self.dir`.

    This method will save latest and best result to
    `self.dir/latest.json` and `self.dir/best.json` respectively.
    """

    results_path = os.path.join(self.dir, "results.json")
    self.save(
        {
            "name": self.name,
            "id": self.id,
            "timestamp": self.timestamp,
            "results": self.results,
        },
        results_path,
        indent=4,
    )
    ret = {"name": self.name, "id": self.id, "timestamp": self.timestamp}
    result = self.latest_result
    if isinstance(result, FlatDict):
        result = result.dict()
    # This is slower but ensure id is the first key
    if result is not None:
        ret.update(result)
    latest_path = os.path.join(self.dir, "latest.json")
    self.save(ret, latest_path, indent=4)
    if self.is_best:
        best_path = os.path.join(self.dir, "best.json")
        shutil.copy(latest_path, best_path)

uuid

Python

uuid() -> UUID

UUID of the config.

Source code in danling/runner/base_runner.py

@cached_property
def uuid(self) -> UUID:
    r"""
    UUID of the config.
    """

    return uuid5(self.run_uuid, self.id)

log_interval

Python

log_interval() -> int

Log interval.

Source code in danling/runner/base_runner.py

@cached_property
def log_interval(self) -> int:
    r"""
    Log interval.
    """

    if self.config.get("log_interval") is not None:
        return self.config.log_interval  # type: ignore[return-value]
    if self.dataloaders:
        return max(ceil(max(len(d) for d in self.dataloaders.values()) / 10), 1)
    return 1

project_root

Python

project_root() -> str

Directory of the project.

Source code in danling/runner/base_runner.py

@cached_ensure_dir
def project_root(self) -> str:
    r"""
    Directory of the project.
    """

    return self.config.get("project_root", "experiments")

dir

Python

dir() -> str

Directory of the run.

Source code in danling/runner/base_runner.py

@cached_ensure_dir
def dir(self) -> str:
    r"""
    Directory of the run.
    """

    if "dir" in self.config:
        return self.config.dir
    return os.path.join(self.project_root, f"{self.name}-{self.id}", self.timestamp)

log_file

Python

log_file() -> str

Path of log file.

Source code in danling/runner/base_runner.py

@cached_ensure_parent_dir
def log_file(self) -> str:
    r"""
    Path of log file.
    """

    if "log_file" in self.config:
        return self.config.log_file
    return os.path.join(self.dir, "run.log")

checkpoint_dir

Python

checkpoint_dir() -> str

Directory of checkpoints.

Source code in danling/runner/base_runner.py

@cached_ensure_dir
def checkpoint_dir(self) -> str:
    r"""
    Directory of checkpoints.
    """

    if "checkpoint_dir" in self.config:
        return self.config.checkpoint_dir
    return os.path.join(self.dir, self.config.checkpoint_dir_name)

Config

Bases: Config

Configuration class for managing and persisting all states of a DanLing Runner.

The Config class provides a hierarchical configuration system that handles:

1. Parameter management: Hyperparameters, model settings, training options
2. Experiment tracking: IDs, names, and other metadata for runs and experiments
3. Serialization: Save/load configurations from files or command line
4. Reproducibility: Tracking seeds and settings for reproducible runs

Config inherits from Config and provides attribute-style access to nested values:

config = Config()

# Attribute-style access (recommended)
config.optim.lr = 1e-3
config.network.type = "resnet50"

# Dictionary-style access (alternative)
config["optim"]["lr"] = 1e-3
config["network"]["type"] = "resnet50"

Config objects support three types of hierarchical attribute access patterns:

1. Direct assignment for simple values:

   config.epochs = 10
   

2. Auto-created nested objects for hierarchical settings:

   # Auto-creates the nested objects
   config.optim.lr = 0.01
   config.optim.weight_decay = 1e-4
   

3. Class-level annotations for typed properties with defaults:

   class MyConfig(Config):
       epochs: int = 10
       learning_rate: float = 0.001
   

Command-line integration is built-in. You can define a configuration and then override values via command line arguments:

config = MyConfig()
config.parse()  # Parse CLI args, e.g., --epochs 20 --optim.lr 0.01

General:

Name	Type	Description
run_name	str	Human-readable name for this run. Defaults to "DanLing".
run_id	str	Unique identifier (hex string) for this run, derived from run_uuid.
run_uuid	(UUID, property)	Unique UUID generated from experiment_uuid and config hash.
experiment_name	str	Human-readable name for the experiment. Defaults to "DanLing".
experiment_id	str	Unique identifier for experiment, typically the git commit hash. Defaults to "xxxxxxxxxxxxxxxx" if not in a git repo or git/gitpython not installed.
experiment_uuid	(UUID, property)	UUID derived from experiment_id. Defaults to UUID('78787878-7878-7878-7878-787878787878') if not in a git repo.

Reproducibility:

Name	Type	Description
seed	int	Random seed for reproducibility. If not set, a random value is generated.
deterministic	bool	Whether to enforce deterministic operations in PyTorch. Defaults to False for better performance. Set to True for exact reproducibility.

Progress:

Name	Type	Description
steps	int	Current training step count.
epochs	int	Current epoch count.
step_begin	int	First step to run (for resuming). Defaults to 0.
epoch_begin	int	First epoch to run (for resuming). Defaults to 0.
step_end	int	Last step to run (optional). Use either this or epoch_end.
epoch_end	int	Last epoch to run (optional). Use either this or step_end.

Model Evaluation:

Name	Type	Description
score_split	str	Dataset split to use for model selection. Defaults to None.
score_name	str	Metric name to use for model selection. Defaults to “loss”.

I/O:

Name	Type	Description
project_root	str	Root directory for experiments. Defaults to "experiments".
checkpoint_dir_name	str	Subdirectory name for checkpoints. Defaults to "checkpoints".
log	bool	Whether to enable file logging. Defaults to True.
tensorboard	bool	Whether to use TensorBoard for visualization. Defaults to False.
log_interval	int	Iterations between log outputs. If None, auto-calculated.
save_interval	int	Epochs between checkpoint saves. If None, only save best/latest.

Examples:

Basic usage:

# Create a config
config = Config()
config.network.type = "resnet18"
config.optim.lr = 0.001
config.epochs = 10

# Use in a runner
runner = Runner(config)

Custom config class with typed attributes:

class TrainingConfig(Config):
    # Type annotations provide auto-completion and validation
    epochs: int = 100
    batch_size: int = 32
    precision: str = "fp16"

    def __init__(self):
        super().__init__()
        # Initialize nested settings
        self.optim.type = "adamw"
        self.optim.lr = 1e-3

    def post(self):
        # Called after parsing CLI args
        super().post()
        # Create derived settings
        self.experiment_name = f"{self.network.type}_{self.optim.lr}"

Command-line integration:

# Override config settings via CLI
python train.py --epochs 50 --optim.lr 0.0005 --network.type resnet50

Note

Always store all parameters needed to reproduce a run in the Config. The Config is automatically saved with checkpoints, enabling exact resumption.

See Also

• Runner: Main runner class that uses this config.
• chanfig.Config: Base config implementation.

Source code in danling/runner/config.py

class Config(chanfig.Config):  # pylint: disable=too-many-instance-attributes
    r"""
    Configuration class for managing and persisting all states of a DanLing Runner.

    The Config class provides a hierarchical configuration system that handles:

    1. **Parameter management**: Hyperparameters, model settings, training options
    2. **Experiment tracking**: IDs, names, and other metadata for runs and experiments
    3. **Serialization**: Save/load configurations from files or command line
    4. **Reproducibility**: Tracking seeds and settings for reproducible runs

    Config inherits from [`Config`][chanfig.Config] and provides attribute-style access to nested values:

    ```python
    config = Config()

    # Attribute-style access (recommended)
    config.optim.lr = 1e-3
    config.network.type = "resnet50"

    # Dictionary-style access (alternative)
    config["optim"]["lr"] = 1e-3
    config["network"]["type"] = "resnet50"
    ```

    Config objects support three types of hierarchical attribute access patterns:

    1. **Direct assignment** for simple values:
       ```python
       config.epochs = 10
       ```

    2. **Auto-created nested objects** for hierarchical settings:
       ```python
       # Auto-creates the nested objects
       config.optim.lr = 0.01
       config.optim.weight_decay = 1e-4
       ```

    3. **Class-level annotations** for typed properties with defaults:
       ```python
       class MyConfig(Config):
           epochs: int = 10
           learning_rate: float = 0.001
       ```

    Command-line integration is built-in. You can define a configuration and
    then override values via command line arguments:

    ```python
    config = MyConfig()
    config.parse()  # Parse CLI args, e.g., --epochs 20 --optim.lr 0.01
    ```

    Attributes: General:
        run_name (str): Human-readable name for this run. Defaults to `"DanLing"`.
        run_id (str): Unique identifier (hex string) for this run, derived from `run_uuid`.
        run_uuid (UUID, property): Unique UUID generated from experiment_uuid and config hash.
        experiment_name (str): Human-readable name for the experiment. Defaults to `"DanLing"`.
        experiment_id (str): Unique identifier for experiment, typically the git commit hash.
            Defaults to `"xxxxxxxxxxxxxxxx"` if not in a git repo or git/gitpython not installed.
        experiment_uuid (UUID, property): UUID derived from `experiment_id`.
            Defaults to `UUID('78787878-7878-7878-7878-787878787878')` if not in a git repo.

    Attributes: Reproducibility:
        seed (int): Random seed for reproducibility. If not set, a random value is generated.
        deterministic (bool): Whether to enforce deterministic operations in PyTorch.
            Defaults to `False` for better performance. Set to `True` for exact reproducibility.

    Attributes: Progress:
        steps (int): Current training step count.
        epochs (int): Current epoch count.
        step_begin (int): First step to run (for resuming). Defaults to 0.
        epoch_begin (int): First epoch to run (for resuming). Defaults to 0.
        step_end (int): Last step to run (optional). Use either this or `epoch_end`.
        epoch_end (int): Last epoch to run (optional). Use either this or `step_end`.

    Attributes: Model Evaluation:
        score_split (str): Dataset split to use for model selection. Defaults to None.
        score_name (str): Metric name to use for model selection. Defaults to "loss".

    Attributes: I/O:
        project_root (str): Root directory for experiments. Defaults to `"experiments"`.
        checkpoint_dir_name (str): Subdirectory name for checkpoints. Defaults to `"checkpoints"`.
        log (bool): Whether to enable file logging. Defaults to `True`.
        tensorboard (bool): Whether to use TensorBoard for visualization. Defaults to `False`.
        log_interval (int): Iterations between log outputs. If None, auto-calculated.
        save_interval (int): Epochs between checkpoint saves. If None, only save best/latest.

    Examples:
        Basic usage:
        ```python
        # Create a config
        config = Config()
        config.network.type = "resnet18"
        config.optim.lr = 0.001
        config.epochs = 10

        # Use in a runner
        runner = Runner(config)
        ```

        Custom config class with typed attributes:
        ```python
        class TrainingConfig(Config):
            # Type annotations provide auto-completion and validation
            epochs: int = 100
            batch_size: int = 32
            precision: str = "fp16"

            def __init__(self):
                super().__init__()
                # Initialize nested settings
                self.optim.type = "adamw"
                self.optim.lr = 1e-3

            def post(self):
                # Called after parsing CLI args
                super().post()
                # Create derived settings
                self.experiment_name = f"{self.network.type}_{self.optim.lr}"
        ```

        Command-line integration:
        ```bash
        # Override config settings via CLI
        python train.py --epochs 50 --optim.lr 0.0005 --network.type resnet50
        ```

    Note:
        Always store all parameters needed to reproduce a run in the Config.
        The Config is automatically saved with checkpoints, enabling exact resumption.

    See Also:
        - [`Runner`][danling.runner.Runner]: Main runner class that uses this config.
        - [`chanfig.Config`](https://github.com/ultmaster/chanfig): Base config implementation.
    """

    # DO NOT set default value in class, as they won't be stored in `__dict__`.

    run_name: str = defaults.RUN_NAME
    run_id: str
    experiment_name: str = defaults.EXPERIMENT_NAME
    experiment_id: str

    seed: Optional[int] = None
    deterministic: bool = False

    steps: int = 0
    epochs: int = 0
    step_begin: int = 0
    epoch_begin: int = 0
    step_end: Optional[int] = None
    epoch_end: Optional[int] = None

    score_split: Optional[str] = None
    score_name: str = "loss"

    project_root: str = "experiments"
    checkpoint_dir_name: str = "checkpoints"
    log: bool = True
    tensorboard: bool = False
    log_interval: Optional[int] = None
    save_interval: Optional[int] = None

    def __post_init__(self):
        if "experiment_id" not in self:
            self.experiment_id = get_git_hash() or defaults.EXPERIMENT_ID
        if "run_id" not in self:
            self.run_id = self.run_uuid.hex
        self.setattr("ignored_keys_in_hash", defaults.IGNORED_NAMES_IN_HASH)

    @property
    def experiment_uuid(self) -> UUID:
        r"""
        UUID of the experiment.
        """

        return UUID(bytes=bytes(self.experiment_id.ljust(16, "x")[:16], encoding="ascii"))

    @property
    def run_uuid(self) -> UUID:
        r"""
        UUID of the run.
        """

        ignored_keys_in_hash = self.getattr("ignored_keys_in_hash", defaults.IGNORED_NAMES_IN_HASH)
        state: chanfig.Config = chanfig.Config({k: v for k, v in self.dict().items() if k not in ignored_keys_in_hash})
        return uuid5(self.experiment_uuid, state.yamls())

    def __hash__(self) -> int:
        return int(self.run_uuid.hex, 16)

experiment_uuid property

Python

experiment_uuid: UUID

UUID of the experiment.

run_uuid property

Python

run_uuid: UUID

UUID of the run.

DeepSpeedRunner

Bases: TorchRunner

DeepSpeed-powered runner for large-scale model operations.

DeepSpeedRunner integrates Microsoft’s DeepSpeed framework to enable efficient execution of large language models and other compute-intensive neural networks with advanced memory optimization techniques.

Key features:

• ZeRO (Zero Redundancy Optimizer) for memory-efficient distributed operations
• Mixed precision execution with reduced memory footprint
• Progressive layer dropping for computational acceleration
• 3D parallelism (data, model, pipeline) for billion+ parameter models
• Optimized checkpoint management for large models

DeepSpeedRunner is particularly useful for:

• Operating large language models efficiently across multiple GPUs
• Running larger models within limited GPU memory constraints
• Achieving near-linear scaling in multi-node distributed settings
• Evaluating very large models that wouldn’t fit in memory with standard approaches

Configuration

DeepSpeed requires a specific configuration format. You can provide a DeepSpeed configuration in the deepspeed key of your runner config, or let DanLing auto-generate a configuration based on your other settings.

Note

For optimal performance, DeepSpeed requires NVMe offloading and specific GPU configurations. See DeepSpeed documentation for details.

See Also

• TorchRunner: PyTorch DDP runner.
• DeepSpeed Documentation: Official DeepSpeed docs.

Source code in danling/runner/deepspeed_runner.py

class DeepSpeedRunner(TorchRunner):
    r"""
    DeepSpeed-powered runner for large-scale model operations.

    DeepSpeedRunner integrates Microsoft's DeepSpeed framework to enable efficient execution
    of large language models and other compute-intensive neural networks with advanced memory
    optimization techniques.

    Key features:

    * ZeRO (Zero Redundancy Optimizer) for memory-efficient distributed operations
    * Mixed precision execution with reduced memory footprint
    * Progressive layer dropping for computational acceleration
    * 3D parallelism (data, model, pipeline) for billion+ parameter models
    * Optimized checkpoint management for large models

    DeepSpeedRunner is particularly useful for:

    * Operating large language models efficiently across multiple GPUs
    * Running larger models within limited GPU memory constraints
    * Achieving near-linear scaling in multi-node distributed settings
    * Evaluating very large models that wouldn't fit in memory with standard approaches

    Configuration:
        DeepSpeed requires a specific configuration format. You can provide a DeepSpeed
        configuration in the `deepspeed` key of your runner config, or let DanLing auto-generate
        a configuration based on your other settings.

    Note:
        For optimal performance, DeepSpeed requires NVMe offloading and specific
        GPU configurations. See DeepSpeed documentation for details.

    See Also:
        - [`TorchRunner`][danling.runner.TorchRunner]: PyTorch DDP runner.
        - [DeepSpeed Documentation](https://www.deepspeed.ai/docs/): Official DeepSpeed docs.
    """

    def __init__(self, config: Config) -> None:
        ds.check()
        super().__init__(config)

    def init_distributed(self) -> None:
        r"""
        Set up distributed training.

        Initialise process group and set up DDP variables.
        """

        backend = self.config.get("backend", os.getenv("BACKEND"))
        init_method = self.config.get("init_method", os.getenv("INIT_METHOD"))
        world_size = int(self.config.get("world_size", os.getenv("WORLD_SIZE", "1")))
        rank = int(self.config.get("rank", os.getenv("RANK", "0")))
        if world_size > 1:
            if torch.cuda.is_available():
                torch.cuda.set_device(self.get_local_rank())
            deepspeed.init_distributed(dist_backend=backend, init_method=init_method, world_size=world_size, rank=rank)
            object_list = [self.id, self.timestamp]
            dist.broadcast_object_list(object_list)
            self.id, self.timestamp = object_list

    def __post_init__(self):
        if self.datasets:
            self.build_dataloaders()
        self.model = self.model.to(self.device)
        if self.ema is not None:
            self.ema = self.ema.to(self.device)
        self.config.deepspeed = self.get_deepspeed_config()
        self.model, self.optimizer, _, self.scheduler = deepspeed.initialize(
            model=self.model,
            optimizer=self.optimizer,
            lr_scheduler=self.scheduler,
            config=self.config.deepspeed,
        )

    def advance(self, loss) -> None:
        self.backward(loss)
        if self.config.get("max_grad_value") is not None:
            clip_grad_value_(self.model.parameters(), self.config["max_grad_value"])
        self.model.step()
        self.steps = self.model.global_steps

    def backward(self, loss: torch.Tensor) -> None:
        return self.model.backward(loss)

    def get_local_rank(self) -> int:
        local_rank = self.config.get("local_rank", os.getenv("LOCAL_RANK"))
        if local_rank is not None:
            return int(local_rank)
        rank = self.config.get("rank", os.getenv("RANK"))
        world_size = self.config.get("world_size", os.getenv("WORLD_SIZE"))
        if world_size is None or rank is None:
            raise ValueError("Please provide either `local_rank` or `world_size` and `rank`")
        return int(world_size) % int(rank)

    def unwrap(self, model: nn.Module) -> nn.Module:
        while isinstance(model, (deepspeed.DeepSpeedEngine, nn.parallel.DistributedDataParallel)):
            model = model.module
        return model

    @property
    def deepspeed(self) -> NestedDict | None:
        if isinstance(self.model, deepspeed.DeepSpeedEngine):
            return self.model.config
        return None

    @catch
    def save_checkpoint(self, name: str = "latest", epoch: int | None = None, save_best: bool = True) -> None:
        r"""
        Save checkpoint to `self.checkpoint_dir`.

        Args:
            name: Name of the checkpoint. Defaults to `"latest"`.
            epoch: Epoch to save. Defaults to `self.epochs`.
            save_best: If `True`, when `self.is_best` is `True`, the checkpoint will also be copied to
                `self.checkpoint_dir/best`.

        If `self.config.save_interval` is positive and `epochs + 1` is a multiple of `save_interval`,
        the checkpoint will also be copied to `self.checkpoint_dir/epoch-{epochs}`.
        """

        epoch = epoch or self.epochs
        save_interval = self.config.get("save_interval", -1)
        latest_path = os.path.join(self.checkpoint_dir, name)
        os.makedirs(latest_path, exist_ok=True)
        self.yaml(os.path.join(latest_path, "runner.yaml"))
        self.model.save_checkpoint(
            self.checkpoint_dir, tag=name, client_state={"runner": self.config.dict()}, save_latest=False
        )
        if self.ema is not None:
            self.save(self.state_dict(), os.path.join(latest_path, "ema.pth"))
        if save_interval > 0 and (epoch + 1) % save_interval == 0:
            save_path = os.path.join(self.checkpoint_dir, f"epoch-{epoch}")
            shutil.copytree(latest_path, save_path, dirs_exist_ok=True)
        if save_best and self.is_best:
            best_path = os.path.join(self.checkpoint_dir, "best")
            shutil.copytree(latest_path, best_path, dirs_exist_ok=True)

    def state_dict(self, cls: Callable = dict) -> Mapping:
        if self.ema is None:
            raise ValueError("It is not necessary to save model in DeepSpeed without EMA")
        return cls(
            runner=self.config.dict(),
            ema=self.ema.state_dict(),
        )

    def load_checkpoint(self, checkpoint: bytes | str | os.PathLike, *args, **kwargs) -> None:  # type: ignore[override]
        """
        Load model, optimizer, and scheduler from checkpoint.

        Args:
            checkpoint: Checkpoint (or its path) to load.
            *args: Additional arguments to pass to `self.load`.
            **kwargs: Additional keyword arguments to pass to `self.load`.

        Raises:
            ValueError: If `model` is not defined.
            ValueError: If `model` is not an instance of `deepspeed.DeepSpeedEngine`.

        See Also:
            [`from_checkpoint`][danling.BaseRunner.from_checkpoint]: Build runner from checkpoint.
            [`load_pretrained`][danling.BaseRunner.load_pretrained]: Load model parameters from pretrained checkpoint.
        """

        if self.model is None:
            raise ValueError("model is not defined")
        if not isinstance(self.model, deepspeed.DeepSpeedEngine):
            raise ValueError("model is not an instance of `deepspeed.DeepSpeedEngine`")

        self.model.load_checkpoint(checkpoint)
        self.config.checkpoint = checkpoint

    def load_pretrained(self, checkpoint: bytes | str | os.PathLike, *args, **kwargs) -> None:  # type: ignore[override]
        """
        Load model from pretrained checkpoint.

        This method only loads the model weights.

        Args:
            checkpoint: Pretrained checkpoint directory.
            *args: Additional arguments to pass to `self.load`.
            **kwargs: Additional keyword arguments to pass to `self.load`.

        Raises:
            ValueError: If `model` is not defined.

        See Also:
            [`load_checkpoint`][danling.BaseRunner.load_checkpoint]: Load model, optimizer, and scheduler from
                checkpoint.
        """

        if self.model is None:
            raise ValueError("model is not defined")

        self.model.load_checkpoint(checkpoint, load_module_only=True)
        self.config.pretrained = checkpoint

    def load_config(
        self, checkpoint: bytes | str | os.PathLike, overwrite: bool = False, *args, **kwargs  # type: ignore[override]
    ) -> None:
        r"""
        Load config from checkpoint.

        Args:
            checkpoint: Checkpoint (or its path) to load.
            overwrite: If `True`, overwrite the current config with the loaded config.
                Defaults to `False`.
            *args: Additional arguments to pass to `self.load`.
            **kwargs: Additional keyword arguments to pass to `self.load`.

        Raises:
            FileNotFoundError: If `checkpoint` does not exists.
        """

        if isinstance(checkpoint, bytes):
            checkpoint = checkpoint.decode()

        config = self.load(os.path.join(checkpoint, "runner.yaml"), *args, **kwargs)
        self.config.merge(config, overwrite=overwrite)
        self.step_begin = config["steps"] + 1
        self.epoch_begin = config["epochs"] + 1

init_distributed

Python

init_distributed() -> None

Set up distributed training.

Initialise process group and set up DDP variables.

Source code in danling/runner/deepspeed_runner.py

def init_distributed(self) -> None:
    r"""
    Set up distributed training.

    Initialise process group and set up DDP variables.
    """

    backend = self.config.get("backend", os.getenv("BACKEND"))
    init_method = self.config.get("init_method", os.getenv("INIT_METHOD"))
    world_size = int(self.config.get("world_size", os.getenv("WORLD_SIZE", "1")))
    rank = int(self.config.get("rank", os.getenv("RANK", "0")))
    if world_size > 1:
        if torch.cuda.is_available():
            torch.cuda.set_device(self.get_local_rank())
        deepspeed.init_distributed(dist_backend=backend, init_method=init_method, world_size=world_size, rank=rank)
        object_list = [self.id, self.timestamp]
        dist.broadcast_object_list(object_list)
        self.id, self.timestamp = object_list

save_checkpoint

Python

save_checkpoint(name: str = 'latest', epoch: int | None = None, save_best: bool = True) -> None

Save checkpoint to self.checkpoint_dir.

Parameters:

Name	Type	Description	Default
name	str	Name of the checkpoint. Defaults to "latest".	'latest'
epoch	int | None	Epoch to save. Defaults to self.epochs.	None
save_best	bool	If True, when self.is_best is True, the checkpoint will also be copied to self.checkpoint_dir/best.	True

If self.config.save_interval is positive and epochs + 1 is a multiple of save_interval, the checkpoint will also be copied to self.checkpoint_dir/epoch-{epochs}.

Source code in danling/runner/deepspeed_runner.py

@catch
def save_checkpoint(self, name: str = "latest", epoch: int | None = None, save_best: bool = True) -> None:
    r"""
    Save checkpoint to `self.checkpoint_dir`.

    Args:
        name: Name of the checkpoint. Defaults to `"latest"`.
        epoch: Epoch to save. Defaults to `self.epochs`.
        save_best: If `True`, when `self.is_best` is `True`, the checkpoint will also be copied to
            `self.checkpoint_dir/best`.

    If `self.config.save_interval` is positive and `epochs + 1` is a multiple of `save_interval`,
    the checkpoint will also be copied to `self.checkpoint_dir/epoch-{epochs}`.
    """

    epoch = epoch or self.epochs
    save_interval = self.config.get("save_interval", -1)
    latest_path = os.path.join(self.checkpoint_dir, name)
    os.makedirs(latest_path, exist_ok=True)
    self.yaml(os.path.join(latest_path, "runner.yaml"))
    self.model.save_checkpoint(
        self.checkpoint_dir, tag=name, client_state={"runner": self.config.dict()}, save_latest=False
    )
    if self.ema is not None:
        self.save(self.state_dict(), os.path.join(latest_path, "ema.pth"))
    if save_interval > 0 and (epoch + 1) % save_interval == 0:
        save_path = os.path.join(self.checkpoint_dir, f"epoch-{epoch}")
        shutil.copytree(latest_path, save_path, dirs_exist_ok=True)
    if save_best and self.is_best:
        best_path = os.path.join(self.checkpoint_dir, "best")
        shutil.copytree(latest_path, best_path, dirs_exist_ok=True)

load_checkpoint

Python

load_checkpoint(checkpoint: bytes | str | PathLike, *args, **kwargs) -> None

Load model, optimizer, and scheduler from checkpoint.

Parameters:

Name	Type	Description	Default
checkpoint	bytes | str | PathLike	Checkpoint (or its path) to load.	required
*args		Additional arguments to pass to self.load.	()
**kwargs		Additional keyword arguments to pass to self.load.	{}

Raises:

Type	Description
ValueError	If model is not defined.
ValueError	If model is not an instance of deepspeed.DeepSpeedEngine.

See Also

from_checkpoint: Build runner from checkpoint. load_pretrained: Load model parameters from pretrained checkpoint.

Source code in danling/runner/deepspeed_runner.py

def load_checkpoint(self, checkpoint: bytes | str | os.PathLike, *args, **kwargs) -> None:  # type: ignore[override]
    """
    Load model, optimizer, and scheduler from checkpoint.

    Args:
        checkpoint: Checkpoint (or its path) to load.
        *args: Additional arguments to pass to `self.load`.
        **kwargs: Additional keyword arguments to pass to `self.load`.

    Raises:
        ValueError: If `model` is not defined.
        ValueError: If `model` is not an instance of `deepspeed.DeepSpeedEngine`.

    See Also:
        [`from_checkpoint`][danling.BaseRunner.from_checkpoint]: Build runner from checkpoint.
        [`load_pretrained`][danling.BaseRunner.load_pretrained]: Load model parameters from pretrained checkpoint.
    """

    if self.model is None:
        raise ValueError("model is not defined")
    if not isinstance(self.model, deepspeed.DeepSpeedEngine):
        raise ValueError("model is not an instance of `deepspeed.DeepSpeedEngine`")

    self.model.load_checkpoint(checkpoint)
    self.config.checkpoint = checkpoint

load_pretrained

Python

load_pretrained(checkpoint: bytes | str | PathLike, *args, **kwargs) -> None

Load model from pretrained checkpoint.

This method only loads the model weights.

Parameters:

Name	Type	Description	Default
checkpoint	bytes | str | PathLike	Pretrained checkpoint directory.	required
*args		Additional arguments to pass to self.load.	()
**kwargs		Additional keyword arguments to pass to self.load.	{}

Raises:

Type	Description
ValueError	If model is not defined.

See Also

load_checkpoint: Load model, optimizer, and scheduler from checkpoint.

Source code in danling/runner/deepspeed_runner.py

def load_pretrained(self, checkpoint: bytes | str | os.PathLike, *args, **kwargs) -> None:  # type: ignore[override]
    """
    Load model from pretrained checkpoint.

    This method only loads the model weights.

    Args:
        checkpoint: Pretrained checkpoint directory.
        *args: Additional arguments to pass to `self.load`.
        **kwargs: Additional keyword arguments to pass to `self.load`.

    Raises:
        ValueError: If `model` is not defined.

    See Also:
        [`load_checkpoint`][danling.BaseRunner.load_checkpoint]: Load model, optimizer, and scheduler from
            checkpoint.
    """

    if self.model is None:
        raise ValueError("model is not defined")

    self.model.load_checkpoint(checkpoint, load_module_only=True)
    self.config.pretrained = checkpoint

load_config

Python

load_config(checkpoint: bytes | str | PathLike, overwrite: bool = False, *args, **kwargs) -> None

Load config from checkpoint.

Parameters:

Name	Type	Description	Default
checkpoint	bytes | str | PathLike	Checkpoint (or its path) to load.	required
overwrite	bool	If True, overwrite the current config with the loaded config. Defaults to False.	False
*args		Additional arguments to pass to self.load.	()
**kwargs		Additional keyword arguments to pass to self.load.	{}

Raises:

Type	Description
FileNotFoundError	If checkpoint does not exists.

Source code in danling/runner/deepspeed_runner.py

def load_config(
    self, checkpoint: bytes | str | os.PathLike, overwrite: bool = False, *args, **kwargs  # type: ignore[override]
) -> None:
    r"""
    Load config from checkpoint.

    Args:
        checkpoint: Checkpoint (or its path) to load.
        overwrite: If `True`, overwrite the current config with the loaded config.
            Defaults to `False`.
        *args: Additional arguments to pass to `self.load`.
        **kwargs: Additional keyword arguments to pass to `self.load`.

    Raises:
        FileNotFoundError: If `checkpoint` does not exists.
    """

    if isinstance(checkpoint, bytes):
        checkpoint = checkpoint.decode()

    config = self.load(os.path.join(checkpoint, "runner.yaml"), *args, **kwargs)
    self.config.merge(config, overwrite=overwrite)
    self.step_begin = config["steps"] + 1
    self.epoch_begin = config["epochs"] + 1

Runner

Bases: BaseRunner

Dynamic platform-selecting runner that serves as the primary entry point for DanLing.

The Runner class automatically selects the most appropriate distributed training platform based on your configuration and available packages. It dynamically modifies its class at initialization to transform into the chosen platform’s runner implementation.

Key features:

• Automatic platform selection: Chooses the best available backend
• Dynamic class transformation: Becomes a TorchRunner, DeepSpeedRunner, or AccelerateRunner
• Unified interface: Provides consistent API across all platforms

Platform selection logic:

1. If config.platform is “auto” (default):
2. Uses DeepSpeed if available
3. Falls back to PyTorch otherwise
4. If config.platform is explicitly set to “torch”, “deepspeed”, or “accelerate”:
5. Uses the specified platform
6. Raises an error if the required packages are not installed

Usage Examples:

# Automatic platform selection
config = Config({"platform": "auto"})
runner = Runner(config)  # Will use DeepSpeed if available, otherwise PyTorch

# Explicit platform selection
config = Config({"platform": "accelerate"})
runner = Runner(config)  # Will use Accelerate

Extension Guidelines:

• DO inherit from Runner when you want to add functionality that should work across all platforms:

  class MyRunner(Runner):
      def __init__(self, config):
          super().__init__(config)
          # Your initialization code
  
      def my_method(self):
          # Custom functionality
          pass
  

• DON’T inherit from a specific platform runner (like TorchRunner) unless you’re implementing a new distributed training framework.

Parameters:

Name	Type	Description	Default
config	Config	Configuration object containing runner settings. The platform key determines which backend implementation will be used.	required

Raises:

Type	Description
ValueError	If an unknown platform is specified or required packages are missing.

See Also

• BaseRunner: Base class with core functionality.
• TorchRunner: PyTorch DDP implementation.
• DeepSpeedRunner: DeepSpeed implementation.
• AccelerateRunner: HuggingFace Accelerate implementation.

Source code in danling/runner/runner.py

class Runner(BaseRunner):
    r"""
    Dynamic platform-selecting runner that serves as the primary entry point for DanLing.

    The Runner class automatically selects the most appropriate distributed training platform
    based on your configuration and available packages. It dynamically modifies its class
    at initialization to transform into the chosen platform's runner implementation.

    Key features:

    * **Automatic platform selection**: Chooses the best available backend
    * **Dynamic class transformation**: Becomes a TorchRunner, DeepSpeedRunner, or AccelerateRunner
    * **Unified interface**: Provides consistent API across all platforms

    Platform selection logic:

    1. If `config.platform` is "auto" (default):
       - Uses DeepSpeed if available
       - Falls back to PyTorch otherwise
    2. If `config.platform` is explicitly set to "torch", "deepspeed", or "accelerate":
       - Uses the specified platform
       - Raises an error if the required packages are not installed

    Usage Examples:

    ```python
    # Automatic platform selection
    config = Config({"platform": "auto"})
    runner = Runner(config)  # Will use DeepSpeed if available, otherwise PyTorch

    # Explicit platform selection
    config = Config({"platform": "accelerate"})
    runner = Runner(config)  # Will use Accelerate
    ```

    Extension Guidelines:

    * **DO inherit from Runner** when you want to add functionality that should work across all platforms:
      ```python
      class MyRunner(Runner):
          def __init__(self, config):
              super().__init__(config)
              # Your initialization code

          def my_method(self):
              # Custom functionality
              pass
      ```

    * **DON'T inherit from a specific platform runner** (like TorchRunner) unless you're implementing
      a new distributed training framework.

    Args:
        config: Configuration object containing runner settings. The `platform` key
               determines which backend implementation will be used.

    Raises:
        ValueError: If an unknown platform is specified or required packages are missing.

    See Also:
        - [`BaseRunner`][danling.runner.BaseRunner]: Base class with core functionality.
        - [`TorchRunner`][danling.runner.TorchRunner]: PyTorch DDP implementation.
        - [`DeepSpeedRunner`][danling.runner.DeepSpeedRunner]: DeepSpeed implementation.
        - [`AccelerateRunner`][danling.runner.AccelerateRunner]: HuggingFace Accelerate implementation.
    """

    def __new__(cls, config: Config) -> Runner:
        platform = config.get("platform", "auto").lower()

        if platform == "auto":
            platform = "deepspeed" if ds.is_successful() else "torch"
        config["platform"] = platform

        if platform == "accelerate":
            ac.check()
            return super().__new__(type("AccelerateRunner", (cls, AccelerateRunner), {}))
        if platform == "deepspeed":
            ds.check()
            return super().__new__(type("DeepSpeedRunner", (cls, DeepSpeedRunner), {}))
        if platform == "torch":
            return super().__new__(type("TorchRunner", (cls, TorchRunner), {}))

        raise ValueError(f"Unknown platform: {platform}. Valid options are: torch, accelerate, deepspeed")

TorchRunner

Bases: BaseRunner

PyTorch-based unified runner for model training, evaluation, and inference.

TorchRunner implements the complete machine learning workflow using PyTorch’s native capabilities, providing a comprehensive solution for the entire model lifecycle.

This runner serves as the core implementation for PyTorch-based workflows, offering:

• Complete workflow with training, evaluation, and inference capabilities
• Native DDP support for efficient multi-GPU/multi-node operations
• Mixed precision execution via torch.cuda.amp
• Gradient accumulation for effective batch size scaling
• Flexible checkpoint management and experiment tracking
• Standardized evaluation protocols and metric collection

TorchRunner is the most flexible backend in DanLing, making it an ideal choice for extending with custom functionality or when maximum compatibility is required.

Note

When running multi-GPU operations with TorchRunner, the environment variables for distributed execution (WORLD_SIZE, RANK, LOCAL_RANK) must be properly set.

See Also

• BaseRunner: Base class for all DanLing runners.
• AccelerateRunner: Runner using HuggingFace Accelerate.
• DeepSpeedRunner: Runner using Microsoft DeepSpeed.

Source code in danling/runner/torch_runner.py

class TorchRunner(BaseRunner):
    r"""
    PyTorch-based unified runner for model training, evaluation, and inference.

    TorchRunner implements the complete machine learning workflow using PyTorch's native
    capabilities, providing a comprehensive solution for the entire model lifecycle.

    This runner serves as the core implementation for PyTorch-based workflows, offering:

    * Complete workflow with training, evaluation, and inference capabilities
    * Native DDP support for efficient multi-GPU/multi-node operations
    * Mixed precision execution via torch.cuda.amp
    * Gradient accumulation for effective batch size scaling
    * Flexible checkpoint management and experiment tracking
    * Standardized evaluation protocols and metric collection

    TorchRunner is the most flexible backend in DanLing, making it an ideal choice for
    extending with custom functionality or when maximum compatibility is required.

    Note:
        When running multi-GPU operations with TorchRunner, the environment variables for distributed
        execution (WORLD_SIZE, RANK, LOCAL_RANK) must be properly set.

    See Also:
        - [`BaseRunner`][danling.runner.BaseRunner]: Base class for all DanLing runners.
        - [`AccelerateRunner`][danling.runner.AccelerateRunner]: Runner using HuggingFace Accelerate.
        - [`DeepSpeedRunner`][danling.runner.DeepSpeedRunner]: Runner using Microsoft DeepSpeed.
    """

    model: nn.Module
    ema: nn.Module | None = None
    criterion: nn.Module
    optimizer: optim.Optimizer
    scheduler: optim.lr_scheduler._LRScheduler

    def __post_init__(self):
        if self.datasets:
            self.build_dataloaders()
        self.model = self.model.to(self.device)
        if self.ema is not None:
            self.ema = self.ema.to(self.device)
        if self.distributed and not isinstance(
            self.model, (nn.parallel.DistributedDataParallel, nn.parallel.DataParallel)
        ):
            self.model = nn.parallel.DistributedDataParallel(self.model)

    def train(self, train_splits: list[str] | None = None, evaluate_splits: list[str] | None = None) -> NestedDict:
        r"""
        Perform training on `split`.

        Args:
            train_splits: list of split to run train.
                Defaults to `["train"]`.
            evaluate_splits: list of split to run evaluate.
                Defaults to `self.dataloaders` except for those in `train_splits`.

        Return:
            NestedDict: train results
        """

        early_stop_counter = 0
        if train_splits is not None:
            self.train_splits = sorted(train_splits)
        if evaluate_splits is not None:
            self.evaluate_splits = sorted(evaluate_splits)
        if not self.train_splits:
            warn("No training split is found. Will only evaluate for one epoch.", stacklevel=2)
            self.epoch_end = self.epoch_begin + 1
        print(f"Begin training from {self.epoch_begin} to {self.epoch_end}")
        print(f"Training splits: {self.train_splits}")
        print(f"Evaluation splits: {self.evaluate_splits}")
        patience = self.config.get("patience", float("inf"))
        for epoch in range(self.epoch_begin, self.epoch_end):  # type: ignore
            self.epochs = epoch
            result = NestedDict()
            result.setattr("convert_mapping", True)
            for split in self.train_splits:
                result[split] = self.train_epoch(split)
            for split in self.evaluate_splits:
                result[split] = self.evaluate_epoch(split)
            self.append_result(result)
            print(self.format_epoch_result(result))
            self.save_result()
            if self.config.save_interval is not None:
                self.save_checkpoint()
            """@nni.report_intermediate_result(self.latest_score)"""
            early_stop_counter = 0 if self.is_best else early_stop_counter + 1
            if early_stop_counter > patience:
                print("early stop")
                break
        """@nni.report_final_result(self.latest_score)"""
        return self.results

    def train_epoch(self, split: str = "train") -> NestedDict:
        r"""
        Train one epoch on `split`.

        Args:
            split (str): split to run train

        Return:
            NestedDict: train result
        """

        self.mode = "train"  # type: ignore
        self.split = split
        loader = self.dataloaders[split]
        length = len(loader) - 1
        last_print_iteration = -1
        self.meters.reset()
        if self.train_metrics is not None:
            self.metrics = self.train_metrics
        if self.metrics is not None:
            self.metrics.reset()
        if hasattr(loader.batch_sampler, "set_epoch"):
            loader.batch_sampler.set_epoch(self.epochs)
        if hasattr(loader.sampler, "set_epoch"):
            loader.sampler.set_epoch(self.epochs)
        is_cuda = self.device == torch.device("cuda")
        batch_time = time()

        for iteration, data in enumerate(loader):
            _, loss = self.train_step(data)

            if self.log_interval > 0 and (iteration > 0 and iteration % self.log_interval == 0 or iteration == length):
                interval = iteration - last_print_iteration
                if is_cuda:
                    torch.cuda.synchronize()
                if self.scheduler is not None:
                    self.meters.lr.update(self.scheduler.get_last_lr()[0])
                self.meters.loss.update(self.reduce(loss).item())
                self.meters.time.update((time() - batch_time) / interval)
                batch_time = time()
                self.step_log(split, iteration, length)
                last_print_iteration = iteration

        result = self.get_epoch_result()
        return result

    def train_step(self, data) -> Tuple[Any, torch.Tensor]:
        data = to_device(data, self.device)
        with self.autocast(), self.accumulate():
            input = data["input"] if isinstance(data, Mapping) else data[0]
            target = data["target"] if isinstance(data, Mapping) else data[1]
            pred = self.model(**input) if isinstance(input, Mapping) else self.model(input)
            loss = self.criterion(pred, target)
            if self.metrics is not None:
                self.metrics.update(pred.squeeze(-1), target)
            self.advance(loss)
        return pred, loss

    def advance(self, loss) -> None:
        r"""
        Backward loss and step optimizer & scheduler.

        Args:
            loss: The loss tensor from which to backpropagate.
        """

        self.backward(loss / self.accum_steps)
        if self.accum_steps <= 1 or self.steps % self.accum_steps == 0:
            if self.config.get("max_grad_value") is not None:
                clip_grad_value_(self.model.parameters(), self.config["max_grad_value"])
            if self.config.get("max_grad_norm") is not None:
                clip_grad_norm_(self.model.parameters(), self.config["max_grad_norm"])
            self.optimizer.step()
            self.optimizer.zero_grad()
            self.steps += 1

    def evaluate(self, evaluate_splits: list[str] | None = None) -> NestedDict:
        r"""
        Perform evaluation on `evaluate_splits`.

        Args:
            evaluate_splits: list of split to run evaluate.
                Defaults to `["evaluate"]`.

        Return:
            NestedDict: evaluation result
        """

        if evaluate_splits is not None:
            self.evaluate_splits = sorted(evaluate_splits)
        if not self.evaluate_splits:
            raise ValueError("No evaluation splits found.")
        print("Begin evaluation")
        print(f"Evaluation splits: {self.evaluate_splits}")
        result = NestedDict()
        result.setattr("convert_mapping", True)
        for split in self.evaluate_splits:
            result[split] = self.evaluate_epoch(split=split)
        print(self.format_epoch_result(result))
        return result

    # torch.inference_mode cause experiments to hang
    # @torch.inference_mode()
    def evaluate_epoch(self, split: str = "val") -> NestedDict:
        r"""
        Evaluate one epoch on `split`.

        Args:
            split (str): split to run evaluate

        Return:
            NestedDict: evaluation result
        """

        self.mode = RunnerMode.evaluate
        self.split = split
        loader = self.dataloaders[split]
        length = len(loader) - 1
        last_print_iteration = -1
        self.meters.reset()
        if self.evaluate_metrics is not None:
            self.metrics = self.evaluate_metrics
        if self.metrics is not None:
            self.metrics.reset()
        batch_time = time()
        is_cuda = self.device == torch.device("cuda")

        for iteration, data in enumerate(loader):
            _, loss = self.evaluate_step(data)

            if self.log_interval > 0 and (iteration > 0 and iteration % self.log_interval == 0 or iteration == length):
                interval = iteration - last_print_iteration
                if is_cuda:
                    torch.cuda.synchronize()
                self.meters.loss.update(self.reduce(loss).item())
                self.meters.time.update((time() - batch_time) / interval)
                batch_time = time()
                self.step_log(split, iteration, length)
                last_print_iteration = iteration

        result = self.get_epoch_result()
        self.write_result(result, split, self.epochs)
        return result

    def evaluate_step(self, data) -> Tuple[Any, torch.Tensor]:
        data = to_device(data, self.device)
        input = data["input"] if isinstance(data, Mapping) else data[0]
        target = data["target"] if isinstance(data, Mapping) else data[1]
        model = self.ema or self.model
        pred = model(**input) if isinstance(input, Mapping) else model(input)
        loss = self.criterion(pred, target)
        if self.metrics is not None:
            self.metrics.update(pred.squeeze(-1), target)
        return pred, loss

    @torch.inference_mode()
    def infer(self, split: str = "infer") -> list[float]:
        r"""
        Perform inference on `split`.

        Args:
            split (str): split to run inference

        Return:
            Tensor: inference outputs
        """

        self.mode = RunnerMode.infer
        loader = self.dataloaders[split]
        output: list[float] = []
        model = self.ema or self.model
        for _, data in tqdm(enumerate(loader), total=len(loader)):
            data = to_device(data, self.device)
            input = data["input"] if isinstance(data, Mapping) else data[0]
            pred = model(**input) if isinstance(input, Mapping) else model(input)
            output.extend(pred.squeeze(-1).tolist())

        if self.distributed:
            torch.cuda.synchronize()
            output = self.gather_for_metrics(output)
        return output

    def backward(self, loss: torch.Tensor) -> None:
        r"""
        Backward loss.

        Args:
            loss: Loss to backward.
        """

        loss.backward()

    def has_nan_inf_grad(self, model: nn.Module | None = None) -> bool:
        r"""
        Check if model has NaN or Inf gradients.

        Args:
            model: Model to check.
                Defaults to `self.model`.

        Return:
            bool: True if NaN or Inf is detected in gradients.
        """
        model = model or self.model
        for name, param in model.named_parameters():
            if param.grad is not None:
                if torch.isnan(param.grad).any():
                    print(f"NaN detected in gradients of parameter: {name}")
                    return True
                if torch.isinf(param.grad).any():
                    print(f"Inf detected in gradients of parameter: {name}")
                    return True
        return False

    def init_distributed(self) -> None:
        r"""
        Set up distributed training.

        Initialise process group and set up DDP variables.
        """

        backend = self.config.get("backend", os.getenv("BACKEND"))
        init_method = self.config.get("init_method", os.getenv("INIT_METHOD"))
        world_size = int(self.config.get("world_size", os.getenv("WORLD_SIZE", "1")))
        rank = int(self.config.get("rank", os.getenv("RANK", "0")))
        if world_size > 1:
            if torch.cuda.is_available():
                torch.cuda.set_device(self.local_rank)
            dist.init_process_group(backend, init_method, world_size=world_size, rank=rank)
            object_list = [self.id, self.timestamp]
            dist.broadcast_object_list(object_list)
            self.id, self.timestamp = object_list

    @on_main_process
    def init_tensorboard(self, *args, **kwargs) -> None:
        r"""
        Set up Tensoraoard SummaryWriter.
        """
        from torch.utils.tensorboard.writer import SummaryWriter  # pylint: disable=C0415

        if "log_dir" not in kwargs:
            kwargs["log_dir"] = self.dir

        self.writer = SummaryWriter(*args, **kwargs)
        self.writer.add_scalar = catch(OSError, verbose=False)(self.writer.add_scalar)

    def set_seed(self, seed: int = None, bias: int = None) -> int:  # type: ignore[assignment]
        r"""
        Set up random seed.

        Args:
            seed: Random seed to set.
                Defaults to `self.config.seed` (`config.seed`).

            bias: Make the seed different for each processes.
                This is used to ensure the data augmentation are applied differently on every processes.
                Defaults to `self.rank`.
                Set to `False` to disable this feature.
        Returns:
            Random seed set.
        """

        seed = seed or self.config.seed  # type: ignore[assignment]
        if seed is None:
            if self.inited:
                seed = random.randint(0, 2**32 - 1)
                if self.distributed:
                    object_list = [seed]
                    dist.broadcast_object_list(object_list)
                    seed = object_list[0]
                self.config.seed = seed
        else:
            seed = defaults.SEED
        bias = bias or self.rank
        if bias:
            seed += bias
        torch.manual_seed(seed)
        torch.cuda.manual_seed(seed)
        if np_random is not None:
            np_random.seed(seed)
        random.seed(seed)
        return seed

    def set_deterministic(self) -> None:
        cudnn.benchmark = False
        cudnn.deterministic = True
        if torch.__version__ >= "1.8.0":
            torch.use_deterministic_algorithms(True)

    def state_dict(self, cls: Callable = dict) -> Mapping:
        if self.model is None:
            raise ValueError("Model must be defined when calling state_dict")
        return cls(
            runner=self.config.dict(),
            model=self.unwrap(self.model).state_dict(),
            ema=self.ema.state_dict() if self.ema else None,
            optimizer=self.optimizer.state_dict() if self.optimizer else None,
            scheduler=self.scheduler.state_dict() if self.scheduler else None,
        )

    def unwrap(self, model: nn.Module) -> nn.Module:
        if isinstance(model, (nn.parallel.DistributedDataParallel, nn.parallel.DataParallel)):
            return model.module
        return model

    def build_dataloaders(self):
        datasets = {k: d for k, d in self.datasets.items() if k not in self.dataloaders}
        default_kwargs = self.config.get("dataloader", NestedDict())
        dataloader_kwargs = NestedDict({k: default_kwargs.pop(k) for k in self.datasets if k in default_kwargs})
        for k, d in datasets.items():
            dataloader_kwargs.setdefault(k, NestedDict())
            dataloader_kwargs[k].merge(default_kwargs, overwrite=False)
            shuffle = dataloader_kwargs[k].pop("shuffle", getattr(d, "train", True))
            if self.distributed:
                sampler = utils.data.distributed.DistributedSampler(d, shuffle=shuffle)
            else:
                sampler = utils.data.RandomSampler(d) if shuffle else utils.data.SequentialSampler(d)
            dataloader_kwargs[k].setdefault("drop_last", not getattr(d, "train", True))
            self.dataloaders[k] = utils.data.DataLoader(
                d, sampler=sampler, collate_fn=self.collate_fn, **dataloader_kwargs[k]
            )

    @staticmethod
    def collate_fn(batch):
        return utils.data.dataloader.default_collate(batch)

    @contextmanager
    def autocast(self):
        if self.config.get("precision") is None:
            yield nullcontext()
        else:
            yield torch.autocast(self.device.type, dtype=get_precision(self.config.precision))

    @contextmanager
    def accumulate(self):
        if self.accum_steps <= 1 or self.steps % self.accum_steps == 0:
            yield nullcontext()
        else:
            yield self.model.no_sync()

    def get_optimizer(self, name: str):
        if name.lower() == "sgd":
            return optim.SGD
        if name.lower() == "asgd":
            return optim.ASGD
        if name.lower() in {"torch_adam", "torch_adamw"}:
            return optim.Adam
        if ds is not None:
            if name.lower() == "adagrad":
                return ds.ops.adagrad.DeepSpeedCPUAdagrad
            if name.lower() in {"adam", "adamw"}:
                if torch.cuda.device_count() > 0:
                    return ds.ops.adam.FusedAdam
                return ds.ops.adam.DeepSpeedCPUAdam
            if name.lower() in {"cpu", "cpu_adam", "cpuadam", "cpu_adamw", "cpuadamw"}:
                return ds.ops.adam.DeepSpeedCPUAdam
            if name.lower() == "lamb":
                if torch.cuda.device_count() > 0:
                    return ds.ops.lamb.FusedLamb
                return ds.ops.lamb.DeepSpeedCPULamb
            if name.lower() in {"cpulamb", "cpu_lamb"}:
                return ds.ops.lamb.DeepSpeedCPULamb
            if name.lower() == "lion":
                if torch.cuda.device_count() > 0:
                    return ds.ops.lion.FusedLion
                return ds.ops.lion.DeepSpeedCPULion
            if name.lower() in {"cpulion", "cpu_lion"}:
                return ds.ops.lion.DeepSpeedCPULion
        if name.lower() in {"adam", "adamw"}:
            return optim.AdamW
        if name.lower() == "adadelta":
            return optim.Adadelta
        if name.lower() == "adafactor":
            return optim.Adafactor
        if name.lower() == "adagrad":
            return optim.Adagrad
        if name.lower() == "adamax":
            return optim.Adamax
        if name.lower() == "lbfgs":
            return optim.LBFGS
        if name.lower() == "nadam":
            return optim.NAdam
        if name.lower() == "radam":
            return optim.RAdam
        if name.lower() == "rmsprop":
            return optim.RMSprop
        if name.lower() == "rprop":
            return optim.Rprop

    def get_deepspeed_config(self, config: NestedDict | str = None) -> NestedDict:  # pylint: disable=R0912, R0915
        r"""
        Preprocess DeepSpeed config.
        """

        if config is None and "deepspeed" in self.config:
            config = self.config.deepspeed
        if isinstance(config, str):
            config = NestedDict(config)
        if config is None:
            config = NestedDict()
        if config.get("steps_per_print", "auto") == "auto":
            config["steps_per_print"] = self.log_interval
        if config.get("train_micro_batch_size_per_gpu", "auto") == "auto":
            config["train_micro_batch_size_per_gpu"] = self.batch_size
        if config.get("gradient_accumulation_steps", "auto") == "auto":
            if self.accum_steps > 1:
                config["gradient_accumulation_steps"] = self.accum_steps
            else:
                config.pop("gradient_accumulation_steps", None)
        if "amp" in config:
            amp = config["amp"]
            if amp.get("enabled", "auto") == "auto":
                amp["enabled"] = "true"
            if amp.get("opt_level", "auto") == "auto":
                amp["opt_level"] = "O1"
        if "zero_optimization" in config:
            zero = config["zero_optimization"]
            if zero.get("allgather_bucket_size") == "auto":
                zero["allgather_bucket_size"] = 1e6
            if zero.get("reduce_bucket_size") == "auto":
                zero["reduce_bucket_size"] = 1e6
            if zero.get("stage3_max_live_parameters") == "auto":
                zero["stage3_max_live_parameters"] = 1e8
            if zero.get("stage3_max_live_gradients") == "auto":
                zero["stage3_max_live_gradients"] = 1e8
            if zero.get("stage3_max_reuse_distance") == "auto":
                zero["stage3_max_reuse_distance"] = 1e8
            if zero.get("stage3_prefetch_bucket_size") == "auto":
                zero["stage3_prefetch_bucket_size"] = 1e6
            if zero.get("stage3_param_persistence_threshold") == "auto":
                zero["stage3_param_persistence_threshold"] = 1e8
            if "amp" in config:
                if "fp16" not in config:
                    config["fp16"] = NestedDict()
                if config["fp16"].get("enabled", "auto"):
                    config["fp16"]["enabled"] = config["amp"]["enabled"]
                warn(
                    f"AMP is not compatible with ZeRO. Automatically set 'fp16' to {config['amp']['enabled']}",
                    stacklevel=2,
                )
                del config["amp"]
        if "optimizer" in config:
            if config["optimizer"].get("type", "auto") == "auto":
                config["optimizer"]["type"] = "Adam"
            if "params" not in config["optimizer"]:
                config["optimizer"]["params"] = NestedDict()
            optimizer = config["optimizer"]["params"]
            if optimizer.get("lr", "auto") == "auto":
                optimizer["lr"] = self.config.get("optim.lr", 1e-3)
            if optimizer.get("weight_decay", "auto") == "auto":
                optimizer["weight_decay"] = self.config.get("optim.weight_decay", 1e-2)
            if optimizer.get("betas") == "auto":
                optimizer["betas"] = (0.9, 0.999)
            if optimizer.get("eps") == "auto":
                optimizer["eps"] = 1e-8
        if "scheduler" in config:
            if config["scheduler"].get("type", "auto") == "auto":
                config["scheduler"]["type"] = "WarmupCosineLR"
            if "params" not in config["scheduler"]:
                config["scheduler"]["params"] = NestedDict()
            scheduler = config["scheduler"]["params"]
            if scheduler.get("total_num_steps", "auto") == "auto":
                scheduler["total_num_steps"] = self.total_steps
            if scheduler.get("warmup_num_steps", "auto") == "auto":
                scheduler["warmup_num_steps"] = scheduler["total_num_steps"] // 20
            if config["scheduler"]["type"] in ("WarmupLR", "WarmupDecayLR"):
                if scheduler.get("warmup_max_lr", "auto") == "auto":
                    if self.optimizer:
                        scheduler["warmup_max_lr"] = self.optimizer.param_groups[0]["lr"]
                    elif "optimizer" in config:
                        scheduler["warmup_max_lr"] = config["optimizer"]["params"]["lr"]
                    else:
                        scheduler["warmup_max_lr"] = self.config.get("optim.lr", 1e-3)
                if scheduler.get("warmup_min_lr", "auto") == "auto":
                    scheduler["warmup_min_lr"] = 1e-9
            else:
                scheduler.pop("warmup_max_lr", None)
                scheduler.pop("warmup_min_lr", None)
        if config.get("gradient_clipping", "auto") == "auto" and self.config.get("max_grad_norm") is not None:
            config["gradient_clipping"] = self.config["max_grad_norm"]
        return config

    @property
    def device(self):
        return torch.device("cuda", self.local_rank) if torch.cuda.is_available() else "cpu"

    @property
    def mode(self) -> RunnerMode:
        return self._mode

    @mode.setter
    def mode(self, mode: str | RunnerMode) -> None:
        if isinstance(mode, str):
            mode = RunnerMode(mode)
        self._mode = mode
        if self.model is not None:
            self.model.train(mode == RunnerMode.train)
        if self.ema is not None:
            self.ema.train(mode == RunnerMode.train)

    @property
    def rank(self) -> int:
        if self.distributed:
            return dist.get_rank()
        return 0

    @property
    def local_rank(self) -> int:
        if local_rank := os.getenv("LOCAL_RANK"):
            return int(local_rank)
        return 0

    @property
    def world_size(self) -> int:
        r"""
        Number of Processes.
        """

        return get_world_size()

    @property
    def distributed(self) -> bool:
        return self.world_size > 1

    @cached_property
    def accum_steps(self) -> int:
        return self.config.get("accum_steps", 1)

    @staticmethod
    def reduce(tensor: torch.Tensor) -> torch.Tensor:
        if torch.distributed.is_available() and torch.distributed.is_initialized():
            dist.all_reduce(tensor)
        return tensor

world_size property

Python

world_size: int

Number of Processes.

train

Python

train(train_splits: list[str] | None = None, evaluate_splits: list[str] | None = None) -> NestedDict

Perform training on split.

Parameters:

Name	Type	Description	Default
train_splits	list[str] | None	list of split to run train. Defaults to ["train"].	None
evaluate_splits	list[str] | None	list of split to run evaluate. Defaults to self.dataloaders except for those in train_splits.	None

Return

Source code in danling/runner/torch_runner.py

def train(self, train_splits: list[str] | None = None, evaluate_splits: list[str] | None = None) -> NestedDict:
    r"""
    Perform training on `split`.

    Args:
        train_splits: list of split to run train.
            Defaults to `["train"]`.
        evaluate_splits: list of split to run evaluate.
            Defaults to `self.dataloaders` except for those in `train_splits`.

    Return:
        NestedDict: train results
    """

    early_stop_counter = 0
    if train_splits is not None:
        self.train_splits = sorted(train_splits)
    if evaluate_splits is not None:
        self.evaluate_splits = sorted(evaluate_splits)
    if not self.train_splits:
        warn("No training split is found. Will only evaluate for one epoch.", stacklevel=2)
        self.epoch_end = self.epoch_begin + 1
    print(f"Begin training from {self.epoch_begin} to {self.epoch_end}")
    print(f"Training splits: {self.train_splits}")
    print(f"Evaluation splits: {self.evaluate_splits}")
    patience = self.config.get("patience", float("inf"))
    for epoch in range(self.epoch_begin, self.epoch_end):  # type: ignore
        self.epochs = epoch
        result = NestedDict()
        result.setattr("convert_mapping", True)
        for split in self.train_splits:
            result[split] = self.train_epoch(split)
        for split in self.evaluate_splits:
            result[split] = self.evaluate_epoch(split)
        self.append_result(result)
        print(self.format_epoch_result(result))
        self.save_result()
        if self.config.save_interval is not None:
            self.save_checkpoint()
        """@nni.report_intermediate_result(self.latest_score)"""
        early_stop_counter = 0 if self.is_best else early_stop_counter + 1
        if early_stop_counter > patience:
            print("early stop")
            break
    """@nni.report_final_result(self.latest_score)"""
    return self.results

train_epoch

Python

train_epoch(split: str = 'train') -> NestedDict

Train one epoch on split.

Parameters:

Name	Type	Description	Default
split	str	split to run train	'train'

Return

Source code in danling/runner/torch_runner.py

def train_epoch(self, split: str = "train") -> NestedDict:
    r"""
    Train one epoch on `split`.

    Args:
        split (str): split to run train

    Return:
        NestedDict: train result
    """

    self.mode = "train"  # type: ignore
    self.split = split
    loader = self.dataloaders[split]
    length = len(loader) - 1
    last_print_iteration = -1
    self.meters.reset()
    if self.train_metrics is not None:
        self.metrics = self.train_metrics
    if self.metrics is not None:
        self.metrics.reset()
    if hasattr(loader.batch_sampler, "set_epoch"):
        loader.batch_sampler.set_epoch(self.epochs)
    if hasattr(loader.sampler, "set_epoch"):
        loader.sampler.set_epoch(self.epochs)
    is_cuda = self.device == torch.device("cuda")
    batch_time = time()

    for iteration, data in enumerate(loader):
        _, loss = self.train_step(data)

        if self.log_interval > 0 and (iteration > 0 and iteration % self.log_interval == 0 or iteration == length):
            interval = iteration - last_print_iteration
            if is_cuda:
                torch.cuda.synchronize()
            if self.scheduler is not None:
                self.meters.lr.update(self.scheduler.get_last_lr()[0])
            self.meters.loss.update(self.reduce(loss).item())
            self.meters.time.update((time() - batch_time) / interval)
            batch_time = time()
            self.step_log(split, iteration, length)
            last_print_iteration = iteration

    result = self.get_epoch_result()
    return result

advance

Python

advance(loss) -> None

Backward loss and step optimizer & scheduler.

Parameters:

Name	Type	Description	Default
loss		The loss tensor from which to backpropagate.	required

Source code in danling/runner/torch_runner.py

def advance(self, loss) -> None:
    r"""
    Backward loss and step optimizer & scheduler.

    Args:
        loss: The loss tensor from which to backpropagate.
    """

    self.backward(loss / self.accum_steps)
    if self.accum_steps <= 1 or self.steps % self.accum_steps == 0:
        if self.config.get("max_grad_value") is not None:
            clip_grad_value_(self.model.parameters(), self.config["max_grad_value"])
        if self.config.get("max_grad_norm") is not None:
            clip_grad_norm_(self.model.parameters(), self.config["max_grad_norm"])
        self.optimizer.step()
        self.optimizer.zero_grad()
        self.steps += 1

evaluate

Python

evaluate(evaluate_splits: list[str] | None = None) -> NestedDict

Perform evaluation on evaluate_splits.

Parameters:

Name	Type	Description	Default
evaluate_splits	list[str] | None	list of split to run evaluate. Defaults to ["evaluate"].	None

Return

Source code in danling/runner/torch_runner.py

def evaluate(self, evaluate_splits: list[str] | None = None) -> NestedDict:
    r"""
    Perform evaluation on `evaluate_splits`.

    Args:
        evaluate_splits: list of split to run evaluate.
            Defaults to `["evaluate"]`.

    Return:
        NestedDict: evaluation result
    """

    if evaluate_splits is not None:
        self.evaluate_splits = sorted(evaluate_splits)
    if not self.evaluate_splits:
        raise ValueError("No evaluation splits found.")
    print("Begin evaluation")
    print(f"Evaluation splits: {self.evaluate_splits}")
    result = NestedDict()
    result.setattr("convert_mapping", True)
    for split in self.evaluate_splits:
        result[split] = self.evaluate_epoch(split=split)
    print(self.format_epoch_result(result))
    return result

evaluate_epoch

Python

evaluate_epoch(split: str = 'val') -> NestedDict

Evaluate one epoch on split.

Parameters:

Name	Type	Description	Default
split	str	split to run evaluate	'val'

Return

Source code in danling/runner/torch_runner.py

def evaluate_epoch(self, split: str = "val") -> NestedDict:
    r"""
    Evaluate one epoch on `split`.

    Args:
        split (str): split to run evaluate

    Return:
        NestedDict: evaluation result
    """

    self.mode = RunnerMode.evaluate
    self.split = split
    loader = self.dataloaders[split]
    length = len(loader) - 1
    last_print_iteration = -1
    self.meters.reset()
    if self.evaluate_metrics is not None:
        self.metrics = self.evaluate_metrics
    if self.metrics is not None:
        self.metrics.reset()
    batch_time = time()
    is_cuda = self.device == torch.device("cuda")

    for iteration, data in enumerate(loader):
        _, loss = self.evaluate_step(data)

        if self.log_interval > 0 and (iteration > 0 and iteration % self.log_interval == 0 or iteration == length):
            interval = iteration - last_print_iteration
            if is_cuda:
                torch.cuda.synchronize()
            self.meters.loss.update(self.reduce(loss).item())
            self.meters.time.update((time() - batch_time) / interval)
            batch_time = time()
            self.step_log(split, iteration, length)
            last_print_iteration = iteration

    result = self.get_epoch_result()
    self.write_result(result, split, self.epochs)
    return result

infer

Python

infer(split: str = 'infer') -> list[float]

Perform inference on split.

Parameters:

Name	Type	Description	Default
split	str	split to run inference	'infer'

Return

Source code in danling/runner/torch_runner.py

@torch.inference_mode()
def infer(self, split: str = "infer") -> list[float]:
    r"""
    Perform inference on `split`.

    Args:
        split (str): split to run inference

    Return:
        Tensor: inference outputs
    """

    self.mode = RunnerMode.infer
    loader = self.dataloaders[split]
    output: list[float] = []
    model = self.ema or self.model
    for _, data in tqdm(enumerate(loader), total=len(loader)):
        data = to_device(data, self.device)
        input = data["input"] if isinstance(data, Mapping) else data[0]
        pred = model(**input) if isinstance(input, Mapping) else model(input)
        output.extend(pred.squeeze(-1).tolist())

    if self.distributed:
        torch.cuda.synchronize()
        output = self.gather_for_metrics(output)
    return output

backward

Python

backward(loss: Tensor) -> None

Backward loss.

Parameters:

Name	Type	Description	Default
loss	Tensor	Loss to backward.	required

Source code in danling/runner/torch_runner.py

def backward(self, loss: torch.Tensor) -> None:
    r"""
    Backward loss.

    Args:
        loss: Loss to backward.
    """

    loss.backward()

has_nan_inf_grad

Python

has_nan_inf_grad(model: Module | None = None) -> bool

Check if model has NaN or Inf gradients.

Parameters:

Name	Type	Description	Default
model	Module | None	Model to check. Defaults to self.model.	None

Return

Source code in danling/runner/torch_runner.py

def has_nan_inf_grad(self, model: nn.Module | None = None) -> bool:
    r"""
    Check if model has NaN or Inf gradients.

    Args:
        model: Model to check.
            Defaults to `self.model`.

    Return:
        bool: True if NaN or Inf is detected in gradients.
    """
    model = model or self.model
    for name, param in model.named_parameters():
        if param.grad is not None:
            if torch.isnan(param.grad).any():
                print(f"NaN detected in gradients of parameter: {name}")
                return True
            if torch.isinf(param.grad).any():
                print(f"Inf detected in gradients of parameter: {name}")
                return True
    return False

init_distributed

Python

init_distributed() -> None

Set up distributed training.

Initialise process group and set up DDP variables.

Source code in danling/runner/torch_runner.py

def init_distributed(self) -> None:
    r"""
    Set up distributed training.

    Initialise process group and set up DDP variables.
    """

    backend = self.config.get("backend", os.getenv("BACKEND"))
    init_method = self.config.get("init_method", os.getenv("INIT_METHOD"))
    world_size = int(self.config.get("world_size", os.getenv("WORLD_SIZE", "1")))
    rank = int(self.config.get("rank", os.getenv("RANK", "0")))
    if world_size > 1:
        if torch.cuda.is_available():
            torch.cuda.set_device(self.local_rank)
        dist.init_process_group(backend, init_method, world_size=world_size, rank=rank)
        object_list = [self.id, self.timestamp]
        dist.broadcast_object_list(object_list)
        self.id, self.timestamp = object_list

init_tensorboard

Python

init_tensorboard(*args, **kwargs) -> None

Set up Tensoraoard SummaryWriter.

Source code in danling/runner/torch_runner.py

@on_main_process
def init_tensorboard(self, *args, **kwargs) -> None:
    r"""
    Set up Tensoraoard SummaryWriter.
    """
    from torch.utils.tensorboard.writer import SummaryWriter  # pylint: disable=C0415

    if "log_dir" not in kwargs:
        kwargs["log_dir"] = self.dir

    self.writer = SummaryWriter(*args, **kwargs)
    self.writer.add_scalar = catch(OSError, verbose=False)(self.writer.add_scalar)

set_seed

Python

set_seed(seed: int = None, bias: int = None) -> int

Set up random seed.

Parameters:

Name	Type	Description	Default
seed	int	Random seed to set. Defaults to self.config.seed (config.seed).	None
bias	int	Make the seed different for each processes. This is used to ensure the data augmentation are applied differently on every processes. Defaults to self.rank. Set to False to disable this feature.	None

Source code in danling/runner/torch_runner.py

def set_seed(self, seed: int = None, bias: int = None) -> int:  # type: ignore[assignment]
    r"""
    Set up random seed.

    Args:
        seed: Random seed to set.
            Defaults to `self.config.seed` (`config.seed`).

        bias: Make the seed different for each processes.
            This is used to ensure the data augmentation are applied differently on every processes.
            Defaults to `self.rank`.
            Set to `False` to disable this feature.
    Returns:
        Random seed set.
    """

    seed = seed or self.config.seed  # type: ignore[assignment]
    if seed is None:
        if self.inited:
            seed = random.randint(0, 2**32 - 1)
            if self.distributed:
                object_list = [seed]
                dist.broadcast_object_list(object_list)
                seed = object_list[0]
            self.config.seed = seed
    else:
        seed = defaults.SEED
    bias = bias or self.rank
    if bias:
        seed += bias
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    if np_random is not None:
        np_random.seed(seed)
    random.seed(seed)
    return seed

get_deepspeed_config

Python

get_deepspeed_config(config: NestedDict | str = None) -> NestedDict

Preprocess DeepSpeed config.

Source code in danling/runner/torch_runner.py

def get_deepspeed_config(self, config: NestedDict | str = None) -> NestedDict:  # pylint: disable=R0912, R0915
    r"""
    Preprocess DeepSpeed config.
    """

    if config is None and "deepspeed" in self.config:
        config = self.config.deepspeed
    if isinstance(config, str):
        config = NestedDict(config)
    if config is None:
        config = NestedDict()
    if config.get("steps_per_print", "auto") == "auto":
        config["steps_per_print"] = self.log_interval
    if config.get("train_micro_batch_size_per_gpu", "auto") == "auto":
        config["train_micro_batch_size_per_gpu"] = self.batch_size
    if config.get("gradient_accumulation_steps", "auto") == "auto":
        if self.accum_steps > 1:
            config["gradient_accumulation_steps"] = self.accum_steps
        else:
            config.pop("gradient_accumulation_steps", None)
    if "amp" in config:
        amp = config["amp"]
        if amp.get("enabled", "auto") == "auto":
            amp["enabled"] = "true"
        if amp.get("opt_level", "auto") == "auto":
            amp["opt_level"] = "O1"
    if "zero_optimization" in config:
        zero = config["zero_optimization"]
        if zero.get("allgather_bucket_size") == "auto":
            zero["allgather_bucket_size"] = 1e6
        if zero.get("reduce_bucket_size") == "auto":
            zero["reduce_bucket_size"] = 1e6
        if zero.get("stage3_max_live_parameters") == "auto":
            zero["stage3_max_live_parameters"] = 1e8
        if zero.get("stage3_max_live_gradients") == "auto":
            zero["stage3_max_live_gradients"] = 1e8
        if zero.get("stage3_max_reuse_distance") == "auto":
            zero["stage3_max_reuse_distance"] = 1e8
        if zero.get("stage3_prefetch_bucket_size") == "auto":
            zero["stage3_prefetch_bucket_size"] = 1e6
        if zero.get("stage3_param_persistence_threshold") == "auto":
            zero["stage3_param_persistence_threshold"] = 1e8
        if "amp" in config:
            if "fp16" not in config:
                config["fp16"] = NestedDict()
            if config["fp16"].get("enabled", "auto"):
                config["fp16"]["enabled"] = config["amp"]["enabled"]
            warn(
                f"AMP is not compatible with ZeRO. Automatically set 'fp16' to {config['amp']['enabled']}",
                stacklevel=2,
            )
            del config["amp"]
    if "optimizer" in config:
        if config["optimizer"].get("type", "auto") == "auto":
            config["optimizer"]["type"] = "Adam"
        if "params" not in config["optimizer"]:
            config["optimizer"]["params"] = NestedDict()
        optimizer = config["optimizer"]["params"]
        if optimizer.get("lr", "auto") == "auto":
            optimizer["lr"] = self.config.get("optim.lr", 1e-3)
        if optimizer.get("weight_decay", "auto") == "auto":
            optimizer["weight_decay"] = self.config.get("optim.weight_decay", 1e-2)
        if optimizer.get("betas") == "auto":
            optimizer["betas"] = (0.9, 0.999)
        if optimizer.get("eps") == "auto":
            optimizer["eps"] = 1e-8
    if "scheduler" in config:
        if config["scheduler"].get("type", "auto") == "auto":
            config["scheduler"]["type"] = "WarmupCosineLR"
        if "params" not in config["scheduler"]:
            config["scheduler"]["params"] = NestedDict()
        scheduler = config["scheduler"]["params"]
        if scheduler.get("total_num_steps", "auto") == "auto":
            scheduler["total_num_steps"] = self.total_steps
        if scheduler.get("warmup_num_steps", "auto") == "auto":
            scheduler["warmup_num_steps"] = scheduler["total_num_steps"] // 20
        if config["scheduler"]["type"] in ("WarmupLR", "WarmupDecayLR"):
            if scheduler.get("warmup_max_lr", "auto") == "auto":
                if self.optimizer:
                    scheduler["warmup_max_lr"] = self.optimizer.param_groups[0]["lr"]
                elif "optimizer" in config:
                    scheduler["warmup_max_lr"] = config["optimizer"]["params"]["lr"]
                else:
                    scheduler["warmup_max_lr"] = self.config.get("optim.lr", 1e-3)
            if scheduler.get("warmup_min_lr", "auto") == "auto":
                scheduler["warmup_min_lr"] = 1e-9
        else:
            scheduler.pop("warmup_max_lr", None)
            scheduler.pop("warmup_min_lr", None)
    if config.get("gradient_clipping", "auto") == "auto" and self.config.get("max_grad_norm") is not None:
        config["gradient_clipping"] = self.config["max_grad_norm"]
    return config

NestedTensor

A container for variable-length tensors that enables efficient batch operations.

NestedTensor solves a fundamental problem in deep learning: handling sequences of different lengths in batch operations. Instead of excessive padding or complex bucketing, NestedTensor provides an elegant solution that maintains both efficiency and usability.

The class provides three main views of the data: - .tensor: A padded tensor with zeros (or other value) in place of missing elements - .mask: A boolean mask indicating which elements are real vs padding - .concat: A flattened tensor containing all elements concatenated (no padding)

When indexing a NestedTensor, the behavior depends on the index type: 1. Integer index (nt[0]): Returns a single tensor without padding 2. Slice index (nt[:]): Returns a tuple of (padded_tensor, mask) 3. Tuple index (nt[:, 1:]): Returns a new NestedTensor with the specified sliced shape

Attributes:

Name	Type	Description
_storage		The sequence of original tensors (internal use)
tensor	Tensor	Padded tensor representing all sequences with padding
mask	Tensor	Boolean mask where True indicates real elements, False indicates padding
concat	Tensor	Concatenated tensor of all sequences without padding
batch_first	bool	Whether the first dimension is the batch dimension (B, N, *) If False, the first dimension is the sequence dimension (N, B, *)
padding_value	SupportsFloat	Value used for padding in the padded tensor
mask_value	bool	Value used in the mask to indicate padding positions (usually False)

Parameters:

Name	Type	Description	Default
*tensors	Iterable[Tensor]	Variable-length tensors or sequences to store	()
batch_first	bool	Whether to use batch-first representation.	True
padding_value	SupportsFloat	Value to use for padding.	0.0
mask_value	bool	Value to use for padding positions in mask.	False

Raises:

Type	Description
ValueError	If tensors is not an iterable or is empty

Examples:

Basic usage:

>>> nested_tensor = NestedTensor(torch.tensor([1, 2, 3]), torch.tensor([4, 5]))
>>> nested_tensor.shape
torch.Size([2, 3])
>>> nested_tensor.tensor  # Padded representation
tensor([[1, 2, 3],
        [4, 5, 0]])
>>> nested_tensor.mask  # Mask showing real vs padding values
tensor([[ True,  True,  True],
        [ True,  True, False]])
>>> nested_tensor.concat  # Concatenated version (no padding)
tensor([1, 2, 3, 4, 5])

>>> nested_tensor[0]  # First tensor (no padding)
tensor([1, 2, 3])
>>> nested_tensor[:2]  # Padded tensor and mask
NestedTensor([[1, 2, 3],
        [4, 5, 0]])
>>> nested_tensor[:, 1:]  # Slice operations return a new NestedTensor
NestedTensor([[2, 3],
        [5, 0]])

Type conversion:

>>> nested_tensor.to(torch.float).tensor
tensor([[1., 2., 3.],
        [4., 5., 0.]])
>>> nested_tensor.half().tensor
tensor([[1., 2., 3.],
        [4., 5., 0.]], dtype=torch.float16)

Conversion to Python types:

>>> nested_tensor.tolist()
[[1, 2, 3], [4, 5]]

Creating from Python lists:

>>> NestedTensor(*[[1, 2, 3], [4, 5]])
NestedTensor([[1, 2, 3],
        [4, 5, 0]])

Source code in danling/tensor/nested_tensor.py

class NestedTensor:
    r"""
    A container for variable-length tensors that enables efficient batch operations.

    `NestedTensor` solves a fundamental problem in deep learning: handling sequences of different lengths
    in batch operations. Instead of excessive padding or complex bucketing, `NestedTensor` provides an
    elegant solution that maintains both efficiency and usability.

    The class provides three main views of the data:
    - `.tensor`: A padded tensor with zeros (or other value) in place of missing elements
    - `.mask`: A boolean mask indicating which elements are real vs padding
    - `.concat`: A flattened tensor containing all elements concatenated (no padding)

    When indexing a `NestedTensor`, the behavior depends on the index type:
    1. Integer index (`nt[0]`): Returns a single tensor without padding
    2. Slice index (`nt[:]`): Returns a tuple of (padded_tensor, mask)
    3. Tuple index (`nt[:, 1:]`): Returns a new `NestedTensor` with the specified sliced shape

    Attributes:
        _storage: The sequence of original tensors (internal use)
        tensor: Padded tensor representing all sequences with padding
        mask: Boolean mask where True indicates real elements, False indicates padding
        concat: Concatenated tensor of all sequences without padding
        batch_first: Whether the first dimension is the batch dimension (B, N, *)
            If `False`, the first dimension is the sequence dimension (N, B, *)
        padding_value: Value used for padding in the padded tensor
        mask_value: Value used in the mask to indicate padding positions (usually False)

    Args:
        *tensors: Variable-length tensors or sequences to store
        batch_first: Whether to use batch-first representation.
        padding_value: Value to use for padding.
        mask_value: Value to use for padding positions in mask.

    Raises:
        ValueError: If `tensors` is not an iterable or is empty

    Examples:
        Basic usage:
        >>> nested_tensor = NestedTensor(torch.tensor([1, 2, 3]), torch.tensor([4, 5]))
        >>> nested_tensor.shape
        torch.Size([2, 3])
        >>> nested_tensor.tensor  # Padded representation
        tensor([[1, 2, 3],
                [4, 5, 0]])
        >>> nested_tensor.mask  # Mask showing real vs padding values
        tensor([[ True,  True,  True],
                [ True,  True, False]])
        >>> nested_tensor.concat  # Concatenated version (no padding)
        tensor([1, 2, 3, 4, 5])

        Indexing:
        >>> nested_tensor[0]  # First tensor (no padding)
        tensor([1, 2, 3])
        >>> nested_tensor[:2]  # Padded tensor and mask
        NestedTensor([[1, 2, 3],
                [4, 5, 0]])
        >>> nested_tensor[:, 1:]  # Slice operations return a new NestedTensor
        NestedTensor([[2, 3],
                [5, 0]])

        Type conversion:
        >>> nested_tensor.to(torch.float).tensor
        tensor([[1., 2., 3.],
                [4., 5., 0.]])
        >>> nested_tensor.half().tensor
        tensor([[1., 2., 3.],
                [4., 5., 0.]], dtype=torch.float16)

        Conversion to Python types:
        >>> nested_tensor.tolist()
        [[1, 2, 3], [4, 5]]

        Creating from Python lists:
        >>> NestedTensor(*[[1, 2, 3], [4, 5]])
        NestedTensor([[1, 2, 3],
                [4, 5, 0]])
    """

    __storage: Tuple[Tensor, ...]

    dtype: torch.dtype | None = None
    device: torch.device | None = None
    requires_grad: bool | None = None
    _pin_memory: bool = False

    batch_first: bool = True
    padding_value: SupportsFloat = 0.0
    mask_value: bool = False

    def __init__(
        self,
        *tensors: Iterable[Tensor],
        dtype: torch.dtype | None = None,
        device: torch.device | None = None,
        requires_grad: bool | None = None,
        pin_memory: bool = False,
        batch_first: bool = True,
        padding_value: SupportsFloat = 0.0,
        mask_value: bool = False,
    ) -> None:
        self.dtype = dtype
        self.device = device
        self.requires_grad = requires_grad
        self._pin_memory = pin_memory
        if len(tensors) == 1 and isinstance(tensors, Sequence):
            tensors = tensors[0]  # type: ignore
        self._storage = tensors
        self.batch_first = batch_first
        self.padding_value = padding_value
        self.mask_value = mask_value

    @property
    def _storage(self):
        return self.__storage

    @_storage.setter
    def _storage(self, tensors: Sequence):
        if not isinstance(tensors, Iterable):
            raise ValueError(f"tensors must be an Iterable, bug got {type(tensors)}.")
        if isinstance(tensors, Tensor) and hasattr(tensors, "unbind"):
            tensors = tensors.unbind()
        tensors_ = []
        for t in tensors:
            if not isinstance(t, Tensor):
                t = torch.tensor(
                    t,
                    dtype=self.dtype,
                    device=self.device,
                    pin_memory=self._pin_memory,
                )
            else:
                t = t.to(self.device, dtype=self.dtype)
            if self.requires_grad is not None:
                t.requires_grad_(self.requires_grad)
            tensors_.append(t)
        if len(tensors_) == 0:
            self.__storage = ()
            return
        tensors = tuple(tensors_)
        self.dtype = tensors[0].dtype
        self.device = tensors[0].device
        self.requires_grad = tensors[0].requires_grad
        # if drop_last=False, the last element is likely not a NestedTensor and has an extra batch dimension
        ndims = {t.ndim for t in tensors[:-1]}
        if len(ndims) == 1:
            (ndim,) = ndims
            if tensors[-1].ndim == ndim + 1 and tensors[-1].size(0) == 1:
                tensors = tensors[:-1] + (tensors[-1].squeeze(0),)
        self.__storage = tensors

    def storage(self):
        return self._storage

    @property
    def tensor_mask(self) -> Tuple[Tensor, Tensor]:
        r"""
        Return a tuple of padded tensor and mask tensor.

        Examples:
            >>> nested_tensor = NestedTensor([torch.tensor([1, 2, 3]), torch.tensor([4, 5])])
            >>> nested_tensor.tensor_mask
            (tensor([[1, 2, 3],
                    [4, 5, 0]]), tensor([[ True,  True,  True],
                    [ True,  True, False]]))
        """

        return self._tensor_mask(
            self._storage, self.batch_first, self.padding_value, self.mask_value, self.requires_grad
        )

    @method_cache(maxsize=1)
    def _tensor_mask(
        self, storage: tuple, batch_first: bool, padding_value: SupportsFloat, mask_value: bool, requires_grad: bool
    ) -> Tensor:
        if storage[0].dim() == 0:
            return torch.stack(storage, dim=0), torch.full(
                (len(storage),), not mask_value, dtype=torch.bool, device=self.device
            )
        return tensor_mask(
            storage,
            size=self.size(),
            batch_first=batch_first,
            padding_value=float(padding_value),
            mask_value=mask_value,
        )

    @property
    def tensor(self) -> Tensor:
        r"""
        Return a single tensor by padding all the tensors.

        Examples:
            >>> nested_tensor = NestedTensor([torch.tensor([1, 2, 3]), torch.tensor([4, 5])])
            >>> nested_tensor.tensor
            tensor([[1, 2, 3],
                    [4, 5, 0]])
        """

        return self._tensor(self._storage, self.batch_first, self.padding_value, self.requires_grad)

    @method_cache(maxsize=1)
    def _tensor(self, storage: tuple, batch_first: bool, padding_value: SupportsFloat, requires_grad: bool) -> Tensor:
        if storage[0].dim() == 0:
            return torch.stack(storage, dim=0)
        return pad_tensor(storage, size=self.size(), batch_first=batch_first, padding_value=float(padding_value))

    @property
    def mask(self) -> Tensor:
        r"""
        Padding mask of `tensor`.

        Examples:
            >>> nested_tensor = NestedTensor([torch.tensor([1, 2, 3]), torch.tensor([4, 5])])
            >>> nested_tensor.mask
            tensor([[ True,  True,  True],
                    [ True,  True, False]])
        """

        return self._mask(self._storage, self.batch_first, self.mask_value, self.requires_grad)

    @method_cache(maxsize=1)
    def _mask(self, storage: tuple, batch_first: bool, mask_value: bool, requires_grad: bool) -> Tensor:
        if storage[0].dim() == 0:
            return torch.full((len(storage),), not mask_value, dtype=torch.bool, device=self.device)
        return mask_tensor(storage, size=self.size(), batch_first=batch_first, mask_value=mask_value)

    @property
    def concat(self) -> Tensor:
        r"""
        Concat `tensor` in padding dim.

        This is particularly useful when calculating loss or passing `Linear` to avoid unnecessary computation.

        Examples:
            >>> nested_tensor = NestedTensor([torch.randn(9, 8), torch.randn(11, 8)])
            >>> nested_tensor.concat.shape
            torch.Size([20, 8])
            >>> nested_tensor = NestedTensor([torch.randn(9, 9, 8), torch.randn(11, 11, 8)])
            >>> nested_tensor.concat.shape
            torch.Size([202, 8])
            >>> nested_tensor = NestedTensor([torch.randn(9, 9, 8, 6), torch.randn(11, 11, 8, 6)])
            >>> nested_tensor.concat.shape
            torch.Size([202, 8, 6])
            >>> nested_tensor = NestedTensor([torch.randn(9, 9, 8, 7), torch.randn(11, 11, 8, 6)])
            >>> nested_tensor.concat.shape
            torch.Size([1293, 8])
            >>> nested_tensor = NestedTensor([torch.randn(1, 9, 9, 5), torch.randn(1, 11, 11, 5)])
        """
        shape = list(self.size())  # type: ignore[arg-type]
        shape = shape[1:] if self.batch_first else shape[0] + shape[2:]
        elem = self._storage[0]
        if elem.shape == shape:
            return torch.cat(self._storage, dim=1 if self.batch_first else 0)
        static_dims = set(range(len(shape)))
        for i, s in enumerate(shape):
            if not all(t.size(i) == s for t in self._storage):
                shape[i] = -1
                static_dims.remove(i)
        target_shape = [-1] + [s for s in shape if s != -1]
        storage = [i.reshape(target_shape) for i in self._storage]
        return torch.cat(storage, dim=0 if self.batch_first else 1)

    @property
    def torch(self) -> Tensor:
        r"""
        Create a `torch.nested.nested_tensor` object from `self`.

        Examples:
            >>> nested_tensor = NestedTensor([[2, 3, 5], [7, 8]])
            >>> nested_tensor.torch
            nested_tensor([
              tensor([2, 3, 5]),
              tensor([7, 8])
            ])
        """
        if nested is None:
            raise ImportError("torch.nested is not available, please install torch with nested support.")
        return nested.nested_tensor(list(self._storage))

    def unbind(self, dim: int = 0) -> Tuple[Tensor, ...]:
        r"""
        Unbind the NestedTensor.
        """
        if dim != 0:
            raise ValueError(f"NestedTensor can only be unbound along dimension 0, got dimension {dim} instead.")
        return self._storage

    @property
    def occupancy(self) -> float:
        r"""
        Occupancy of the NestedTensor.

        Examples:
            >>> nested_tensor = NestedTensor([torch.tensor([1, 2, 3, 4]), torch.tensor([5, 6])])
            >>> nested_tensor.occupancy
            0.75
        """

        return self.numel() / self.shape.numel()  # type: ignore[union-attr]

    @classmethod
    def from_tensor_mask(cls, tensor: Tensor, mask: Tensor, **kwargs):
        r"""
        Build a `NestedTensor` object from a padded `Tensor` and corresponding mask `Tensor`.

        Args:
            tensor: Padded Tensor.
            mask: Tensor Mask.

        Examples:
            >>> padded_tensor = torch.tensor([[1, 2, 3, 0, 0],
            ...                                [4, 5, 0, 0, 0],
            ...                                [6, 7, 8, 9, 0]])
            >>> mask_tensor = torch.tensor([[1, 1, 1, 0, 0],
            ...                             [1, 1, 0, 0, 0],
            ...                             [1, 1, 1, 1, 0]])
            >>> nested_tensor = NestedTensor.from_tensor_mask(padded_tensor, mask_tensor)
            >>> nested_tensor
            NestedTensor([[1, 2, 3, 0],
                    [4, 5, 0, 0],
                    [6, 7, 8, 9]])
        """

        if mask.ndim == 1:
            return cls(tensor, **kwargs)
        if mask.ndim == 2:
            return cls((t[slice(0, m.sum())] for t, m in zip(tensor, mask)), **kwargs)
        return cls(
            (
                t[[slice(0, (m.sum(dim=dim) > 0).sum().item()) for dim in reversed(range(m.dim()))]]
                for t, m in zip(tensor, mask)
            ),
            **kwargs,
        )

    def nested_like(self, tensor: Tensor, strict: bool = True) -> Self:
        r"""
        Create a new `NestedTensor` from a `Tensor`.
        The newly created `NestedTensor` will have the same shape as current `NestedTensor`.

        Args:
            tensor: The tensor to be converted to `NestedTensor`.
            strict: Check if the shape of `tensor` is the same as the current `NestedTensor`.

        Examples:
            >>> nested_tensor = NestedTensor([torch.tensor([1, 2, 3]), torch.tensor([4, 5])])
            >>> (nested_tensor == nested_tensor.nested_like(nested_tensor)).all()
            tensor(True)
            >>> tensor = nested_tensor.tensor
            >>> (nested_tensor == nested_tensor.nested_like(tensor)).all()
            tensor(True)
            >>> f = nested_tensor.nested_like(torch.randn(2, 2))
            Traceback (most recent call last):
            ValueError: The shape of NestedTensor and input tensor does not match, torch.Size([2, 3]) != torch.Size([2, 2])
            >>> p = nested_tensor.nested_like(torch.randn(2, 2), False)
            >>> p = nested_tensor.nested_like(torch.randn(3, 3), False)
            Traceback (most recent call last):
            ValueError: The batch size of NestedTensor and input tensor does not match, 2 != 3
        """  # noqa: E501

        if isinstance(tensor, NestedTensor):
            return tensor.clone()

        if strict and self.shape != tensor.shape:
            raise ValueError(
                f"The shape of NestedTensor and input tensor does not match, {self.shape} != {tensor.shape}"
            )
        if self.size(0) != tensor.size(0):
            raise ValueError(
                f"The batch size of NestedTensor and input tensor does not match, {self.size(0)} != {tensor.size(0)}"
            )
        return self.__class__([o[tuple(slice(0, dim) for dim in t.shape)] for t, o in zip(self._storage, tensor)])

    @classmethod
    def __torch_function__(cls, func, types, args=(), kwargs=None) -> Self:
        if kwargs is None:
            kwargs = {}
        if func in NestedTensorFuncRegistry:
            return NestedTensorFuncRegistry[func](*args, **kwargs)
        args = [a.tensor if hasattr(a, "tensor") else a for a in args]
        for k, v in kwargs.items():
            if hasattr(v, "tensor"):
                kwargs[k] = v.tensor
        output = func(*args, **kwargs)
        if isinstance(output, (Tensor, NestedTensor)):
            return output
        return cls(output)

    @classmethod
    def __torch_dispatch__(cls, func, types, args=(), kwargs=None) -> Self:
        args = [a.tensor if hasattr(a, "tensor") else a for a in args]
        for k, v in kwargs.items():
            if hasattr(v, "tensor"):
                kwargs[k] = v.tensor
        output = func(*args, **kwargs)
        if isinstance(output, (Tensor, NestedTensor)):
            return output
        return cls(output)

    def __getitem__(self, index: int | slice | list | tuple | Tensor | NestedTensor) -> Tensor | NestedTensor:
        if isinstance(index, int):
            return self._storage[index]
        if isinstance(index, (slice, list)):
            storage = tuple(self._storage[index] if isinstance(index, slice) else [self._storage[i] for i in index])
            return NestedTensor(storage, **self._state)
        if isinstance(index, tuple):
            return NestedTensor([t[index[0]][index[1:]] for t in self._storage])
        if isinstance(index, Tensor):
            index = self.nested_like(index, strict=False)
        if isinstance(index, NestedTensor):
            return NestedTensor([t[i] for t, i in zip(self._storage, index._storage)])
        raise ValueError(f"Unsupported index type {type(index)}")

    def __setitem__(self, index: int | slice | list | tuple, value: Tensor | NestedTensor) -> None:
        """
        Set values in the NestedTensor at the specified index.

        Args:
            index: The index to modify. Can be an integer, slice, list, or tuple.
            value: The new value to set. Can be a Tensor or NestedTensor.

        Examples:
            >>> nested_tensor = NestedTensor([torch.tensor([1, 2, 3]), torch.tensor([4, 5])])
            >>> nested_tensor[0] = torch.tensor([6, 7, 8])
            >>> nested_tensor[0]
            tensor([6, 7, 8])
            >>> nested_tensor[1] = torch.tensor([9, 10, 11, 12])
            >>> nested_tensor.shape
            torch.Size([2, 4])
        """
        if isinstance(index, int):
            if isinstance(value, NestedTensor):
                if len(value._storage) != 1:
                    raise ValueError(
                        f"When setting with an integer index, value must have a single tensor, but got {len(value)}"
                    )
                value = value._storage[0]
            if not isinstance(value, Tensor):
                value = torch.tensor(value, device=self.device)
            # Create a new list of tensors to modify
            storage_list = list(self._storage)
            storage_list[index] = value
            self._storage = tuple(storage_list)
        elif isinstance(index, (slice, list)):
            if isinstance(value, Tensor):
                # Convert tensor to NestedTensor if it's a regular tensor
                if value.dim() > 1 and value.size(0) > 1:
                    value = NestedTensor(value.unbind(0))
                else:
                    value = NestedTensor([value])

            if isinstance(index, slice):
                start, stop, step = index.indices(len(self._storage))
                indices = range(start, stop, step)
            else:
                indices = index  # type: ignore[assignment]

            if len(indices) != len(value._storage):
                raise ValueError(
                    f"Size mismatch: tried to assign {len(value._storage)} values to {len(indices)} indices"
                )

            storage_list = list(self._storage)
            for i, idx in enumerate(indices):
                storage_list[idx] = value._storage[i]
            self._storage = tuple(storage_list)
        elif isinstance(index, tuple):
            if len(index) < 2:
                raise ValueError("Tuple index must have at least two elements")

            first_idx, rest_idx = index[0], index[1:]

            if isinstance(first_idx, int):
                # Handle single tensor modification
                storage_list = list(self._storage)
                tensor = storage_list[first_idx]
                tensor[rest_idx] = value
                storage_list[first_idx] = tensor
                self._storage = tuple(storage_list)
            elif isinstance(first_idx, (slice, list)):
                # Handle multiple tensor modification
                if isinstance(first_idx, slice):
                    start, stop, step = first_idx.indices(len(self._storage))
                    indices = range(start, stop, step)
                else:
                    indices = first_idx  # type: ignore[assignment]

                storage_list = list(self._storage)
                for idx in indices:
                    tensor = storage_list[idx]
                    tensor[rest_idx] = value
                    storage_list[idx] = tensor
                self._storage = tuple(storage_list)
            else:
                raise ValueError(f"Unsupported first index type {type(first_idx)}")
        else:
            raise ValueError(f"Unsupported index type {type(index)}")

    def __getattr__(self, name: str) -> Any:
        if not self._storage:
            raise ValueError(f"Unable to get {name} from an empty {self.__class__.__name__}")
        ret = [getattr(i, name) for i in self._storage]
        elem = ret[0]
        if isinstance(elem, Tensor):
            return NestedTensor(ret, **self._state)
        if callable(elem):
            return NestedTensorFuncWrapper(ret, state=self._state)
        if elem.__hash__ is not None and len(set(ret)) == 1:
            return elem
        return ret

    def __iter__(self):
        return iter(self._storage)

    @property
    def _state(self) -> Mapping:
        return self.__state__(return_dtype=False, return_device=True, return_requires_grad=False)

    def __state__(
        self, return_dtype: bool = True, return_device: bool = True, return_requires_grad: bool = False
    ) -> Mapping:
        state = {k: v for k, v in self.__dict__.items() if not (k.startswith("_") or k.endswith("_"))}
        if not return_dtype:
            state.pop("dtype", None)
        if not return_device:
            state.pop("device", None)
        if not return_requires_grad:
            state.pop("requires_grad", None)
        return state

    def __setstate__(self, state: Mapping) -> None:
        self.__dict__.update(state)

    def __len__(self) -> int:
        return len(self._storage)

    def __repr__(self):
        if not self._storage:
            return self.__class__.__name__ + "()"
        return self.__class__.__name__ + repr(self.tensor)[len(self.tensor.__class__.__name__) :]  # noqa: E203

    def __bool__(self) -> int:
        return all(bool(x) for x in self._storage)

    def __gt__(  # type: ignore[override]
        self, other: Tensor | NestedTensor | SupportsFloat
    ) -> bool | Tensor | NestedTensor:
        if isinstance(other, Tensor) and self.shape == other.shape:
            other = self.nested_like(other)
        if isinstance(other, NestedTensor):
            return NestedTensor(i > j for i, j in zip(self._storage, other._storage))
        if isinstance(other, (int, float, Tensor)):
            return NestedTensor([x > other for x in self._storage], **self._state)
        raise TypeError(f"> not supported between instances of '{type(self)}' and '{type(other)}'")

    def __ge__(  # type: ignore[override]
        self, other: Tensor | NestedTensor | SupportsFloat
    ) -> bool | Tensor | NestedTensor:
        if isinstance(other, Tensor) and self.shape == other.shape:
            other = self.nested_like(other)
        if isinstance(other, NestedTensor):
            return NestedTensor(i >= j for i, j in zip(self._storage, other._storage))
        if isinstance(other, (int, float, Tensor)):
            return NestedTensor([x >= other for x in self._storage], **self._state)
        raise TypeError(f">= not supported between instances of '{type(self)}' and '{type(other)}'")

    def __eq__(  # type: ignore[override]
        self, other: Tensor | NestedTensor | SupportsFloat
    ) -> bool | Tensor | NestedTensor:
        if isinstance(other, Tensor) and self.shape == other.shape:
            other = self.nested_like(other)
        if isinstance(other, NestedTensor):
            return NestedTensor(i == j for i, j in zip(self._storage, other._storage))
        if isinstance(other, (int, float, Tensor)):
            return NestedTensor([x == other for x in self._storage], **self._state)
        return False

    def __ne__(  # type: ignore[override]
        self, other: Tensor | NestedTensor | SupportsFloat
    ) -> bool | Tensor | NestedTensor:
        if isinstance(other, Tensor) and self.shape == other.shape:
            other = self.nested_like(other)
        if isinstance(other, NestedTensor):
            return NestedTensor(i != j for i, j in zip(self._storage, other._storage))
        if isinstance(other, (int, float, Tensor)):
            return NestedTensor([x != other for x in self._storage], **self._state)
        return True

    def __le__(  # type: ignore[override]
        self, other: Tensor | NestedTensor | SupportsFloat
    ) -> bool | Tensor | NestedTensor:
        if isinstance(other, Tensor) and self.shape == other.shape:
            other = self.nested_like(other)
        if isinstance(other, NestedTensor):
            return NestedTensor(i <= j for i, j in zip(self._storage, other._storage))
        if isinstance(other, (int, float, Tensor)):
            return NestedTensor([x <= other for x in self._storage], **self._state)
        raise TypeError(f"<= not supported between instances of '{type(self)}' and '{type(other)}'")

    def __lt__(  # type: ignore[override]
        self, other: Tensor | NestedTensor | SupportsFloat
    ) -> bool | Tensor | NestedTensor:
        if isinstance(other, Tensor) and self.shape == other.shape:
            other = self.nested_like(other)
        if isinstance(other, NestedTensor):
            return NestedTensor(i < j for i, j in zip(self._storage, other._storage))
        if isinstance(other, (int, float, Tensor)):
            return NestedTensor([x < other for x in self._storage], **self._state)
        raise TypeError(f"< not supported between instances of '{type(self)}' and '{type(other)}'")

    def __abs__(self):
        return NestedTensor([abs(value) for value in self._storage], **self._state)

    def __add__(self, other: Tensor | NestedTensor | SupportsFloat):
        if isinstance(other, Tensor) and self.shape == other.shape:
            other = self.nested_like(other)
        if isinstance(other, NestedTensor):
            return NestedTensor([x + y for x, y in zip(self._storage, other._storage)], **self._state)
        return NestedTensor([value + other for value in self._storage], **self._state)

    def __radd__(self, other: Tensor | NestedTensor | SupportsFloat):
        if isinstance(other, Tensor) and self.shape == other.shape:
            other = self.nested_like(other)
        if isinstance(other, NestedTensor):
            return NestedTensor([y + x for x, y in zip(self._storage, other._storage)], **self._state)
        return NestedTensor([other + value for value in self._storage], **self._state)

    def __iadd__(self, other: Tensor | NestedTensor | SupportsFloat):
        if isinstance(other, Tensor) and self.shape == other.shape:
            other = self.nested_like(other)
        if hasattr(other, "to"):
            other = other.to(self.dtype)
        if isinstance(other, NestedTensor):
            for x, y in zip(self._storage, other._storage):
                x += y
        else:
            for value in self._storage:
                value += other
        return self

    def __sub__(self, other: Tensor | NestedTensor | SupportsFloat):
        if isinstance(other, Tensor) and self.shape == other.shape:
            other = self.nested_like(other)
        if isinstance(other, NestedTensor):
            return NestedTensor([x - y for x, y in zip(self._storage, other._storage)], **self._state)
        return NestedTensor([value - other for value in self._storage], **self._state)

    def __rsub__(self, other: Tensor | NestedTensor | SupportsFloat):
        if isinstance(other, Tensor) and self.shape == other.shape:
            other = self.nested_like(other)
        if isinstance(other, NestedTensor):
            return NestedTensor([y - x for x, y in zip(self._storage, other._storage)], **self._state)
        return NestedTensor([other - value for value in self._storage], **self._state)

    def __isub__(self, other: Tensor | NestedTensor | SupportsFloat):
        if isinstance(other, Tensor) and self.shape == other.shape:
            other = self.nested_like(other)
        if hasattr(other, "to"):
            other = other.to(self.dtype)
        if isinstance(other, NestedTensor):
            for x, y in zip(self._storage, other._storage):
                x -= y
        else:
            for value in self._storage:
                value -= other
        return self

    def __pos__(self):
        return NestedTensor([+x for x in self._storage])

    def __neg__(self):
        return NestedTensor([-x for x in self._storage])

    def __invert__(self):
        return NestedTensor([~x for x in self._storage])

    def __mul__(self, other: Tensor | NestedTensor | SupportsFloat):
        if isinstance(other, Tensor) and self.shape == other.shape:
            other = self.nested_like(other)
        if isinstance(other, NestedTensor):
            return NestedTensor([x * y for x, y in zip(self._storage, other._storage)], **self._state)
        return NestedTensor([value * other for value in self._storage], **self._state)

    def __rmul__(self, other: Tensor | NestedTensor | SupportsFloat):
        if isinstance(other, Tensor) and self.shape == other.shape:
            other = self.nested_like(other)
        if isinstance(other, NestedTensor):
            return NestedTensor([y * x for x, y in zip(self._storage, other._storage)], **self._state)
        return NestedTensor([other * value for value in self._storage], **self._state)

    def __imul__(self, other: Tensor | NestedTensor | SupportsFloat):
        if isinstance(other, Tensor) and self.shape == other.shape:
            other = self.nested_like(other)
        if hasattr(other, "to"):
            other = other.to(self.dtype)
        if isinstance(other, NestedTensor):
            for x, y in zip(self._storage, other._storage):
                x *= y
        else:
            for value in self._storage:
                value *= other
        return self

    def __pow__(self, other: Tensor | NestedTensor | SupportsFloat):
        if isinstance(other, Tensor) and self.shape == other.shape:
            other = self.nested_like(other)
        if isinstance(other, NestedTensor):
            return NestedTensor([x**y for x, y in zip(self._storage, other._storage)], **self._state)
        return NestedTensor([value**other for value in self._storage], **self._state)

    def __rpow__(self, other: Tensor | NestedTensor | SupportsFloat):
        if isinstance(other, Tensor) and self.shape == other.shape:
            other = self.nested_like(other)
        if isinstance(other, NestedTensor):
            return NestedTensor([y**x for x, y in zip(self._storage, other._storage)], **self._state)
        return NestedTensor([other**value for value in self._storage], **self._state)

    def __ipow__(self, other: Tensor | NestedTensor | SupportsFloat):
        if isinstance(other, Tensor) and self.shape == other.shape:
            other = self.nested_like(other)
        if hasattr(other, "to"):
            other = other.to(self.dtype)
        if isinstance(other, NestedTensor):
            for x, y in zip(self._storage, other._storage):
                x **= y
        else:
            for value in self._storage:
                value **= other
        return self

    def __matmul__(self, other: Tensor | NestedTensor | SupportsFloat):
        if isinstance(other, Tensor) and self.shape == other.shape:
            other = self.nested_like(other)
        if isinstance(other, NestedTensor):
            return NestedTensor([x @ y for x, y in zip(self._storage, other._storage)], **self._state)
        return NestedTensor([value @ other for value in self._storage], **self._state)

    def __rmatmul__(self, other: Tensor | NestedTensor | SupportsFloat):
        if isinstance(other, Tensor) and self.shape == other.shape:
            other = self.nested_like(other)
        if isinstance(other, NestedTensor):
            return NestedTensor([y @ x for x, y in zip(self._storage, other._storage)], **self._state)
        return NestedTensor([other @ value for value in self._storage], **self._state)

    def __imatmul__(self, other: Tensor | NestedTensor | SupportsFloat):
        if isinstance(other, Tensor) and self.shape == other.shape:
            other = self.nested_like(other)
        if hasattr(other, "to"):
            other = other.to(self.dtype)
        if isinstance(other, NestedTensor):
            for x, y in zip(self._storage, other._storage):
                x @= y
        else:
            for value in self._storage:
                value @= other
        return self

    def __truediv__(self, other: Tensor | NestedTensor | SupportsFloat):
        if isinstance(other, Tensor) and self.shape == other.shape:
            other = self.nested_like(other)
        if isinstance(other, NestedTensor):
            return NestedTensor([x / y for x, y in zip(self._storage, other._storage)], **self._state)
        return NestedTensor([value / other for value in self._storage], **self._state)

    def __rtruediv__(self, other: Tensor | NestedTensor | SupportsFloat):
        if isinstance(other, Tensor) and self.shape == other.shape:
            other = self.nested_like(other)
        if isinstance(other, NestedTensor):
            return NestedTensor([y / x for x, y in zip(self._storage, other._storage)], **self._state)
        return NestedTensor([other / value for value in self._storage], **self._state)

    def __itruediv__(self, other: Tensor | NestedTensor | SupportsFloat):
        if isinstance(other, Tensor) and self.shape == other.shape:
            other = self.nested_like(other)
        if hasattr(other, "to"):
            other = other.to(self.dtype)
        if isinstance(other, NestedTensor):
            for x, y in zip(self._storage, other._storage):
                x /= y
        else:
            for value in self._storage:
                value /= other
        return self

    def __floordiv__(self, other: Tensor | NestedTensor | SupportsFloat):
        if isinstance(other, Tensor) and self.shape == other.shape:
            other = self.nested_like(other)
        if isinstance(other, NestedTensor):
            return NestedTensor([x // y for x, y in zip(self._storage, other._storage)], **self._state)
        return NestedTensor([value // other for value in self._storage], **self._state)

    def __rfloordiv__(self, other: Tensor | NestedTensor | SupportsFloat):
        if isinstance(other, Tensor) and self.shape == other.shape:
            other = self.nested_like(other)
        if isinstance(other, NestedTensor):
            return NestedTensor([y // x for x, y in zip(self._storage, other._storage)], **self._state)
        return NestedTensor([other // value for value in self._storage], **self._state)

    def __ifloordiv__(self, other: Tensor | NestedTensor | SupportsFloat):
        if isinstance(other, Tensor) and self.shape == other.shape:
            other = self.nested_like(other)
        if hasattr(other, "to"):
            other = other.to(self.dtype)
        if isinstance(other, NestedTensor):
            for x, y in zip(self._storage, other._storage):
                x //= y
        else:
            for value in self._storage:
                value //= other
        return self

    def __mod__(self, other: Tensor | NestedTensor | SupportsFloat):
        if isinstance(other, Tensor) and self.shape == other.shape:
            other = self.nested_like(other)
        if isinstance(other, NestedTensor):
            return NestedTensor([x % y for x, y in zip(self._storage, other._storage)], **self._state)
        return NestedTensor([value % other for value in self._storage], **self._state)

    def __rmod__(self, other: Tensor | NestedTensor | SupportsFloat):
        if isinstance(other, Tensor) and self.shape == other.shape:
            other = self.nested_like(other)
        if isinstance(other, NestedTensor):
            return NestedTensor([y % x for x, y in zip(self._storage, other._storage)], **self._state)
        return NestedTensor([other % value for value in self._storage], **self._state)

    def __imod__(self, other: Tensor | NestedTensor | SupportsFloat):
        if isinstance(other, Tensor) and self.shape == other.shape:
            other = self.nested_like(other)
        if hasattr(other, "to"):
            other = other.to(self.dtype)
        if isinstance(other, NestedTensor):
            for x, y in zip(self._storage, other._storage):
                x %= y
        else:
            for value in self._storage:
                value %= other
        return self

    def __and__(self, other: Tensor | NestedTensor | SupportsFloat):
        if isinstance(other, Tensor) and self.shape == other.shape:
            other = self.nested_like(other)
        if isinstance(other, NestedTensor):
            return NestedTensor([x & y for x, y in zip(self._storage, other._storage)], **self._state)
        return NestedTensor([value & other for value in self._storage], **self._state)

    def __rand__(self, other: Tensor | NestedTensor | SupportsFloat):
        if isinstance(other, Tensor) and self.shape == other.shape:
            other = self.nested_like(other)
        if isinstance(other, NestedTensor):
            return NestedTensor([y & x for x, y in zip(self._storage, other._storage)], **self._state)
        return NestedTensor([other & value for value in self._storage], **self._state)

    def __iand__(self, other: Tensor | NestedTensor | SupportsFloat):
        if isinstance(other, Tensor) and self.shape == other.shape:
            other = self.nested_like(other)
        if hasattr(other, "to"):
            other = other.to(self.dtype)
        if isinstance(other, NestedTensor):
            for x, y in zip(self._storage, other._storage):
                x &= y
        else:
            for value in self._storage:
                value &= other
        return self

    def __or__(self, other: Tensor | NestedTensor | SupportsFloat):
        if isinstance(other, Tensor) and self.shape == other.shape:
            other = self.nested_like(other)
        if isinstance(other, NestedTensor):
            return NestedTensor([x | y for x, y in zip(self._storage, other._storage)], **self._state)
        return NestedTensor([value | other for value in self._storage], **self._state)

    def __ror__(self, other: Tensor | NestedTensor | SupportsFloat):
        if isinstance(other, Tensor) and self.shape == other.shape:
            other = self.nested_like(other)
        if isinstance(other, NestedTensor):
            return NestedTensor([y | x for x, y in zip(self._storage, other._storage)], **self._state)
        return NestedTensor([other | value for value in self._storage], **self._state)

    def __ior__(self, other: Tensor | NestedTensor | SupportsFloat):
        if isinstance(other, Tensor) and self.shape == other.shape:
            other = self.nested_like(other)
        if hasattr(other, "to"):
            other = other.to(self.dtype)
        if isinstance(other, NestedTensor):
            for x, y in zip(self._storage, other._storage):
                x |= y
        else:
            for value in self._storage:
                value |= other
        return self

    def __xor__(self, other: Tensor | NestedTensor | SupportsFloat):
        if isinstance(other, Tensor) and self.shape == other.shape:
            other = self.nested_like(other)
        if isinstance(other, NestedTensor):
            return NestedTensor([x ^ y for x, y in zip(self._storage, other._storage)], **self._state)
        return NestedTensor([value ^ other for value in self._storage], **self._state)

    def __rxor__(self, other: Tensor | NestedTensor | SupportsFloat):
        if isinstance(other, Tensor) and self.shape == other.shape:
            other = self.nested_like(other)
        if isinstance(other, NestedTensor):
            return NestedTensor([y ^ x for x, y in zip(self._storage, other._storage)], **self._state)
        return NestedTensor([other ^ value for value in self._storage], **self._state)

    def __ixor__(self, other: Tensor | NestedTensor | SupportsFloat):
        if isinstance(other, Tensor) and self.shape == other.shape:
            other = self.nested_like(other)
        if hasattr(other, "to"):
            other = other.to(self.dtype)
        if isinstance(other, NestedTensor):
            for x, y in zip(self._storage, other._storage):
                x ^= y
        else:
            for value in self._storage:
                value ^= other
        return self

    def __lshift__(self, other: Tensor | NestedTensor | SupportsFloat):
        if isinstance(other, Tensor) and self.shape == other.shape:
            other = self.nested_like(other)
        if isinstance(other, NestedTensor):
            return NestedTensor([x << y for x, y in zip(self._storage, other._storage)], **self._state)
        return NestedTensor([value << other for value in self._storage], **self._state)

    def __rlshift__(self, other: Tensor | NestedTensor | SupportsFloat):
        if isinstance(other, Tensor) and self.shape == other.shape:
            other = self.nested_like(other)
        if isinstance(other, NestedTensor):
            return NestedTensor([y << x for x, y in zip(self._storage, other._storage)], **self._state)
        return NestedTensor([other << value for value in self._storage], **self._state)

    def __ilshift__(self, other: Tensor | NestedTensor | SupportsFloat):
        if isinstance(other, Tensor) and self.shape == other.shape:
            other = self.nested_like(other)
        if hasattr(other, "to"):
            other = other.to(self.dtype)
        if isinstance(other, NestedTensor):
            for x, y in zip(self._storage, other._storage):
                x <<= y
        else:
            for value in self._storage:
                value <<= other
        return self

    def __rshift__(self, other: Tensor | NestedTensor | SupportsFloat):
        if isinstance(other, Tensor) and self.shape == other.shape:
            other = self.nested_like(other)
        if isinstance(other, NestedTensor):
            return NestedTensor([x >> y for x, y in zip(self._storage, other._storage)], **self._state)
        return NestedTensor([value >> other for value in self._storage], **self._state)

    def __rrshift__(self, other: Tensor | NestedTensor | SupportsFloat):
        if isinstance(other, Tensor) and self.shape == other.shape:
            other = self.nested_like(other)
        if isinstance(other, NestedTensor):
            return NestedTensor([y >> x for x, y in zip(self._storage, other._storage)], **self._state)
        return NestedTensor([other >> value for value in self._storage], **self._state)

    def __irshift__(self, other: Tensor | NestedTensor | SupportsFloat):
        if isinstance(other, Tensor) and self.shape == other.shape:
            other = self.nested_like(other)
        if hasattr(other, "to"):
            other = other.to(self.dtype)
        if isinstance(other, NestedTensor):
            for x, y in zip(self._storage, other._storage):
                x >>= y
        else:
            for value in self._storage:
                value >>= other
        return self

    def all(self, dim: int | None = None, keepdim: bool = False) -> bool | Tensor | NestedTensor:
        r"""
        Tests if all elements in NestedTensor evaluate to True.

        Examples:
            >>> nested_tensor = NestedTensor([torch.ones(2, 4, dtype=torch.bool), torch.ones(3, 5, dtype=torch.bool)])
            >>> nested_tensor.all()
            tensor(True)
            >>> nested_tensor.all(dim=0)
            tensor([True, True])
            >>> nested_tensor.all(dim=0, keepdim=True)
            tensor([[True, True]])
            >>> nested_tensor.all(dim=1)
            NestedTensor([[ True,  True,  True,  True, False],
                    [ True,  True,  True,  True,  True]])
            >>> nested_tensor.all(dim=1, keepdim=True)
            NestedTensor([[[ True,  True,  True,  True, False]],
            <BLANKLINE>
                    [[ True,  True,  True,  True,  True]]])
            >>> nested_tensor.batch_first = False
            >>> nested_tensor.all(dim=1)
            tensor([True, True])
            >>> nested_tensor.all(dim=0)
            NestedTensor([[ True,  True],
                    [ True,  True],
                    [ True,  True],
                    [ True,  True],
                    [False,  True]])
            >>> nested_tensor.all(dim=-2)
            tensor([True, True])
        """

        if dim is None:
            return torch.tensor(all(i.all() for i in self._storage))

        if dim < 0:
            dim += self.dim()

        if (self.batch_first and dim == 0) or (not self.batch_first and dim == 1):
            if keepdim:
                return torch.tensor([i.all() for i in self._storage]).unsqueeze(0 if self.batch_first else 1)
            return torch.tensor([i.all() for i in self._storage])

        if self.batch_first or dim != 0:
            dim -= 1

        ret = [i.all(dim=dim, keepdim=keepdim) for i in self._storage]
        try:
            return torch.stack(ret)
        except (RuntimeError, ValueError):
            return NestedTensor(ret, **self._state)

    def dim(self) -> int:
        r"""
        Number of dimension of the NestedTensor.

        Examples:
            >>> nested_tensor = NestedTensor([torch.tensor([1, 2, 3]), torch.tensor([4, 5])])
            >>> nested_tensor.dim()
            2
        """

        return self._dim(self._storage)

    @method_cache(maxsize=1)
    def _dim(self, storage: Tuple[Tensor, ...]) -> int:  # type: ignore[name-defined]
        return max(t.dim() for t in storage) + 1

    def max(self, dim: int | None = None, keepdim: bool = False) -> Tensor | NestedTensor:
        if dim is None:
            return torch.stack([t.max() for t in self._storage]).max()

        if dim < 0:
            dim += self.dim()

        if (self.batch_first and dim == 0) or (not self.batch_first and dim == 1):
            if keepdim:
                return torch.stack([t.max() for t in self._storage]).unsqueeze(0 if self.batch_first else 1)
            return torch.stack([t.max() for t in self._storage])

        if self.batch_first or dim != 0:
            dim -= 1

        ret = [i.max(dim=dim, keepdim=keepdim) for i in self._storage]
        try:
            return torch.stack(ret)
        except (RuntimeError, ValueError):
            return NestedTensor(ret, **self._state)

    def mean(
        self,
        dim: int | None = None,
        keepdim: bool = False,
        *,
        dtype: torch.dtype | None = None,  # type: ignore[name-defined]
    ) -> Tensor | NestedTensor:
        if dim is None:
            return sum([t.sum(dtype=dtype) for t in self._storage]) / self.numel()

        if dim < 0:
            dim += self.dim()

        if (self.batch_first and dim == 0) or (not self.batch_first and dim == 1):
            if keepdim:
                return torch.stack([t.mean(dtype=dtype) for t in self._storage]).unsqueeze(0 if self.batch_first else 1)
            return torch.stack([t.mean(dtype=dtype) for t in self._storage])

        if self.batch_first or dim != 0:
            dim -= 1

        ret = [i.mean(dim=dim, keepdim=keepdim, dtype=dtype) for i in self._storage]
        try:
            return torch.stack(ret)
        except (RuntimeError, ValueError):
            return NestedTensor(ret, **self._state)

    def min(self, dim: int | None = None, keepdim: bool = False) -> Tensor | NestedTensor:
        if dim is None:
            return torch.stack([t.min() for t in self._storage]).min()

        if dim < 0:
            dim += self.dim()

        if (self.batch_first and dim == 0) or (not self.batch_first and dim == 1):
            if keepdim:
                return torch.stack([t.min() for t in self._storage]).unsqueeze(0 if self.batch_first else 1)
            return torch.stack([t.min() for t in self._storage])

        if self.batch_first or dim != 0:
            dim -= 1

        ret = [i.min(dim=dim, keepdim=keepdim) for i in self._storage]
        try:
            return torch.stack(ret)
        except (RuntimeError, ValueError):
            return NestedTensor(ret, **self._state)

    @property
    def ndim(self) -> int:
        r"""
        Alias for `dim()`.
        """

        return self.dim()

    def numel(self) -> int:
        r"""
        Number of elements in the NestedTensor.

        Examples:
            >>> nested_tensor = NestedTensor([torch.tensor([1, 2, 3]), torch.tensor([4, 5])])
            >>> nested_tensor.numel()
            5
        """

        return sum(t.numel() for t in self._storage)

    def requires_grad_(self, requires_grad: bool = True):
        self.requires_grad = requires_grad
        for t in self._storage:
            t.requires_grad = requires_grad
        return self

    def reshape(self, *shape) -> Tensor:
        r"""
        Returns a torch tensor with a different shape.

        Note:
            since NestedTensor is a collection of tensors, the reshape operation is ambiguous.

            Therefore, it is converted to a tensor and then reshaped.

        Args:
            shape: The desired size of each dimension.

        Examples:
            >>> nested_tensor = NestedTensor([torch.tensor([1, 2, 3]), torch.tensor([4, 5])])
            >>> nested_tensor.reshape(3, 2)
            tensor([[1, 2],
                    [3, 4],
                    [5, 0]])
            >>> nested_tensor.reshape(2, 3)
            tensor([[1, 2, 3],
                    [4, 5, 0]])
        """

        return self.tensor.reshape(*shape)

    @property
    def shape(self) -> torch.Size | int:  # type: ignore[name-defined]
        r"""
        Alias for `size()`.
        """

        return self.size()

    def size(self, dim: int | None = None) -> torch.Size | int:  # type: ignore[name-defined]
        r"""
        Returns the size of the self `NestedTensor`.

        Args:
            dim: If not specified, the returned value is a `torch.Size`, a subclass of `tuple`.
                If specified, returns an `int` holding the size of that dimension.
                Defaults to `None`.

        Examples:
            >>> nested_tensor = NestedTensor([torch.tensor([1, 2, 3]), torch.tensor([4, 5])])
            >>> nested_tensor.size()
            torch.Size([2, 3])
            >>> nested_tensor.size(0)
            2
            >>> nested_tensor[1] = torch.tensor([4, 5, 6, 7])
            >>> nested_tensor.shape
            torch.Size([2, 4])
            >>> nested_tensor.size(1)
            4
        """

        return self._size(self._storage, dim, self.batch_first)

    @method_cache(maxsize=1)
    def _size(
        self,
        storage: Tuple[Tensor, ...],
        dim: int | None = None,
        batch_first: bool = True,
    ) -> torch.Size | int:  # type: ignore[name-defined]
        if dim is not None:
            if dim == 0:
                return len(storage)
            return max(t.size(dim - 1) for t in storage)
        if max(t.dim() for t in storage) == 0:
            return torch.Size((len(storage),))
        ndim = max(t.dim() for t in storage)
        size = [max(t.shape[i] if i < len(t.shape) else 0 for t in storage) for i in range(ndim)]
        size.insert(0 if batch_first else 1, len(storage))
        return torch.Size(size)

    def sum(self, dim: int | None = None, keepdim: bool = False) -> Tensor | NestedTensor:
        if dim is None:
            return torch.stack([t.sum() for t in self._storage]).sum()

        if dim < 0:
            dim += self.dim()

        if (self.batch_first and dim == 0) or (not self.batch_first and dim == 1):
            if keepdim:
                return torch.stack([t.sum() for t in self._storage]).unsqueeze(0 if self.batch_first else 1)
            return torch.stack([t.sum() for t in self._storage])

        if self.batch_first or dim != 0:
            dim -= 1

        ret = [i.sum(dim=dim, keepdim=keepdim) for i in self._storage]
        try:
            return torch.stack(ret)
        except (RuntimeError, ValueError):
            return NestedTensor(ret, **self._state)

    def to(self, *args, **kwargs):
        return NestedTensor(
            tuple(t.to(*args, **kwargs) for t in self._storage),
            **self.__state__(return_dtype=False, return_device=False),
        )

    def tolist(self) -> list:
        r"""
        Convert a NestedTensor to a list of lists of values.

        Examples:
            >>> nested_tensor = NestedTensor([torch.tensor([1, 2, 3]), torch.tensor([4, 5])])
            >>> nested_tensor.tolist()
            [[1, 2, 3], [4, 5]]
        """

        return [t.tolist() for t in self._storage]

    def view(self, *shape) -> Tensor:
        r"""
        Returns a torch tensor with a different shape.

        Note:
            since NestedTensor is a collection of tensors, the view operation is ambiguous.

            Therefore, it is converted to a tensor and then reshaped.

        Args:
            shape: The desired size of each dimension.

        Examples:
            >>> nested_tensor = NestedTensor([torch.tensor([1, 2, 3]), torch.tensor([4, 5])])
            >>> nested_tensor.view(3, 2)
            tensor([[1, 2],
                    [3, 4],
                    [5, 0]])
            >>> nested_tensor.view(2, 3)
            tensor([[1, 2, 3],
                    [4, 5, 0]])
        """

        return self.tensor.view(*shape)

    def where(self, condition: Tensor | NestedTensor, other: Tensor | NestedTensor | SupportsFloat) -> Self:
        r"""
        Return a NestedTensor of elements selected from either self or other, depending on condition.

        Examples:
            >>> nested_tensor = NestedTensor([torch.tensor([1, 2, 3]), torch.tensor([4, 5])])
            >>> nested_tensor.where(nested_tensor > 2, torch.tensor([[6, 5, 4], [3, 2, 1]]))
            NestedTensor([[6, 5, 3],
                    [4, 5, 0]])
            >>> nested_tensor.where(nested_tensor > 2, NestedTensor([[6, 5, 4], [3, 2]]))
            NestedTensor([[6, 5, 3],
                    [4, 5, 0]])
            >>> nested_tensor.where(torch.tensor(True), NestedTensor([[6, 5, 4], [3, 2]]))
            NestedTensor([[1, 2, 3],
                    [4, 5, 0]])
        """

        if isinstance(condition, Tensor) and self.shape == condition.shape:
            condition = self.nested_like(condition)
        if isinstance(other, Tensor) and self.shape == other.shape:
            other = self.nested_like(other)
        if isinstance(condition, NestedTensor) and isinstance(other, NestedTensor):
            return self.__class__(
                [x.where(c, y) for x, c, y in zip(self._storage, condition._storage, other._storage)], **self._state
            )
        if isinstance(condition, NestedTensor):
            return self.__class__([x.where(c, other) for x, c in zip(self._storage, condition._storage)], **self._state)
        if isinstance(other, NestedTensor):
            return self.__class__([x.where(condition, y) for x, y in zip(self._storage, other._storage)], **self._state)
        return self.__class__(x.where(condition, other) for x in self._storage)

tensor_mask property

Python

tensor_mask: Tuple[Tensor, Tensor]

Return a tuple of padded tensor and mask tensor.

Examples:

>>> nested_tensor = NestedTensor([torch.tensor([1, 2, 3]), torch.tensor([4, 5])])
>>> nested_tensor.tensor_mask
(tensor([[1, 2, 3],
        [4, 5, 0]]), tensor([[ True,  True,  True],
        [ True,  True, False]]))

tensor property

Python

tensor: Tensor

Return a single tensor by padding all the tensors.

Examples:

>>> nested_tensor = NestedTensor([torch.tensor([1, 2, 3]), torch.tensor([4, 5])])
>>> nested_tensor.tensor
tensor([[1, 2, 3],
        [4, 5, 0]])

mask property

Python

mask: Tensor

Padding mask of tensor.

Examples:

>>> nested_tensor = NestedTensor([torch.tensor([1, 2, 3]), torch.tensor([4, 5])])
>>> nested_tensor.mask
tensor([[ True,  True,  True],
        [ True,  True, False]])

concat property

Python

concat: Tensor

Concat tensor in padding dim.

This is particularly useful when calculating loss or passing Linear to avoid unnecessary computation.

Examples:

>>> nested_tensor = NestedTensor([torch.randn(9, 8), torch.randn(11, 8)])
>>> nested_tensor.concat.shape
torch.Size([20, 8])
>>> nested_tensor = NestedTensor([torch.randn(9, 9, 8), torch.randn(11, 11, 8)])
>>> nested_tensor.concat.shape
torch.Size([202, 8])
>>> nested_tensor = NestedTensor([torch.randn(9, 9, 8, 6), torch.randn(11, 11, 8, 6)])
>>> nested_tensor.concat.shape
torch.Size([202, 8, 6])
>>> nested_tensor = NestedTensor([torch.randn(9, 9, 8, 7), torch.randn(11, 11, 8, 6)])
>>> nested_tensor.concat.shape
torch.Size([1293, 8])
>>> nested_tensor = NestedTensor([torch.randn(1, 9, 9, 5), torch.randn(1, 11, 11, 5)])

torch property

Python

torch: Tensor

Create a torch.nested.nested_tensor object from self.

Examples:

>>> nested_tensor = NestedTensor([[2, 3, 5], [7, 8]])
>>> nested_tensor.torch
nested_tensor([
  tensor([2, 3, 5]),
  tensor([7, 8])
])

occupancy property

Python

occupancy: float

Occupancy of the NestedTensor.

Examples:

>>> nested_tensor = NestedTensor([torch.tensor([1, 2, 3, 4]), torch.tensor([5, 6])])
>>> nested_tensor.occupancy
0.75

ndim property

Python

ndim: int

Alias for dim().

shape property

Python

shape: Size | int

Alias for size().

unbind

Python

unbind(dim: int = 0) -> Tuple[Tensor, ...]

Unbind the NestedTensor.

Source code in danling/tensor/nested_tensor.py

def unbind(self, dim: int = 0) -> Tuple[Tensor, ...]:
    r"""
    Unbind the NestedTensor.
    """
    if dim != 0:
        raise ValueError(f"NestedTensor can only be unbound along dimension 0, got dimension {dim} instead.")
    return self._storage

from_tensor_mask classmethod

Python

from_tensor_mask(tensor: Tensor, mask: Tensor, **kwargs)

Build a NestedTensor object from a padded Tensor and corresponding mask Tensor.

Parameters:

Name	Type	Description	Default
tensor	Tensor	Padded Tensor.	required
mask	Tensor	Tensor Mask.	required

Examples:

>>> padded_tensor = torch.tensor([[1, 2, 3, 0, 0],
...                                [4, 5, 0, 0, 0],
...                                [6, 7, 8, 9, 0]])
>>> mask_tensor = torch.tensor([[1, 1, 1, 0, 0],
...                             [1, 1, 0, 0, 0],
...                             [1, 1, 1, 1, 0]])
>>> nested_tensor = NestedTensor.from_tensor_mask(padded_tensor, mask_tensor)
>>> nested_tensor
NestedTensor([[1, 2, 3, 0],
        [4, 5, 0, 0],
        [6, 7, 8, 9]])

Source code in danling/tensor/nested_tensor.py

@classmethod
def from_tensor_mask(cls, tensor: Tensor, mask: Tensor, **kwargs):
    r"""
    Build a `NestedTensor` object from a padded `Tensor` and corresponding mask `Tensor`.

    Args:
        tensor: Padded Tensor.
        mask: Tensor Mask.

    Examples:
        >>> padded_tensor = torch.tensor([[1, 2, 3, 0, 0],
        ...                                [4, 5, 0, 0, 0],
        ...                                [6, 7, 8, 9, 0]])
        >>> mask_tensor = torch.tensor([[1, 1, 1, 0, 0],
        ...                             [1, 1, 0, 0, 0],
        ...                             [1, 1, 1, 1, 0]])
        >>> nested_tensor = NestedTensor.from_tensor_mask(padded_tensor, mask_tensor)
        >>> nested_tensor
        NestedTensor([[1, 2, 3, 0],
                [4, 5, 0, 0],
                [6, 7, 8, 9]])
    """

    if mask.ndim == 1:
        return cls(tensor, **kwargs)
    if mask.ndim == 2:
        return cls((t[slice(0, m.sum())] for t, m in zip(tensor, mask)), **kwargs)
    return cls(
        (
            t[[slice(0, (m.sum(dim=dim) > 0).sum().item()) for dim in reversed(range(m.dim()))]]
            for t, m in zip(tensor, mask)
        ),
        **kwargs,
    )

nested_like

Python

nested_like(tensor: Tensor, strict: bool = True) -> Self

Create a new NestedTensor from a Tensor. The newly created NestedTensor will have the same shape as current NestedTensor.

Parameters:

Name	Type	Description	Default
tensor	Tensor	The tensor to be converted to NestedTensor.	required
strict	bool	Check if the shape of tensor is the same as the current NestedTensor.	True

Examples:

>>> nested_tensor = NestedTensor([torch.tensor([1, 2, 3]), torch.tensor([4, 5])])
>>> (nested_tensor == nested_tensor.nested_like(nested_tensor)).all()
tensor(True)
>>> tensor = nested_tensor.tensor
>>> (nested_tensor == nested_tensor.nested_like(tensor)).all()
tensor(True)
>>> f = nested_tensor.nested_like(torch.randn(2, 2))
Traceback (most recent call last):
ValueError: The shape of NestedTensor and input tensor does not match, torch.Size([2, 3]) != torch.Size([2, 2])
>>> p = nested_tensor.nested_like(torch.randn(2, 2), False)
>>> p = nested_tensor.nested_like(torch.randn(3, 3), False)
Traceback (most recent call last):
ValueError: The batch size of NestedTensor and input tensor does not match, 2 != 3

Source code in danling/tensor/nested_tensor.py

def nested_like(self, tensor: Tensor, strict: bool = True) -> Self:
    r"""
    Create a new `NestedTensor` from a `Tensor`.
    The newly created `NestedTensor` will have the same shape as current `NestedTensor`.

    Args:
        tensor: The tensor to be converted to `NestedTensor`.
        strict: Check if the shape of `tensor` is the same as the current `NestedTensor`.

    Examples:
        >>> nested_tensor = NestedTensor([torch.tensor([1, 2, 3]), torch.tensor([4, 5])])
        >>> (nested_tensor == nested_tensor.nested_like(nested_tensor)).all()
        tensor(True)
        >>> tensor = nested_tensor.tensor
        >>> (nested_tensor == nested_tensor.nested_like(tensor)).all()
        tensor(True)
        >>> f = nested_tensor.nested_like(torch.randn(2, 2))
        Traceback (most recent call last):
        ValueError: The shape of NestedTensor and input tensor does not match, torch.Size([2, 3]) != torch.Size([2, 2])
        >>> p = nested_tensor.nested_like(torch.randn(2, 2), False)
        >>> p = nested_tensor.nested_like(torch.randn(3, 3), False)
        Traceback (most recent call last):
        ValueError: The batch size of NestedTensor and input tensor does not match, 2 != 3
    """  # noqa: E501

    if isinstance(tensor, NestedTensor):
        return tensor.clone()

    if strict and self.shape != tensor.shape:
        raise ValueError(
            f"The shape of NestedTensor and input tensor does not match, {self.shape} != {tensor.shape}"
        )
    if self.size(0) != tensor.size(0):
        raise ValueError(
            f"The batch size of NestedTensor and input tensor does not match, {self.size(0)} != {tensor.size(0)}"
        )
    return self.__class__([o[tuple(slice(0, dim) for dim in t.shape)] for t, o in zip(self._storage, tensor)])

__setitem__

Python

__setitem__(index: int | slice | list | tuple, value: Tensor | NestedTensor) -> None

Set values in the NestedTensor at the specified index.

Parameters:

Name	Type	Description	Default
index	int | slice | list | tuple	The index to modify. Can be an integer, slice, list, or tuple.	required
value	Tensor | NestedTensor	The new value to set. Can be a Tensor or NestedTensor.	required

Examples:

>>> nested_tensor = NestedTensor([torch.tensor([1, 2, 3]), torch.tensor([4, 5])])
>>> nested_tensor[0] = torch.tensor([6, 7, 8])
>>> nested_tensor[0]
tensor([6, 7, 8])
>>> nested_tensor[1] = torch.tensor([9, 10, 11, 12])
>>> nested_tensor.shape
torch.Size([2, 4])

Source code in danling/tensor/nested_tensor.py

def __setitem__(self, index: int | slice | list | tuple, value: Tensor | NestedTensor) -> None:
    """
    Set values in the NestedTensor at the specified index.

    Args:
        index: The index to modify. Can be an integer, slice, list, or tuple.
        value: The new value to set. Can be a Tensor or NestedTensor.

    Examples:
        >>> nested_tensor = NestedTensor([torch.tensor([1, 2, 3]), torch.tensor([4, 5])])
        >>> nested_tensor[0] = torch.tensor([6, 7, 8])
        >>> nested_tensor[0]
        tensor([6, 7, 8])
        >>> nested_tensor[1] = torch.tensor([9, 10, 11, 12])
        >>> nested_tensor.shape
        torch.Size([2, 4])
    """
    if isinstance(index, int):
        if isinstance(value, NestedTensor):
            if len(value._storage) != 1:
                raise ValueError(
                    f"When setting with an integer index, value must have a single tensor, but got {len(value)}"
                )
            value = value._storage[0]
        if not isinstance(value, Tensor):
            value = torch.tensor(value, device=self.device)
        # Create a new list of tensors to modify
        storage_list = list(self._storage)
        storage_list[index] = value
        self._storage = tuple(storage_list)
    elif isinstance(index, (slice, list)):
        if isinstance(value, Tensor):
            # Convert tensor to NestedTensor if it's a regular tensor
            if value.dim() > 1 and value.size(0) > 1:
                value = NestedTensor(value.unbind(0))
            else:
                value = NestedTensor([value])

        if isinstance(index, slice):
            start, stop, step = index.indices(len(self._storage))
            indices = range(start, stop, step)
        else:
            indices = index  # type: ignore[assignment]

        if len(indices) != len(value._storage):
            raise ValueError(
                f"Size mismatch: tried to assign {len(value._storage)} values to {len(indices)} indices"
            )

        storage_list = list(self._storage)
        for i, idx in enumerate(indices):
            storage_list[idx] = value._storage[i]
        self._storage = tuple(storage_list)
    elif isinstance(index, tuple):
        if len(index) < 2:
            raise ValueError("Tuple index must have at least two elements")

        first_idx, rest_idx = index[0], index[1:]

        if isinstance(first_idx, int):
            # Handle single tensor modification
            storage_list = list(self._storage)
            tensor = storage_list[first_idx]
            tensor[rest_idx] = value
            storage_list[first_idx] = tensor
            self._storage = tuple(storage_list)
        elif isinstance(first_idx, (slice, list)):
            # Handle multiple tensor modification
            if isinstance(first_idx, slice):
                start, stop, step = first_idx.indices(len(self._storage))
                indices = range(start, stop, step)
            else:
                indices = first_idx  # type: ignore[assignment]

            storage_list = list(self._storage)
            for idx in indices:
                tensor = storage_list[idx]
                tensor[rest_idx] = value
                storage_list[idx] = tensor
            self._storage = tuple(storage_list)
        else:
            raise ValueError(f"Unsupported first index type {type(first_idx)}")
    else:
        raise ValueError(f"Unsupported index type {type(index)}")

all

Python

all(dim: int | None = None, keepdim: bool = False) -> bool | Tensor | NestedTensor

Tests if all elements in NestedTensor evaluate to True.

Examples:

>>> nested_tensor = NestedTensor([torch.ones(2, 4, dtype=torch.bool), torch.ones(3, 5, dtype=torch.bool)])
>>> nested_tensor.all()
tensor(True)
>>> nested_tensor.all(dim=0)
tensor([True, True])
>>> nested_tensor.all(dim=0, keepdim=True)
tensor([[True, True]])
>>> nested_tensor.all(dim=1)
NestedTensor([[ True,  True,  True,  True, False],
        [ True,  True,  True,  True,  True]])
>>> nested_tensor.all(dim=1, keepdim=True)
NestedTensor([[[ True,  True,  True,  True, False]],

        [[ True,  True,  True,  True,  True]]])
>>> nested_tensor.batch_first = False
>>> nested_tensor.all(dim=1)
tensor([True, True])
>>> nested_tensor.all(dim=0)
NestedTensor([[ True,  True],
        [ True,  True],
        [ True,  True],
        [ True,  True],
        [False,  True]])
>>> nested_tensor.all(dim=-2)
tensor([True, True])

Source code in danling/tensor/nested_tensor.py

def all(self, dim: int | None = None, keepdim: bool = False) -> bool | Tensor | NestedTensor:
    r"""
    Tests if all elements in NestedTensor evaluate to True.

    Examples:
        >>> nested_tensor = NestedTensor([torch.ones(2, 4, dtype=torch.bool), torch.ones(3, 5, dtype=torch.bool)])
        >>> nested_tensor.all()
        tensor(True)
        >>> nested_tensor.all(dim=0)
        tensor([True, True])
        >>> nested_tensor.all(dim=0, keepdim=True)
        tensor([[True, True]])
        >>> nested_tensor.all(dim=1)
        NestedTensor([[ True,  True,  True,  True, False],
                [ True,  True,  True,  True,  True]])
        >>> nested_tensor.all(dim=1, keepdim=True)
        NestedTensor([[[ True,  True,  True,  True, False]],
        <BLANKLINE>
                [[ True,  True,  True,  True,  True]]])
        >>> nested_tensor.batch_first = False
        >>> nested_tensor.all(dim=1)
        tensor([True, True])
        >>> nested_tensor.all(dim=0)
        NestedTensor([[ True,  True],
                [ True,  True],
                [ True,  True],
                [ True,  True],
                [False,  True]])
        >>> nested_tensor.all(dim=-2)
        tensor([True, True])
    """

    if dim is None:
        return torch.tensor(all(i.all() for i in self._storage))

    if dim < 0:
        dim += self.dim()

    if (self.batch_first and dim == 0) or (not self.batch_first and dim == 1):
        if keepdim:
            return torch.tensor([i.all() for i in self._storage]).unsqueeze(0 if self.batch_first else 1)
        return torch.tensor([i.all() for i in self._storage])

    if self.batch_first or dim != 0:
        dim -= 1

    ret = [i.all(dim=dim, keepdim=keepdim) for i in self._storage]
    try:
        return torch.stack(ret)
    except (RuntimeError, ValueError):
        return NestedTensor(ret, **self._state)

dim

Python

dim() -> int

Number of dimension of the NestedTensor.

Examples:

>>> nested_tensor = NestedTensor([torch.tensor([1, 2, 3]), torch.tensor([4, 5])])
>>> nested_tensor.dim()
2

Source code in danling/tensor/nested_tensor.py

def dim(self) -> int:
    r"""
    Number of dimension of the NestedTensor.

    Examples:
        >>> nested_tensor = NestedTensor([torch.tensor([1, 2, 3]), torch.tensor([4, 5])])
        >>> nested_tensor.dim()
        2
    """

    return self._dim(self._storage)

numel

Python

numel() -> int

Number of elements in the NestedTensor.

Examples:

>>> nested_tensor = NestedTensor([torch.tensor([1, 2, 3]), torch.tensor([4, 5])])
>>> nested_tensor.numel()
5

Source code in danling/tensor/nested_tensor.py

def numel(self) -> int:
    r"""
    Number of elements in the NestedTensor.

    Examples:
        >>> nested_tensor = NestedTensor([torch.tensor([1, 2, 3]), torch.tensor([4, 5])])
        >>> nested_tensor.numel()
        5
    """

    return sum(t.numel() for t in self._storage)

reshape

Python

reshape(*shape) -> Tensor

Returns a torch tensor with a different shape.

Note

since NestedTensor is a collection of tensors, the reshape operation is ambiguous.

Therefore, it is converted to a tensor and then reshaped.

Parameters:

Name	Type	Description	Default
shape		The desired size of each dimension.	()

Examples:

>>> nested_tensor = NestedTensor([torch.tensor([1, 2, 3]), torch.tensor([4, 5])])
>>> nested_tensor.reshape(3, 2)
tensor([[1, 2],
        [3, 4],
        [5, 0]])
>>> nested_tensor.reshape(2, 3)
tensor([[1, 2, 3],
        [4, 5, 0]])

Source code in danling/tensor/nested_tensor.py

def reshape(self, *shape) -> Tensor:
    r"""
    Returns a torch tensor with a different shape.

    Note:
        since NestedTensor is a collection of tensors, the reshape operation is ambiguous.

        Therefore, it is converted to a tensor and then reshaped.

    Args:
        shape: The desired size of each dimension.

    Examples:
        >>> nested_tensor = NestedTensor([torch.tensor([1, 2, 3]), torch.tensor([4, 5])])
        >>> nested_tensor.reshape(3, 2)
        tensor([[1, 2],
                [3, 4],
                [5, 0]])
        >>> nested_tensor.reshape(2, 3)
        tensor([[1, 2, 3],
                [4, 5, 0]])
    """

    return self.tensor.reshape(*shape)

size

Python

size(dim: int | None = None) -> Size | int

Returns the size of the self NestedTensor.

Parameters:

Name	Type	Description	Default
dim	int | None	If not specified, the returned value is a torch.Size, a subclass of tuple. If specified, returns an int holding the size of that dimension. Defaults to None.	None

Examples:

>>> nested_tensor = NestedTensor([torch.tensor([1, 2, 3]), torch.tensor([4, 5])])
>>> nested_tensor.size()
torch.Size([2, 3])
>>> nested_tensor.size(0)
2
>>> nested_tensor[1] = torch.tensor([4, 5, 6, 7])
>>> nested_tensor.shape
torch.Size([2, 4])
>>> nested_tensor.size(1)
4

Source code in danling/tensor/nested_tensor.py

def size(self, dim: int | None = None) -> torch.Size | int:  # type: ignore[name-defined]
    r"""
    Returns the size of the self `NestedTensor`.

    Args:
        dim: If not specified, the returned value is a `torch.Size`, a subclass of `tuple`.
            If specified, returns an `int` holding the size of that dimension.
            Defaults to `None`.

    Examples:
        >>> nested_tensor = NestedTensor([torch.tensor([1, 2, 3]), torch.tensor([4, 5])])
        >>> nested_tensor.size()
        torch.Size([2, 3])
        >>> nested_tensor.size(0)
        2
        >>> nested_tensor[1] = torch.tensor([4, 5, 6, 7])
        >>> nested_tensor.shape
        torch.Size([2, 4])
        >>> nested_tensor.size(1)
        4
    """

    return self._size(self._storage, dim, self.batch_first)

tolist

Python

tolist() -> list

Convert a NestedTensor to a list of lists of values.

Examples:

>>> nested_tensor = NestedTensor([torch.tensor([1, 2, 3]), torch.tensor([4, 5])])
>>> nested_tensor.tolist()
[[1, 2, 3], [4, 5]]

Source code in danling/tensor/nested_tensor.py

def tolist(self) -> list:
    r"""
    Convert a NestedTensor to a list of lists of values.

    Examples:
        >>> nested_tensor = NestedTensor([torch.tensor([1, 2, 3]), torch.tensor([4, 5])])
        >>> nested_tensor.tolist()
        [[1, 2, 3], [4, 5]]
    """

    return [t.tolist() for t in self._storage]

view

Python

view(*shape) -> Tensor

Returns a torch tensor with a different shape.

Note

since NestedTensor is a collection of tensors, the view operation is ambiguous.

Therefore, it is converted to a tensor and then reshaped.

Parameters:

Name	Type	Description	Default
shape		The desired size of each dimension.	()

Examples:

>>> nested_tensor = NestedTensor([torch.tensor([1, 2, 3]), torch.tensor([4, 5])])
>>> nested_tensor.view(3, 2)
tensor([[1, 2],
        [3, 4],
        [5, 0]])
>>> nested_tensor.view(2, 3)
tensor([[1, 2, 3],
        [4, 5, 0]])

Source code in danling/tensor/nested_tensor.py

def view(self, *shape) -> Tensor:
    r"""
    Returns a torch tensor with a different shape.

    Note:
        since NestedTensor is a collection of tensors, the view operation is ambiguous.

        Therefore, it is converted to a tensor and then reshaped.

    Args:
        shape: The desired size of each dimension.

    Examples:
        >>> nested_tensor = NestedTensor([torch.tensor([1, 2, 3]), torch.tensor([4, 5])])
        >>> nested_tensor.view(3, 2)
        tensor([[1, 2],
                [3, 4],
                [5, 0]])
        >>> nested_tensor.view(2, 3)
        tensor([[1, 2, 3],
                [4, 5, 0]])
    """

    return self.tensor.view(*shape)

where

Python

where(condition: Tensor | NestedTensor, other: Tensor | NestedTensor | SupportsFloat) -> Self

Return a NestedTensor of elements selected from either self or other, depending on condition.

Examples:

>>> nested_tensor = NestedTensor([torch.tensor([1, 2, 3]), torch.tensor([4, 5])])
>>> nested_tensor.where(nested_tensor > 2, torch.tensor([[6, 5, 4], [3, 2, 1]]))
NestedTensor([[6, 5, 3],
        [4, 5, 0]])
>>> nested_tensor.where(nested_tensor > 2, NestedTensor([[6, 5, 4], [3, 2]]))
NestedTensor([[6, 5, 3],
        [4, 5, 0]])
>>> nested_tensor.where(torch.tensor(True), NestedTensor([[6, 5, 4], [3, 2]]))
NestedTensor([[1, 2, 3],
        [4, 5, 0]])

Source code in danling/tensor/nested_tensor.py

def where(self, condition: Tensor | NestedTensor, other: Tensor | NestedTensor | SupportsFloat) -> Self:
    r"""
    Return a NestedTensor of elements selected from either self or other, depending on condition.

    Examples:
        >>> nested_tensor = NestedTensor([torch.tensor([1, 2, 3]), torch.tensor([4, 5])])
        >>> nested_tensor.where(nested_tensor > 2, torch.tensor([[6, 5, 4], [3, 2, 1]]))
        NestedTensor([[6, 5, 3],
                [4, 5, 0]])
        >>> nested_tensor.where(nested_tensor > 2, NestedTensor([[6, 5, 4], [3, 2]]))
        NestedTensor([[6, 5, 3],
                [4, 5, 0]])
        >>> nested_tensor.where(torch.tensor(True), NestedTensor([[6, 5, 4], [3, 2]]))
        NestedTensor([[1, 2, 3],
                [4, 5, 0]])
    """

    if isinstance(condition, Tensor) and self.shape == condition.shape:
        condition = self.nested_like(condition)
    if isinstance(other, Tensor) and self.shape == other.shape:
        other = self.nested_like(other)
    if isinstance(condition, NestedTensor) and isinstance(other, NestedTensor):
        return self.__class__(
            [x.where(c, y) for x, c, y in zip(self._storage, condition._storage, other._storage)], **self._state
        )
    if isinstance(condition, NestedTensor):
        return self.__class__([x.where(c, other) for x, c in zip(self._storage, condition._storage)], **self._state)
    if isinstance(other, NestedTensor):
        return self.__class__([x.where(condition, y) for x, y in zip(self._storage, other._storage)], **self._state)
    return self.__class__(x.where(condition, other) for x in self._storage)

PNTensor

Bases: Tensor

A tensor wrapper that enables automatic collation into NestedTensor with PyTorch DataLoader.

PNTensor (Potential Nested Tensor) seamlessly bridges the gap between individual tensors and batched NestedTensor objects in PyTorch workflows. It’s designed specifically to work with PyTorch’s DataLoader collation mechanism, allowing datasets to return variable-length tensors that will automatically be combined into a NestedTensor when batched.

The class provides three properties that mirror those of NestedTensor: - .tensor: The tensor itself (self) - .mask: A tensor of ones with the same shape as self - .concat: The tensor itself (self)

Examples:

Basic usage with PyTorch DataLoader:

>>> from torch.utils.data import Dataset, DataLoader
>>> from danling.tensor import PNTensor
>>> class VariableLengthDataset(Dataset):
...     def __init__(self, data):
...         self.data = data
...     def __len__(self):
...         return len(self.data)
...     def __getitem__(self, idx):
...         return PNTensor(self.data[idx])
>>> # Create a dataset with variable-length sequences
>>> dataset = VariableLengthDataset([[1, 2, 3], [4, 5], [6, 7, 8, 9]])
>>> dataloader = DataLoader(dataset, batch_size=3)
>>> # The DataLoader automatically produces NestedTensor batches
>>> batch = next(iter(dataloader))
>>> batch
NestedTensor([[1., 2., 3., 0.],
        [4., 5., 0., 0.],
        [6., 7., 8., 9.]])

Using PNTensor directly:

>>> tensor = PNTensor([1, 2, 3])
>>> tensor
PNTensor([1., 2., 3.])
>>> tensor.tensor
PNTensor([1., 2., 3.])
>>> tensor.mask
PNTensor([True, True, True])

Source code in danling/tensor/tensor.py

class PNTensor(Tensor):
    r"""
    A tensor wrapper that enables automatic collation into NestedTensor with PyTorch DataLoader.

    `PNTensor` (Potential Nested Tensor) seamlessly bridges the gap between individual tensors
    and batched `NestedTensor` objects in PyTorch workflows. It's designed specifically to work
    with PyTorch's DataLoader collation mechanism, allowing datasets to return variable-length
    tensors that will automatically be combined into a `NestedTensor` when batched.

    The class provides three properties that mirror those of NestedTensor:
    - `.tensor`: The tensor itself (self)
    - `.mask`: A tensor of ones with the same shape as self
    - `.concat`: The tensor itself (self)

    Attributes:
        Inherits all attributes from torch.Tensor

    Methods:
        Inherits all methods from torch.Tensor

    Examples:
        Basic usage with PyTorch DataLoader:

        >>> from torch.utils.data import Dataset, DataLoader
        >>> from danling.tensor import PNTensor
        >>> class VariableLengthDataset(Dataset):
        ...     def __init__(self, data):
        ...         self.data = data
        ...     def __len__(self):
        ...         return len(self.data)
        ...     def __getitem__(self, idx):
        ...         return PNTensor(self.data[idx])
        >>> # Create a dataset with variable-length sequences
        >>> dataset = VariableLengthDataset([[1, 2, 3], [4, 5], [6, 7, 8, 9]])
        >>> dataloader = DataLoader(dataset, batch_size=3)
        >>> # The DataLoader automatically produces NestedTensor batches
        >>> batch = next(iter(dataloader))
        >>> batch
        NestedTensor([[1., 2., 3., 0.],
                [4., 5., 0., 0.],
                [6., 7., 8., 9.]])

        Using PNTensor directly:

        >>> tensor = PNTensor([1, 2, 3])
        >>> tensor
        PNTensor([1., 2., 3.])
        >>> tensor.tensor
        PNTensor([1., 2., 3.])
        >>> tensor.mask
        PNTensor([True, True, True])
    """

    @property
    def tensor(self) -> Tensor:
        r"""
        Identical to `self`.

        Returns:
            (torch.Tensor):

        Examples:
            >>> tensor = torch.tensor([1, 2, 3])
            >>> pn_tensor = PNTensor([1, 2, 3])
            >>> bool((tensor == pn_tensor).all())
            True
            >>> bool((tensor == pn_tensor.tensor).all())
            True
        """

        return self

    @property
    def mask(self) -> Tensor:
        r"""
        Identical to `torch.ones_like(self)`.

        Returns:
            (torch.Tensor):

        Examples:
            >>> tensor = torch.tensor([1, 2, 3])
            >>> pn_tensor = PNTensor([1, 2, 3])
            >>> bool((pn_tensor.mask == torch.ones_like(pn_tensor)).all().item())
            True
        """

        return torch.ones_like(self, dtype=torch.bool)

    @property
    def contact(self) -> Tensor:
        r"""
        Identical to `self`.

        Returns:
            (torch.Tensor):

        Examples:
            >>> tensor = torch.tensor([1, 2, 3])
            >>> pn_tensor = PNTensor([1, 2, 3])
            >>> bool((tensor == pn_tensor).all())
            True
            >>> bool((tensor == pn_tensor.contact).all())
            True
        """

        return self

    def new_empty(self, *args, **kwargs):
        return PNTensor(super().new_empty(*args, **kwargs))

tensor property

Python

tensor: Tensor

Identical to self.

Returns:

Type	Description
Tensor

Examples:

>>> tensor = torch.tensor([1, 2, 3])
>>> pn_tensor = PNTensor([1, 2, 3])
>>> bool((tensor == pn_tensor).all())
True
>>> bool((tensor == pn_tensor.tensor).all())
True

mask property

Python

mask: Tensor

Identical to torch.ones_like(self).

Returns:

Type	Description
Tensor

Examples:

>>> tensor = torch.tensor([1, 2, 3])
>>> pn_tensor = PNTensor([1, 2, 3])
>>> bool((pn_tensor.mask == torch.ones_like(pn_tensor)).all().item())
True

contact property

Python

contact: Tensor

Identical to self.

Returns:

Type	Description
Tensor

Examples:

>>> tensor = torch.tensor([1, 2, 3])
>>> pn_tensor = PNTensor([1, 2, 3])
>>> bool((tensor == pn_tensor).all())
True
>>> bool((tensor == pn_tensor.contact).all())
True

ensure_dir

Bases: property

Ensure a directory property exists.

Examples:

>>> @ensure_dir
... def dir(self) -> str:
...     return os.path.join("path", "to", "dir")

Source code in danling/utils/descriptors.py

class ensure_dir(property):
    r"""
    Ensure a directory property exists.

    Examples:
        >>> @ensure_dir
        ... def dir(self) -> str:
        ...     return os.path.join("path", "to", "dir")
    """

    def __get__(self, instance, owner=None):
        val = super().__get__(instance, owner)
        makedirs(val, exist_ok=True)
        return val

Metrics

Bases: Metric

A comprehensive metric tracking system that maintains the complete history of predictions and labels.

Metrics is designed for computing evaluations that require access to the entire dataset history, such as AUROC, Pearson correlation, or other metrics that cannot be meaningfully averaged batch-by-batch.

Attributes:

Name	Type	Description
metrics	FlatDict[str, Callable]	A dictionary of metric functions to be computed
preprocess	Callable	Optional preprocessing function to apply to inputs and targets
val	RoundDict[str, float | flist]	Metrics computed on the current batch only
avg	RoundDict[str, float | flist]	Metrics computed on all accumulated data
input	Tensor	The input tensor from the latest batch
target	Tensor	The target tensor from the latest batch
inputs	Tensor	Concatenation of all input tensors seen so far
targets	Tensor	Concatenation of all target tensors seen so far

Parameters:

Name	Type	Description	Default
*args		A single mapping of metrics or callable metric functions	()
device	device | None	Device to store tensors on	None
preprocess	Callable	Function to preprocess inputs before computing metrics	base_preprocess
**metrics	Callable	Named metric functions to compute	{}

Examples:

>>> from danling.metric.functional import auroc, auprc
>>> from danling.metric.preprocess import preprocess_binary
>>> metrics = Metrics(auroc=auroc, auprc=auprc)
>>> metrics
Metrics('auroc', 'auprc')
>>> metrics.update([0.2, 0.3, 0.5, 0.7], [0, 1, 0, 1])
>>> metrics.input  # predicted values of current batch
tensor([0.2000, 0.3000, 0.5000, 0.7000])
>>> metrics.target  # ground truth of current batch
tensor([0, 1, 0, 1])
>>> metrics.inputs  # predicted values of all data
tensor([0.2000, 0.3000, 0.5000, 0.7000])
>>> metrics.targets  # ground truth of all data
tensor([0, 1, 0, 1])
>>> metrics.val  # Metrics of current batch on current device
RoundDict(
  ('auroc'): 0.75
  ('auprc'): 0.8333333730697632
)
>>> metrics.avg  # Metrics of all data on all devices
RoundDict(
  ('auroc'): 0.75
  ('auprc'): 0.8333333730697632
)
>>> metrics.update([0.1, 0.4, 0.6, 0.8], [0, 0, 1, 0])
>>> metrics.input  # predicted values of current batch
tensor([0.1000, 0.4000, 0.6000, 0.8000])
>>> metrics.target  # ground truth of current batch
tensor([0, 0, 1, 0])
>>> metrics.inputs  # predicted values of all data
tensor([0.2000, 0.3000, 0.5000, 0.7000, 0.1000, 0.4000, 0.6000, 0.8000])
>>> metrics.targets  # ground truth of all data
tensor([0, 1, 0, 1, 0, 0, 1, 0])
>>> metrics.val.round(4)  # Metrics of current batch on current device
RoundDict(
  ('auroc'): 0.6667
  ('auprc'): 0.5
)
>>> metrics.avg.round(4)  # Metrics of all data on all devices
RoundDict(
  ('auroc'): 0.6667
  ('auprc'): 0.5556
)
>>> f"{metrics:.4f}"
'auroc: 0.6667 (0.6667)\tauprc: 0.5000 (0.5556)'
>>> metrics = Metrics(auroc=auroc, auprc=auprc, preprocess=preprocess_binary)
>>> metrics.update([[0.1, 0.4, 0.6, 0.8], [0.1, 0.4, 0.6]], [[0, -100, 1, 0], [0, -100, 1]])
>>> metrics.input, metrics.target
(tensor([0.1000, 0.6000, 0.8000, 0.1000, 0.6000]), tensor([0, 1, 0, 0, 1]))

Notes

• Metrics stores the complete prediction and target history, which is memory-intensive but necessary for metrics like AUROC that operate on the entire dataset.
• For metrics that can be meaningfully averaged batch-by-batch (like accuracy), consider using MetricMeter for better memory efficiency.
• All metrics are synchronized across devices in distributed training environments.

See Also

• MetricMeters: Memory-efficient metric tracker that averages multiple metrics batch-by-batch.

Source code in danling/metric/metrics.py

class Metrics(Metric):
    r"""
    A comprehensive metric tracking system that maintains the complete history of predictions and labels.

    Metrics is designed for computing evaluations that require access to the entire dataset history,
    such as AUROC, Pearson correlation, or other metrics that cannot be meaningfully averaged batch-by-batch.

    Attributes:
        metrics: A dictionary of metric functions to be computed
        preprocess: Optional preprocessing function to apply to inputs and targets
        val: Metrics computed on the current batch only
        avg: Metrics computed on all accumulated data
        input: The input tensor from the latest batch
        target: The target tensor from the latest batch
        inputs: Concatenation of all input tensors seen so far
        targets: Concatenation of all target tensors seen so far

    Args:
        *args: A single mapping of metrics or callable metric functions
        device: Device to store tensors on
        preprocess: Function to preprocess inputs before computing metrics
        **metrics: Named metric functions to compute

    Examples:
        >>> from danling.metric.functional import auroc, auprc
        >>> from danling.metric.preprocess import preprocess_binary
        >>> metrics = Metrics(auroc=auroc, auprc=auprc)
        >>> metrics
        Metrics('auroc', 'auprc')
        >>> metrics.update([0.2, 0.3, 0.5, 0.7], [0, 1, 0, 1])
        >>> metrics.input  # predicted values of current batch
        tensor([0.2000, 0.3000, 0.5000, 0.7000])
        >>> metrics.target  # ground truth of current batch
        tensor([0, 1, 0, 1])
        >>> metrics.inputs  # predicted values of all data
        tensor([0.2000, 0.3000, 0.5000, 0.7000])
        >>> metrics.targets  # ground truth of all data
        tensor([0, 1, 0, 1])
        >>> metrics.val  # Metrics of current batch on current device
        RoundDict(
          ('auroc'): 0.75
          ('auprc'): 0.8333333730697632
        )
        >>> metrics.avg  # Metrics of all data on all devices
        RoundDict(
          ('auroc'): 0.75
          ('auprc'): 0.8333333730697632
        )
        >>> metrics.update([0.1, 0.4, 0.6, 0.8], [0, 0, 1, 0])
        >>> metrics.input  # predicted values of current batch
        tensor([0.1000, 0.4000, 0.6000, 0.8000])
        >>> metrics.target  # ground truth of current batch
        tensor([0, 0, 1, 0])
        >>> metrics.inputs  # predicted values of all data
        tensor([0.2000, 0.3000, 0.5000, 0.7000, 0.1000, 0.4000, 0.6000, 0.8000])
        >>> metrics.targets  # ground truth of all data
        tensor([0, 1, 0, 1, 0, 0, 1, 0])
        >>> metrics.val.round(4)  # Metrics of current batch on current device
        RoundDict(
          ('auroc'): 0.6667
          ('auprc'): 0.5
        )
        >>> metrics.avg.round(4)  # Metrics of all data on all devices
        RoundDict(
          ('auroc'): 0.6667
          ('auprc'): 0.5556
        )
        >>> f"{metrics:.4f}"
        'auroc: 0.6667 (0.6667)\tauprc: 0.5000 (0.5556)'
        >>> metrics = Metrics(auroc=auroc, auprc=auprc, preprocess=preprocess_binary)
        >>> metrics.update([[0.1, 0.4, 0.6, 0.8], [0.1, 0.4, 0.6]], [[0, -100, 1, 0], [0, -100, 1]])
        >>> metrics.input, metrics.target
        (tensor([0.1000, 0.6000, 0.8000, 0.1000, 0.6000]), tensor([0, 1, 0, 0, 1]))

    Notes:
        - `Metrics` stores the complete prediction and target history, which is memory-intensive
          but necessary for metrics like AUROC that operate on the entire dataset.
        - For metrics that can be meaningfully averaged batch-by-batch (like accuracy),
          consider using [`MetricMeter`][danling.metric.metric_meter.MetricMeter] for better memory efficiency.
        - All metrics are synchronized across devices in distributed training environments.

    See Also:
        - [`MetricMeters`][danling.metric.metric_meter.MetricMeters]:
            Memory-efficient metric tracker that averages multiple metrics batch-by-batch.
    """

    metrics: FlatDict[str, Callable]
    preprocess: Callable = base_preprocess
    _input: Tensor
    _target: Tensor
    _inputs: Tensor
    _targets: Tensor

    def __init__(
        self,
        *args,
        device: torch.device | None = None,
        preprocess: Callable = base_preprocess,
        **metrics: Callable,
    ):
        super().__init__(device=device)
        self._add_state("_input", torch.empty(0))
        self._add_state("_target", torch.empty(0))
        self._add_state("_inputs", torch.empty(0))
        self._add_state("_targets", torch.empty(0))
        if args:
            from .metric_meter import MetricMeters

            if len(args) == 1 and isinstance(args[0], MetricMeters):
                meters = args[0]
                for name, meter in meters.items():
                    metrics.setdefault(name, meter.metric)
                if preprocess is base_preprocess:
                    preprocess = meters.getattr("preprocess")
            else:
                for metric in args:
                    if not callable(metric):
                        raise ValueError(f"Expected metric to be callable, but got {type(metric)}")
                    metrics.setdefault(metric.__name__, metric)
        self.metrics = FlatDict(**metrics)
        self.preprocess = preprocess

    def update(self, input: Tensor | NestedTensor | Sequence, target: Tensor | NestedTensor | Sequence) -> None:
        input, target = self.preprocess(input, target)
        world_size = get_world_size()
        if world_size > 1:
            input, target = self._sync(input, world_size), self._sync(target, world_size)
        if (
            isinstance(input, (Tensor, NestedTensor))
            and isinstance(target, (Tensor, NestedTensor))
            and input.ndim == target.ndim + 1
        ):
            input = input.squeeze(-1)
        if isinstance(input, (Tensor, NestedTensor)):
            input = input.detach().to(self.device)
        if isinstance(target, (Tensor, NestedTensor)):
            target = target.detach().to(self.device)
        self._input = input
        self._target = target
        if self._inputs.numel() == 0 or self._targets.numel() == 0:
            self._inputs = input
            self._targets = target
        else:
            self._inputs = torch.cat([self._inputs, input])
            self._targets = torch.cat([self._targets, target])

    def value(self) -> RoundDict[str, float | flist]:
        return self.calculate(self.input, self.target)

    def average(self) -> RoundDict[str, float | flist]:
        return self.calculate(self.inputs, self.targets)

    def compute(self) -> RoundDict[str, float | flist]:
        return self.average()

    @property
    def val(self) -> RoundDict[str, float | flist]:
        return self.value()

    @property
    def avg(self) -> RoundDict[str, float | flist]:
        return self.average()

    @torch.inference_mode()
    def calculate(self, input: Tensor, target: Tensor) -> RoundDict[str, flist | float]:
        if (
            isinstance(input, (Tensor, NestedTensor))
            and input.numel() == 0 == target.numel()
            or isinstance(input, (list, dict))
            and len(input) == 0 == len(target)
        ):
            return RoundDict({name: nan for name in self.metrics.keys()})
        ret = RoundDict()
        for name, metric in self.metrics.items():
            score = self._calculate(metric, input, target)
            if isinstance(score, Mapping):
                ret.merge(score)
            else:
                ret[name] = score
        return ret

    @torch.inference_mode()
    def _calculate(self, metric, input: Tensor, target: Tensor) -> flist | float:
        score = metric(input, target)
        if isinstance(score, Tensor):
            return score.item() if score.numel() == 1 else flist(score.tolist())
        return score

    @torch.inference_mode()
    def merge_state(self, metrics: Iterable):
        raise NotImplementedError()

    def _sync(self, tensor: Tensor, world_size: int):
        local_size = torch.tensor([tensor.shape[0]], dtype=torch.int64, device=tensor.device)
        size_list = [torch.zeros_like(local_size) for _ in range(world_size)]
        dist.all_gather(size_list, local_size)
        sizes = torch.cat(size_list)
        max_size = sizes.max()

        padded_tensor = torch.empty((max_size, *tensor.shape[1:]), dtype=tensor.dtype, device=tensor.device)
        padded_tensor[: tensor.shape[0]] = tensor
        gathered_tensors = [torch.empty_like(padded_tensor) for _ in range(world_size)]
        dist.all_gather(gathered_tensors, padded_tensor)
        slices = [gathered_tensors[i][: sizes[i]] for i in range(world_size) if sizes[i] > 0]
        return torch.cat(slices, dim=0)

    @property
    def input(self) -> Tensor:
        return self._input

    @property
    def target(self) -> Tensor:
        return self._target

    @property
    def inputs(self) -> Tensor:
        return self._inputs

    @property
    def targets(self) -> Tensor:
        return self._targets

    def __repr__(self):
        keys = tuple(i for i in self.metrics.keys())
        return f"{self.__class__.__name__}{keys}"

    def __format__(self, format_spec):
        val, avg = self.value(), self.average()
        return "\t".join(
            [f"{key}: {val[key].__format__(format_spec)} ({avg[key].__format__(format_spec)})" for key in self.metrics]
        )

    def reset(self: Self) -> Self:  # pragma: no cover
        r"""
        Reset the metric state variables to their default value.
        The tensors in the default values are also moved to the device of
        the last ``self.to(device)`` call.
        """
        for state_name, default in self._state_name_to_default.items():
            if isinstance(default, Tensor):
                setattr(self, state_name, default.clone().to(self.device))
            elif isinstance(default, list):
                setattr(
                    self,
                    state_name,
                    flist(tensor.clone().to(self.device) for tensor in default),
                )
            elif isinstance(default, dict):
                setattr(
                    self,
                    state_name,
                    DefaultDict(
                        lambda: torch.tensor(0.0, device=self.device),
                        {key: tensor.clone().to(self.device) for key, tensor in default.items()},
                    ),
                )
            elif isinstance(default, (int, float)):
                setattr(self, state_name, default)
            else:
                raise TypeError(
                    f"Invalid type for default value for {state_name}. Received {type(default)},"
                    "but expected ``Tensor``, a list of ``Tensor``,"
                    "a dictionary with ``Tensor``, int, or float."
                )
        return self

reset

Python

reset() -> Self

Reset the metric state variables to their default value. The tensors in the default values are also moved to the device of the last self.to(device) call.

Source code in danling/metric/metrics.py

def reset(self: Self) -> Self:  # pragma: no cover
    r"""
    Reset the metric state variables to their default value.
    The tensors in the default values are also moved to the device of
    the last ``self.to(device)`` call.
    """
    for state_name, default in self._state_name_to_default.items():
        if isinstance(default, Tensor):
            setattr(self, state_name, default.clone().to(self.device))
        elif isinstance(default, list):
            setattr(
                self,
                state_name,
                flist(tensor.clone().to(self.device) for tensor in default),
            )
        elif isinstance(default, dict):
            setattr(
                self,
                state_name,
                DefaultDict(
                    lambda: torch.tensor(0.0, device=self.device),
                    {key: tensor.clone().to(self.device) for key, tensor in default.items()},
                ),
            )
        elif isinstance(default, (int, float)):
            setattr(self, state_name, default)
        else:
            raise TypeError(
                f"Invalid type for default value for {state_name}. Received {type(default)},"
                "but expected ``Tensor``, a list of ``Tensor``,"
                "a dictionary with ``Tensor``, int, or float."
            )
    return self

MultiTaskMetrics

Bases: MultiTaskDict

Examples:

>>> from danling.metric.functional import accuracy
>>> metrics = MultiTaskMetrics()
>>> metrics.dataset1.cls = Metrics(acc=accuracy)
>>> metrics.dataset2 = MetricMeters(acc=accuracy)
>>> metrics
MultiTaskMetrics(<class 'danling.metric.multitask.MultiTaskMetrics'>,
  ('dataset1'): MultiTaskMetrics(<class 'danling.metric.multitask.MultiTaskMetrics'>,
    ('cls'): Metrics('acc',)
  )
  ('dataset2'): MetricMeters('acc',)
)
>>> metrics.update({"dataset1.cls": {"input": [0.2, 0.4, 0.6, 0.7], "target": [0, 1, 0, 1]}, "dataset2": ([0.1, 0.4, 0.6, 0.8], [1, 0, 0, 0])})
>>> f"{metrics:.4f}"
'dataset1.cls: acc: 0.5000 (0.5000)\ndataset2: acc: 0.2500 (0.2500)'
>>> metrics.setattr("return_average", True)
>>> metrics.update({"dataset1.cls": [[0.1, 0.4, 0.6, 0.8], [0, 0, 1, 0]], "dataset2": {"input": [0.2, 0.3, 0.6, 0.7], "target": [0, 0, 0, 1]}})
>>> f"{metrics:.4f}"
'dataset1.cls: acc: 0.7500 (0.6250)\ndataset2: acc: 0.7500 (0.5000)'
>>> metrics.update(dict(loss=""))
Traceback (most recent call last):
ValueError: Metric loss not found in ...

Notes

• MultiTaskMetrics manages nested hierarchies of MetricMeters for multiple tasks/datasets
• Supports hierarchical access using dot notation or dictionary-style access
• All metrics are updated simultaneously with a single update() call
• Provides a structured way to track metrics across different tasks or model components

See Also

• [MultiTaskMetrics][danling.metric.metrics.MultiTaskMetrics]: Metric tracker that stores the complete prediction and target history for multi-task learning.

Source code in danling/metric/multitask.py

class MultiTaskMetrics(MultiTaskDict):
    r"""
    Examples:
        >>> from danling.metric.functional import accuracy
        >>> metrics = MultiTaskMetrics()
        >>> metrics.dataset1.cls = Metrics(acc=accuracy)
        >>> metrics.dataset2 = MetricMeters(acc=accuracy)
        >>> metrics
        MultiTaskMetrics(<class 'danling.metric.multitask.MultiTaskMetrics'>,
          ('dataset1'): MultiTaskMetrics(<class 'danling.metric.multitask.MultiTaskMetrics'>,
            ('cls'): Metrics('acc',)
          )
          ('dataset2'): MetricMeters('acc',)
        )
        >>> metrics.update({"dataset1.cls": {"input": [0.2, 0.4, 0.6, 0.7], "target": [0, 1, 0, 1]}, "dataset2": ([0.1, 0.4, 0.6, 0.8], [1, 0, 0, 0])})
        >>> f"{metrics:.4f}"
        'dataset1.cls: acc: 0.5000 (0.5000)\ndataset2: acc: 0.2500 (0.2500)'
        >>> metrics.setattr("return_average", True)
        >>> metrics.update({"dataset1.cls": [[0.1, 0.4, 0.6, 0.8], [0, 0, 1, 0]], "dataset2": {"input": [0.2, 0.3, 0.6, 0.7], "target": [0, 0, 0, 1]}})
        >>> f"{metrics:.4f}"
        'dataset1.cls: acc: 0.7500 (0.6250)\ndataset2: acc: 0.7500 (0.5000)'
        >>> metrics.update(dict(loss=""))  # doctest: +ELLIPSIS
        Traceback (most recent call last):
        ValueError: Metric loss not found in ...

    Notes:
        - `MultiTaskMetrics` manages nested hierarchies of MetricMeters for multiple tasks/datasets
        - Supports hierarchical access using dot notation or dictionary-style access
        - All metrics are updated simultaneously with a single `update()` call
        - Provides a structured way to track metrics across different tasks or model components

    See Also:
        - [`MultiTaskMetrics`][danling.metric.metrics.MultiTaskMetrics]:
            Metric tracker that stores the complete prediction and target history for multi-task learning.
    """  # noqa: E501

    def __init__(self, *args, **kwargs):
        super().__init__(*args, default_factory=MultiTaskMetrics, **kwargs)

    def update(  # type: ignore[override] # pylint: disable=W0221
        self,
        values: Mapping[str, Mapping[str, Tensor | NestedTensor | Sequence]],
    ) -> None:
        r"""
        Updates the average and current value in all metrics.

        Args:
            values: A dictionary mapping metric names to their values.
        """

        for metric, value in values.items():
            if metric not in self:
                raise ValueError(f"Metric {metric} not found in {self}")
            if isinstance(self[metric], MultiTaskMetrics):
                for name, met in self[metric].items():
                    if name in value:
                        val = value[name]
                        if isinstance(value, Mapping):
                            met.update(**val)
                        elif isinstance(value, Sequence):
                            met.update(*val)
                        else:
                            raise ValueError(f"Expected value to be a Mapping or Sequence, but got {type(value)}")
            elif isinstance(self[metric], (Metrics, MetricMeters, MetricMeter)):
                if isinstance(value, Mapping):
                    self[metric].update(**value)
                elif isinstance(value, Sequence):
                    self[metric].update(*value)
                else:
                    raise ValueError(f"Expected value to be a Mapping or Sequence, but got {type(value)}")
            else:
                raise ValueError(
                    f"Expected {metric} to be an instance of MultiTaskMetrics, Metrics, MetricMeters, or MetricMeter, "
                    f"but got {type(self[metric])}"
                )

    def set(  # pylint: disable=W0237
        self,
        name: str,
        metric: MultiTaskMetrics | MetricMeter | MetricMeters | Callable,  # type: ignore[override]
    ) -> None:
        if not isinstance(metric, (MultiTaskMetrics, Metrics, MetricMeters, MetricMeter)):
            raise ValueError(
                f"Expected {metric} to be an instance of MultiTaskMetrics, Metrics, MetricMeters, or MetricMeter, "
                f"but got {type(self[metric])}"
            )
        super().set(name, metric)

update

Python

update(values: Mapping[str, Mapping[str, Tensor | NestedTensor | Sequence]]) -> None

Updates the average and current value in all metrics.

Parameters:

Name	Type	Description	Default
values	Mapping[str, Mapping[str, Tensor | NestedTensor | Sequence]]	A dictionary mapping metric names to their values.	required

Source code in danling/metric/multitask.py

def update(  # type: ignore[override] # pylint: disable=W0221
    self,
    values: Mapping[str, Mapping[str, Tensor | NestedTensor | Sequence]],
) -> None:
    r"""
    Updates the average and current value in all metrics.

    Args:
        values: A dictionary mapping metric names to their values.
    """

    for metric, value in values.items():
        if metric not in self:
            raise ValueError(f"Metric {metric} not found in {self}")
        if isinstance(self[metric], MultiTaskMetrics):
            for name, met in self[metric].items():
                if name in value:
                    val = value[name]
                    if isinstance(value, Mapping):
                        met.update(**val)
                    elif isinstance(value, Sequence):
                        met.update(*val)
                    else:
                        raise ValueError(f"Expected value to be a Mapping or Sequence, but got {type(value)}")
        elif isinstance(self[metric], (Metrics, MetricMeters, MetricMeter)):
            if isinstance(value, Mapping):
                self[metric].update(**value)
            elif isinstance(value, Sequence):
                self[metric].update(*value)
            else:
                raise ValueError(f"Expected value to be a Mapping or Sequence, but got {type(value)}")
        else:
            raise ValueError(
                f"Expected {metric} to be an instance of MultiTaskMetrics, Metrics, MetricMeters, or MetricMeter, "
                f"but got {type(self[metric])}"
            )

to_device

Python

to_device(data: Any, device: device)

Move data to device.

Source code in danling/data/utils.py

def to_device(data: Any, device: torch.device):
    r"""Move data to device."""
    if isinstance(data, list):
        return [to_device(i, device) for i in data]
    if isinstance(data, tuple):
        return tuple(to_device(i, device) for i in data)
    if isinstance(data, NestedDict):
        return NestedDict({k: to_device(v, device) for k, v in data.all_items()})
    if isinstance(data, dict):
        return FlatDict({k: to_device(v, device) for k, v in data.items()})
    if hasattr(data, "to"):
        return data.to(device)
    return data

catch

Python

catch(error: Exceptions = Exception, exclude: Exceptions | None = None, callback: Callable = print_exc, *callback_args, **callback_kwargs)

Decorator to catch error except for exclude. Detailed traceback will be printed to stderr.

catch is extremely useful for unfatal errors. For example, Runner saves checkpoint regularly, however, this might break running if the space is full. Decorating save method with catch will allow you to catch these errors and continue your running.

Parameters:

Name	Type	Description	Default
error	Exceptions	Exceptions to be caught.	Exception
exclude	Exceptions | None	Exceptions to be excluded.	None
callback	Callable	Callback to be called when an error occurs. The first four arguments to callback are exc, func, args, kwargs. Additional arguments should be passed with *callback_args and **callback_kwargs.	print_exc
callback_args		Arguments to be passed to callback.	()
callback_kwargs		Keyword arguments to be passed to callback.	{}

Examples:

>>> def file_not_found(*args, **kwargs):
...     raise FileNotFoundError
>>> func = file_not_found
>>> func()
Traceback (most recent call last):
FileNotFoundError
>>> func = catch(OSError)(file_not_found)
>>> func()
>>> func = catch(IOError)(file_not_found)
>>> func()
>>> func = catch(ZeroDivisionError)(file_not_found)
>>> func()
Traceback (most recent call last):
FileNotFoundError

Source code in danling/utils/decorators.py

@flexible_decorator
def catch(  # pylint: disable=keyword-arg-before-vararg
    error: Exceptions = Exception,
    exclude: Exceptions | None = None,
    callback: Callable = print_exc,
    *callback_args,
    **callback_kwargs,
):
    r"""
    Decorator to catch `error` except for `exclude`.
    Detailed traceback will be printed to `stderr`.

    `catch` is extremely useful for unfatal errors.
    For example, `Runner` saves checkpoint regularly, however, this might break running if the space is full.
    Decorating `save` method with `catch` will allow you to catch these errors and continue your running.

    Args:
        error: Exceptions to be caught.
        exclude: Exceptions to be excluded.
        callback: Callback to be called when an error occurs.
            The first four arguments to `callback` are `exc`, `func`, `args`, `kwargs`.
            Additional arguments should be passed with `*callback_args` and `**callback_kwargs`.
        callback_args: Arguments to be passed to `callback`.
        callback_kwargs: Keyword arguments to be passed to `callback`.

    Examples:
        >>> def file_not_found(*args, **kwargs):
        ...     raise FileNotFoundError
        >>> func = file_not_found
        >>> func()
        Traceback (most recent call last):
        FileNotFoundError
        >>> func = catch(OSError)(file_not_found)
        >>> func()
        >>> func = catch(IOError)(file_not_found)
        >>> func()
        >>> func = catch(ZeroDivisionError)(file_not_found)
        >>> func()
        Traceback (most recent call last):
        FileNotFoundError
    """

    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):  # pylint: disable=inconsistent-return-statements
            try:
                return func(*args, **kwargs)
            except error as exc:  # pylint: disable=broad-exception-caught
                if exclude is not None and isinstance(exc, exclude):
                    raise exc
                callback(exc, func, args, kwargs, *callback_args, **callback_kwargs)

        return wrapper

    decorator.__doc__ = catch.__doc__

    return decorator

debug

Python

debug(enable: bool = True, error: Exceptions = Exception, exclude: Optional[Exceptions] = None)

Contextmanager to enter debug mode on error except for exclude.

debug is intended to be used to catch the error and enter debug mode. Since it is mainly for development purposed, we include an enable args so that it can be deactivated.

Parameters:

Name	Type	Description	Default
enable	bool	Whether to enable the contextmanager. Defaults to True.	True
error	Exceptions	The error to catch. Defaults to Exception.	Exception
exclude	Optional[Exceptions]	The error to exclude. Defaults to None.	None

Source code in danling/utils/context_managers.py

@contextmanager
def debug(
    enable: bool = True,
    error: Exceptions = Exception,
    exclude: Optional[Exceptions] = None,
):
    """
    Contextmanager to enter debug mode on `error` except for `exclude`.

    `debug` is intended to be used to catch the error and enter debug mode.
    Since it is mainly for development purposed, we include an `enable` args so that it can be deactivated.

    Args:
        enable: Whether to enable the contextmanager.
            Defaults to `True`.
        error: The error to catch.
            Defaults to `Exception`.
        exclude: The error to exclude.
            Defaults to `None`.
    """

    if not enable:
        yield
        return
    try:
        yield
    except error as exc:  # pylint: disable=broad-exception-caught
        if exclude is not None and isinstance(exc, exclude):
            raise exc
        _, m, tb = sys.exc_info()
        print(repr(m), file=sys.stderr)
        pdb.post_mortem(tb)
    finally:
        pass

flexible_decorator

Python

flexible_decorator(maybe_decorator: Optional[Callable] = None)

Meta decorator to allow bracket-less decorator when no arguments are passed.

Examples:

For decorator defined as follows:

>>> @flexible_decorator
... def decorator(*args, **kwargs):
...     def wrapper(func, *args, **kwargs):
...         pass
...     return wrapper

The following two are equivalent:

>>> @decorator
... def func(*args, **kwargs):
...     pass

>>> @decorator()
... def func(*args, **kwargs):
...     pass

Source code in danling/utils/decorators.py

def flexible_decorator(maybe_decorator: Optional[Callable] = None):
    r"""
    Meta decorator to allow bracket-less decorator when no arguments are passed.

    Examples:
        For decorator defined as follows:

        >>> @flexible_decorator
        ... def decorator(*args, **kwargs):
        ...     def wrapper(func, *args, **kwargs):
        ...         pass
        ...     return wrapper

        The following two are equivalent:

        >>> @decorator
        ... def func(*args, **kwargs):
        ...     pass

        >>> @decorator()
        ... def func(*args, **kwargs):
        ...     pass
    """

    def decorator(func: Callable):
        @wraps(func)
        def wrapper(*args, **kwargs):
            if len(args) == 1 and isfunction(args[0]):
                return func(**kwargs)(args[0])
            return func(*args, **kwargs)

        return wrapper

    if maybe_decorator is None:
        return decorator
    return decorator(maybe_decorator)

is_json_serializable

Python

is_json_serializable(obj: Any) -> bool

Check if obj is JSON serializable.

Source code in danling/utils/io.py

def is_json_serializable(obj: Any) -> bool:
    r"""
    Check if `obj` is JSON serializable.
    """
    try:
        json.dumps(obj)
        return True
    except (TypeError, OverflowError):
        return False

load

Python

load(file: PathStr, *args: List[Any], **kwargs: Dict[str, Any]) -> Any

Load any file with supported extensions.

Source code in danling/utils/io.py

def load(file: PathStr, *args: List[Any], **kwargs: Dict[str, Any]) -> Any:
    r"""
    Load any file with supported extensions.
    """
    if not os.path.isfile(file):
        raise ValueError(f"Trying to load {file!r} but it is not a file.")
    extension = os.path.splitext(file)[-1].lower()[1:]
    if extension in PYTORCH:
        if not TORCH_AVAILABLE:
            raise ImportError(f"Trying to load {file!r} but torch is not installed.")
        return torch.load(file, *args, **kwargs)
    if extension in NUMPY:
        if not NUMPY_AVAILABLE:
            raise ImportError(f"Trying to load {file!r} but numpy is not installed.")
        return numpy.load(file, *args, **kwargs)
    if extension in JSON:
        with open(file) as fp:
            return json.load(fp, *args, **kwargs)  # type: ignore[arg-type]
    if extension in YAML:
        with open(file) as fp:
            kwargs.setdefault("Loader", yaml.FullLoader)  # type: ignore[arg-type]
            return yaml.load(fp, *args, **kwargs)  # type: ignore[arg-type]
    if extension in PICKLE:
        with open(file, "rb") as fp:
            return pickle.load(fp, *args, **kwargs)  # type: ignore[arg-type]
    if extension in PANDAS_SUPPORTED:
        return load_pandas(file, *args, **kwargs)
    raise ValueError(f"Tying to load {file!r} with unsupported extension={extension!r}")

load_pandas

Python

load_pandas(file: PathStr, *args: List[Any], **kwargs: Dict[str, Any]) -> Any

Load any pandas data file with supported extensions.

Source code in danling/utils/io.py

def load_pandas(file: PathStr, *args: List[Any], **kwargs: Dict[str, Any]) -> Any:
    r"""
    Load any pandas data file with supported extensions.
    """
    if not PANDAS_AVAILABLE:
        raise ImportError(f"Trying to load {file!r} but pandas is not installed.")
    if not os.path.isfile(file):
        raise ValueError(f"Trying to load {file!r} but it is not a file.")
    extension = os.path.splitext(file)[-1].lower()[1:]
    if extension in PANDAS or extension in PICKLE:
        return pandas.read_pickle(file, *args, **kwargs)
    if extension in PARQUET:
        return pandas.read_parquet(file, *args, **kwargs)
    if extension in H5:
        return pandas.read_hdf(file, *args, **kwargs)
    if extension in CSV:
        return pandas.read_csv(file, *args, **kwargs)
    if extension in JSON:
        return pandas.read_json(file, *args, **kwargs)
    if extension in EXCEL:
        return pandas.read_excel(file, *args, **kwargs)
    if extension in XML:
        return pandas.read_xml(file, *args, **kwargs)
    if extension in SQL:
        return pandas.read_sql(file, *args, **kwargs)
    raise ValueError(f"Tying to load {file!r} with unsupported extension={extension!r}")

method_cache

Python

method_cache(maxsize: int | None = 128, typed: bool = False)

Decorator to cache the result of an instance method.

functools.lru_cache uses a strong reference to the instance, which will make the instance immortal and break the garbage collection.

method_cache uses a weak reference to the instance to resolve this issue.

See Also

Source code in danling/utils/decorators.py

@flexible_decorator
def method_cache(maxsize: int | None = 128, typed: bool = False):
    r"""
    Decorator to cache the result of an instance method.

    `functools.lru_cache` uses a strong reference to the instance,
    which will make the instance immortal and break the garbage collection.

    `method_cache` uses a weak reference to the instance to resolve this issue.

    See Also:
        https://rednafi.github.io/reflections/dont-wrap-instance-methods-with-functoolslru_cache-decorator-in-python.html
    """

    def decorator(func):
        @wraps(func)
        def wrapper(self, *args, **kwargs):
            self_ref = ref(self)

            @wraps(func)
            @lru_cache(maxsize=maxsize, typed=typed)
            def cached_method(*args, **kwargs):
                return func(self_ref(), *args, **kwargs)

            setattr(self, func.__name__, cached_method)
            return cached_method(*args, **kwargs)

        return wrapper

    return decorator

save

Python

save(obj: Any, file: PathStr, *args: List[Any], **kwargs: Dict[str, Any]) -> File

Save any file with supported extensions.

Source code in danling/utils/io.py

def save(obj: Any, file: PathStr, *args: List[Any], **kwargs: Dict[str, Any]) -> File:
    r"""
    Save any file with supported extensions.
    """
    extension = os.path.splitext(file)[-1].lower()[1:]
    if extension in PYTORCH:
        if not TORCH_AVAILABLE:
            raise ImportError(f"Trying to save {obj} to {file!r} but torch is not installed.")
        torch.save(obj, file, *args, **kwargs)
    elif extension in NUMPY:
        if not NUMPY_AVAILABLE:
            raise ImportError(f"Trying to save {obj} to {file!r} but numpy is not installed.")
        numpy.save(file, obj, *args, **kwargs)
    elif extension in PANDAS:
        if not PANDAS_AVAILABLE:
            raise ImportError(f"Trying to save {obj} to {file!r} but pandas is not installed.")
        pandas.to_pickle(obj, file, *args, **kwargs)
    elif extension in PARQUET:
        if isinstance(obj, pandas.DataFrame):
            obj.to_parquet(file, *args, **kwargs)
        elif not PYARROW_AVAILABLE:
            raise ImportError(f"Trying to save {obj} to {file!r} but pyarrow is not installed.")
        else:
            pyarrow.parquet.write_table(obj, file, *args, **kwargs)
    elif extension in CSV:
        if isinstance(obj, pandas.DataFrame):
            obj.to_csv(file, *args, **kwargs)
        else:
            raise NotImplementedError(f"Trying to save {obj} to {file!r} but is not supported")
    elif extension in JSON:
        if isinstance(obj, FlatDict):
            obj.json(file)
        else:
            with open(file, "w") as fp:
                json.dump(obj, fp, *args, **kwargs)  # type: ignore[arg-type]
    elif extension in YAML:
        if isinstance(obj, FlatDict):
            obj.yaml(file)
        else:
            with open(file, "w") as fp:
                yaml.dump(obj, fp, *args, **kwargs)  # type: ignore[arg-type, call-overload]
    elif extension in PICKLE:
        with open(file, "wb") as fp:
            pickle.dump(obj, fp, *args, **kwargs)  # type: ignore[arg-type]
    else:
        raise ValueError(f"Tying to save {obj} to {file!r} with unsupported extension={extension!r}")
    return file