Add gpu support to tracincp rand projection #969
Conversation
NarineK
force-pushed
the
add_gpu_support_to_tracincp
branch
from
c846d1d
to
b652271
Compare
tests/influence/_utils/common.py
Outdated
| return ( | ||
| wrap_model_in_dataparallel(net) if use_gpu else net, | ||
| dataset, | ||
| [s.cuda() for s in test_samples] |
There was a problem hiding this comment.
nit: Maybe a helper function for moving to GPU for either single tensor or list ([elem.cuda() for elem in val] if isinstance(val, list) else val.cuda()) might be helpful to make the if elses here more concise as well as the use-cases above?
| @@ -568,10 +563,28 @@ def _basic_computation_tracincp_fast( | |||
| targets (tensor): If computing influence scores on a loss function, | |||
| these are the labels corresponding to the batch `inputs`. | |||
| """ | |||
| global layer_inputs | |||
| layer_inputs = [] | |||
| layer_inputs: Dict[Module, Dict[device, Tuple[Tensor, ...]]] = defaultdict(dict) | |||
There was a problem hiding this comment.
nit: Looks like we only use this dictionary for results of one module, could instead just make this Dict[device, Tuple[Tensor, ...]] ?
Seems like extract_device_ids expects this Dict of Dict structure, but can probably just be replaced with getting the device_id attribute of the model if available.
|
@NarineK has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator. |
NarineK
force-pushed
the
add_gpu_support_to_tracincp
branch
from
b4bc0fa
to
a1bcfba
Compare
|
@NarineK has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator. |
|
@NarineK has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator. |
|
@NarineK has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator. |
Adds gpu support to tracincp rand projection.
Cleaned up un-passed args to _load_flexible_state_dict