add influence gpu tests not using DataParallel

tests/influence/_core/test_tracin_k_most_influential.py

@@ -1,5 +1,5 @@
 import tempfile
-from typing import Callable
+from typing import Callable, Union
 
 import torch
 import torch.nn as nn
@@ -18,7 +18,7 @@
 class TestTracInGetKMostInfluential(BaseTest):
 
     use_gpu_list = (
-        [True, False]
+        [False, "cuda", "cuda_data_parallel"]
         if torch.cuda.is_available() and torch.cuda.device_count() != 0
         else [False]
     )
@@ -48,7 +48,9 @@ class TestTracInGetKMostInfluential(BaseTest):
                             DataInfluenceConstructor(
                                 TracInCP,
                                 name="linear2",
-                                layers=["module.linear2"] if use_gpu else ["linear2"],
+                                layers=["module.linear2"]
+                                if use_gpu == "cuda_data_parallel"
+                                else ["linear2"],
                             ),
                             False,
                         ),
@@ -83,7 +85,7 @@ def test_tracin_k_most_influential(
         proponents: bool,
         batch_size: int,
         k: int,
-        use_gpu: bool,
+        use_gpu: Union[bool, str],
         aggregate: bool,
     ) -> None:
         """

tests/influence/_core/test_tracin_self_influence.py

@@ -1,5 +1,5 @@
 import tempfile
-from typing import Callable
+from typing import Callable, Union
 
 import torch
 import torch.nn as nn
@@ -19,7 +19,7 @@
 class TestTracInSelfInfluence(BaseTest):
 
     use_gpu_list = (
-        [True, False]
+        [False, "cuda", "cuda_data_parallel"]
         if torch.cuda.is_available() and torch.cuda.device_count() != 0
         else [False]
     )
@@ -37,7 +37,9 @@ class TestTracInSelfInfluence(BaseTest):
                     DataInfluenceConstructor(
                         TracInCP,
                         name="TracInCP_linear1",
-                        layers=["module.linear1"] if use_gpu else ["linear1"],
+                        layers=["module.linear1"]
+                        if use_gpu == "cuda_data_parallel"
+                        else ["linear1"],
                     ),
                 ),
                 (
@@ -46,7 +48,7 @@ class TestTracInSelfInfluence(BaseTest):
                         TracInCP,
                         name="TracInCP_linear1_linear2",
                         layers=["module.linear1", "module.linear2"]
-                        if use_gpu
+                        if use_gpu == "cuda_data_parallel"
                         else ["linear1", "linear2"],
                     ),
                 ),
@@ -87,7 +89,7 @@ def test_tracin_self_influence(
         reduction: str,
         tracin_constructor: Callable,
         unpack_inputs: bool,
-        use_gpu: bool,
+        use_gpu: Union[bool, str],
     ) -> None:
         with tempfile.TemporaryDirectory() as tmpdir:
             (net, train_dataset,) = get_random_model_and_data(

tests/influence/_utils/common.py

@@ -183,6 +183,16 @@ def forward(self, *inputs):
 def get_random_model_and_data(
     tmpdir, unpack_inputs, return_test_data=True, use_gpu=False
 ):
+    """
+    `use_gpu` can either be
+    - `False`: returned model is on cpu
+    - `'cuda'`: returned model is on gpu
+    - `'cuda_data_parallel``: returned model is a `DataParallel` model, and on cpu
+    The need to differentiate between `'cuda'` and `'cuda_data_parallel'`
+    is that sometimes we may want to test a model that is on cpu, but is *not*
+    wrapped in `DataParallel`.
+    """
+    assert use_gpu in [False, "cuda", "cuda_data_parallel"]
 
     in_features, hidden_nodes, out_features = 5, 4, 3
     num_inputs = 2
@@ -209,7 +219,11 @@ def get_random_model_and_data(
         if hasattr(net, "pre"):
             net.pre.weight.data = net.pre.weight.data.double()
         checkpoint_name = "-".join(["checkpoint-reg", str(i + 1) + ".pt"])
-        net_adjusted = _wrap_model_in_dataparallel(net) if use_gpu else net
+        net_adjusted = (
+            _wrap_model_in_dataparallel(net)
+            if use_gpu == "cuda_data_parallel"
+            else (net.to(device="cuda") if use_gpu == "cuda" else net)
+        )
         torch.save(net_adjusted.state_dict(), os.path.join(tmpdir, checkpoint_name))
 
     num_samples = 50
@@ -238,7 +252,9 @@ def get_random_model_and_data(
 
     if return_test_data:
         return (
-            _wrap_model_in_dataparallel(net) if use_gpu else net,
+            _wrap_model_in_dataparallel(net)
+            if use_gpu == "cuda_data_parallel"
+            else (net.to(device="cuda") if use_gpu == "cuda" else net),
             dataset,
             _move_sample_to_cuda(test_samples)
             if isinstance(test_samples, list) and use_gpu
@@ -248,7 +264,12 @@ def get_random_model_and_data(
             test_labels.cuda() if use_gpu else test_labels,
         )
     else:
-        return _wrap_model_in_dataparallel(net) if use_gpu else net, dataset
+        return (
+            _wrap_model_in_dataparallel(net)
+            if use_gpu == "cuda_data_parallel"
+            else (net.to(device="cuda") if use_gpu == "cuda" else net),
+            dataset,
+        )
 
 
 class DataInfluenceConstructor: