Captum: CUDA OOM Error

[zwacke]

Hi,

I am currently integrating Captum into my deep learning tool kit, thx for providing this lib.

When I try to run IntegratedGradients on a standard densenet201 model that is on a cuda device (11GB vram), I am getting an out-of-memory error even for one input image.

Just a quick check: Is this normal behaviour?

[vivekmig]

Hi @zwacke , it depends on what you've set for n_steps. Integrated Gradients needs to compute gradients at each of the n_steps points, so one input image expands to a batch of n_step evaluations, which is likely causing the out of memory issue.

You can try either reducing n_steps to a smaller value or otherwise using the internal_batch_size argument in IG, which splits the evaluations into batches with at most internal_batch_size examples. Setting internal_batch_size to a smaller value (corresponding to a batch size that can fit in memory) while keeping n_steps the same should also work. Let us know if this resolves the issue!

[NarineK]

Hi @zwacke, did Vivek's suggestion help ? If so can we close the issue ?

[zwacke]

Yes, thank you it did. Defaulting to a working internal_batch_size seems to be the recommendable approach. If I get it right, this especially helps when wrapping an attribution method in NoiseTunnel, or any method that introduces further expansion of the overall input batch size.

[NarineK]

Nice! Glad that it helped. Feel free to close the issue if it's fixed!

[tranvnhan]

I encountered the same problem today. By reducing n_steps to a smaller value, it works now.