hyrax.pytorch_ignite ==================== .. py:module:: hyrax.pytorch_ignite Attributes ---------- .. autoapisummary:: hyrax.pytorch_ignite.logger Classes ------- .. autoapisummary:: hyrax.pytorch_ignite.SubsetSequentialSampler hyrax.pytorch_ignite.HyraxEvents Functions --------- .. autoapisummary:: hyrax.pytorch_ignite.setup_dataset hyrax.pytorch_ignite.setup_model hyrax.pytorch_ignite.dist_data_loader hyrax.pytorch_ignite.create_splits hyrax.pytorch_ignite.create_splits_from_fractions hyrax.pytorch_ignite._inner_loop hyrax.pytorch_ignite._create_process_func hyrax.pytorch_ignite.create_engine hyrax.pytorch_ignite.extract_model_method hyrax.pytorch_ignite.create_evaluator hyrax.pytorch_ignite.create_validator hyrax.pytorch_ignite.create_tester hyrax.pytorch_ignite.attach_best_checkpoint hyrax.pytorch_ignite.create_trainer hyrax.pytorch_ignite.create_save_batch_callback hyrax.pytorch_ignite.fixup_engine Module Contents --------------- .. py:data:: logger .. py:class:: SubsetSequentialSampler(indices: collections.abc.Sequence[int], generator=None) Bases: :py:obj:`torch.utils.data.Sampler`\ [\ :py:obj:`int`\ ] Samples elements sequentially from a given list of indices, without replacement. :param indices: sequence a sequence of indices .. py:attribute:: indices :type: collections.abc.Sequence[int] .. py:attribute:: generator :value: None .. py:method:: __iter__() -> collections.abc.Iterator[int] .. py:method:: __len__() -> int .. py:function:: setup_dataset(config: dict, *, splits: tuple[str, Ellipsis] | None = None, shuffle: bool = True) -> dict[str, hyrax.datasets.data_provider.DataProvider] This function creates an instance of the requested dataset(s) specified in the runtime configuration for the given splits (data_groups). It will create an instance of a DataProvider, and return that as the dataset. :param config: The runtime configuration :type config: dict :param splits: When provided, only create DataProvider instances for the groups whose names appear in *splits*. Groups present in the data_request but not listed here are silently skipped. When ``None`` (the default) every group in the data_request is loaded — preserving backward compatibility. :type splits: tuple[str, ...] | None, optional :param shuffle: Whether to shuffle indices when computing ``split_fraction``-based partitions via :func:`create_splits_from_fractions`. Defaults to ``True``. Set to ``False`` for inference / test verbs where deterministic ordering is required. :type shuffle: bool, optional :returns: A dictionary mapping data group names to DataProvider instances. :rtype: dict[str, DataProvider] .. py:function:: setup_model(config: dict, dataset: hyrax.datasets.data_provider.DataProvider) -> torch.nn.Module Create a model object based on the configuration. :param config: The runtime configuration :type config: dict :param dataset: The dataset object that will provide data to the model for training or inference. Here it is only used to provide a data sample to the model so that it can resize itself at runtime if necessary. :type dataset: DataProvider :returns: An instance of the model class specified in the configuration :rtype: torch.nn.Module .. py:function:: dist_data_loader(dataset: torch.utils.data.Dataset, config: dict, split: Union[str, list[str], bool] = False) Create Pytorch Ignite distributed data loaders It is recommended that each verb needing dataloaders only call this function once. :param dataset: A Hyrax dataset instance :type dataset: hyrax.datasets.dataset_registry.HyraxDataset :param config: Hyrax runtime configuration :type config: dict :param split: The name(s) of the split we want to use from the data set. If this is false or not passed, then a single data loader is returned that corresponds to the entire dataset. :type split: Union[str, list[str]], Optional :returns: * *Dataloader (or an ignite-wrapped equivalent)* -- This is the distributed dataloader, formed by calling ignite.distributed.auto_dataloader * *For multiple splits, we return a dictionary where the keys are the names of the splits* * *and the value is either a Dataloader as described above or the value None if the split* * *was not configured.* .. py:function:: create_splits(data_set: torch.utils.data.Dataset, config: dict) Returns train, test, and validation indexes constructed to be used with the passed in dataset. The allocation of indexes in the underlying dataset to samplers depends on the data_set section of the config dict. .. deprecated:: This function and the associated configuration style using ``config["data_set"]["train_size"]``, ``config["data_set"]["validate_size"]``, and ``config["data_set"]["test_size"]`` is deprecated and will be removed in a future release. Please migrate to defining separate dataset groups in ``[data_request]`` with ``split_fraction`` for each group. :param data_set: The data set to use :type data_set: Dataset :param config: Configuration that defines dataset splits :type config: dict :param split: Name of the split to use. :type split: str .. py:function:: create_splits_from_fractions(dataset_providers: dict[str, Any], config: dict, *, shuffle: bool = True) -> dict[str, list[int]] Partition a shared set of indices across dataset groups using the ``split_fraction`` defined on each ``DataProvider``. All providers in *dataset_providers* are expected to wrap the **same underlying data source** (same ``data_location``). The full index range ``[0, len)`` of the first provider is shuffled deterministically (when *shuffle* is ``True``) using ``config["data_set"]["seed"]``, then sliced into contiguous, non-overlapping segments proportional to each provider's ``split_fraction``. :param dataset_providers: Mapping of group name (e.g. ``"train"``, ``"validate"``) to a ``DataProvider`` instance whose ``split_fraction`` is set. :type dataset_providers: dict[str, Any] :param config: The Hyrax runtime configuration. Only ``config["data_set"]["seed"]`` is used here. :type config: dict :param shuffle: Whether to shuffle the index array before slicing. Defaults to ``True``. Set to ``False`` for inference / test workloads where deterministic sequential ordering is required. :type shuffle: bool, optional :returns: Mapping of group name → list of indices assigned to that group. :rtype: dict[str, list[int]] :raises RuntimeError: If any provider is missing a ``split_fraction``, if the fractions sum to more than 1.0, or if providers have mismatched lengths. .. py:function:: _inner_loop(func, prepare_inputs, device, config, engine, batch) This wraps a model-specific function (func) to move data to the appropriate device. .. py:function:: _create_process_func(funcname, device, model, config) .. py:function:: create_engine(funcname: str, device: torch.device, model: torch.nn.Module, config: dict) -> ignite.engine.Engine Unified creation of the pytorch engine object for either an evaluator or trainer. This function will automatically unwrap a distributed model to find the necessary function, and construct the necessary functions to transfer data to the device on every batch, so model code can be the same no matter where the model is being run. :param funcname: The function name on the model that we will call in the core of the engine loop, and be called once per batch :type funcname: str :param device: The device the engine will run the model on :type device: torch.device :param model: The Model the engine will be using :type model: torch.nn.Module :param config: The runtime config in use :type config: dict .. py:function:: extract_model_method(model, method_name) Extract a method from a model, which may be wrapped in a DistributedDataParallel or DataParallel object. For instance, method_name could be `train_batch` or `infer_batch`. :param model: The model to extract the method from :type model: nn.Module, DistributedDataParallel, or DataParallel :param method_name: Name of the method to extract :type method_name: str :returns: The method extracted from the model :rtype: Callable .. py:function:: create_evaluator(model: torch.nn.Module, save_function: collections.abc.Callable[[torch.Tensor, torch.Tensor], Any], config: dict) -> ignite.engine.Engine Creates an evaluator engine Primary purpose of this function is to attach the appropriate handlers to an evaluator engine :param model: The model to evaluate :type model: torch.nn.Module :param save_function: A function which will receive Engine.state.output at the end of each iteration. The intent is for the results of evaluation to be saved. :type save_function: Callable[[torch.Tensor], Any] :param config: The runtime config in use :type config: dict :returns: Engine object which when run will evaluate the model. :rtype: pytorch-ignite.Engine .. py:function:: create_validator(model: torch.nn.Module, config: dict, validation_data_loader: torch.utils.data.DataLoader, trainer: ignite.engine.Engine) -> ignite.engine.Engine This function creates a Pytorch Ignite engine object that will be used to validate the model. :param model: The model to train :type model: torch.nn.Module :param config: Hyrax runtime configuration :type config: dict :param validation_data_loader: The data loader for the validation data :type validation_data_loader: DataLoader :param trainer: The engine object that will be used to train the model. We will use specific hooks in the trainer to determine when to run the validation engine. :type trainer: pytorch-ignite.Engine :returns: Engine object that will be used to train the model. :rtype: pytorch-ignite.Engine .. py:function:: create_tester(model: torch.nn.Module, config: dict) -> ignite.engine.Engine This function creates a Pytorch Ignite engine object that will be used to test the model and compute metrics without updating model weights. :param model: The model to test :type model: torch.nn.Module :param config: Hyrax runtime configuration :type config: dict :returns: Engine object that will be used to test the model and compute metrics. :rtype: pytorch-ignite.Engine .. py:function:: attach_best_checkpoint(engine: ignite.engine.Engine, model: torch.nn.Module, trainer: ignite.engine.Engine, results_directory: pathlib.Path) -> None Attach a best-checkpoint handler to ``engine``, scored on ``engine.state.output["loss"]``. Call this function *after* both ``create_trainer`` and (optionally) ``create_validator`` have been called so that handler registration order is correct. When a validator is available, pass it as ``engine`` so that checkpointing is driven by validation loss. When no validator is available, pass the trainer as ``engine`` so that checkpointing falls back to training loss — preserving the previous behaviour. The saved checkpoint format is identical to the one produced by ``create_trainer``, so existing resume logic is fully backward-compatible. :param engine: The engine whose ``output["loss"]`` is used as the checkpoint score. Pass the validator when one exists; otherwise pass the trainer. If the engine has a ``hyrax_label`` attribute, it will be included in the checkpoint filename. :type engine: pytorch-ignite.Engine :param model: The model being trained. Must expose ``model.optimizer`` and optionally ``model.scheduler``. :type model: torch.nn.Module :param trainer: The training engine. Used to derive the global step counter and to attach the end-of-training log handler. :type trainer: pytorch-ignite.Engine :param results_directory: Directory where checkpoint files are written. :type results_directory: Path .. py:function:: create_trainer(model: torch.nn.Module, config: dict, results_directory: pathlib.Path) -> ignite.engine.Engine This function is originally copied from here: https://github.com/pytorch-ignite/examples/blob/main/tutorials/intermediate/cifar10-distributed.py#L164 It was substantially trimmed down to make it easier to understand. :param model: The model to train :type model: torch.nn.Module :param config: Hyrax runtime configuration :type config: dict :param results_directory: The directory where training results will be saved :type results_directory: Path :returns: Engine object that will be used to train the model. :rtype: pytorch-ignite.Engine .. py:function:: create_save_batch_callback(results_dir) Create a callback function for saving batch results during inference or testing. This factory function creates a closure that captures the output directory, then returns a callback that can be used by pytorch_ignite engines to save model outputs batch by batch. :param results_dir: Directory where results should be saved :type results_dir: Path :returns: A callback function with signature (batch, batch_results) that saves results :rtype: callable .. py:class:: HyraxEvents(value: str, event_filter: collections.abc.Callable | None = None, name: str | None = None) Bases: :py:obj:`ignite.engine.EventEnum` Workaround event for a pytorch ignite bug. See fixup_engine for details .. py:attribute:: HYRAX_EPOCH_COMPLETED :value: 'HyraxEpochCompleted' .. py:function:: fixup_engine(engine: ignite.engine.Engine) Workaround for this pytorch ignite bug (https://github.com/pytorch/ignite/issues/3372) where engine.state.output is not available at EPOCH_COMPLETED or later times (COMPLETED, etc) We create a new event HYRAX_EPOCH_COMPLETED which triggers at ITERATION_COMPLETED, but only on the final iteration. This is just before the erronious state reset. This hack relies on pytorch ignite internal state, but can be removed as soon as our fix is mainlined (https://github.com/pytorch/ignite/pull/3373) in version 0.6.0 estimated August 2025