drytorch 0.1.0__py3-none-any.whl

drytorch/__init__.py

@@ -0,0 +1,165 @@
+"""Package initialization and default tracker setup.
+
+It automatically initializes some trackers with sets of settings (modes) that
+work well together. The mode can be set as an environmental variable
+DRYTORCH_INIT_MODE before loading the package or explicitly reset after.
+The DRYTORCH_STRICT environment variable turns reproducibility warnings into
+errors when set to 1, requiring a readable commit.
+
+Available modes:
+    1) standard: log to stderr, preferring tqdm over the built-in logger.
+    2) hydra: log to stdout and accommodate default Hydra settings.
+    3) minimal: reduce output and avoid dumping metadata.
+    4) none: no tracker is added to the default ones.
+
+Attributes:
+    INIT_MODE: the mode the trackers will be initialized with at the start.
+        If DRYTORCH_INIT_MODE is not present, it defaults to "standard".
+"""
+
+import importlib.metadata
+import logging
+import os
+import sys
+import warnings
+
+from typing import Literal, TypeGuard
+
+from drytorch.core import exceptions
+from drytorch.core.exceptions import FailedOptionalImportWarning
+from drytorch.core.experimenting import Experiment
+from drytorch.core.tracking import (
+    Tracker,
+    extend_default_trackers,
+    remove_all_default_trackers,
+)
+from drytorch.lib.evaluations import Evaluation, Test, Validation
+from drytorch.lib.learning import LearningSchema
+from drytorch.lib.loading import DataLoader
+from drytorch.lib.models import Model
+from drytorch.lib.objectives import Loss, Metric
+from drytorch.lib.training import Trainer
+from drytorch.trackers import logging as builtin_logging
+from drytorch.trackers.logging import INFO_LEVELS
+
+
+__version__ = importlib.metadata.version('drytorch')
+
+__all__ = [
+    'DataLoader',
+    'Evaluation',
+    'Experiment',
+    'LearningSchema',
+    'Loss',
+    'Metric',
+    'Model',
+    'Test',
+    'Trainer',
+    'Validation',
+    'extend_default_trackers',
+    'init_trackers',
+    'remove_all_default_trackers',
+]
+
+_InitMode = Literal['standard', 'hydra', 'minimal']
+
+logger = logging.getLogger('drytorch')
+
+
+def init_trackers(mode: _InitMode = 'standard') -> None:
+    """Initialize trackers used by default during the experiment.
+
+    Three initializations are available:
+        1) standard: log to stderr, preferring tqdm over the built-in logger.
+        2) hydra: log to stdout and accommodate default Hydra settings.
+        3) minimal: reduce output and avoid dumping metadata.
+
+    Args:
+        mode: one of the suggested initialization modes.
+
+    Raises:
+        ValueError if mode is not available.
+    """
+    remove_all_default_trackers()
+    tracker_list: list[Tracker] = [builtin_logging.BuiltinLogger()]
+    is_tqdm_installed = _add_tqdm(tracker_list, mode=mode)
+    if mode == 'hydra':
+        # hydra logs to stdout by default
+        builtin_logging.enable_default_handler(sys.stdout)
+        builtin_logging.enable_propagation()
+        _add_yaml(tracker_list)
+        verbosity = builtin_logging.INFO_LEVELS.metrics
+
+    elif mode == 'minimal':
+        if is_tqdm_installed:
+            verbosity = builtin_logging.INFO_LEVELS.training
+        else:
+            verbosity = builtin_logging.INFO_LEVELS.epoch
+            builtin_logging.set_formatter('progress')
+
+    elif mode == 'standard':
+        _add_yaml(tracker_list)
+        if is_tqdm_installed:
+            verbosity = builtin_logging.INFO_LEVELS.epoch
+        else:
+            verbosity = builtin_logging.INFO_LEVELS.metrics
+
+    else:
+        raise ValueError('Mode {mode} not available.')
+
+    extend_default_trackers(tracker_list)
+    builtin_logging.set_verbosity(verbosity)
+    return
+
+
+def _add_tqdm(tracker_list: list[Tracker], mode: _InitMode) -> bool:
+    try:
+        from drytorch.trackers import tqdm
+    except (ImportError, ModuleNotFoundError):
+        warnings.warn(FailedOptionalImportWarning('tqdm'), stacklevel=2)
+        return False
+
+    if mode == 'hydra':
+        # progress bar disappears leaving only log metrics.
+        tqdm_logger = tqdm.TqdmLogger(leave=False)
+    elif mode == 'minimal':
+        # double bar replaces most logs.
+        tqdm_logger = tqdm.TqdmLogger(enable_training_bar=True)
+    else:
+        # console metrics from the progress bar
+        tqdm_logger = tqdm.TqdmLogger()
+
+    tracker_list.append(tqdm_logger)
+    return True
+
+
+def _add_yaml(tracker_list: list[Tracker]) -> bool:
+    try:
+        from drytorch.trackers import yaml
+    except (ImportError, ModuleNotFoundError):
+        warnings.warn(FailedOptionalImportWarning('yaml'), stacklevel=2)
+        return False
+
+    tracker_list.append(yaml.MetadataDumper())
+    return True
+
+
+def _check_mode_is_valid(
+    mode: str,
+) -> TypeGuard[Literal['standard', 'hydra', 'minimal']]:
+    return mode in ('standard', 'hydra', 'minimal')
+
+
+INIT_MODE = os.getenv('DRYTORCH_INIT_MODE', 'standard')
+if _check_mode_is_valid(INIT_MODE):
+    logger.log(INFO_LEVELS.internal, 'Initializing %s mode.', INIT_MODE)
+    init_trackers(INIT_MODE)
+elif INIT_MODE != 'none':
+    raise ValueError(f'DRYTORCH_INIT_MODE: {INIT_MODE} not a valid setting.')
+
+warnings.filterwarnings('ignore', category=exceptions.GitUnavailableWarning)
+_strict = os.getenv('DRYTORCH_STRICT', '0')
+if _strict == '1':
+    warnings.filterwarnings('error', category=exceptions.ReproducibilityWarning)
+elif _strict != '0':
+    raise ValueError(f'DRYTORCH_STRICT: {_strict} not a valid setting.')

