Merge pull request #250 from jrzaurin/multi-gpu

Added Support for multiple GPUs

Author: jrzaurin

pytorch_widedeep/models/wide_deep.py

@@ -297,9 +297,13 @@ def _forward_wide(self, X: Dict[str, Union[Tensor, List[Tensor]]]) -> Tensor:
             first_model_mode = list(X.keys())[0]
             if isinstance(X[first_model_mode], list):
                 batch_size = X[first_model_mode][0].size(0)
+                # Get device from input tensor
+                device = X[first_model_mode][0].device
             else:
                 batch_size = X[first_model_mode].size(0)  # type: ignore[union-attr]
-            out = torch.zeros(batch_size, self.pred_dim).to(self.wd_device)
+                # Get device from input tensor
+                device = X[first_model_mode].device  # type: ignore[union-attr]
+            out = torch.zeros(batch_size, self.pred_dim, device=device)
 
         return out
 
@@ -331,7 +335,9 @@ def _forward_deep(
     def _forward_deephead(
         self, X: Dict[str, Union[Tensor, List[Tensor]]], wide_out: Tensor
     ) -> Union[Tensor, Tuple[Tensor, Tensor]]:
-        deepside = torch.FloatTensor().to(self.wd_device)
+        # Get device from wide_out
+        device = wide_out.device
+        deepside = torch.FloatTensor().to(device)
 
         if self.deeptabular is not None:
             if self.is_tabnet:

pytorch_widedeep/training/_base_trainer.py

@@ -78,10 +78,22 @@ def __init__(
         self.seed = seed
 
         self.model = to_device_model(model, self.device)
-        if self.model.is_tabnet:
+
+        self.is_model_tabnet = model.is_tabnet
+        if self.is_model_tabnet:
             self.lambda_sparse = kwargs.get("lambda_sparse", 1e-3)
+
+        # Simply we need this attribute
         self.model.wd_device = self.device
 
+        use_multi_gpu = kwargs.get("use_multi_gpu", False) and self.device.startswith(
+            "cuda"
+        )
+        if use_multi_gpu and torch.cuda.device_count() > 1:
+            if self.verbose:
+                print(f"Using {torch.cuda.device_count()} GPUs for training")
+            self.model = torch.nn.DataParallel(self.model)
+
         self.objective = objective
         self.method: str = _ObjectiveToMethod.get(objective)  # type: ignore
 
@@ -444,6 +456,14 @@ def _set_device_and_num_workers(**kwargs) -> Tuple[str, int]:
         num_workers = kwargs.get("num_workers", default_num_workers)
         default_device = setup_device()
         device = kwargs.get("device", default_device)
+
+        # Check for multi-GPU setup
+        use_cuda = device.startswith("cuda")
+        use_multi_gpu = use_cuda and kwargs.get("use_multi_gpu", False)
+
+        if use_multi_gpu and torch.cuda.device_count() > 1:
+            device = f"cuda:{torch.cuda.current_device()}"
+
         return device, num_workers
 
     def __repr__(self) -> str:  # noqa: C901

pytorch_widedeep/training/trainer.py

@@ -160,12 +160,21 @@ class Trainer(BaseTrainer):
         - **num_workers**: `int`<br/>
             number of workers to be used internally by the data loaders
 
+        - **use_multi_gpu**: `bool`<br/>
+            If True, the model will be trained on multiple GPUs. This is
+            only supported for the `deeptabular` component.
+
+            NOTE: this is an experimental feature and might not work as expected
+            in some cases. In the particular case of the `Trainer` class, it has
+            been extensively tested.
+
         - **lambda_sparse**: `float`<br/>
             lambda sparse parameter in case the `deeptabular` component is `TabNet`
 
         - **class_weight**: `List[float]`<br/>
             This is the `weight` or `pos_weight` parameter in
             `CrossEntropyLoss` and `BCEWithLogitsLoss`, depending on whether
+
         - **reducelronplateau_criterion**: `str`
             This sets the criterion that will be used by the lr scheduler to
             take a step: One of _'loss'_ or _'metric'_. The ReduceLROnPlateau
@@ -834,22 +843,58 @@ def _do_finetune(
         r"""
         Simple wrap-up to individually fine-tune model components
         """
-        if self.model.deephead is not None:
+
+        if isinstance(self.model, torch.nn.DataParallel):
+            wide_component = (
+                torch.nn.DataParallel(self.model.module.wide)
+                if self.model.module.wide
+                else None
+            )
+            deeptabular_component = (
+                torch.nn.DataParallel(self.model.module.deeptabular)
+                if self.model.module.deeptabular
+                else None
+            )
+            deeptext_component = (
+                torch.nn.DataParallel(self.model.module.deeptext)
+                if self.model.module.deeptext
+                else None
+            )
+            deepimage_component = (
+                torch.nn.DataParallel(self.model.module.deepimage)
+                if self.model.module.deepimage
+                else None
+            )
+            deephead_component = (
+                torch.nn.DataParallel(self.model.module.deephead)
+                if self.model.module.deephead
+                else None
+            )
+        else:
+            wide_component = self.model.wide if self.model.wide else None
+            deeptabular_component = (
+                self.model.deeptabular if self.model.deeptabular else None
+            )
+            deeptext_component = self.model.deeptext if self.model.deeptext else None
+            deepimage_component = self.model.deepimage if self.model.deepimage else None
+            deephead_component = self.model.deephead if self.model.deephead else None
+
+        if deephead_component is not None:
             raise ValueError(
                 "Currently warming up is only supported without a fully connected 'DeepHead'"
             )
 
         finetuner = FineTune(self.loss_fn, self.metric, self.method, self.verbose)  # type: ignore[arg-type]
-        if self.model.wide:
-            finetuner.finetune_all(self.model.wide, "wide", loader, n_epochs, max_lr)
+        if wide_component:
+            finetuner.finetune_all(wide_component, "wide", loader, n_epochs, max_lr)
 
-        if self.model.deeptabular:
+        if deeptabular_component:
             if deeptabular_gradual:
                 assert (
                     deeptabular_layers is not None
                 ), "deeptabular_layers must be passed if deeptabular_gradual=True"
                 finetuner.finetune_gradual(
-                    self.model.deeptabular,
+                    deeptabular_component,
                     "deeptabular",
                     loader,
                     deeptabular_max_lr,
@@ -858,16 +903,16 @@ def _do_finetune(
                 )
             else:
                 finetuner.finetune_all(
-                    self.model.deeptabular, "deeptabular", loader, n_epochs, max_lr
+                    deeptabular_component, "deeptabular", loader, n_epochs, max_lr
                 )
 
-        if self.model.deeptext:
+        if deeptext_component:
             if deeptext_gradual:
                 assert (
                     deeptext_layers is not None
                 ), "deeptext_layers must be passed if deeptext_gradual=True"
                 finetuner.finetune_gradual(
-                    self.model.deeptext,
+                    deeptext_component,
                     "deeptext",
                     loader,
                     deeptext_max_lr,
@@ -876,16 +921,16 @@ def _do_finetune(
                 )
             else:
                 finetuner.finetune_all(
-                    self.model.deeptext, "deeptext", loader, n_epochs, max_lr
+                    deeptext_component, "deeptext", loader, n_epochs, max_lr
                 )
 
-        if self.model.deepimage:
+        if deepimage_component:
             if deepimage_gradual:
                 assert (
                     deepimage_layers is not None
                 ), "deepimage_layers must be passed if deepimage_gradual=True"
                 finetuner.finetune_gradual(
-                    self.model.deepimage,
+                    deepimage_component,
                     "deepimage",
                     loader,
                     deepimage_max_lr,
@@ -894,7 +939,7 @@ def _do_finetune(
                 )
             else:
                 finetuner.finetune_all(
-                    self.model.deepimage, "deepimage", loader, n_epochs, max_lr
+                    deepimage_component, "deepimage", loader, n_epochs, max_lr
                 )
 
     def _train_epoch(
@@ -944,7 +989,7 @@ def _train_step(
 
         y_pred = self.model(X)
 
-        if self.model.is_tabnet:
+        if self.is_model_tabnet:
             loss = self.loss_fn(y_pred[0], y) - self.lambda_sparse * y_pred[1]
             score = self._get_score(y_pred[0], y, is_train=True)
         else:
@@ -1008,7 +1053,7 @@ def _eval_step(
             y = to_device(y, self.device)
 
             y_pred = self.model(X)
-            if self.model.is_tabnet:
+            if self.is_model_tabnet:
                 loss = self.loss_fn(y_pred[0], y) - self.lambda_sparse * y_pred[1]
                 score = self._get_score(y_pred[0], y, is_train=False)
             else:
@@ -1119,7 +1164,7 @@ def _predict(  # type: ignore[override, return]  # noqa: C901
                                     X[k] = to_device(v, self.device)
                             preds = (
                                 self.model(X)
-                                if not self.model.is_tabnet
+                                if not self.is_model_tabnet
                                 else self.model(X)[0]
                             )
                             if self.method == "binary":
@@ -1170,6 +1215,7 @@ def _extract_kwargs(kwargs):
             "prefetch_factor",
             "persistent_workers",
             "oversample_mul",
+            "pin_memory",
         ]
         finetune_params = [
             "n_epochs",

pytorch_widedeep/training/trainer_from_folder.py

@@ -144,7 +144,17 @@ class TrainerFromFolder(Trainer):
          - **num_workers**: `int`<br/>
              number of workers to be used internally by the data loaders
 
-         - **lambda_sparse**: `float`<br/>
+        - **use_multi_gpu**: `bool`<br/>
+            If True, the model will be trained on multiple GPUs. This is
+            only supported for the `deeptabular` component.
+
+            NOTE: this is an experimental feature and might not work as expected
+            in some cases. While for the `Trainer` class, it has been extensively
+            tested, for the `TrainerFromFolder` class, it has not been tested
+            that thoroughly (in principle the `TrainerFromFolder` inherits from
+            the `Trainer` class, so it should work).
+
+        - **lambda_sparse**: `float`<br/>
              lambda sparse parameter in case the `deeptabular` component is `TabNet`
 
          - **class_weight**: `List[float]`<br/>

pytorch_widedeep/utils/general_utils.py

@@ -6,7 +6,7 @@
 
 def setup_device() -> str:
     if torch.cuda.is_available():
-        return "cuda"
+        return f"cuda:{torch.cuda.current_device()}"
     elif torch.backends.mps.is_available():
         return "mps"
     else:
@@ -24,11 +24,10 @@ def to_device_model(model, device: str):  # noqa: C901
     # insistent transformation since it some cases overall approaches such as
     # model.to('mps') do not work
 
-    if device in ["cpu", "cuda"]:
+    if device == "cpu" or (device.startswith("cuda") and torch.cuda.is_available()):
         return model.to(device)
 
     if device == "mps":
-
         try:
             return model.to(device)
         except (RuntimeError, TypeError):