def 462

Author: R-Palazzo

ctgan/synthesizers/_utils.py

@@ -0,0 +1,53 @@
+import sys
+import warnings
+
+import torch
+
+
+def _validate_gpu_parameters(enable_gpu, cuda):
+    """Validate both the `enable_gpu` and `cuda` parameters.
+
+    The logic here is to:
+    - Raise a warning if `cuda` is set because it's deprecated.
+    - Raise an error if both parameters are set in a conflicting way.
+    - Return the resolved `enable_gpu` value.
+    """
+    if cuda is not None:
+        warnings.warn(
+            '`cuda` parameter is deprecated and will be removed in a future release. '
+            'Please use `enable_gpu` instead.',
+            FutureWarning,
+        )
+        if not enable_gpu:
+            raise ValueError(
+                'Cannot resolve the provided values of `cuda` and `enable_gpu` parameters. '
+                'Please use only `enable_gpu`.'
+            )
+
+        enable_gpu = cuda
+
+    return enable_gpu
+
+
+def _set_device(enable_gpu, device=None):
+    """Set the torch device based on the `enable_gpu` parameter and system capabilities."""
+    if device:
+        return torch.device(device)
+
+    if enable_gpu:
+        if sys.platform == 'darwin':  # macOS
+            if getattr(torch.backends, 'mps', None) and torch.backends.mps.is_available():
+                device = 'mps'
+            else:
+                device = 'cpu'
+        else:  # Linux/Windows
+            device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    else:
+        device = 'cpu'
+
+    return torch.device(device)
+
+
+def validate_and_set_device(enable_gpu, cuda):
+    enable_gpu = _validate_gpu_parameters(enable_gpu, cuda)
+    return _set_device(enable_gpu)

ctgan/synthesizers/base.py

@@ -5,6 +5,8 @@
 import numpy as np
 import torch
 
+from ctgan.synthesizers._utils import _set_device
+
 
 @contextlib.contextmanager
 def set_random_states(random_state, set_model_random_state):
@@ -118,7 +120,7 @@ def save(self, path):
     @classmethod
     def load(cls, path):
         """Load the model stored in the passed `path`."""
-        device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
+        device = _set_device(enable_gpu=True)
         model = torch.load(path, weights_only=False)
         model.set_device(device)
         return model

ctgan/synthesizers/ctgan.py

@@ -12,6 +12,7 @@
 from ctgan.data_sampler import DataSampler
 from ctgan.data_transformer import DataTransformer
 from ctgan.errors import InvalidDataError
+from ctgan.synthesizers._utils import _set_device, validate_and_set_device
 from ctgan.synthesizers.base import BaseSynthesizer, random_state
 
 
@@ -138,8 +139,11 @@ class CTGAN(BaseSynthesizer):
         pac (int):
             Number of samples to group together when applying the discriminator.
             Defaults to 10.
+        enable_gpu (bool):
+            Whether to attempt to use GPU for computation.
+            Defaults to ``True``.
         cuda (bool):
-            Whether to attempt to use cuda for GPU computation.
+            ** Deprecated ** Whether to attempt to use cuda for GPU computation.
             If this is False or CUDA is not available, CPU will be used.
             Defaults to ``True``.
     """
@@ -159,7 +163,8 @@ def __init__(
         verbose=False,
         epochs=300,
         pac=10,
-        cuda=True,
+        enable_gpu=True,
+        cuda=None,
     ):
         assert batch_size % 2 == 0
 
@@ -178,16 +183,8 @@ def __init__(
         self._verbose = verbose
         self._epochs = epochs
         self.pac = pac
-
-        if not cuda or not torch.cuda.is_available():
-            device = 'cpu'
-        elif isinstance(cuda, str):
-            device = cuda
-        else:
-            device = 'cuda'
-
-        self._device = torch.device(device)
-
+        self._device = validate_and_set_device(enable_gpu, cuda)
+        self._enable_gpu = cuda if cuda is not None else enable_gpu
         self._transformer = None
         self._data_sampler = None
         self._generator = None
@@ -544,6 +541,6 @@ def sample(self, n, condition_column=None, condition_value=None):
 
     def set_device(self, device):
         """Set the `device` to be used ('GPU' or 'CPU)."""
-        self._device = device
+        self._device = _set_device(self._enable_gpu, device)
         if self._generator is not None:
             self._generator.to(self._device)

ctgan/synthesizers/tvae.py

@@ -10,6 +10,7 @@
 from tqdm import tqdm
 
 from ctgan.data_transformer import DataTransformer
+from ctgan.synthesizers._utils import _set_device, validate_and_set_device
 from ctgan.synthesizers.base import BaseSynthesizer, random_state
 
 
@@ -114,8 +115,9 @@ def __init__(
         batch_size=500,
         epochs=300,
         loss_factor=2,
-        cuda=True,
+        enable_gpu=True,
         verbose=False,
+        cuda=None,
     ):
         self.embedding_dim = embedding_dim
         self.compress_dims = compress_dims
@@ -127,15 +129,8 @@ def __init__(
         self.epochs = epochs
         self.loss_values = pd.DataFrame(columns=['Epoch', 'Batch', 'Loss'])
         self.verbose = verbose
-
-        if not cuda or not torch.cuda.is_available():
-            device = 'cpu'
-        elif isinstance(cuda, str):
-            device = cuda
-        else:
-            device = 'cuda'
-
-        self._device = torch.device(device)
+        self._device = validate_and_set_device(enable_gpu, cuda)
+        self._enable_gpu = cuda if cuda is not None else enable_gpu
 
     @random_state
     def fit(self, train_data, discrete_columns=()):
@@ -241,6 +236,6 @@ def sample(self, samples):
         return self.transformer.inverse_transform(data, sigmas.detach().cpu().numpy())
 
     def set_device(self, device):
-        """Set the `device` to be used ('GPU' or 'CPU)."""
-        self._device = device
+        """Set the `device` to be used ('GPU' or 'CPU')."""
+        self._device = _set_device(self._enable_gpu, device)
         self.decoder.to(self._device)