tink/nerfstudio
4
0
Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-11 10:34:29 +08:00
Code

Speedup eval and training with better data management (#492)

Browse Source 
move to tensors to gpu earlier
This commit is contained in:
Matthew Tancik 2022-09-19 18:36:12 -07:00 committed by GitHub
parent 96e1f2a698
commit a43776d76d
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
5 changed files with 13 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
nerfactory/cameras/cameras.py
View File

@ -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,

2
nerfactory/cameras/rays.py
View File

@ -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"]

15
nerfactory/fields/compound_field.py
View File

@ -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:

1
nerfactory/models/modules/ray_generator.py
View File

@ -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

2
nerfactory/utils/math.py
View File

@ -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]}"