add GPU number and lazy load img to GPU

[sword4869]

hello,

1. I add a gpu number argument in ModelParams class.
2. In order to save GPU memory, I cancelled the loading of img to GPU when creating the camera, and delayed loading img to GPU during training. This greatly saves GPU memory. I tested it in RTX2080ti and less than 2GB during initial training, even if trained to 30k, less than 5GB.

[pablospe]

This is a great improvement in memory usage, this should be merged into main branch. @Snosixtyboo

[pablospe]

Maybe equivalent?

if lazy_load:
    data_device = "cpu"

try:
    self.data_device = torch.device(data_device)
....

If this is correct, perhaps lazy_load can be replaced directly with data_device (and it needs to be exposed in train.py). Perhaps adding a comment in the README, there is a question about the memory and the different approaches when there is no 24GB.

[sword4869]

After your reminder, I realized that other tensors were directly migrated through . cuda, and data_device is only responsible for images, which can indeed function as a lazy load.

There is no need to expose data_device in train.py as it is already specified in the arguments/__init__.py.

All in all, nothing needs to be modified. If possible, I suggest modifying the "24 GB VRAM" in README, as it can easily make people mistakenly think that this is the minimum configuration.

[pablospe]

I am not sure to understand your reply. What do you mean by . cuda? I think the lazy_load was a good option to add, so you don't need to modify every time the __init__.py. Perhaps, could you expand how you would modify the README to clarify that data_device can be used for lazy loading?

[sword4869]

"other tensors were directly migrated through . cuda" precisely means that they are fixed to the cuda device, e.g. create gaussian.

There is no need to need to modify data_device every time the __init__.py. The train.py uses argparse.ArgumentParser to analyse the command arguments. And data_device is can be called by python train.py --data_device cpu -s <path to COLMAP or NeRF Synthetic dataset>.

# scene/cameras.py
# if here data_device is specified as `cpu`, self.original_image is on cpu
try:
    self.data_device = torch.device(data_device)
except Exception as e:
    print(e)
    print(f"[Warning] Custom device {data_device} failed, fallback to default cuda device" )
    self.data_device = torch.device("cuda")

self.original_image = image.clamp(0.0, 1.0).to(self.data_device)
self.image_width = self.original_image.shape[2]
self.image_height = self.original_image.shape[1]

if gt_alpha_mask is not None:
    self.original_image *= gt_alpha_mask.to(self.data_device)
else:
    self.original_image *= torch.ones((1, self.image_height, self.image_width), device=self.data_device)

# train.py
# original_image is sent to gpu while iterating
gt_image = viewpoint_cam.original_image.cuda()

data_device in README

[sword4869]

I modify REAMDE to clarify lazy load and recommit it.

[pablospe]

thanks for the clarifications!

[NiklasVoigt]

typo, should be device with a c --data_device cpu

[sword4869]

typo, should be device with a c --data_device cpu

tks for ur careful check

[Ky1eYang]

maybe we can reduce more usage of memory by lazy call function loadCam.

def lazy_call(f, *args, **kwargs):
    return lambda: f(*args, **kwargs)

def cameraList_from_camInfos(cam_infos, resolution_scale, args, lazy_load):
    camera_list = []
    # pdb.set_trace()
    for id, c in enumerate(cam_infos):
        if lazy_load:
            camera_list.append(lazy_call(loadCam, args, id, c, resolution_scale))
        else:
            camera_list.append(loadCam(args, id, c, resolution_scale))
    return camera_list`

In function 'cameraList_from_camInfos', we can delay executing loadCam and execute it in 'training'.

# Pick a random Camera
if not viewpoint_stack:
    viewpoint_stack = scene.getTrainCameras().copy()
viewpoint_cam = viewpoint_stack.pop(randint(0, len(viewpoint_stack)-1))
if scene.lazy_load:
    viewpoint_cam = viewpoint_cam()