search

MadryLab / trak

A fast, effective data attribution method for neural networks in PyTorch
https://trak.csail.mit.edu/
MIT License
180 stars 25 forks source link

[Bug Fix] fix `is_grads_dict` for `IterativeGradientComputer` #58

Closed TheaperDeng closed 10 months ago

TheaperDeng commented 11 months ago

is_grads_dict is not set properly, which makes cases where grad is other instance (e.g., tensor) fails.

In FunctionalGradientComputer, it may works properly since grads is a dictionary.

https://github.com/MadryLab/trak/blob/39bf22ac0a803dec6dab7f3cd29c168ecdc069a7/trak/gradient_computers.py#L126-L160

While in IterativeGradientComputer, the return value could be a tensor.

https://github.com/MadryLab/trak/blob/39bf22ac0a803dec6dab7f3cd29c168ecdc069a7/trak/gradient_computers.py#L215-L237

kristian-georgiev commented 11 months ago

Great catch! In fact, is_grads_dict as it exists now violates the abstraction set by AbstractProjector. Therefore, setting is_grads_dict whenever the projector is different from CudaProjector will fail. I'll fix the abstraction violation, add a commit here, and merge. Thanks!

kristian-georgiev commented 10 months ago

Fixed (https://github.com/MadryLab/trak/commit/ead7aa4bde0e7b9c00971543e65b5597583e616d). Ended up adding a few minor docs/tests changes to the commit, as well as bringing back an example of an "iterative" model output function.

kristian-georgiev commented 10 months ago

Closing this PR and adding you as co-author of the other commit (https://github.com/MadryLab/trak/commit/ead7aa4bde0e7b9c00971543e65b5597583e616d) :)