drytorch/contrib/__init__.py

@@ -0,0 +1 @@
+"""Extensions and contributions."""

drytorch/contrib/optuna.py

@@ -0,0 +1,265 @@
+"""Support for optuna."""
+
+import collections
+import string
+
+from collections.abc import Callable, Sequence
+from typing import Any, Final, TypeVar
+
+import omegaconf
+import optuna
+
+from typing_extensions import override
+
+from drytorch.core import exceptions
+from drytorch.core import protocols as p
+from drytorch.lib import hooks, monitoring
+
+
+__all__ = [
+    'OptunaError',
+    'TrialCallback',
+    'get_final_value',
+    'suggest_overrides',
+]
+
+_Target = TypeVar('_Target', bound=p.TargetType)
+_Output = TypeVar('_Output', bound=p.OutputType)
+
+
+class OptunaError(exceptions.DryTorchError):
+    """Base class for Optuna errors."""
+
+    _template = 'Optuna: {}'
+
+
+class TrialCallback(hooks.TrainerHook):
+    """Report a metric to an optuna trial and prune the training.
+
+    Attributes:
+        reported: the values reported at each epoch.
+        trial: ``optuna`` trial.
+        watch: what the callback reports.
+        when: decides whether the hook runs when the trainer calls it.
+    """
+
+    reported: Final[dict[int, float]]
+    trial: Final[optuna.Trial]
+    watch: Final[monitoring.Watch]
+    _monitor: monitoring.Monitor | None
+
+    def __init__(
+        self,
+        trial: optuna.Trial,
+        *,
+        watch: monitoring.Watch | None = None,
+    ) -> None:
+        """Initialize.
+
+        Args:
+            trial: ``optuna`` trial.
+            watch: what to report. Default reports the loss of the trainer.
+        """
+        super().__init__()
+        self.reported = {}
+        self.trial = trial
+        self.watch = monitoring.Watch() if watch is None else watch
+        self._monitor = None
+        return
+
+    @property
+    def monitor(self) -> monitoring.Monitor:
+        """The record of the reported metric.
+
+        Raises:
+            ResultNotAvailableError: if no trainer bound the callback yet.
+        """
+        if self._monitor is None:
+            raise exceptions.ResultNotAvailableError()
+
+        return self._monitor
+
+    @override
+    def bind_trainer(self, trainer: p.TrainerProtocol[Any, Any, Any]) -> None:
+        """Create the record of the metric for the trainer."""
+        super().bind_trainer(trainer)
+        self._monitor = self.watch.resolve(trainer)
+        return
+
+    @override
+    def run(self) -> None:
+        """Report the value to the trial and prune the training.
+
+        Raises:
+            optuna.TrialPruned: if the trial should be pruned.
+        """
+        epoch = self.trainer.model.epoch
+        self.monitor.record(epoch)
+        if epoch not in self.monitor.values:
+            return
+
+        value = self.monitor.value
+        self.trial.report(value, epoch)
+        self.reported[epoch] = value
+        if self.trial.should_prune():
+            msg = f'Optuna pruning while monitoring {self.monitor.metric_name}'
+            self.trainer.terminate_training(msg)
+            raise optuna.TrialPruned()
+
+        return
+
+
+def suggest_overrides(
+    tune_cfg: omegaconf.DictConfig,
+    trial: optuna.Trial,
+    *,
+    use_full_name: bool = False,
+) -> list[str]:
+    """Suggest values for a trial from structured configurations.
+
+    This function helps integrate optuna into hydra by specifying trial
+    parameters present in the hydra run configuration.
+
+    The configuration file (loadable with hydra) should follow this structure:
+
+    .. code-block:: yaml
+
+        tune:
+          params:
+            param_name:
+              suggest: "suggest_float"  # or other optuna suggest method
+              settings:
+                low: 0.0
+                high: 1.0
+            list_param:
+              suggest: "suggest_list"
+              settings:
+                min_length: 1
+                max_length: 5
+                suggest: "suggest_float"  # method for sampling list elements
+                settings:
+                  low: 0.0
+                  high: 1.0
+        overrides: []  # additional static overrides
+
+    For 'suggest_list' configurations, the settings must specify:
+    - min_length and max_length: bounds for the size of the list.
+    - nested suggest and settings: used to sample each list element.
+
+    The resulting values can be used with hydra.initialize and hydra.compose.
+    Example usage:
+
+    .. code-block:: python
+
+        import hydra
+
+        with hydra.initialize(version_base=None, config_path='path/to/config'):
+            overrides = suggest_overrides(tune_cfg, trial)
+            dict_cfg = hydra.compose(config_name='config', overrides=overrides)
+
+    Here, "your_hydra_config" is the name of the configuration file that
+    includes the configuration parameters to override.
+
+    Args:
+        tune_cfg: a structure specifying how to sample new parameter values.
+        trial: the optuna trial related to the sampled parameters.
+        use_full_name: use the fully qualified setting name. Default to a
+            human-readable name.
+
+    Returns:
+        A list of strings for hydra configuration overrides.
+
+    Raises:
+        OptunaError: if the suggested configuration is invalid.
+    """
+
+    def _get_short_name(name: str) -> str:
+        *prefix_parts, suffix = name.rsplit('.', maxsplit=2)
+        if suffix.isdigit() and prefix_parts:
+            suffix = f'{prefix_parts[-1]} {suffix}'
+
+        return string.capwords(suffix.replace('_', ' '))
+
+    short_names = {k: _get_short_name(k) for k in tune_cfg.tune.params}
+    name_counts = collections.Counter(short_names.values())
+
+    all_overrides: list[str] = [*tune_cfg.overrides]
+    for setting_name, param_value in tune_cfg.tune.params.items():
+        param_name = short_names[setting_name]
+        if use_full_name or name_counts[param_name] > 1:
+            param_name = setting_name
+
+        if param_value.suggest == 'suggest_list':
+            new_value: list[Any] = []
+            for i in range(
+                trial.suggest_int(
+                    name=f'{param_name} #',
+                    low=param_value.settings.min_length,
+                    high=param_value.settings.max_length,
+                )
+            ):
+                try:
+                    bound_suggest = getattr(trial, param_value.settings.suggest)
+                except AttributeError as ae:
+                    msg = f'Invalid suggest configuration: {ae}.'
+                    raise OptunaError(msg) from ae
+                new_value.append(
+                    bound_suggest(
+                        f'{param_name} {i}', **param_value.settings.settings
+                    )
+                )
+        else:
+            try:
+                bound_suggest = getattr(trial, param_value.suggest)
+            except AttributeError as ae:
+                msg = f'Invalid suggest configuration: {ae}.'
+                raise OptunaError(msg) from ae
+
+            new_value = bound_suggest(param_name, **param_value.settings)
+        all_overrides.append(f'{setting_name}={new_value}')
+
+    return all_overrides
+
+
+def get_final_value(
+    trial: optuna.Trial,
+    *,
+    filter_fn: Callable[[Sequence[float]], float] | None = None,
+) -> float:
+    """Calculates a trial's final value from its intermediate reported values.
+
+    This function aggregates the intermediate values reported during trial
+    optimization using trial.report().
+
+    Important: This function will not work with trials created using study.ask()
+    as these don't populate the intermediate values in the corresponding
+    FrozenTrial.
+
+    Args:
+        trial: the completed ``optuna`` trial to evaluate.
+        filter_fn: function to aggregate the trial's intermediate values.
+            Defaults to min or max depending on the study direction.
+
+    Returns:
+        The aggregated final value for the trial.
+
+    Raises:
+        OptunaError: if the trial has no reported values, or if there's
+            a trial number mismatch.
+    """
+    current_study = trial.study
+    if filter_fn is None:
+        filter_fn = min if current_study.direction.name == 'MINIMIZE' else max
+
+    for frozen_trial in reversed(current_study.trials):
+        if frozen_trial.number == trial.number:
+            break
+    else:
+        raise OptunaError('trial number mismatch.')
+
+    reported_values = list(frozen_trial.intermediate_values.values())
+    if not reported_values:
+        msg = 'trial has no reported values. Did you use study.optimize?'
+        raise OptunaError(msg)
+
+    return filter_fn(reported_values)

