Optim-wip: Add JIT support to all transforms & some image parameterizations

[ProGamerGov]

• Added JIT support to all transforms.

• Added JIT tests for all applicable transforms.

• There's currently a bug that prevents CenterCrop from working with JIT, so I added the @torch.jit.ignore decorator to it's forward function to avoid the bug. This is so that it won't throw an error if someone uses torch.jit.script on it. JIT is overriding variables of type List[int] to type Tuple[int, int] for an unknown reason during function calls pytorch#69938

• JIT is also told to ignore NChannelsToRGB due to it's unsupported inner functions.

• The ToRGB didn't support JIT due to the usage of named dimensions, so I implemented a second forward function that is only used when using JIT.

• Add JIT support NaturalImage, FFTImage, & PixelImage.

• Improved NaturalImage & FFTImage tests, and test coverage.

• Added ImageParameterization instance support for NaturalImage. This improvement should make it easier to use parameterization enhancements like SharedImage, and will be helpful for custom parameterizations that don't use the standard input variable set (size, channels, batch, & init).

• Added asserts to verify NaturalImage parameterization inputs are instances or types of ImageParameterization.

• Some transform type hints had to be changed due to a bug in JIT's Union support: JIT RuntimeError: 'Union[Tensor, List[float], List[int]]' object is not subscriptable pytorch#69434

• Fixed bug with RandomCrop transform.

• Replaced default RandomScale transform with Interpolation based variant. Renamed the old variant to RandomScaleAffine.

• CenterCrop / center_crop now adds padding if the crop size is larger than the input dimensions.

• Add distributions support to both versions of RandomScale. Ludwig wanted this based off of his original PR.

• Changed NumSeqOrTensorType hint to NumSeqOrTensorOrProbDistType.

• Added more comprehensive testing to applicable transforms. Test coverage reported by pytest should now be 100% minus the version specific tests.

• Transform torch version checks required for module forward functions are performed in their __init__ function so that they work with JIT.

• Add TransformationRobustness() transform.

  • Similar to Lucid's standard_transforms, the TransformationRobustness() is a convenient helper for applying the standard set of transforms in the correct order. The default values have also been tuned to optimal values.
  • It performs the following transforms: A padding transform, one or more RandomSpatialJitter transforms, a RandomScale transform, a RandomRotation transform, a single RandomSpatialJitter transform, and then a final optional CenterCrop or Pad To transform.
  • There's a pretty significant improvement to visualization quality when using TransformationRobustness(), and it makes it easier to use transform robustness properly..
  • Links to Lucid's versions: https://github.com/tensorflow/lucid/blob/master/lucid/optvis/transform.py#L164, https://github.com/tensorflow/lucid/blob/master/lucid/recipes/activation_atlas/render.py#L127
  • Ludwig's original Captum PR had the same transform planned: https://github.com/ludwigschubert/captum/blob/f1fd0729dece59564a7c10b7b397617d8a09a247/captum/optim/param/transform.py#L88
  • It was also used for feature visualization with OpenAI's CLIP model: https://github.com/openai/CLIP-featurevis/blob/master/example_facets.py#L61

    pad_module = torch.nn.ConstantPad2d(2, value=0.5)
    transforms = opt.transforms.TransformationRobustness(padding_transform=pad_module)

JIT support for the InceptionV1 model was added to #655

This PR only changes / adds around 711 lines of code, with the other 1306 lines being from tests that feature a lot of redundant / simple code.

[ProGamerGov]

Looks like there was an error with the Insights install in the tests. Everything optim related is working fine however.

[ProGamerGov]

I've been using the TransformationRobustness in all of the newer notebooks due to it's ease of use and because it dramatically improves visualization quality.

The optimizing with transparency tutorial notebook requires this PR for the TransformationRobustness transform and the fixes to RandomCrop.