utils.py

# Copyright (c) MONAI Consortium
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#     http://www.apache.org/licenses/LICENSE-2.0
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import annotations

from abc import ABC, abstractmethod
from collections.abc import Callable, Mapping, Sequence
from typing import TYPE_CHECKING, Any, cast

import torch
import torch.nn as nn

from monai.transforms import apply_transform
from monai.utils import IgniteInfo, ensure_tuple, min_version, optional_import
from monai.utils.enums import CommonKeys, GanKeys

if TYPE_CHECKING:
    from ignite.engine import EventEnum
else:
    EventEnum, _ = optional_import(
        "ignite.engine", IgniteInfo.OPT_IMPORT_VERSION, min_version, "EventEnum", as_type="base"
    )

__all__ = [
    "IterationEvents",
    "get_devices_spec",
    "default_prepare_batch",
    "PrepareBatch",
    "PrepareBatchDefault",
    "PrepareBatchExtraInput",
    "DiffusionPrepareBatch",
    "VPredictionPrepareBatch",
    "default_make_latent",
    "engine_apply_transform",
    "default_metric_cmp_fn",
    "GradientAccumulation",
]


class IterationEvents(EventEnum):
    """
    Additional Events engine can register and trigger in the iteration process.
    Refer to the example in ignite: https://pytorch.org/ignite/generated/ignite.engine.events.EventEnum.html.
    These Events can be triggered during training iteration:
    `FORWARD_COMPLETED` is the Event when `network(image, label)` completed.
    `LOSS_COMPLETED` is the Event when `loss(pred, label)` completed.
    `BACKWARD_COMPLETED` is the Event when `loss.backward()` completed.
    `MODEL_COMPLETED` is the Event when all the model related operations completed.
    `INNER_ITERATION_STARTED` is the Event when the iteration has an inner loop and the loop is started.
    `INNER_ITERATION_COMPLETED` is the Event when the iteration has an inner loop and the loop is completed.
    """

    FORWARD_COMPLETED = "forward_completed"
    LOSS_COMPLETED = "loss_completed"
    BACKWARD_COMPLETED = "backward_completed"
    MODEL_COMPLETED = "model_completed"
    INNER_ITERATION_STARTED = "inner_iteration_started"
    INNER_ITERATION_COMPLETED = "inner_iteration_completed"


def get_devices_spec(devices: Sequence[torch.device | str] | None = None) -> list[torch.device]:
    """
    Get a valid specification for one or more devices. If `devices` is None get devices for all CUDA devices available.
    If `devices` is and zero-length structure a single CPU compute device is returned. In any other cases `devices` is
    returned unchanged.

    Args:
        devices: list of devices to request, None for all GPU devices, [] for CPU.

    Raises:
        RuntimeError: When all GPUs are selected (``devices=None``) but no GPUs are available.

    Returns:
        list of torch.device: list of devices.

    """
    if devices is None:
        devices = [torch.device(f"cuda:{d:d}") for d in range(torch.cuda.device_count())]

        if not devices:
            raise RuntimeError("No GPU devices available.")

    elif len(devices) == 0:
        devices = [torch.device("cpu")]

    else:
        devices = list(devices)

    devices = [torch.device(d) if isinstance(d, str) else d for d in devices]
    return devices  # type: ignore


def default_prepare_batch(
    batchdata: dict[str, torch.Tensor] | torch.Tensor | Sequence[torch.Tensor],
    device: str | torch.device | None = None,
    non_blocking: bool = False,
    **kwargs: Any,
) -> tuple[torch.Tensor, torch.Tensor | None] | torch.Tensor:
    """
    Default function to prepare the data for current iteration.

    The input `batchdata` is either a single tensor, a pair of tensors, or a dictionary of data. In the first case the
    return value is the tensor and None, in the second case the return value is the two tensors, and in the dictionary
    case the return value depends on what keys are present. if `CommonKeys.IMAGE` and `CommonKeys.LABEL` are present
    then the tensors they key to are returned, if only `CommonKeys.IMAGE` is present that tensor and None is returned.
    If `CommonKeys.REALS` is present this is returned with None. All returned tensors are moved to the given device
    using the given non-blocking argument before being returned.

    This function implements the expected API for a `prepare_batch` callable in Ignite:
    https://pytorch.org/ignite/v0.4.8/generated/ignite.engine.create_supervised_trainer.html

    Args:
        batchdata: input batch data which is either a single tensor, a pair, or a dictionary
        device: device to move every returned tensor to
        non_blocking: equivalent argument for `Tensor.to`
        kwargs: further arguments for `Tensor.to`

    Returns:
        image, label(optional).
    """
    if not isinstance(batchdata, dict):
        if isinstance(batchdata, torch.Tensor):
            return batchdata.to(device=device, non_blocking=non_blocking, **kwargs), None
        elif len(batchdata) == 2:
            image, label = batchdata
            return (
                image.to(device=device, non_blocking=non_blocking, **kwargs),
                label.to(device=device, non_blocking=non_blocking, **kwargs),
            )

        raise AssertionError("Default prepare_batch expects a single tensor, a tensor pair, or dictionary input data.")

    if isinstance(batchdata.get(CommonKeys.LABEL), torch.Tensor):
        return (
            batchdata[CommonKeys.IMAGE].to(device=device, non_blocking=non_blocking, **kwargs),
            batchdata[CommonKeys.LABEL].to(device=device, non_blocking=non_blocking, **kwargs),
        )

    if GanKeys.REALS in batchdata:
        return batchdata[GanKeys.REALS].to(device=device, non_blocking=non_blocking, **kwargs)

    return batchdata[CommonKeys.IMAGE].to(device=device, non_blocking=non_blocking, **kwargs), None