drytorch/contrib/torcheval.py

@@ -0,0 +1,88 @@
+"""Support for torcheval syncing recommended for distributed training."""
+
+from __future__ import annotations
+
+from typing import Final, TypeVar
+
+import torch
+
+from torcheval import metrics
+from torcheval.metrics import toolkit
+
+from drytorch.core import exceptions
+from drytorch.core import protocols as p
+
+
+__all__ = [
+    'from_torcheval',
+]
+
+
+_Tensor = torch.Tensor
+
+_ComputeReturn = TypeVar('_ComputeReturn', bound=_Tensor | dict[str, _Tensor])
+
+
+def from_torcheval(
+    torch_eval: metrics.Metric[_ComputeReturn],
+    /,
+    *,
+    higher_is_better: bool | None = None,
+    name: str | None = None,
+) -> p.SyncedObjectiveProtocol[_Tensor, _Tensor]:
+    """Returns a wrapper of a Metric from torcheval with a sync method.
+
+    Args:
+        torch_eval: the torcheval metric to wrap.
+        higher_is_better: whether higher values are better.
+        name: optional name for the metric. Defaults to the class name.
+
+    Raises:
+        exceptions.ComputedMetricsTypeError: If the metric computes to an
+            invalid type.
+    """
+
+    class _TorchEvalWithSync:
+        higher_is_better: Final[bool | None]
+        metric: Final[metrics.Metric[_ComputeReturn]]
+        name: Final[str]
+        _synced_value: _ComputeReturn | None
+
+        def __init__(self, _metric: metrics.Metric[_ComputeReturn]) -> None:
+            self.higher_is_better = higher_is_better
+            self.metric = _metric
+            name_str = name if name is not None else _metric.__class__.__name__
+            self.name = name_str
+            self._synced_value = None
+            return
+
+        def compute(self) -> dict[str, _Tensor]:
+            if self._synced_value is None:
+                val = self.metric.compute()
+            else:
+                val = self._synced_value
+
+            if isinstance(val, dict):
+                return val
+
+            if not isinstance(val, _Tensor):
+                raise exceptions.ComputedMetricsTypeError(type(val))
+
+            return {self.name: val}
+
+        def reset(self) -> None:
+            self.metric.reset()
+            self._synced_value = None
+            return
+
+        def sync(self) -> None:
+            """Use torcheval toolkit to synchronize and compute metrics."""
+            self._synced_value = toolkit.sync_and_compute(self.metric)
+            return
+
+        def update(self, outputs: _Tensor, targets: _Tensor) -> None:
+            self.metric.update(outputs, targets)
+            self._synced_value = None
+            return
+
+    return _TorchEvalWithSync(torch_eval)

