gaussian-splatting icon indicating copy to clipboard operation
gaussian-splatting copied to clipboard

Assign image tensors to `data_device` immediately on creation.

Open GaneshBannur opened this issue 6 months ago • 0 comments

The tensors which are created from PIL images are first created on the CPU. https://github.com/graphdeco-inria/gaussian-splatting/blob/d9fad7b3450bf4bd29316315032d57157e23a515/utils/general_utils.py#L23 If data_device is "cuda" they are later moved to the GPU. Normally, unreferenced tensors on the CPU should be released but PyTorch doesn't seem to do this. This results in high CPU RAM consumption for the entire training duration even when data_device is "cuda".

Moving the tensors to data_device immediately on creation results in a dramatic decrease in CPU RAM consumption when data_device is "cuda". When training on a T4 instance on Colab with 200 images, CPU RAM consumption went from 10GB down to 2GB. The GPU vRAM consumption doesn't increase as tensors are eventually moved to the GPU anyway.

It might help to move all tensors to data_device immediately on creation since PyTorch doesn't seem to deallocate RAM for CPU tensors.

GaneshBannur avatar Feb 21 '24 11:02 GaneshBannur