tink/nerfstudio
4
0
Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-11 02:24:38 +08:00
Code

Use torch.where for faster contract (#1132)

Browse Source 
use where for faster contract
This commit is contained in:
Penn 2022-12-15 13:24:31 -08:00 committed by GitHub
parent d51718ec2b
commit bc690632bb
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 2 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
nerfstudio/field_components/spatial_distortions.py
View File

@ -65,12 +65,8 @@ class SceneContraction(SpatialDistortion):
def forward(self, positions):
def contract(x):
mag = torch.linalg.norm(x, ord=self.order, dim=-1)
mask = mag >= 1
x_new = x.clone()
x_new[mask] = (2 - (1 / mag[mask][..., None])) * (x[mask] / mag[mask][..., None])
return x_new
mag = torch.linalg.norm(x, ord=self.order, dim=-1)[..., None]
return torch.where(mag < 1, x, (2 - (1 / mag)) * (x / mag))
if isinstance(positions, Gaussians):
means = contract(positions.mean.clone())