drytorch/contrib/torchmetrics.py

@@ -0,0 +1,118 @@
+"""Utilities to ensure compatibility with torchmetrics."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Final
+
+import torch
+
+from drytorch.core import exceptions
+from drytorch.core import protocols as p
+
+
+__all__ = [
+    'from_torchmetrics',
+]
+
+if TYPE_CHECKING:
+    from torchmetrics import metric
+
+_Tensor = torch.Tensor
+
+
+def from_torchmetrics(
+    torch_metric: metric.CompositionalMetric,
+    /,
+) -> p.SyncedLossProtocol[_Tensor, _Tensor]:
+    """Returns a wrapper of a CompositionalMetric for integration.
+
+    Args:
+        torch_metric: the composed torchmetrics metric to wrap.
+
+    Returns:
+        The wrapped metric, usable as a loss.
+
+    Raises:
+        RepeatedMetricsError: if two components of the same class are
+            present in the composition. Subclass the metric so each
+            configuration has its own class name.
+    """
+
+    class _TorchMetricCompositionalMetric:
+        """Wrapper of Compositional Metric reporting internal calculations.
+
+        Note that this CompositionalMetric can be used as loss.
+        """
+
+        higher_is_better: Final[bool | None]
+        metric: Final[metric.CompositionalMetric]
+        name = 'Loss'
+        _leaf_metrics: list[metric.Metric]
+
+        def __init__(self, _metric: metric.CompositionalMetric) -> None:
+            self.higher_is_better = _metric.higher_is_better
+            self.metric = _metric
+            self.metric.sync_on_compute = False
+            self.metric.dist_sync_on_step = False
+            self._leaf_metrics = _get_leaf_metrics(self.metric)
+            _disable_leaf_syncing(self._leaf_metrics)
+            return
+
+        def compute(self) -> dict[str, _Tensor]:
+            """Output a dictionary of metric values for each component."""
+            dict_output = dict[str, _Tensor]({self.name: self.metric.compute()})
+            for m in self._leaf_metrics:
+                if isinstance(value := m.compute(), _Tensor):
+                    dict_output[m.__class__.__name__] = value
+
+            return dict_output
+
+        def forward(self, outputs: _Tensor, targets: _Tensor) -> _Tensor:
+            return self.metric(outputs, targets)
+
+        def reset(self) -> None:
+            self.metric.reset()
+            return
+
+        def sync(self) -> None:
+            for m in self._leaf_metrics:
+                m.sync()
+
+            return
+
+        def update(self, outputs: _Tensor, targets: _Tensor) -> None:
+            self.metric.update(outputs, targets)
+            return
+
+    return _TorchMetricCompositionalMetric(torch_metric)
+
+
+def _get_leaf_metrics(
+    comp_metric: metric.CompositionalMetric,
+) -> list[metric.Metric]:
+    leaves: list[metric.Metric] = []
+    stack: list[metric.Metric | float | int | _Tensor | None] = [comp_metric]
+    names_seen: set[str] = set()
+    while stack:
+        node = stack.pop()
+        if isinstance(node, comp_metric.__class__):
+            stack.extend([node.metric_b, node.metric_a])
+        elif isinstance(node, float | int | _Tensor) or node is None:
+            continue
+        else:
+            leaves.append(node)
+            name = node.__class__.__name__
+            if name in names_seen:
+                raise exceptions.RepeatedMetricsError([name])
+
+            names_seen.add(name)
+
+    return leaves
+
+
+def _disable_leaf_syncing(leaf_metrics: list[metric.Metric]) -> None:
+    for m in leaf_metrics:
+        m.sync_on_compute = False
+        m.dist_sync_on_step = False
+
+    return

drytorch/core/__init__.py

@@ -0,0 +1 @@
+"""Core modules setting up and defining the framework."""