From 0bfb364d2521707b966a99fea28350fcbb06c891 Mon Sep 17 00:00:00 2001 From: Florian Prill <63-m300196@users.noreply.gitlab.dkrz.de> Date: Fri, 27 Sep 2024 14:07:53 +0200 Subject: [PATCH 1/4] Make pin_memory of the Dataloader configurable --- .../training/config/dataloader/native_grid.yaml | 1 + src/anemoi/training/data/datamodule.py | 12 ++++++++---- 2 files changed, 9 insertions(+), 4 deletions(-) diff --git a/src/anemoi/training/config/dataloader/native_grid.yaml b/src/anemoi/training/config/dataloader/native_grid.yaml index 35ed7e35..e6d50801 100644 --- a/src/anemoi/training/config/dataloader/native_grid.yaml +++ b/src/anemoi/training/config/dataloader/native_grid.yaml @@ -1,4 +1,5 @@ prefetch_factor: 2 +pin_memory: True num_workers: training: 8 diff --git a/src/anemoi/training/data/datamodule.py b/src/anemoi/training/data/datamodule.py index 1e119892..70b15506 100644 --- a/src/anemoi/training/data/datamodule.py +++ b/src/anemoi/training/data/datamodule.py @@ -1,9 +1,10 @@ # (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts. +# (C) Copyright 2024 Deutscher Wetterdienst. # This software is licensed under the terms of the Apache Licence Version 2.0 # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. -# In applying this licence, ECMWF does not waive the privileges and immunities -# granted to it by virtue of its status as an intergovernmental organisation -# nor does it submit to any jurisdiction. +# In applying this licence, the above institution do not waive the privileges +# and immunities granted to it by virtue of its status as an intergovernmental +# organisation nor does it submit to any jurisdiction. import logging import os @@ -96,6 +97,9 @@ def __init__(self, config: DictConfig) -> None: ) self.config.dataloader.training.end = self.config.dataloader.validation.start - 1 + if not self.config.dataloader.pin_memory: + LOGGER.info("Data loader memory pinning disabled.") + def _check_resolution(self, resolution: str) -> None: assert ( self.config.data.resolution.lower() == resolution.lower() @@ -185,7 +189,7 @@ def _get_dataloader(self, ds: NativeGridDataset, stage: str) -> DataLoader: num_workers=self.config.dataloader.num_workers[stage], # use of pinned memory can speed up CPU-to-GPU data transfers # see https://pytorch.org/docs/stable/notes/cuda.html#cuda-memory-pinning - pin_memory=True, + pin_memory=self.config.dataloader.pin_memory, # worker initializer worker_init_fn=worker_init_func, # prefetch batches From b007defb71296dec368cad46b0a2f17cff6ecd72 Mon Sep 17 00:00:00 2001 From: Marek Jacob <1129-b380572@users.noreply.gitlab.dkrz.de> Date: Tue, 1 Oct 2024 16:07:21 +0200 Subject: [PATCH 2/4] make pin_memory setting optional --- src/anemoi/training/data/datamodule.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/anemoi/training/data/datamodule.py b/src/anemoi/training/data/datamodule.py index 70b15506..a887bf71 100644 --- a/src/anemoi/training/data/datamodule.py +++ b/src/anemoi/training/data/datamodule.py @@ -97,7 +97,7 @@ def __init__(self, config: DictConfig) -> None: ) self.config.dataloader.training.end = self.config.dataloader.validation.start - 1 - if not self.config.dataloader.pin_memory: + if not self.config.dataloader.get("pin_memory", True): LOGGER.info("Data loader memory pinning disabled.") def _check_resolution(self, resolution: str) -> None: @@ -189,7 +189,7 @@ def _get_dataloader(self, ds: NativeGridDataset, stage: str) -> DataLoader: num_workers=self.config.dataloader.num_workers[stage], # use of pinned memory can speed up CPU-to-GPU data transfers # see https://pytorch.org/docs/stable/notes/cuda.html#cuda-memory-pinning - pin_memory=self.config.dataloader.pin_memory, + pin_memory=self.config.dataloader.get("pin_memory", True), # worker initializer worker_init_fn=worker_init_func, # prefetch batches From 0d66511732d84b518d051d43988e3c7006d61d94 Mon Sep 17 00:00:00 2001 From: Marek Jacob <1129-b380572@users.noreply.gitlab.dkrz.de> Date: Wed, 16 Oct 2024 12:23:14 +0200 Subject: [PATCH 3/4] revert copyright changes --- src/anemoi/training/data/datamodule.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/src/anemoi/training/data/datamodule.py b/src/anemoi/training/data/datamodule.py index a887bf71..b714c965 100644 --- a/src/anemoi/training/data/datamodule.py +++ b/src/anemoi/training/data/datamodule.py @@ -1,10 +1,9 @@ # (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts. -# (C) Copyright 2024 Deutscher Wetterdienst. # This software is licensed under the terms of the Apache Licence Version 2.0 # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. -# In applying this licence, the above institution do not waive the privileges -# and immunities granted to it by virtue of its status as an intergovernmental -# organisation nor does it submit to any jurisdiction. +# In applying this licence, ECMWF does not waive the privileges and immunities +# granted to it by virtue of its status as an intergovernmental organisation +# nor does it submit to any jurisdiction. import logging import os From 9c3985da2d3fceb493b6fe651f8862aeaa4789bb Mon Sep 17 00:00:00 2001 From: Marek Jacob <1129-b380572@users.noreply.gitlab.dkrz.de> Date: Wed, 16 Oct 2024 12:26:43 +0200 Subject: [PATCH 4/4] update CHANGELOG.md --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 3b233a77..a98c7e52 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,7 @@ Please add your functional changes to the appropriate section in the PR. Keep it human-readable, your future self will thank you! ## [Unreleased](https://github.com/ecmwf/anemoi-training/compare/0.1.0...HEAD) +- Make pin_memory of the Dataloader configurable (#64) ### Added