Move only per-chunk data to the GPU to support rendering large images

[decrispell]

Despite chunking of the rays for evaluation, there is a lot of full-image per-pixel data (inputs and outputs) still being stored on the GPU, which can be problematic when rendering very large images.

This change allows passing a CPU camera_ray_bundle to get_outputs_for_camera_ray_bundle(). The per-chunk rays are moved to the GPU, and the per-chunk outputs are moved back to the device of the original camera_ray_bundle.

Behavior of existing code shouldn't change since inputs and outputs will still be stored completely on the GPU - it's just no longer a requirement.

[drewgilliam]

@brentyi & @tancik - hoping you could help out again, this PR is a blocker for our team. Is there more information or more changes needed to get this PR merged? Thanks!

[brentyi]

Sorry for the delay, this seems fine to me! Just left one minor comment.

[decrispell]

I agree that `output.to(input_device)` is cleaner (and it does work), but I changed it to this form to make Pyright happy. (See the errors on the initial commit of this PR).

…

On Fri, Oct 20, 2023 at 9:42 PM Brent Yi ***@***.***> wrote: ***@***.**** approved this pull request. Sorry for the delay, this seems fine to me! Just left one minor comment. ------------------------------ In nerfstudio/models/base_model.py <#2481 (comment)> : > outputs = self.forward(ray_bundle=ray_bundle) for output_name, output in outputs.items(): # type: ignore if not torch.is_tensor(output): # TODO: handle lists of tensors as well continue - outputs_lists[output_name].append(output) + # move the chunk outputs from the model device back to the device of the inputs. + outputs_lists[output_name].append(torch.as_tensor(output, device=input_device)) Does outputs_lists[output_name].append(output.to(input_device)) work? (it seems equivalent, but semantically torch.as_tensor() makes me think that output could be a numpy array) — Reply to this email directly, view it on GitHub <#2481 (review)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAKTNQ4AAJJBXTOA64ZEWCLYAMSAHAVCNFSM6AAAAAA5MXSQ6KVHI2DSMVQWIX3LMV43YUDVNRWFEZLROVSXG5CSMV3GSZLXHMYTMOJRGA2DMNRTGM> . You are receiving this because you authored the thread.Message ID: ***@***.***>

[brentyi]

Ok! I pushed a fix: the error seemed correct, since output has a list[Tensor] | Tensor type and is_tensor(output) isn't annotated by the PyTorch team as a TypeGuard. Switching to isinstance() makes .to() pass the checks.

[decrispell]

That seems like a much better solution, Thanks @brentyi !