class PrepareBatch(ABC):
    """
    Interface of customized prepare_batch in the trainer or evaluator workflows.
    It takes the data of current batch, target device and non_blocking flag as input.
    Args `batchdata`, `device`, `non_blocking` refer to the ignite API:
    https://pytorch.org/ignite/v0.4.8/generated/ignite.engine.create_supervised_trainer.html.
    `kwargs` supports other args for `Tensor.to()` API.
    """

    @abstractmethod
    def __call__(
        self,
        batchdata: dict[str, torch.Tensor],
        device: str | torch.device | None = None,
        non_blocking: bool = False,
        **kwargs: Any,
    ) -> Any:
        raise NotImplementedError(f"Subclass {self.__class__.__name__} must implement this method.")


class PrepareBatchDefault(PrepareBatch):
    """
    This wraps `default_prepare_batch` to return `image` and `label` only, so is consistent with its API.
    """

    def __call__(
        self,
        batchdata: dict[str, torch.Tensor] | torch.Tensor | Sequence[torch.Tensor],
        device: str | torch.device | None = None,
        non_blocking: bool = False,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None] | torch.Tensor:
        """
        Args `batchdata`, `device`, `non_blocking` refer to the ignite API:
        https://pytorch.org/ignite/v0.4.8/generated/ignite.engine.create_supervised_trainer.html.
        `kwargs` supports other args for `Tensor.to()` API.

        """
        return default_prepare_batch(batchdata, device, non_blocking, **kwargs)


class PrepareBatchExtraInput(PrepareBatch):
    """
    Customized prepare batch callable for trainers or evaluators which support extra input data for the network.
    Extra items are specified by the `extra_keys` parameter and are extracted from the input dictionary (ie. the batch).
    This uses `default_prepare_batch` but requires dictionary inputs.

    Args:
        extra_keys: If a string or sequence of strings is provided, values from the input dictionary are extracted from
            those keys and passed to the network as extra positional arguments. If a dictionary is provided, every pair
            `(k, v)` in that dictionary will become a new keyword argument assigning to `k` the value in the input
            dictionary keyed to `v`.
    """

    def __init__(self, extra_keys: str | Sequence[str] | dict[str, str]) -> None:
        self.extra_keys = extra_keys

    def __call__(
        self,
        batchdata: dict[str, torch.Tensor],
        device: str | torch.device | None = None,
        non_blocking: bool = False,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor, tuple, dict]:
        """
        Args `batchdata`, `device`, `non_blocking` refer to the ignite API:
        https://pytorch.org/ignite/v0.4.8/generated/ignite.engine.create_supervised_trainer.html.
        `kwargs` supports other args for `Tensor.to()` API.
        """
        image, label = default_prepare_batch(batchdata, device, non_blocking, **kwargs)
        args_ = list()
        kwargs_ = dict()

        def _get_data(key: str) -> torch.Tensor:
            data = batchdata[key]

            if isinstance(data, torch.Tensor):
                data = data.to(device=device, non_blocking=non_blocking, **kwargs)

            return data

        if isinstance(self.extra_keys, (str, list, tuple)):
            for k in ensure_tuple(self.extra_keys):
                args_.append(_get_data(k))
        elif isinstance(self.extra_keys, dict):
            for k, v in self.extra_keys.items():
                kwargs_.update({k: _get_data(v)})

        return cast(torch.Tensor, image), cast(torch.Tensor, label), tuple(args_), kwargs_


