diff --git a/RELEASES.md b/RELEASES.md index 526f6409e..42cb2bddb 100644 --- a/RELEASES.md +++ b/RELEASES.md @@ -8,6 +8,7 @@ #### Closed issues +- Build the CUDA generator and the CUDA entries of `TorchBackend.__type_list__` lazily, so that using POT with CPU-only torch tensors no longer initialises a CUDA context and claims device memory (PR #847, Issue #612) - Fix the sign issue in updates of the previous transport plan in `ot.batch.proximal_bregman_log_plan_batch` (Issue #842) - Load triton before TensorFlow in `ot.backend` so that building a torch optimizer no longer segfaults the interpreter, and remove the `torch<2.12` pin from the doctest and documentation requirements (PR #839, Issue #816) - Fix mean centering in `ot.dr.fda` and `ot.dr.wda`: `np.mean(X)` returned a scalar instead of the per-feature mean, so `proj` did not center the data as documented. In `ot.dr.fda` the same pattern in the class means made the between-class scatter matrix independent of which features separate the classes, and FDA returned a non-discriminant direction (PR #840) diff --git a/ot/backend.py b/ot/backend.py index af622734d..7d562ec34 100644 --- a/ot/backend.py +++ b/ot/backend.py @@ -2066,7 +2066,7 @@ class TorchBackend(Backend): __name__ = "torch" __type__ = torch_type - __type_list__ = None + # __type_list__ is a property below: its CUDA entries are built lazily. rng_ = None @@ -2074,22 +2074,17 @@ def __init__(self): self.rng_ = torch.Generator("cpu") self.rng_.seed() - self.__type_list__ = [ + # The CUDA generator and the CUDA entries of the type list are built the + # first time something asks for them. Building either here initialises a + # CUDA context, which claims device memory and wakes the GPU even when + # the computation stays entirely on the CPU. + self._type_list_cpu = [ torch.tensor(1, dtype=torch.float32), torch.tensor(1, dtype=torch.float64), ] - - if torch.cuda.is_available(): - self.rng_cuda_ = torch.Generator("cuda") - self.rng_cuda_.seed() - self.__type_list__.append( - torch.tensor(1, dtype=torch.float32, device="cuda") - ) - self.__type_list__.append( - torch.tensor(1, dtype=torch.float64, device="cuda") - ) - else: - self.rng_cuda_ = torch.Generator("cpu") + self._type_list_cuda = None + self._rng_cuda = None + self._cuda_seed = None from torch.autograd import Function from torch.autograd.function import once_differentiable @@ -2133,6 +2128,35 @@ def backward(ctx, g): self.ValFunction = ValFunction self.MatrixSqrtFunction = MatrixSqrtFunction + @property + def __type_list__(self): + if self._type_list_cuda is None: + if torch.cuda.is_available(): + self._type_list_cuda = [ + torch.tensor(1, dtype=torch.float32, device="cuda"), + torch.tensor(1, dtype=torch.float64, device="cuda"), + ] + else: + self._type_list_cuda = [] + return self._type_list_cpu + self._type_list_cuda + + @property + def rng_cuda_(self): + if self._rng_cuda is None: + if torch.cuda.is_available(): + self._rng_cuda = torch.Generator("cuda") + if self._cuda_seed is None: + self._rng_cuda.seed() + else: + self._rng_cuda.manual_seed(self._cuda_seed) + else: + self._rng_cuda = torch.Generator("cpu") + return self._rng_cuda + + @rng_cuda_.setter + def rng_cuda_(self, generator): + self._rng_cuda = generator + def _to_numpy(self, a): if isinstance(a, float) or isinstance(a, int) or isinstance(a, np.ndarray): return np.array(a) @@ -2412,7 +2436,11 @@ def seed(self, seed=None): pass elif isinstance(seed, int): self.rng_.manual_seed(seed) - self.rng_cuda_.manual_seed(seed) + # Remember the seed rather than forcing the CUDA generator into + # existence; it is applied when that generator is first needed. + self._cuda_seed = seed + if self._rng_cuda is not None: + self._rng_cuda.manual_seed(seed) elif isinstance(seed, torch.Generator): if self.device_type(seed) == "GPU": self.rng_cuda_ = seed diff --git a/test/test_backend.py b/test/test_backend.py index 4df918140..e3a060ed0 100644 --- a/test/test_backend.py +++ b/test/test_backend.py @@ -942,3 +942,31 @@ def test_torch_optimizer_after_tensorflow_import(): f"interpreter died with returncode {result.returncode}: " f"{result.stderr.decode(errors='replace')[-2000:]}" ) + + +@pytest.mark.skipif( + not torch or not torch.cuda.is_available(), + reason="Requires torch with CUDA available", +) +def test_no_cuda_context_for_cpu_only_work(): + """Non-regression test for issue #612. + + Building the torch backend used to create a CUDA generator and the CUDA + entries of the type list straight away. That initialises a CUDA context, so + device memory is claimed and the GPU wakes up for a computation that stays + entirely on the CPU. The check needs an interpreter that has not touched + CUDA yet, so it runs in a subprocess. + """ + code = ( + "import torch\n" + "import ot\n" + "x = torch.randn(64, 2)\n" + "ot.dist(x, x)\n" + "assert not torch.cuda.is_initialized(), 'a CUDA context was created'\n" + "assert torch.cuda.memory_allocated() == 0, 'device memory was claimed'\n" + ) + result = subprocess.run([sys.executable, "-c", code], capture_output=True) + assert result.returncode == 0, ( + f"interpreter exited with returncode {result.returncode}: " + f"{result.stderr.decode(errors='replace')[-2000:]}" + )