From f362eae60832c1f41d316b19d636cb3bfd34f7a1 Mon Sep 17 00:00:00 2001 From: Daniel Crispell Date: Fri, 20 Oct 2023 22:41:19 -0400 Subject: [PATCH] Move only per-chunk data to the GPU to support rendering large images (#2481) * Move only per-chunk data to the GPU to support large images * make type checking happy * Use isinstance check torch.is_tensor() doesn't have a TypeGuard annotation --------- Co-authored-by: Brent Yi --- nerfstudio/models/base_model.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/nerfstudio/models/base_model.py b/nerfstudio/models/base_model.py index ebda710b3..d7a221744 100644 --- a/nerfstudio/models/base_model.py +++ b/nerfstudio/models/base_model.py @@ -168,6 +168,7 @@ class Model(nn.Module): Args: camera_ray_bundle: ray bundle to calculate outputs over """ + input_device = camera_ray_bundle.directions.device num_rays_per_chunk = self.config.eval_num_rays_per_chunk image_height, image_width = camera_ray_bundle.origins.shape[:2] num_rays = len(camera_ray_bundle) @@ -176,12 +177,15 @@ class Model(nn.Module): start_idx = i end_idx = i + num_rays_per_chunk ray_bundle = camera_ray_bundle.get_row_major_sliced_ray_bundle(start_idx, end_idx) + # move the chunk inputs to the model device + ray_bundle = ray_bundle.to(self.device) outputs = self.forward(ray_bundle=ray_bundle) for output_name, output in outputs.items(): # type: ignore - if not torch.is_tensor(output): + if not isinstance(output, torch.Tensor): # 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(output.to(input_device)) outputs = {} for output_name, outputs_list in outputs_lists.items(): outputs[output_name] = torch.cat(outputs_list).view(image_height, image_width, -1) # type: ignore