Initial version of async attribution with torch.futures (#1295)

Summary:
Pull Request resolved: #1295

Currently Captum doesn't support async forward functions. Ads R&P team would like this feature in order to replace their custom variant (D56655643) of Feature Ablation with Captum and maintain similar performance.

PyTorch introduce future concepts ([link](https://pytorch.org/docs/stable/futures.html)) so we can adopt it for feature_ablation.py as the first step.

Details:
- Initial evaluation returns a future, save it.
- Each evaluation for each feature for each input will returns an attribution result (plus corresponding weight if applicable), save all those result separately since futures cannot be added up directly.
- When all futures above are done. we can add up the evaluation result to the final outcome as one Tensor per input.
- Since common._run_forward is used by other attribution methods, need to do some type hacking over there. But if users attempt to use those methods async, they will end up in failure before Captum support async for those methods.

TODO: Extend FeatureAttributor to support `torch.futures`

Reviewed By: vivekmig

Differential Revision: D56764316

fbshipit-source-id: 33661a76380dc009f4c9c60323b3c584e5152cda

Author: yucu

captum/_utils/common.py

@@ -15,6 +15,8 @@
     TupleOrTensorOrBoolGeneric,
 )
 from torch import device, Tensor
+
+from torch.futures import Future
 from torch.nn import Module
 
 
@@ -514,7 +516,7 @@ def _run_forward(
     inputs: Any,
     target: TargetType = None,
     additional_forward_args: Any = None,
-) -> Tensor:
+) -> Union[Tensor, Future[Tensor]]:
     forward_func_args = signature(forward_func).parameters
     if len(forward_func_args) == 0:
         output = forward_func()
@@ -532,6 +534,8 @@ def _run_forward(
             else inputs
         )
     )
+    if isinstance(output, torch.futures.Future):
+        return output.then(lambda x: _select_targets(x.value(), target))
     return _select_targets(output, target)
 
 

captum/_utils/gradient.py

@@ -112,6 +112,10 @@ def compute_gradients(
     with torch.autograd.set_grad_enabled(True):
         # runs forward pass
         outputs = _run_forward(forward_fn, inputs, target_ind, additional_forward_args)
+        # _run_forward may return future of Tensor,
+        # but we don't support it here now
+        # And it will fail before here.
+        outputs = cast(Tensor, outputs)
         assert outputs[0].numel() == 1, (
             "Target not provided when necessary, cannot"
             " take gradient with respect to multiple outputs."
@@ -297,6 +301,10 @@ def forward_hook(module, inp, out=None):
             target=target_ind,
             additional_forward_args=additional_forward_args,
         )
+        # _run_forward may return future of Tensor,
+        # but we don't support it here now
+        # And it will fail before here.
+        output = cast(Tensor, output)
     finally:
         for hook in all_hooks:
             hook.remove()