class DiffusionPrepareBatch(PrepareBatch):
    """
    This class is used as a callable for the `prepare_batch` parameter of engine classes for diffusion training.

    Assuming a supervised training process, it will generate a noise field using `get_noise` for an input image, and
    return the image and noise field as the image/target pair plus the noise field the kwargs under the key "noise".
    This assumes the inferer being used in conjunction with this class expects a "noise" parameter to be provided.

    If the `condition_name` is provided, this must refer to a key in the input dictionary containing the condition
    field to be passed to the inferer. This will appear in the keyword arguments under the key "condition".

    """

    def __init__(self, num_train_timesteps: int, condition_name: str | None = None) -> None:
        self.condition_name = condition_name
        self.num_train_timesteps = num_train_timesteps

    def get_noise(self, images: torch.Tensor) -> torch.Tensor:
        """Returns the noise tensor for input tensor `images`, override this for different noise distributions."""
        return torch.randn_like(images)

    def get_timesteps(self, images: torch.Tensor) -> torch.Tensor:
        """Get a timestep, by default this is a random integer between 0 and `self.num_train_timesteps`."""
        return torch.randint(0, self.num_train_timesteps, (images.shape[0],), device=images.device).long()

    def get_target(self, images: torch.Tensor, noise: torch.Tensor, timesteps: torch.Tensor) -> torch.Tensor:
        """Return the target for the loss function, this is the `noise` value by default."""
        return noise

    def __call__(
        self,
        batchdata: dict[str, torch.Tensor],
        device: str | torch.device | None = None,
        non_blocking: bool = False,
        **kwargs: Any,
    ) -> tuple[torch.Tensor, torch.Tensor, tuple, dict]:
        images, _ = default_prepare_batch(batchdata, device, non_blocking, **kwargs)
        noise = self.get_noise(images).to(device, non_blocking=non_blocking, **kwargs)
        timesteps = self.get_timesteps(images).to(device, non_blocking=non_blocking, **kwargs)

        target = self.get_target(images, noise, timesteps).to(device, non_blocking=non_blocking, **kwargs)
        infer_kwargs = {"noise": noise, "timesteps": timesteps}

        if self.condition_name is not None and isinstance(batchdata, Mapping):
            infer_kwargs["condition"] = batchdata[self.condition_name].to(device, non_blocking=non_blocking, **kwargs)

        # return input, target, arguments, and keyword arguments where noise is the target and also a keyword value
        return images, target, (), infer_kwargs


class VPredictionPrepareBatch(DiffusionPrepareBatch):
    """
    This class is used as a callable for the `prepare_batch` parameter of engine classes for diffusion training.

    Assuming a supervised training process, it will generate a noise field using `get_noise` for an input image, and
    from this compute the velocity using the provided scheduler. This value is used as the target in place of the
    noise field itself although the noise is field is in the kwargs under the key "noise". This assumes the inferer
    being used in conjunction with this class expects a "noise" parameter to be provided.

    If the `condition_name` is provided, this must refer to a key in the input dictionary containing the condition
    field to be passed to the inferer. This will appear in the keyword arguments under the key "condition".

    """

    def __init__(self, scheduler: nn.Module, num_train_timesteps: int, condition_name: str | None = None) -> None:
        super().__init__(num_train_timesteps=num_train_timesteps, condition_name=condition_name)
        self.scheduler = scheduler

    def get_target(self, images, noise, timesteps):
        return self.scheduler.get_velocity(images, noise, timesteps)  # type: ignore[operator]


def default_make_latent(
    num_latents: int,
    latent_size: int,
    device: str | torch.device | None = None,
    non_blocking: bool = False,
    **kwargs: Any,
) -> torch.Tensor:
    return torch.randn(num_latents, latent_size).to(device=device, non_blocking=non_blocking, **kwargs)


def engine_apply_transform(batch: Any, output: Any, transform: Callable[..., dict]) -> tuple[Any, Any]:
    """
    Apply transform on `batch` and `output`.
    If `batch` and `output` are dictionaries, temporarily combine them for the transform,
    otherwise, apply the transform for `output` data only.

    """
    if isinstance(batch, dict) and isinstance(output, dict):
        data = dict(batch)
        data.update(output)
        transformed_data = apply_transform(transform, data)

        if not isinstance(transformed_data, dict):
            raise AssertionError("With a dict supplied to apply_transform a single dict return is expected.")

        for k, v in transformed_data.items():
            # split the output data of post transforms into `output` and `batch`,
            # `batch` should be read-only, so save the generated key-value into `output`
            if k in output or k not in batch:
                output[k] = v
            else:
                batch[k] = v
    else:
        output = apply_transform(transform, output)

    return batch, output


def default_metric_cmp_fn(current_metric: float, prev_best: float) -> bool:
    """
    The default function to compare metric values between current metric and previous best metric.

    Args:
        current_metric: metric value of current round computation.
        prev_best: the best metric value of previous rounds to compare with.

    """
    return current_metric > prev_best


def _noop(*args: Any, **kwargs: Any) -> None:
    """No-op callable used to suppress optimizer/scaler methods during gradient accumulation."""


