remove barf

This commit is contained in:
Justin Kerr committed 2023-08-02 09:54:20 -07:00
1 parent 483a6c2eae
commit dc1b809211
3 files changed
+2 -64

No files matched your search

+1 -1
View File
@@ -121,7 +121,7 @@ method_configs["nerfacto"] = TrainerConfig(
},
},
viewer=ViewerConfig(num_rays_per_chunk=1 << 15),
vis="viewer_beta",
vis="viewer",
)
method_configs["nerfacto-big"] = TrainerConfig(
+1 -50
View File
@@ -65,8 +65,6 @@ class NerfactoField(Field):
use_pred_normals: whether to use predicted normals
use_average_appearance_embedding: whether to use average appearance embedding or zeros for inference
spatial_distortion: spatial distortion to apply to the scene
bundle_adjust: whether to use bundle adjustment
coarse_to_fine_iters: iterations (percentage of max iters) at which bundle adjustment is active
"""
aabb: Tensor
@@ -92,14 +90,11 @@ class NerfactoField(Field):
use_transient_embedding: bool = False,
use_semantics: bool = False,
num_semantic_classes: int = 100,
freq_warmup: int = 0,
pass_semantic_gradients: bool = False,
use_pred_normals: bool = False,
use_average_appearance_embedding: bool = False,
spatial_distortion: Optional[SpatialDistortion] = None,
implementation: Literal["tcnn", "torch"] = "tcnn",
bundle_adjust: bool = False,
coarse_to_fine_iters: Optional[Tuple[float, float]] = None,
) -> None:
super().__init__()
@@ -120,9 +115,6 @@ class NerfactoField(Field):
self.use_pred_normals = use_pred_normals
self.pass_semantic_gradients = pass_semantic_gradients
self.base_res = base_res
self.freq_warmup = freq_warmup
self.bundle_adjust = bundle_adjust
self.coarse_to_fine_iters = coarse_to_fine_iters
self.step = 0
self.direction_encoding = SHEncoding(
@@ -208,44 +200,6 @@ class NerfactoField(Field):
implementation=implementation,
)
def step_cb(self, step):
self.step = step
def get_freq_mask(self, B, N, D, device):
import numpy as np
if self.bundle_adjust:
# set masks for hash grid levels
self.max_iters: int = GLOBAL_BUFFER.get("max_iter")
assert self.coarse_to_fine_iters is not None
start, end = self.coarse_to_fine_iters
assert (
start >= 0 and end >= 0
), f"start and end iterations for bundle adjustment have to be positive, got start = {start} and end = {end}"
L = N
# From https://arxiv.org/pdf/2104.06405.pdf equation 14
alpha = ((self.step / self.max_iters) - start) / (end - start) * L
k = torch.arange(L, dtype=torch.float32, device=device)
mask_vals = (1 - (alpha - k).clamp_(min=0, max=1).mul_(np.pi).cos_()) / 2
else:
ts = np.linspace(self.step * N, self.step, N)
mask_vals = np.interp(ts, [0, self.freq_warmup], [0, 1])
mask_vals = torch.from_numpy(mask_vals)
m = mask_vals[None, ..., None].to(device)
return m.repeat((B, 1, D))
def mask_freqs(self, hashgrid_outputs):
"""
hashgrid_outputs is a B x N*D array of the hashgrid values
"""
# N is number of levels (usually 16)
# D is feature dim (usually 2)
N = self.mlp_base_grid.num_levels
D = self.mlp_base_grid.features_per_level
hashgrid_outputs = hashgrid_outputs.view(-1, N, D)
masked = self.get_freq_mask(hashgrid_outputs.shape[0], N, D, hashgrid_outputs.device) * hashgrid_outputs
return masked.view(masked.shape[0], -1)
def get_density(self, ray_samples: RaySamples) -> Tuple[Tensor, Tensor]:
"""Computes and returns the densities."""
if self.spatial_distortion is not None:
@@ -261,10 +215,7 @@ class NerfactoField(Field):
if not self._sample_locations.requires_grad:
self._sample_locations.requires_grad = True
positions_flat = positions.view(-1, 3)
hashgrid_vecs = self.mlp_base_grid(positions_flat)
hashgrid_vecs = self.mask_freqs(hashgrid_vecs)
h = self.mlp_base_mlp(hashgrid_vecs).view(*ray_samples.frustums.shape, -1)
# h = self.mlp_base(positions_flat).view(*ray_samples.frustums.shape, -1)
h = self.mlp_base(positions_flat).view(*ray_samples.frustums.shape, -1)
density_before_activation, base_mlp_out = torch.split(h, [1, self.geo_feat_dim], dim=-1)
self._density_before_activation = density_before_activation
-13
View File
@@ -131,11 +131,6 @@ class NerfactoModelConfig(ModelConfig):
"""Dimension of the appearance embedding."""
camera_optimizer: CameraOptimizerConfig = CameraOptimizerConfig(mode="SO3xR3")
"""Config of the camera optimizer to use"""
bundle_adjust: bool = False
"""Whether to bundle adjust (BARF)"""
coarse_to_fine_iters: Optional[Tuple[float, float]] = (0.0, 0.1)
"""Iterations (as a percentage of total iterations) at which coarse to fine hash grid optimization starts and ends.
Linear interpolation between (start, end) and full activation of hash grid from end onwards."""
class NerfactoModel(Model):
@@ -171,8 +166,6 @@ class NerfactoModel(Model):
use_average_appearance_embedding=self.config.use_average_appearance_embedding,
appearance_embedding_dim=self.config.appearance_embed_dim,
implementation=self.config.implementation,
bundle_adjust=self.config.bundle_adjust,
coarse_to_fine_iters=self.config.coarse_to_fine_iters,
)
self.camera_optimizer: CameraOptimizer = self.config.camera_optimizer.setup(
@@ -287,12 +280,6 @@ class NerfactoModel(Model):
func=self.proposal_sampler.step_cb,
)
)
callbacks.append(
TrainingCallback(
where_to_run=[TrainingCallbackLocation.AFTER_TRAIN_ITERATION],
update_every_num_iters=1,
func=self.field.step_cb,
)
)
return callbacks