Option to process images on GPU for speed at the expense of memory

[AX-I]

Analogous to #2110. This fix speeds up training by over 3x when using multi-gpu DDP.

The observation that prompted this PR: training with --num-devices=2 is significantly slower than training with single gpu, with extremely high cpu usage (~2600% on a 28-core cpu).

Profiling with --logging.profiler=pytorch reveals that collate_image_dataset_batch() in pixel_samplers.py takes 70ms with --num-devices=2 and 30ms with single gpu. The batch image tensors are stored on the cpu, so the indexing operation is cpu-intensive.

This PR adds an option to remove "image" from exclude_batch_keys_from_device to keep the image tensor on the gpu. Using the option with suitable hardware reduces cpu usage from 2600% to 200%, and speeds up 2-gpu training from ~70 Krays/s to ~250 Krays/s. Single-gpu training is also sped up from 100 to 140 Krays/s. VRAM usage is increased from 2gb to 5gb.

Profiling the 2-gpu training with this fix confirms that the pixel sampler is no longer a bottleneck.

---

Training configuration: nerfacto default params, on F2Nerf grass scene
Machine: 2x Xeon Gold 5120 (disabled HT), 2x Titan RTX

[AX-I]

Might also help with #1638

[tancik]

The "expense of memory" is the tricky part. Many users have 8gb cards which nerfacto has been designed to fit onto. When the dataset is moved to GPU, the memory is no longer predictable.

[AX-I]

The "expense of memory" is the tricky part. Many users have 8gb cards which nerfacto has been designed to fit onto. When the dataset is moved to GPU, the memory is no longer predictable.

Converted this to a config option! That should allow for all use cases.
--pipeline.datamanager.images-on-gpu True

[AX-I]

Changed optional to bool and rebased

[tancik]

LGTM