class GradientAccumulation:
    """
    Callable class implementing gradient accumulation for use with ``SupervisedTrainer``.

    Gradients are accumulated over ``accumulation_steps`` mini-batches before calling
    ``optimizer.step()``, simulating a larger effective batch size on memory-constrained
    hardware.

    Pass an instance as ``iteration_update`` when constructing ``SupervisedTrainer``::

        trainer = SupervisedTrainer(
            ...,
            iteration_update=GradientAccumulation(accumulation_steps=4),
        )

    All ``IterationEvents`` (``FORWARD_COMPLETED``, ``LOSS_COMPLETED``,
    ``BACKWARD_COMPLETED``, ``MODEL_COMPLETED``) still fire on every mini-batch, so
    existing handlers (checkpoint savers, metric loggers, etc.) are unaffected.

    When ``epoch_length`` is known, the optimizer is flushed at the end of each epoch
    even if ``epoch_length % accumulation_steps != 0``, so no gradients are silently
    discarded. For iterable datasets (``epoch_length=None``) this flush does not apply.

    The loss stored in ``engine.state.output[Keys.LOSS]`` is the **unscaled**
    original loss value, so metrics and loggers report the true loss.  Internally
    the loss is divided by ``accumulation_steps`` for the backward pass only.

    Args:
        accumulation_steps: number of mini-batches to accumulate before updating
            weights. Must be a positive integer. Default: 2.

    Raises:
        ValueError: when ``accumulation_steps`` is not a positive integer.
    """

    def __init__(self, accumulation_steps: int = 2) -> None:
        if not isinstance(accumulation_steps, int) or accumulation_steps < 1:
            raise ValueError(f"`accumulation_steps` must be a positive integer, got {accumulation_steps!r}.")
        self.accumulation_steps = accumulation_steps

    def __repr__(self) -> str:
        return f"GradientAccumulation(accumulation_steps={self.accumulation_steps})"

    def __call__(self, engine: Any, batchdata: dict[str, Any]) -> dict:
        """
        Execute one iteration with gradient accumulation.

        Args:
            engine: the Ignite engine (usually ``SupervisedTrainer``).
            batchdata: batch data for this iteration.

        Returns:
            the output dict from ``engine._iteration()``.
        """
        acc = self.accumulation_steps

        result: dict

        if acc == 1:
            result = engine._iteration(engine, batchdata)
            return result

        # engine.state.iteration is 1-indexed and already incremented before __call__
        epoch_length = engine.state.epoch_length  # None for iterable datasets
        if epoch_length is not None:
            local_iter = (engine.state.iteration - 1) % epoch_length  # 0-indexed within epoch
            should_zero_grad = local_iter % acc == 0
            should_step = (local_iter + 1) % acc == 0 or (local_iter + 1) == epoch_length
        else:
            local_iter = engine.state.iteration - 1  # 0-indexed global
            should_zero_grad = local_iter % acc == 0
            should_step = (local_iter + 1) % acc == 0

        # Save and conditionally suppress zero_grad. Only clear gradients at the start of an accumulation cycle.
        original_zero_grad = engine.optimizer.zero_grad
        if not should_zero_grad:
            engine.optimizer.zero_grad = _noop

        # Save and wrap loss_function to scale by 1/accumulation_steps. This ensures the per-mini-batch
        # gradient contribution is correct: the scaled loss will be backpropagated, and accumulated gradients
        # will average to the same value they would with the full batch.
        original_loss_fn = engine.loss_function
        engine.loss_function = lambda *args, **kwargs: original_loss_fn(*args, **kwargs) / acc

        # Save and conditionally suppress optimizer.step. Only update weights at the end of an accumulation cycle.
        # Also patch GradScaler.step and GradScaler.update when step is suppressed, for mixed-precision training.
        original_step = engine.optimizer.step
        original_scaler_step = None
        original_scaler_update = None
        if not should_step:
            engine.optimizer.step = _noop
            if hasattr(engine, "scaler") and engine.scaler is not None:
                original_scaler_step = engine.scaler.step
                original_scaler_update = engine.scaler.update
                engine.scaler.step = _noop
                engine.scaler.update = _noop

        try:
            result = engine._iteration(engine, batchdata)
        finally:
            engine.optimizer.zero_grad = original_zero_grad
            engine.loss_function = original_loss_fn
            engine.optimizer.step = original_step
            if original_scaler_step is not None:
                engine.scaler.step = original_scaler_step
                engine.scaler.update = original_scaler_update

        # Restore the unscaled loss for logging and metrics.  The backward pass
        # already used the scaled value, so this only affects what handlers see.
        if CommonKeys.LOSS in result:
            result[CommonKeys.LOSS] = result[CommonKeys.LOSS] * acc

        return result