Optim-wip: Add new StackImage parameterization & JIT support for SharedImage

[ProGamerGov]

• Added SimpleTensorParameterization as a workaround for JIT not supporting nn.ParameterList. It also helps StackImage support tensor inputs. Edit: Later versions of PyTorch now support nn.ParameterList, but the current solution supports all torch versions supported by Captum.
• Added JIT support for SharedImage. I had to separate the interpolation operations due to a bug with JIT's union support: JIT RuntimeError: 'Union[Tensor, List[float], List[int]]' object is not subscriptable pytorch#69434
• Added new parameterization called StackImage, that stacks multiple parameterizations (that are can be on different devices) along the batch dimension. Lucid uses this paramterization a few times, and I thought that I could improve on it with the multi-device functionality.
• Added tests for InputParameterization and ImageParameterization

[ProGamerGov]

The test_py37_conda test failed due to a stochastic error, and thus it can be ignored.

[NarineK]

@ProGamerGov, are there any specific benefits of adding jit support ? I thought that the models that we work with aren't jit-scripted. I see that we already merged a large PR without reviewing related to jit. I was thinking that it could be added later because we need to do additional reviews for jit and @vivekmig would be the best POC for it.

[ProGamerGov]

@NarineK The image parameterizations and transforms can very much benefit from being JIT scripted, without having to JIT script the models themselves. The other JIT PR covered all of the image parameterizations except the final ones in this PR, so that it's size wouldn't be too large.

Though, we can skip this PR for now if you want.

[NarineK]

@ProGamerGov, what kind of benefit do we get from JIT for transformations and image parametrization ? Is it runtime performance ? Did we measure the benefit ?

[ProGamerGov]

@NarineK In my limited testing, JIT scripting the image parameterizations and transforms can shave off small amount of time for standard rendering, though it will vary based on the parameters / settings used. I haven't really done an in-depth dive into the performance changes though, so the results could vary.

[ProGamerGov]

@NarineK Upon testing, there is a clear difference in terms of speed when using JIT scripted parameterizations and transforms:

# https://docs.python.org/3/library/timeit.html
import timeit

setup="""
import captum.optim as opt
import torch
test_transform = torch.jit.script(opt.transforms.TransformationRobustness())
image = torch.jit.script(opt.images.NaturalImage((224, 224), batch=4).cuda())
"""

code="""
_ = test_transform(image())
"""

for _ in range(10):
    output = timeit.timeit(stmt=code, setup=setup, number=5000)
    print(output)

Remove the torch.jit.script calls in the above code to test without JIT.

When I tested the above code on Colab, I got the following results:

# With JIT
10.923049781000032
11.118382118
11.57444840200003
10.89497750999999
10.775770954999984
10.848362547999955
10.786963311000022
10.801879264000036
10.77277677799998
10.866098210000018

# Without JIT
14.257542533999981
14.208807915999955
14.269374133000042
14.300639500999978
14.169171004999953
14.305170376999968
14.284609974000091
14.262493259999928
14.304441671000063
14.142993052999941

[NarineK]

Thank you, @ProGamerGov ! I'm rerunning the tests. Hopefully the conda issue will get fixed.

[ProGamerGov]

@NarineK Looks like the Conda test failed again due to a stochastic error

[NarineK]

It's not clear to me that it is a flaky error caused by a stochastic behavior. It is probably a version issue. Some people were able to fix it by downgrading conda version. Let's leave it as is and see if we are seeing the same issue on master. I don't see it in master PRs.