diff --git a/ignite/contrib/engines/common.py b/ignite/contrib/engines/common.py index 94ea715f992a..95bd6e29b457 100644 --- a/ignite/contrib/engines/common.py +++ b/ignite/contrib/engines/common.py @@ -47,6 +47,8 @@ from ignite.metrics.metric import RunningBatchWise from ignite.utils import deprecated +import gc + def setup_common_training_handlers( trainer: Engine, @@ -83,7 +85,7 @@ def setup_common_training_handlers( save_every_iters: saving interval. By default, `to_save` objects are stored each 1000 iterations. output_path: output path to indicate where `to_save` objects are stored. Argument is mutually - exclusive with ``save_handler``. + exclusive with `save_handler`. lr_scheduler: learning rate scheduler as native torch LRScheduler or ignite's parameter scheduler. with_gpu_stats: if True, :class:`~ignite.metrics.GpuInfo` is attached to the @@ -100,8 +102,8 @@ def setup_common_training_handlers( clear_cuda_cache: if True, `torch.cuda.empty_cache()` is called every end of epoch. Default, True. save_handler: Method or callable - class to use to store ``to_save``. See :class:`~ignite.handlers.checkpoint.Checkpoint` for more details. - Argument is mutually exclusive with ``output_path``. + class to use to store `to_save`. See :class:`~ignite.handlers.checkpoint.Checkpoint` for more details. + Argument is mutually exclusive with `output_path`. kwargs: optional keyword args to be passed to construct :class:`~ignite.handlers.checkpoint.Checkpoint`. """ @@ -184,6 +186,8 @@ def _setup_common_training_handlers( if torch.cuda.is_available() and clear_cuda_cache: trainer.add_event_handler(Events.EPOCH_COMPLETED, empty_cuda_cache) + if hasattr(torch, "mps") and hasattr(torch.mps, "empty_cache") and clear_cuda_cache: + trainer.add_event_handler(Events.EPOCH_COMPLETED, empty_mps_cache) if to_save is not None: if output_path is None and save_handler is None: @@ -276,10 +280,17 @@ def _setup_common_distrib_training_handlers( @trainer.on(Events.EPOCH_STARTED) def distrib_set_epoch(engine: Engine) -> None: - # pyrefly: ignore [missing-attribute] + # pyrely: ignore [missing-attribute] train_sampler.set_epoch(engine.state.epoch - 1) +def empty_mps_cache(_: Engine) -> None: + if hasattr(torch, "mps") and hasattr(torch.mps, "empty_cache"): + torch.mps.empty_cache() + import gc + gc.collect() + + def empty_cuda_cache(_: Engine) -> None: torch.cuda.empty_cache() import gc @@ -318,21 +329,21 @@ def gen_save_best_models_by_val_score( score_sign: float = 1.0, **kwargs: Any, ) -> Checkpoint: - """Method adds a handler to ``evaluator`` to save ``n_saved`` of best models based on the metric - (named by ``metric_name``) provided by ``evaluator`` (i.e. ``evaluator.state.metrics[metric_name]``). + """Method adds a handler to `evaluator` to save `n_saved` of best models based on the metric + (named by `metric_name`) provided by `evaluator` (i.e. `evaluator.state.metrics[metric_name]`). Models with highest metric value will be retained. The logic of how to store objects is delegated to - ``save_handler``. + `save_handler`. Args: save_handler: Method or callable class to use to save engine and other provided objects. Function receives two objects: checkpoint as a dictionary - and filename. If ``save_handler`` is callable class, it can - inherit of :class:`~ignite.handlers.checkpoint.BaseSaveHandler` and optionally implement ``remove`` method + and filename. If `save_handler` is callable class, it can + inherit of :class:`~ignite.handlers.checkpoint.BaseSaveHandler` and optionally implement `remove` method to keep a fixed number of saved checkpoints. In case if user needs to save engine's checkpoint on a disk, - ``save_handler`` can be defined with :class:`~ignite.handlers.DiskSaver`. + `save_handler` can be defined with :class:`~ignite.handlers.DiskSaver`. evaluator: evaluation engine used to provide the score models: model or dictionary with the object to save. Objects should have - implemented ``state_dict`` and ``load_state_dict`` methods. + implemented `state_dict` and `load_state_dict` methods. metric_name: metric name to use for score evaluation. This metric should be present in `evaluator.state.metrics`. n_saved: number of best models to store @@ -380,8 +391,8 @@ def save_best_model_by_val_score( score_sign: float = 1.0, **kwargs: Any, ) -> Checkpoint: - """Method adds a handler to ``evaluator`` to save on a disk ``n_saved`` of best models based on the metric - (named by ``metric_name``) provided by ``evaluator`` (i.e. ``evaluator.state.metrics[metric_name]``). + """Method adds a handler to `evaluator` to save on a disk `n_saved` of best models based on the metric + (named by `metric_name`) provided by `evaluator` (i.e. `evaluator.state.metrics[metric_name]`). Models with highest metric value will be retained. Args: diff --git a/ignite/distributed/auto.py b/ignite/distributed/auto.py index 501e57fc762a..f8f00a319c1b 100644 --- a/ignite/distributed/auto.py +++ b/ignite/distributed/auto.py @@ -31,7 +31,7 @@ def auto_dataloader(dataset: Dataset, **kwargs: Any) -> DataLoader | _MpDeviceLo - if no sampler provided by user, a `torch DistributedSampler`_ is setup. - if a `torch DistributedSampler`_ is provided by user, it is used without wrapping it. - if another sampler is provided, it is wrapped by :class:`~ignite.distributed.auto.DistributedProxySampler`. - - if the default device is 'cuda', `pin_memory` is automatically set to `True`. + - if the default device is 'cuda' or 'mps', ``pin_memory`` is automatically set to ``True``. .. warning:: @@ -50,7 +50,7 @@ def auto_dataloader(dataset: Dataset, **kwargs: Any) -> DataLoader | _MpDeviceLo Examples: .. code-block:: python - import ignite.distribted as idist + import ignite.distributed as idist train_loader = idist.auto_dataloader( train_dataset, @@ -64,8 +64,7 @@ def auto_dataloader(dataset: Dataset, **kwargs: Any) -> DataLoader | _MpDeviceLo .. _torch DataLoader: https://pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader .. _XLA MpDeviceLoader: https://pytorch.org/xla/release/2.0/index.html#running-on-multiple-xla-devices-with-multi-processing - .. _torch DistributedSampler: - https://pytorch.org/docs/stable/data.html#torch.utils.data.distributed.DistributedSampler + .. _torch DistributedSampler: https://pytorch.org/docs/stable/data.html#torch.utils.data.distributed.DistributedSampler .. _torch IterableDataset: https://pytorch.org/docs/stable/data.html#torch.utils.data.IterableDataset """ rank = idist.get_rank() @@ -76,9 +75,9 @@ def auto_dataloader(dataset: Dataset, **kwargs: Any) -> DataLoader | _MpDeviceLo if "batch_size" in kwargs and kwargs["batch_size"] >= world_size: kwargs["batch_size"] //= world_size - nproc = idist.get_nproc_per_node() - if "num_workers" in kwargs and kwargs["num_workers"] >= nproc: - kwargs["num_workers"] = (kwargs["num_workers"] + nproc - 1) // nproc + nprocs = idist.get_nprocs_per_node() + if "num_workers" in kwargs and kwargs["num_workers"] >= nprocs: + kwargs["num_workers"] = (kwargs["num_workers"] + nprocs - 1) // nprocs if "batch_sampler" not in kwargs: if isinstance(dataset, IterableDataset): @@ -118,7 +117,7 @@ def auto_dataloader(dataset: Dataset, **kwargs: Any) -> DataLoader | _MpDeviceLo ) kwargs["pin_memory"] = False else: - kwargs["pin_memory"] = kwargs.get("pin_memory", "cuda" in idist.device().type) + kwargs["pin_memory"] = kwargs.get("pin_memory", "cuda" in idist.device().type or "mps" in idist.device().type) logger.info(f"Use data loader kwargs for dataset '{repr(dataset)[:20].strip()}': \n\t{kwargs}") dataloader = DataLoader(dataset, **kwargs) @@ -155,7 +154,7 @@ def auto_model(model: nn.Module, sync_bn: bool = False, **kwargs: Any) -> nn.Mod Args: model: model to adapt. sync_bn: if True, applies `torch convert_sync_batchnorm`_ to the model for native torch - distributed only. Default, False. Note, if using Nvidia/Apex, batchnorm conversion should be + distributed only. Default, False. Note, if using Nvidia/Apex, batch norm conversion should be applied before calling ``amp.initialize``. kwargs: kwargs to model's wrapping class: `torch DistributedDataParallel`_ or `torch DataParallel`_ if applicable. Please, make sure to use acceptable kwargs for given backend. @@ -166,15 +165,15 @@ def auto_model(model: nn.Module, sync_bn: bool = False, **kwargs: Any) -> nn.Mod Examples: .. code-block:: python - import ignite.distribted as idist + import ignite.distributed as idist model = idist.auto_model(model) - In addition with NVidia/Apex, it can be used in the following way: + In addition with Nvidia/Apex, it can be used in the following way: .. code-block:: python - import ignite.distribted as idist + import ignite.distributed as idist model, optimizer = amp.initialize(model, optimizer, opt_level=opt_level) model = idist.auto_model(model) @@ -242,7 +241,7 @@ def auto_optim(optimizer: Optimizer, **kwargs: Any) -> Optimizer: Internally, this method is no-op for non-distributed and torch native distributed configuration. For XLA distributed configuration, we create a new class that inherits from provided optimizer. - The goal is to override the `step()` method with specific `xm.optimizer_step`_ implementation. + The goal is to override the ``step()`` method with specific `xm.optimizer_step`_ implementation. For Horovod distributed configuration, optimizer is wrapped with Horovod Distributed Optimizer and its state is broadcasted from rank 0 to all other processes. @@ -339,7 +338,7 @@ class _MpDeviceLoader: # From pytorch/xla if `torch_xla.distributed.parallel_loader.MpDeviceLoader` is not available def __init__(self, loader: Any, device: torch.device, **kwargs: Any) -> None: self._loader = loader - # pyrefly: ignore [read-only] + # pyrely: ignore [read-only] self._device = device self._parallel_loader_kwargs = kwargs @@ -356,4 +355,4 @@ def __init__(self, optimizer: Optimizer) -> None: self.wrapped_optimizer = optimizer def step(self, closure: Any = None) -> Any: - xm.optimizer_step(self.wrapped_optimizer, barrier=True) + xm.optimizer_step(self.wrapped_optimizer, barrier=True) \ No newline at end of file