Compute importance measure in mini-batches

[AndreasMadsen]

This should work. I've tested batch_size=8 aganist batch_size=1 and get identical results.

What is left is to run this on compute-canada and make sure there is no out-of-memory issues. Also, maybe compute_batch_size should just be base_dataset.batch_size. I will test that too.

[AndreasMadsen]

@ncmeade @vaibhavad Can you review this. This appears to work. It is still very slow. I'm speculating that the code does not run on the GPU when you simply call model(x). At least I can not find any code in the PyTorchLightning that would cause it do run on the GPU. For the training Trainer.fit is the one to call model.cuda():

[vaibhavad]

I can confirm that the code currently does not use GPU. PID will not be visible on nvidia-smi. I put in .cuda() for base_model and batch, and observed significant time improvement. However, to do it properly within the Lightning framework, the importance measure computation will need to be a LightningModule, which will be run from a Trainer.

[AndreasMadsen]

However, to do it properly within the Lightning framework, the importance measure computation will need to be a LightningModule, which will be run from a Trainer.

That should be fine. I would anyway do something similar for TorchScript support. But how do I get the prediction output using Trainer?

In any case, let's enable GPU in a seperate PR. This is just for mini-batch support.

[vaibhavad]

Looks good to me, I'll follow this to have an input-level batched version for IG as well.

Let's discuss GPU in a separate PR as you suggested. For the time being, just adding .cuda() will work because we are using a single GPU for all our current experiments and model sizes are not that big.

[AndreasMadsen]

Thanks @vaibhavad. I fixed the squeeze issue and merged.