Merge pull request #900 from evansakmar/899_mem_inference_bb_gpu

899 updated MembershipInferenceBlackBox#infer to have CUDA support

Author: beat-buesser

art/attacks/inference/membership_inference/black_box.py

@@ -198,13 +198,7 @@ def fit(self, x: np.ndarray, y: np.ndarray, test_x: np.ndarray, test_y: np.ndarr
             import torch.nn as nn  # lgtm [py/repeated-import]
             import torch.optim as optim  # lgtm [py/repeated-import]
             from torch.utils.data import DataLoader  # lgtm [py/repeated-import]
-
-            use_cuda = torch.cuda.is_available()
-
-            def to_cuda(x):
-                if use_cuda:
-                    x = x.cuda()
-                return x
+            from art.utils import to_cuda
 
             loss_fn = nn.BCELoss()
             optimizer = optim.Adam(self.attack_model.parameters(), lr=self.learning_rate)
@@ -261,14 +255,18 @@ def infer(self, x: np.ndarray, y: Optional[np.ndarray] = None, **kwargs) -> np.n
         if self.default_model and self.attack_model_type == "nn":
             import torch  # lgtm [py/repeated-import]
             from torch.utils.data import DataLoader  # lgtm [py/repeated-import]
+            from art.utils import to_cuda, from_cuda
 
             self.attack_model.eval()
             inferred = None
             test_set = self._get_attack_dataset(f_1=features, f_2=y)
             test_loader = DataLoader(test_set, batch_size=self.batch_size, shuffle=True, num_workers=0)
             for input1, input2, _ in test_loader:
+                input1, input2 = to_cuda(input1), to_cuda(input2)
                 outputs = self.attack_model(input1, input2)
                 predicted = torch.round(outputs)
+                predicted = from_cuda(predicted)
+
                 if inferred is None:
                     inferred = predicted.detach().numpy()
                 else:

art/utils.py

@@ -36,6 +36,7 @@
 from scipy.special import gammainc
 import six
 from tqdm.auto import tqdm
+from torch import Tensor
 
 from art import config
 
@@ -1235,3 +1236,34 @@ def pad_sequence_input(x: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
         x_padded[i, : len(x_i)] = x_i
         x_mask[i, : len(x_i)] = 1
     return x_padded, x_mask
+
+
+# -------------------------------------------------------------------------------------------------------- CUDA SUPPORT
+
+
+def to_cuda(x: Tensor) -> Tensor:
+    """
+    Move the tensor from the CPU to the GPU if a GPU is available.
+
+    :param x: CPU Tensor to move to GPU if available.
+    :return: The CPU Tensor moved to a GPU Tensor.
+    """
+    from torch.cuda import is_available
+    use_cuda = is_available()
+    if use_cuda:
+        x = x.cuda()
+    return x
+
+
+def from_cuda(x: Tensor) -> Tensor:
+    """
+    Move the tensor from the GPU to the CPU if a GPU is available.
+
+    :param x: GPU Tensor to move to CPU if available.
+    :return: The GPU Tensor moved to a CPU Tensor.
+    """
+    from torch.cuda import is_available
+    use_cuda = is_available()
+    if use_cuda:
+        x = x.cpu()
+    return x