From a43776d76d7f227174141e8d52e8fb48d8feaa87 Mon Sep 17 00:00:00 2001
From: Matthew Tancik <tancik@berkeley.edu>
Date: Mon, 19 Sep 2022 18:36:12 -0700
Subject: [PATCH] Speedup eval and training with better data management (#492)

move to tensors to gpu earlier
---
 nerfactory/cameras/cameras.py              |  5 ++++-
 nerfactory/cameras/rays.py                 |  2 +-
 nerfactory/fields/compound_field.py        | 15 +++++++--------
 nerfactory/models/modules/ray_generator.py |  1 -
 nerfactory/utils/math.py                   |  2 +-
 5 files changed, 13 insertions(+), 12 deletions(-)

diff --git a/nerfactory/cameras/cameras.py b/nerfactory/cameras/cameras.py
index 217994fb9..66669640a 100644
--- a/nerfactory/cameras/cameras.py
+++ b/nerfactory/cameras/cameras.py
@@ -171,6 +171,9 @@ class Cameras:
             Rays for the given camera indices and coords.
         """
 
+        if isinstance(camera_indices, torch.Tensor):
+            camera_indices = camera_indices.to(self.device)
+
         if coords is None:
             coords = self.get_image_coords().to(self.device)
 
@@ -237,7 +240,7 @@ class Cameras:
         if not isinstance(camera_indices, torch.Tensor):
             ray_bundle_camera_indices = torch.Tensor([camera_indices]).broadcast_to(pixel_area.shape).to(self.device)
         else:
-            ray_bundle_camera_indices = camera_indices
+            ray_bundle_camera_indices = camera_indices.view(pixel_area.shape)
 
         return RayBundle(
             origins=origins,
diff --git a/nerfactory/cameras/rays.py b/nerfactory/cameras/rays.py
index 51f85e7aa..bb10b6e7c 100644
--- a/nerfactory/cameras/rays.py
+++ b/nerfactory/cameras/rays.py
@@ -113,7 +113,7 @@ class RaySamples(TensorDataclass):
         # mip-nerf version of transmittance calculation:
         transmittance = torch.cumsum(delta_density[..., :-1, :], dim=-2)
         transmittance = torch.cat(
-            [torch.zeros((*transmittance.shape[:1], 1, 1)).to(densities.device), transmittance], dim=-2
+            [torch.zeros((*transmittance.shape[:1], 1, 1), device=densities.device), transmittance], dim=-2
         )
         transmittance = torch.exp(-transmittance)  # [..., "num_samples"]
 
diff --git a/nerfactory/fields/compound_field.py b/nerfactory/fields/compound_field.py
index 9cb26a8ff..712eb15e8 100644
--- a/nerfactory/fields/compound_field.py
+++ b/nerfactory/fields/compound_field.py
@@ -153,11 +153,9 @@ class TCNNCompoundField(Field):
         return density, base_mlp_out
 
     def get_outputs(self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None):
-        # TODO: add valid_mask masking!
-        # tcnn requires directions in the range [0,1]
         if ray_samples.camera_indices is None:
             raise AttributeError("Camera indices are not provided.")
-        camera_indices = ray_samples.camera_indices.squeeze().to(ray_samples.frustums.origins.device)
+        camera_indices = ray_samples.camera_indices.squeeze()
         directions = get_normalized_directions(ray_samples.frustums.directions)
         directions_flat = directions.view(-1, 3)
         d = self.direction_encoding(directions_flat)
@@ -165,8 +163,8 @@ class TCNNCompoundField(Field):
         if self.training:
             embedded_appearance = self.embedding_appearance(camera_indices)
         else:
-            embedded_appearance = torch.zeros((*directions.shape[:-1], self.appearance_embedding_dim)).to(
-                directions.device
+            embedded_appearance = torch.zeros(
+                (*directions.shape[:-1], self.appearance_embedding_dim), device=directions.device
             )
 
         if density_embedding is None:
@@ -250,13 +248,14 @@ class TorchCompoundField(Field):
 
         if ray_samples.camera_indices is None:
             raise AttributeError("Camera indices are not provided.")
-        camera_indices = ray_samples.camera_indices.squeeze().to(ray_samples.frustums.origins.device)
+        camera_indices = ray_samples.camera_indices.squeeze()
         if self.training:
             embedded_appearance = self.embedding_appearance(camera_indices)
         else:
             embedded_appearance = torch.zeros(
-                (*ray_samples.frustums.directions.shape[:-1], self.appearance_embedding_dim)
-            ).to(ray_samples.frustums.directions.device)
+                (*ray_samples.frustums.directions.shape[:-1], self.appearance_embedding_dim),
+                device=ray_samples.frustums.directions.device,
+            )
 
         outputs = {}
         for field_head in self.field_heads:
diff --git a/nerfactory/models/modules/ray_generator.py b/nerfactory/models/modules/ray_generator.py
index c30f40f00..5ac58807f 100644
--- a/nerfactory/models/modules/ray_generator.py
+++ b/nerfactory/models/modules/ray_generator.py
@@ -50,5 +50,4 @@ class RayGenerator(nn.Module):
             camera_indices=c,
             coords=coords,
         )
-        ray_bundle.camera_indices = c[..., None]  # ["num_rays",1]
         return ray_bundle
diff --git a/nerfactory/utils/math.py b/nerfactory/utils/math.py
index 68d163cf7..334c6107c 100644
--- a/nerfactory/utils/math.py
+++ b/nerfactory/utils/math.py
@@ -29,7 +29,7 @@ def components_from_spherical_harmonics(levels: int, directions: TensorType[...,
         directions: Spherical hamonic coefficients
     """
     num_components = levels**2
-    components = torch.zeros((*directions.shape[:-1], num_components)).to(directions.device)
+    components = torch.zeros((*directions.shape[:-1], num_components), device=directions.device)
 
     assert 1 <= levels <= 5, f"SH levels must be in [1,4], got {levels}"
     assert directions.shape[-1] == 3, f"Direction input should have three dimensions. Got {directions.shape[-1]}"