Support TCAV on cuda (#725)

Summary:
Addresses the issues: #721 #719 #720

Pull Request resolved: #725

Reviewed By: bilalsal

Differential Revision: D30356015

Pulled By: NarineK

fbshipit-source-id: 010a5263bdfc33e8c4d3f9de523d9d3ba3969f49

Author: NarineK

captum/_utils/models/linear_model/train.py

@@ -343,9 +343,15 @@ def sklearn_train_linear_model(
         if hasattr(sklearn_model, "classes_")
         else None
     )
+
+    # extract model device
+    device = model.device if hasattr(model, "device") else "cpu"
+
     num_outputs = sklearn_model.coef_.shape[0] if sklearn_model.coef_.ndim > 1 else 1
-    weight_values = torch.FloatTensor(sklearn_model.coef_)  # type: ignore
-    bias_values = torch.FloatTensor([sklearn_model.intercept_])  # type: ignore
+    weight_values = torch.FloatTensor(sklearn_model.coef_).to(device)  # type: ignore
+    bias_values = torch.FloatTensor([sklearn_model.intercept_]).to(  # type: ignore
+        device  # type: ignore
+    )  # type: ignore
     model._construct_model_params(
         norm_type=None,
         weight_values=weight_values.view(num_outputs, -1),

captum/concept/_core/concept.py

@@ -1,7 +1,8 @@
 #!/usr/bin/env python3
 
-from typing import Callable, Iterable, Union
+from typing import Callable, Union
 
+import torch
 from torch.nn import Module
 
 
@@ -18,15 +19,19 @@ class Concept:
     adjectives and words that convey happiness.
     """
 
-    def __init__(self, id: int, name: str, data_iter: Union[None, Iterable]) -> None:
+    def __init__(
+        self, id: int, name: str, data_iter: Union[None, torch.utils.data.DataLoader]
+    ) -> None:
 
         r"""
         Args:
             id (int):   The unique identifier of the concept.
             name (str): A unique name of the concept.
-            data_iter (iter): A pytorch Dataloader object. Combines a dataset
+            data_iter (DataLoader): A pytorch DataLoader object that combines a dataset
                         and a sampler, and provides an iterable over a given
-                        dataset. For more information, please check:
+                        dataset. Only the input batches are provided by `data_iter`.
+                        Concept ids can be used as labels if necessary.
+                        For more information, please check:
                         https://pytorch.org/docs/stable/data.html
 
         Example::

captum/concept/_core/tcav.py

@@ -640,7 +640,9 @@ def _tcav_sub_computation(
 
             # sort classes / concepts in the order specified in concept_keys
             concept_ord = {concept.id: ci for ci, concept in enumerate(concepts)}
-            new_ord = torch.tensor([concept_ord[cls] for cls in cls_set])
+            new_ord = torch.tensor(
+                [concept_ord[cls] for cls in cls_set], device=tcav_score.device
+            )
 
             # sort based on classes
             scores[concepts_key][layer] = {

captum/concept/_utils/classifier.py

@@ -170,16 +170,17 @@ def train_and_eval(
             inputs.append(input)
             labels.append(label)
 
+        device = "cpu" if input is None else input.device
         x_train, x_test, y_train, y_test = _train_test_split(
             torch.cat(inputs), torch.cat(labels), test_split=test_split_ratio
         )
-
+        self.lm.device = device
         self.lm.fit(DataLoader(TensorDataset(x_train, y_train)))
 
         predict = self.lm(x_test)
 
         predict = self.lm.classes()[torch.argmax(predict, dim=1)]
-        score = predict.long() == y_test.long()
+        score = predict.long() == y_test.long().cpu()
 
         accs = score.float().mean()