import math from dataclasses import dataclass from typing import Literal, Optional, Any import torch import torch.nn.functional as F import torch.utils.checkpoint import torchvision.transforms as T from einops import rearrange from torch import nn, Tensor from optgs.dataset.data_types import BatchedViews from optgs.geometry.projection import project, sample_image_grid from optgs.misc.general_utils import SkipBatchException from optgs.misc.io import FrequencyScheduler from optgs.model.decoder.decoder import Decoder from optgs.model.backbones.layer import ResNetFeatureWarpper from optgs.model.types import Gaussians from optgs.scene_trainer.common.gaussian_adapter import build_covariance from optgs.scene_trainer.initializer import InitializerCfg, InitializerColmapCfg, InitializerEdgsCfg, \ InitializerRandomCfg, InitializerPointcloudCfg from optgs.scene_trainer.initializer import InitializerPlyCfg from optgs.scene_trainer.initializer.initializer_resplat import ResplatInitializerCfg from optgs.scene_trainer.optimizer.optimizer import OptimizerInput, LearnedOptimizer, OptimizerOutput, OptimizerState, \ OptimizerPreviousOutput, OptimizerCfg from optgs.scene_trainer.optimizer.optimizer_utils import Number3DGSCfg, Bool3DGSCfg from optgs.scene_trainer.optimizer.optimizer_utils import unpack_gaussians, \ get_visibility_contribution_from_gaussian_obj try: from optgs.model.backbones.point_transformer.layer import (PlainPointTransformer, PointLinearWrapper, MultViewLowresAttn) except: pass from optgs.scene_trainer.optimizer.layer import CustomGroupNorm, AdamInputSmoothing, SlicedG3RNorm, AdamState from optgs.scene_trainer.initializer.initializer import InitializerOutput from optgs.scene_trainer.optimizer.time_embed import get_embedder, TimeEncodingWrapper from optgs.loss.loss_depth_smooth import get_smooth_loss from optgs.scene_trainer.optimizer.optimizer_utils import ( inner_loss_for_input_gradients, chunk_index_iter, split_grads, get_gaussian_param_slices, get_gaussian_param_sizes, pack_gaussians, ) _IMAGENET_NORM = T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) # Gaussian parameter groups the delta head predicts, paired with their key in # get_gaussian_param_sizes. Order matches pack_gaussians / split_delta_gaussians. _GAUSSIAN_GROUPS = [ ("means", "means"), ("scales", "scales"), ("rotations", "quats"), ("opacities", "opacities"), ("shs", "shs"), ] @dataclass class KnnBasedOptimizerCfg(OptimizerCfg): name: Literal["knn_based", "resplat_v1", "resplat_v2", "clogs", "l2s"] # TODO (release) remove clogs # iterative refine no_render_error: bool input_error_shallow_resnet_feature: bool input_error_resnet_feature_layers: int num_basic_blocks: int num_blocks: int concat_init_state: bool # always concat init state during updates replace_init_state: bool # always use the init state during updates state_channels: int block_rmsnorm: bool block_layernorm: bool pt_qk_norm: bool norm_pt_block: bool delta_gaussian_multiple: int # predict more gaussian residuals based on the previous gaussian center residual_init_state: bool # add residual connection in the prediction head to the initial state clamp_max_scale: float clamp_min_scale: float | int clamp_min_raw_scales: float | int clamp_max_raw_scales: float | int clamp_min_raw_opacities: float | int clamp_max_raw_opacities: float | int clamp_min_sh0: float | int clamp_max_sh0: float | int clamp_min_shs: float | int clamp_max_shs: float | int clamp_shs_soft: bool # TODO (release): remove gaussian_head_multiple (kept for next-project research; unused here, always 1) gaussian_head_multiple: int # use multiple non-weight sharing heads to predict multiple gaussians residual_grad_scale: float | int input_gradient_with_ssim_loss: bool attn_proj_channels: int | None no_knn_attn: bool no_tran_block_norm: bool tran_block_act: str | None multi_gaussian_scale_smaller: bool condition_pt_feature: bool same_num_points: bool # when init_gaussian_multiple > 1, refine directly works on it instead of subsampling points knn_samples: int max_active_gaussians: int # KNN use_fused_attn: bool knn_idx_update_every: int # point transformer pt_heads: int # inputs input_alpha: bool input_depth: bool input_depth_smooth_error: bool # input error input_error: bool # render error as input to the refine head input_error_rgb_no_shuffle: bool # sample single pixel instead of pixel unshuffling input_error_add_rgb_feature: bool # resnet input_error_resnet_feature: bool input_error_cache_resnet_feature: bool input_error_no_freeze_resnet_feature: bool # number of views for render error input_error_num_views: int input_error_additional_cross_attn: bool input_error_num_intermediate_views: int # render error with remaining context views input_error_remain_context: bool input_error_merge_remain_context: bool input_error_warp_remain_context: bool input_error_random_num_remain_context: bool input_error_num_remain_context_test: int # render error mv attn input_error_mv_attn: bool input_error_mv_attn_blocks: int # refine global attention with_mv_attn: bool with_mv_attn_lowres: bool no_mv_attn: bool # remove only the attn mv_attn_conv_with_norm: bool # unet-attn conv with norm mv_shuffle_attn: bool # use pixel shuffle to save computation instead of unet mv_attn_with_pos_enc: bool shuffle_attn_no_norm: bool mv_unimatch_attn: bool # input gradients input_gradient: bool input_gradient_log: bool input_gradient_log_clip_deltas: float | int input_gradient_scale: float | int input_gradient_same_loss: bool # use the same loss as the gaussian update input_gradient_loss_reduction: str scale_residual_grads: bool # time encoding use_time_encoding: bool time_encoding_max_steps: int train_global_only: bool # random size refine # update more for low resolution, less for high random_step_with_size: bool # amp use_amp: bool pt_head_amp: bool pt_update_amp: bool use_checkpointing: bool recurrent_use_checkpointing: bool # Normalizing input input_gradient_normalize: bool input_gradient_normalize_type: str input_normalize_state: bool input_normalize_gaussians: bool # State scaling predict_state_scale: bool predict_state_scale_norm: bool # whether to normalize the state before scaling # Use optimizer without condition features init_state_wo_features: bool init_state_type: Literal["random", "constant"] init_state_scale: float | int opt_scales_before_act: bool # optimize scale before activation (raw -> exp -> scale -> log -> raw) # Preprocessing the init gaussians scale_initial_opacities: float | int # TODO (release): remove experimental delta override cluster experimental_run: bool experimental_update: Bool3DGSCfg experimental_use_grads: bool experimental_use_norm_grads: Bool3DGSCfg experimental_lr: Number3DGSCfg # Deactivate gaussians # TODO (release): wire in or remove (only used in experimental_get_visible_gaussian_mask) local_prune_zero_radii: bool local_prune_low_weights: bool local_prune_low_weights_thresh: float | int update_only_nonzero_grad: bool # update learn residual state residual_state: bool # Update head delta_head_layer_num: int delta_head_concat_img: bool delta_head_act: str | None # delta_head activation to predict the deltas delta_head_final_act: str | None # final activation in the delta_head delta_head_hidden_dim_matches: str # rebuttal or submission version delta_head_scalar_scale: bool # predict deltas as scalar * delta / norm(delta) delta_head_scalar_scale_act: str # activation for the scalar scale output # Per-parameter-group update head (Feature A) delta_head_per_param_heads: bool # separate heads per param group, each with own normalize+scale delta_head_per_param_hidden_dim: int # hidden dim for per-param heads (SH head gets 2x) # Per-parameter scalar scales (Feature B) — requires delta_head_scalar_scale=true delta_head_per_param_scales: bool # per-group scalar scales instead of one global scalar # Config from initializer sh_d: int | None init_sh_d: int | None = None # Fow initialization from feed forward, gaussians are aligned with pixels. init_gaussian_multiple: int | None = None latent_downsample: int | None = None @property def scales_clamp_lims(self) -> tuple: """Scale clamp limits in the space scales are refined in. With opt_scales_before_act scales are refined in log space, so the raw (log-space) limits apply; otherwise the activation-space limits apply. """ if self.opt_scales_before_act: return self.clamp_min_raw_scales, self.clamp_max_raw_scales return self.clamp_min_scale, self.clamp_max_scale def update(self, initializer_cfg: InitializerCfg): """ Update the optimizer config based on the initializer config""" # General settings self.init_sh_d = initializer_cfg.get_sh_d() if self.sh_d is None: # get sh_d from initializer if not set self.sh_d = initializer_cfg.get_sh_d() # Settings specific to DepthSplat initializer if isinstance(initializer_cfg, ResplatInitializerCfg): self.latent_downsample = initializer_cfg.latent_downsample self.init_gaussian_multiple = initializer_cfg.init_gaussian_multiple # update proj channels if self.condition_pt_feature: self.condition_channels = initializer_cfg.gaussian_regressor_channels else: self.condition_channels = initializer_cfg.get_pt_in_channels() # Settings specific to Colmap initializer elif isinstance(initializer_cfg, (InitializerPlyCfg, InitializerColmapCfg, InitializerEdgsCfg, InitializerRandomCfg, InitializerPointcloudCfg)): # Since pixels and gaussians are not aligned, we cannot use pixel attributes assert not self.input_error, "The error calculation assumes per pixel gaussians" assert not self.delta_head_concat_img assert not self.input_alpha assert self.init_state_wo_features, "Colmap initializer does not have point features, init_state_wo_features must be set to True" self.init_gaussian_multiple = 1 self.latent_downsample = 1 else: raise ValueError(f"Unsupported initializer config type: {type(initializer_cfg)}") @dataclass class KnnBasedOptimizerState(OptimizerState): """OptimizerState subclass for KNN-based optimizers. Holds the per-Gaussian learned state vectors and adds ADC mutation methods (called by adc/base.py helpers) that keep them in sync with clone/split/prune. """ state: torch.Tensor | None = None init_state: torch.Tensor | None = None # scene-start state; None when concat/replace/residual flags are all off adam_state: AdamState | None = None def clone(self, clone_mask: torch.Tensor, zero_t: bool) -> None: if self.state is not None: cloned = torch.zeros_like(self.state[clone_mask]) if zero_t else self.state[clone_mask] self.state = torch.cat([self.state, cloned], dim=0) if self.init_state is not None: # NOTE: cloned Gaussians inherit the parent's init_state (or zeros when zero_t). # Whether a clone should keep, reset, or re-derive init_state is an open design question. cloned_init = torch.zeros_like(self.init_state[clone_mask]) if zero_t else self.init_state[clone_mask] self.init_state = torch.cat([self.init_state, cloned_init], dim=0) def split(self, split_mask: torch.Tensor, num_splits: int, zero_t: bool) -> None: if self.state is not None: chunks = self.state[split_mask].chunk(num_splits, dim=0) new_states = [torch.zeros_like(c) if zero_t else c for c in chunks] self.state = torch.cat([self.state, *new_states], dim=0) if self.init_state is not None: # NOTE: split Gaussians inherit the parent's init_state (or zeros when zero_t). # Whether a split should keep, reset, or re-derive init_state is an open design question. init_chunks = self.init_state[split_mask].chunk(num_splits, dim=0) new_init = [torch.zeros_like(c) if zero_t else c for c in init_chunks] self.init_state = torch.cat([self.init_state, *new_init], dim=0) def replace(self, from_indices: torch.Tensor, dest_indices: torch.Tensor, zero_t: bool) -> None: if self.state is not None: if zero_t: self.state[dest_indices] = 0.0 else: self.state[dest_indices] = self.state[from_indices] if self.init_state is not None: if zero_t: self.init_state[dest_indices] = 0.0 else: self.init_state[dest_indices] = self.init_state[from_indices] def prune(self, prune_mask: torch.Tensor) -> None: if self.state is not None: self.state = self.state[~prune_mask] if self.init_state is not None: self.init_state = self.init_state[~prune_mask] def add(self, num_new: int) -> None: if num_new <= 0: return if self.state is not None: zeros = torch.zeros((num_new, *self.state.shape[1:]), device=self.state.device, dtype=self.state.dtype) self.state = torch.cat([self.state, zeros], dim=0) # NOTE: newly added Gaussians get a zero init_state (they have no parent to inherit from). # Whether they should instead re-derive init_state is an open design question. if self.init_state is not None: zeros_init = torch.zeros((num_new, *self.init_state.shape[1:]), device=self.init_state.device, dtype=self.init_state.dtype) self.init_state = torch.cat([self.init_state, zeros_init], dim=0) class Abs(nn.Module): def forward(self, x): return torch.abs(x) def get_activation_cls(activation: Optional[str] = None): if activation in ['none', None, 'identity']: return nn.Identity elif activation == 'tanh': return nn.Tanh elif activation == "gelu": return nn.GELU elif activation == 'sigmoid': return nn.Sigmoid elif activation == 'relu': return nn.ReLU elif activation == "softplus": return nn.Softplus elif activation == "abs": return Abs else: raise ValueError(f"Unsupported activation: {activation}") class KnnBasedOptimizer(LearnedOptimizer[KnnBasedOptimizerCfg]): """Learned optimizer that refines a scene's Gaussians with a point-transformer (the core method). Replaces hand-tuned gradient descent (Adam + heuristics) with a network that predicts, at each iteration, how every Gaussian should change. The scene is treated as a point cloud (one point per Gaussian) and the network reasons over local neighborhoods via KNN attention, so an update to one Gaussian is informed by its spatial neighbors. Each Gaussian carries a learned latent state vector alongside its parameters (means, scales, rotations, opacity, SH). The state is what the network reads and writes across iterations; the Gaussian parameters are updated by adding predicted deltas. One optimization step (`_apply_step`, run per iteration by `_forward_impl`): 1. Input signal (`prepare_input_signal`): render the context views and measure how well they match the ground-truth inputs. The signal is the per-Gaussian rendering gradients (`_calc_input_gradients` — a backprop of the reconstruction loss, used only as a feature, never to update weights) and/or the render error in feature space. This tells the network "what is wrong" with the current Gaussians. 2. Update (point transformer): KNN attention over the point cloud mixes each Gaussian's state with its neighbors', conditioned on the input signal, producing an updated per-Gaussian state. 3. Delta head: maps the updated state to per-parameter deltas, which are added to the Gaussians (with clamping/activation handling per parameter group). 4. Optional adaptive density control (clone/split/prune), keeping the state in sync with the changing Gaussian count. `optimizer_preprocessing` runs once at scene start to seed the latent state from the initializer's Gaussians. This class implements both ReSplat and Learn2Splat through different settings in `KnnBasedOptimizerCfg`: ReSplat (`resplat_v1.yaml`) drives the update with feature-space render error and multi-view attention, while Learn2Splat (`learn2splat.yaml`) drives it with normalized per-Gaussian gradients. The many `_build_*` / `_calc_*` helpers wire up these settings. """ OPTIMIZER_NAME = "knn_based" OPTIMIZER_NAME_ALIASES: tuple[str, ...] = () def __init__(self, cfg: KnnBasedOptimizerCfg, save_every: Optional[FrequencyScheduler] = None) -> None: valid = {self.OPTIMIZER_NAME, *self.OPTIMIZER_NAME_ALIASES} assert cfg.name in valid, f"Expected optimizer name {valid}, got {cfg.name}" super().__init__(cfg, save_every) if self.cfg.residual_state: assert not self.cfg.residual_init_state # State channel self.state_channels = self.cfg.state_channels # time embedder if self.cfg.use_time_encoding: self.time_encoder_fn, self.time_embedding_dim = get_embedder(multires=6) else: self.time_encoder_fn = None self.time_embedding_dim = 0 # state_proj if not self.cfg.init_state_wo_features: self.state_proj = nn.Conv2d(self.cfg.condition_channels, self.state_channels, 1) channels, in_channels, update_gaussian_param_num, out_channels, error_features_channels = ( self._define_channels()) self.error_features_channels = error_features_channels self.gaussian_param_num = out_channels if self.cfg.input_error: self.error_feature_extractor = self.get_input_error_feature_extractor() if self.cfg.input_error_add_rgb_feature: if self.cfg.init_gaussian_multiple == 4: # re10k self.rgb_error_proj = nn.Sequential( nn.Linear(3, error_features_channels), nn.LayerNorm(error_features_channels) ) else: self.rgb_error_proj = nn.Sequential( nn.Linear(3 * self.cfg.latent_downsample ** 2, error_features_channels), nn.LayerNorm(error_features_channels) ) self.input_norm = self._build_input_norm(in_channels) self.point_transformer = self._build_point_transformer(channels, in_channels) # predict multiple gaussians out_channels = out_channels * self.cfg.delta_gaussian_multiple if not self.cfg.same_num_points: out_channels = out_channels * self.cfg.init_gaussian_multiple # make sure the input size of the gaussian head is updated accordingly if self.cfg.use_time_encoding: channels += self.time_embedding_dim # Compute per-param group dims (needed by per_param_heads and per_param_scales) if self.cfg.delta_head_per_param_heads or self.cfg.delta_head_per_param_scales: self._per_param_group_dims = self._compute_per_param_group_dims(out_channels) # Scaling state for update head if self.cfg.predict_state_scale: self.state_scale_head = self._build_state_scale_head(in_channels) self.delta_head = self._build_delta_head(in_channels, channels, out_channels) # multiple gaussian heads to predict multiple gaussians if self.cfg.gaussian_head_multiple > 1: self.extra_gaussian_heads = self._build_extra_gaussian_heads(channels, out_channels) # Define error calculation # add global attention to the render error if self.cfg.input_error and self.cfg.input_error_mv_attn: assert self.cfg.input_error_resnet_feature self.error_mv_attn = nn.ModuleList([ MultViewLowresAttn(error_features_channels) for _ in range(self.cfg.input_error_mv_attn_blocks) ]) self.param_slices = get_gaussian_param_slices(self.cfg.sh_d) def _reset_knn_caches(self) -> None: """Invalidate cached KNN indices on all point-transformer sub-modules. Must be called whenever the number of Gaussians changes (e.g. after add_new) so the next forward recomputes KNN from scratch instead of using stale indices that index out-of-bounds into the grown point cloud. """ for module in self.modules(): if hasattr(module, "cache_knn_idx"): module.cache_knn_idx = None @property def adc_object_dict_to_adjust(self): if self.cfg.any_adc: object_dict: dict[str, Any] = {"optimizer_state": None} # For ADC if self.cfg.input_gradient_normalize and self.cfg.input_gradient_normalize_type == "adam": object_dict.update(self.input_norm.subgroups_view(self.param_slices)) else: return None return object_dict def _frozen_param_groups(self) -> set: """Gaussian parameter groups the delta head does not predict (frozen). Any combination of freeze_mean/scale/rotation/opacity is supported. The knn optimizer treats SH as a single group, so SH is frozen only when both freeze_sh0 and freeze_shN are set — a partial SH freeze is not expressible here. """ frozen = set() if self.cfg.freeze_mean: frozen.add("means") if self.cfg.freeze_scale: frozen.add("scales") if self.cfg.freeze_rotation: frozen.add("rotations") if self.cfg.freeze_opacity: frozen.add("opacities") assert self.cfg.freeze_sh0 == self.cfg.freeze_shN, \ "knn optimizer treats SH as one group; freeze_sh0 and freeze_shN must match" if self.cfg.freeze_sh0 and self.cfg.freeze_shN: frozen.add("shs") return frozen def _compute_per_param_group_dims(self, out_channels): """Compute per-parameter-group output dimensions from total out_channels. Returns a dict {group_name: dim} in the same order as split_delta_gaussians, omitting frozen groups, with the delta/init-gaussian multipliers folded in. """ p = get_gaussian_param_sizes(self.cfg.sh_d) excluded = self._frozen_param_groups() multiplier = self.cfg.delta_gaussian_multiple if not self.cfg.same_num_points: multiplier *= self.cfg.init_gaussian_multiple group_dims = {name: p[key] * multiplier for name, key in _GAUSSIAN_GROUPS if name not in excluded} assert sum(group_dims.values()) == out_channels, ( f"Per-param group dims {dict(group_dims)} sum={sum(group_dims.values())} != out_channels={out_channels}" ) return group_dims def _build_per_param_heads(self, channels, out_channels): """Build per-parameter-group heads (Feature A). Each head: Linear(channels, hidden) -> act -> Linear(hidden, dim+1) The +1 is a per-group scalar scale. Each head independently normalizes + scales. """ act_cls = get_activation_cls(self.cfg.delta_head_act) hidden_dim = self.cfg.delta_head_per_param_hidden_dim # Set up scale activation (shared across all per-param heads) scale_act_name = self.cfg.delta_head_scalar_scale_act init_bias_map = {'softplus': -1, 'relu': 1e-8, 'abs': 1e-8} if scale_act_name not in init_bias_map: raise ValueError(f"Unsupported scalar_scale_act: {scale_act_name}") act_class = get_activation_cls(scale_act_name) self.scale_act = act_class(beta=1) if scale_act_name == 'softplus' else act_class() heads = nn.ModuleDict() for name, dim in self._per_param_group_dims.items(): # SH head gets 2x hidden dim (more outputs to predict) h = hidden_dim * 2 if name == "shs" else hidden_dim layers = [nn.Linear(channels, h), act_cls()] for _ in range(self.cfg.delta_head_layer_num - 2): layers += [nn.Linear(h, h), act_cls()] layers.append(nn.Linear(h, dim + 1)) # +1 for scalar scale head = nn.Sequential(*layers) # Zero-init last layer (deltas start at 0) nn.init.zeros_(head[-1].weight) nn.init.zeros_(head[-1].bias) # Init scale bias nn.init.constant_(head[-1].bias[-1], init_bias_map[scale_act_name]) heads[name] = head return heads def _build_delta_head(self, in_channels, channels, out_channels): update_head_activation_cls = get_activation_cls(self.cfg.delta_head_act) final_head_activation_cls = get_activation_cls(self.cfg.delta_head_final_act) # skip connection to the image color if self.cfg.delta_head_concat_img: channels += 3 * (self.cfg.latent_downsample ** 2) # Feature A: per-parameter-group heads (early return — builds ModuleDict instead of Sequential) if self.cfg.delta_head_per_param_heads: assert not self.cfg.delta_head_per_param_scales, "per_param_heads already includes per-group scales" return self._build_per_param_heads(channels, out_channels) if self.cfg.delta_head_scalar_scale: if self.cfg.delta_head_per_param_scales: # Feature B: one scalar scale per parameter group out_channels = out_channels + len(self._per_param_group_dims) else: out_channels = out_channels + 1 # Determine hidden layer size # TODO (release): update_head_hidden_dim_source should be "output" (out_channels). # Using "input" currently as default to reproduce rebuttal results. if self.cfg.delta_head_hidden_dim_matches == "input": hidden_dim = channels # cvpr rebuttal version (current default) else: hidden_dim = out_channels # cvpr submitted version # Build update head layers_list = [ nn.Linear(channels, hidden_dim), update_head_activation_cls() ] for i in range(self.cfg.delta_head_layer_num - 2): layers_list += [ nn.Linear(hidden_dim, hidden_dim), update_head_activation_cls(), ] layers_list += [ nn.Linear(hidden_dim, out_channels), final_head_activation_cls() ] delta_head = nn.Sequential(*layers_list) # init the delta as 0 nn.init.zeros_(delta_head[-2].weight) if final_head_activation_cls == torch.nn.Sigmoid: desired_init_delta = 0.005 bias = math.log(desired_init_delta / (1 - desired_init_delta)) # ~= -4.6 nn.init.constant_(delta_head[-2].bias, bias) else: nn.init.zeros_(delta_head[-2].bias) # Scalar scale output if self.cfg.delta_head_scalar_scale: # Set the initial scale to very low number, to get the gradients flow init_bias_map = { 'softplus': -1, 'relu': 1e-8, 'abs': 1e-8, } act_name = self.cfg.delta_head_scalar_scale_act if act_name not in init_bias_map: raise ValueError(f"Unsupported scalar_scale_out_act: {act_name}") # Initialize bias for scale output(s) if self.cfg.delta_head_per_param_scales: num_groups = len(self._per_param_group_dims) for i in range(num_groups): nn.init.constant_(delta_head[-2].bias[-(num_groups - i)], init_bias_map[act_name]) else: nn.init.constant_(delta_head[-2].bias[-1], init_bias_map[act_name]) # Create activation act_class = get_activation_cls(act_name) self.scale_act = act_class(beta=1) if act_name == 'softplus' else act_class() return delta_head def _build_extra_gaussian_heads(self, channels, out_channels): update_head_activation = get_activation_cls(self.cfg.delta_head_act) final_head_activation = get_activation_cls(self.cfg.delta_head_final_act) extra_gaussian_heads = nn.ModuleList() for i in range(self.cfg.gaussian_head_multiple - 1): extra_gaussian_heads.append( nn.Sequential( nn.Linear(channels, channels), update_head_activation(), nn.Linear(channels, out_channels), final_head_activation() ) ) # init the delta as 0 nn.init.zeros_(extra_gaussian_heads[i][-2].weight) nn.init.zeros_(extra_gaussian_heads[i][-2].bias) return extra_gaussian_heads def _build_input_norm(self, in_channels): if self.cfg.input_gradient_normalize: assert self.cfg.input_gradient, "for now we only normalize when using gradient as input" if self.cfg.input_gradient_normalize_type == 'layer': return nn.LayerNorm(in_channels) elif self.cfg.input_gradient_normalize_type == 'group': return CustomGroupNorm([self.gaussian_param_num, self.state_channels, self.gaussian_param_num]) elif self.cfg.input_gradient_normalize_type == 'batch': return nn.BatchNorm1d(in_channels, affine=False) elif self.cfg.input_gradient_normalize_type == 'g3r': return SlicedG3RNorm(in_channels, slice(-self.gaussian_param_num, None)) elif self.cfg.input_gradient_normalize_type == 'adam': assert not self.cfg.input_gradient_log and self.cfg.input_gradient_scale == 1 return AdamInputSmoothing(input_slice=slice(-self.gaussian_param_num, None)) else: raise ValueError(f"normalization type not supported {self.cfg.input_gradient_normalize_type}") else: return nn.Identity() def _build_point_transformer(self, channels, in_channels): point_transformer = nn.Sequential( PointLinearWrapper(in_channels, channels), PlainPointTransformer(channels, self.cfg.knn_samples, num_blocks=self.cfg.num_basic_blocks, qk_norm=self.cfg.pt_qk_norm, norm_pt_block=self.cfg.norm_pt_block, num_heads=self.cfg.pt_heads, no_rpe=True, no_attn=self.cfg.no_knn_attn, no_norm=self.cfg.no_tran_block_norm, act=self.cfg.tran_block_act, attn_proj_channels=self.cfg.attn_proj_channels, with_mv_attn=self.cfg.with_mv_attn, with_mv_attn_lowres=self.cfg.with_mv_attn_lowres, no_mv_attn=self.cfg.no_mv_attn, conv_with_norm=self.cfg.mv_attn_conv_with_norm, mv_shuffle_attn=self.cfg.mv_shuffle_attn, with_pos_enc=self.cfg.mv_attn_with_pos_enc, shuffle_attn_no_norm=self.cfg.shuffle_attn_no_norm, mv_unimatch_attn=self.cfg.mv_unimatch_attn, use_checkpointing=self.cfg.use_checkpointing, use_fused_attn=self.cfg.use_fused_attn, knn_idx_update_every=self.cfg.knn_idx_update_every ) ) # Init normalization layers if self.cfg.input_normalize_state: for block in point_transformer[1].blocks: nn.init.zeros_(block.norm1.bias) nn.init.zeros_(block.norm2.bias) nn.init.ones_(block.norm1.weight) nn.init.ones_(block.norm2.weight) return point_transformer @staticmethod def _build_state_scale_head(in_channels): state_scale_head = nn.Sequential( nn.Linear(in_channels, in_channels // 2), nn.ReLU(), nn.Linear(in_channels // 2, 1), nn.ReLU() ) # Init the scale to 1 # nn.init.zeros_(state_scale_head[-2].weight) nn.init.ones_(state_scale_head[-2].bias) return state_scale_head def _define_channels(self): # The optimizer refines a fixed per-Gaussian representation, independent of the initializer: # means/3D position(3), scales(3), rotation(4), opacity(1), SH(3*sh_d). Whatever position # encoding the initializer used (e.g. resplat's 2D pixel offset) has already been unprojected # to a 3D point by the time gaussians reach the optimizer. Sized from the same per-group sizes # the freeze logic uses below, so the total and the per-group counts cannot drift. param_sizes = get_gaussian_param_sizes(self.cfg.sh_d) gaussian_param_num = sum(param_sizes.values()) # Get error channels if self.cfg.input_error: error_channels, error_feature_channels = self.define_error_channels() else: error_channels, error_feature_channels = 0, 0 # Get gradient channels if self.cfg.input_gradient: gradient_channels = gaussian_param_num * self.cfg.init_gaussian_multiple else: gradient_channels = 0 # final input channels input_signal_channels = gradient_channels + error_channels if self.cfg.same_num_points: in_channels = (gaussian_param_num + self.state_channels + input_signal_channels) else: in_channels = (gaussian_param_num * self.cfg.init_gaussian_multiple + self.state_channels + input_signal_channels) if self.cfg.concat_init_state: in_channels += self.state_channels out_channels = gaussian_param_num for name, key in _GAUSSIAN_GROUPS: if name in self._frozen_param_groups(): out_channels -= param_sizes[key] channels = self.state_channels if self.cfg.input_alpha: # pixel shuffle the alpha channel to the latent resolution in_channels += self.cfg.latent_downsample ** 2 # alpha if self.cfg.input_depth or self.cfg.input_depth_smooth_error: # pixel shuffle the depth channel to the latent resolution in_channels += self.cfg.latent_downsample ** 2 # depth return channels, in_channels, gaussian_param_num, out_channels, error_feature_channels def define_error_channels(self): if self.cfg.no_render_error: error_channels = 0 else: if self.cfg.input_error_rgb_no_shuffle: error_channels = 3 else: error_channels = 3 * self.cfg.latent_downsample ** 2 if self.cfg.input_error_resnet_feature: # 3 scales: 1/2, 1/4, 1/8, channels: 64, 64, 128 if self.cfg.input_error_resnet_feature_layers in (18, 34): error_feature_channels = 64 + 64 if self.cfg.input_error_shallow_resnet_feature else 64 + 64 + 128 elif self.cfg.input_error_resnet_feature_layers == 50: error_feature_channels = 64 + 256 + 512 else: raise NotImplementedError error_channels = error_feature_channels else: error_feature_channels = 256 return error_channels, error_feature_channels def optimizer_preprocessing(self, optimizer_input: OptimizerInput, from_init: bool) -> None: """Prepare the initializer's Gaussians and seed the per-Gaussian latent state at scene start. On a fresh scene (from_init) scales the initial opacities, pads/truncates SH to sh_d, and builds the latent state vector; on a buffer resume it only refreshes the state, leaving init_state intact. """ if self.cfg.input_error_remain_context or self.cfg.input_error_merge_remain_context: assert self.cfg.input_error_cache_resnet_feature # Image dimensions context = optimizer_input.context b, v, _, h, w = context["image"].shape # Prepare Gaussians if from_init: # Scale initial opacities (in normal scale) opacities = optimizer_input.prev_output.gaussians.opacities # post activation, in [0, 1] scaled_opacities = opacities * self.cfg.scale_initial_opacities # default to 1.0 optimizer_input.prev_output.gaussians.opacities = scaled_opacities # Process shs shs = optimizer_input.prev_output.gaussians.harmonics # [B, N, 3, init_sh_d] init_sh_d = shs.shape[-1] if init_sh_d != self.cfg.sh_d: if init_sh_d > self.cfg.sh_d: shs = shs[:, :, :, :self.cfg.sh_d] # truncate [B, N, 3, sh_d] else: pad = self.cfg.sh_d - init_sh_d shs = F.pad(shs, (0, pad), "constant", 0) optimizer_input.prev_output.gaussians.harmonics = shs # Prepare state # Gaussians dimensions n = optimizer_input.prev_output.gaussians.means.shape[1] vector_state = self.get_vector_state(b, v, n, optimizer_input, from_init) if from_init: # Set everything so that the optimizer isn't aware whether it's a new scene # Convert InitializerOutput to OptimizerPreviousOutput optimizer_input.prev_output = OptimizerPreviousOutput(gaussians=optimizer_input.prev_output.gaussians, state=KnnBasedOptimizerState()) optimizer_input.prev_output.state.state = vector_state # init_state captures the scene-start state used by some experiments; # only set it on a fresh scene so replay-buffer resumes preserve the original value. # Left None when none of the three init_state flags are active. if from_init and (self.cfg.concat_init_state or self.cfg.replace_init_state or self.cfg.residual_init_state): optimizer_input.prev_output.state.init_state = vector_state def _forward_impl( self, i: int, optimizer_input: OptimizerInput, optimizer_output: OptimizerOutput, full_context: BatchedViews, full_target: BatchedViews, **kwargs ) -> OptimizerOutput: """Run one optimization iteration: apply the update step, then ADC, then save renders/stats.""" with self.benchmarker.time("iter"): # Unpack iter_context: BatchedViews = optimizer_input.context target: BatchedViews = optimizer_input.target renderer: Decoder = optimizer_input.renderer b, v, _, h, w = iter_context["image"].shape assert b == 1, "Batch size > 1 not supported for post-processing" # Log number of gaussians self.benchmarker.record("gaussians", optimizer_input.prev_output.gaussians.means.shape[1]) # One optimization step res = self._apply_step( i, optimizer_input, optimizer_output ) updated_gaussians: Gaussians = res[0] state: Tensor = res[1] meta_for_adc: dict = res[2] updates: dict[str, Tensor] = res[3] grads_raw: Tensor | None = res[4] normalized_grads: Tensor | None = res[5] updated_scaled_state: Tensor | None = res[6] gaussians_sel: Tensor | None = res[7] # Log stats if grads_raw is not None: nonzero_grads = (grads_raw != 0).any(-1) # [B, G] assert nonzero_grads.shape[0] == 1 self.benchmarker.record("nonzero_grads", nonzero_grads[0].sum().item()) # Densification and Pruning if self.cfg.any_adc: n_before_adc = updated_gaussians.means.shape[1] # Prepare objects to adjust during ADC. # Write the current state back so ADC mutation methods see it. optimizer_input.prev_output.state.state = state object_dict = self.adc_object_dict_to_adjust object_dict["optimizer_state"] = optimizer_input.prev_output.state # Apply ADC — mutates state (and init_state if set) in place self.apply_adc( i=i, v=v, h=h, w=w, adc_state=optimizer_input.prev_output.state.adc_state, gaussians=updated_gaussians, meta=meta_for_adc, object_dict_to_adjust=object_dict ) # Read back the (possibly grown/pruned) state tensor state = optimizer_input.prev_output.state.state del object_dict["optimizer_state"] if self.cfg.input_gradient_normalize and self.cfg.input_gradient_normalize_type == "adam": self.input_norm.aggregate_from_subgroups(object_dict, self.param_slices) # If N changed (add_new grew the population), stale KNN caches in the # point transformer modules would index out-of-bounds on the next forward # pass → CUDA illegal memory access. Reset them so they are recomputed. if updated_gaussians.means.shape[1] != n_before_adc: self._reset_knn_caches() # Save updated gaussians and state optimizer_input.prev_output.gaussians = updated_gaussians optimizer_input.prev_output.state.state = state if self.cfg.input_gradient_normalize_type == "adam": optimizer_input.prev_output.state.adam_state = self.input_norm.get_state() if self.training: optimizer_output.gaussian_list.append(updated_gaussians) # Save per-iteration debug info (test time only) if not self.training and self.save_every(i + 1, tag="info"): self._save_info_stats(optimizer_output, updated_gaussians, updates, grads_raw, normalized_grads, updated_scaled_state, gaussians_sel, state) # Post-update context + target renders self._save_post_update_renders( i, optimizer_input, optimizer_output, updated_gaussians, full_context, full_target, ) # Optimizer output is being changed in place, but for clarity we return it return optimizer_output def _save_info_stats(self, optimizer_output, updated_gaussians, updates, grads_raw, normalized_grads, updated_scaled_state, gaussians_sel, state): """Prepare per-iteration debug stats (deltas, gradients, state norms) and save them. knn-specific prep: unpack shs into sh0s/shNs, split the flat gradient vectors into per-param dicts, and restore zeroed-out (excluded) Gaussians; then hand the already-CPU values to the shared Optimizer._append_info. """ # unpack shs into sh0s / shNs shs = updates.pop("shs") # [1, N, 3*sh_d] assert shs.shape[0] == 1, "Batch size > 1 not supported" shs = rearrange(shs.squeeze(0), "n (c x) -> n c x", c=3, x=self.cfg.sh_d) # [N, 3, sh_d] updates["sh0s"] = shs[..., 0:1] updates["shNs"] = shs[..., 1:] if self.cfg.sh_d > 1 else None info = {"deltas": {k: v.squeeze(0).cpu() if v is not None else None for k, v in updates.items()}} if grads_raw is not None: if gaussians_sel is not None: # Restore the zeroed-out (excluded) Gaussians so the logged tensors are full-size b, g_valid, d = grads_raw.shape g = state.shape[0] grads_raw_full = torch.zeros((b, g, d)) normalized_grads_full = torch.zeros((b, g, d)) grads_raw_full[:, gaussians_sel, :] = grads_raw.cpu() normalized_grads_full[:, gaussians_sel, :] = normalized_grads.cpu() grads_raw = grads_raw_full normalized_grads = normalized_grads_full info["grads"] = split_grads(grads_raw.cpu(), self.cfg) if normalized_grads is not None: info["normalized_grads"] = split_grads(normalized_grads.cpu(), self.cfg) assert info["grads"]["means"].shape == info["normalized_grads"]["means"].shape, \ f"Shape mismatch between grads and normalized_grads: " \ f"{info['grads']['means'].shape} vs {info['normalized_grads']['means'].shape}" else: info["normalized_grads"] = None if updated_scaled_state is not None: info["states_norms"] = torch.norm(updated_scaled_state, dim=-1).cpu() # [B, N] self._append_info(optimizer_output, updated_gaussians, **info) def _apply_step( self, i, optimizer_input: OptimizerInput, optimizer_output: OptimizerOutput ) -> tuple[Gaussians, Tensor, dict, dict[str, Tensor], Tensor | None, Tensor | None, Tensor | None, Tensor | None]: """Compute one update: input signal -> point-transformer state update -> delta head -> new Gaussians. Returns the updated Gaussians and latent state plus the intermediates ADC and the debug logging need. """ # Unpacking context = optimizer_input.context debug_dict = optimizer_input.debug_dict gaussians = optimizer_input.prev_output.gaussians # Gaussian object of [B, N, C] state = optimizer_input.prev_output.state.state # [N, C] init_state = optimizer_input.prev_output.state.init_state # [N, C] # Get input signal for the optimizer model (errors/gradients) with self.benchmarker.time("decoder"): input_signal, gaussian_grads_raw, gaussian_grads, grad_sign, context_render_output, means2d_grads = ( self.prepare_input_signal(context, gaussians, optimizer_input.renderer) ) # Preparing meta for ADC if means2d_grads is not None: means2d_grads = means2d_grads.detach() # [B, V, N, 2] meta_for_adc = { "visibility_filter": context_render_output.visibility_filter.detach(), # [B, V, N] "radii": context_render_output.radii.detach(), # [B, V, N, 1] "means_2d_grads": means2d_grads, # [B, V, N, 2] } # Some initializers pad the scene with extra, invisible (zero-opacity) Gaussians so that every # GPU optimizes the same number of Gaussians (ColmapInitializer, train_fixed_gaussians_num; the # true count is gaussians.nr_valid). The optimizer predicts an update for every Gaussian, so # updating one of these dummies would raise its opacity and turn it into a real Gaussian. Because # the dummies are invisible their gradient is zero, so updating only nonzero-gradient Gaussians # (update_only_nonzero_grad=True, done by filter_gaussians_for_optimizer_step) leaves them as they # are. Every training config does this; the check below rejects the combination we don't support. has_padding = 0 <= gaussians.nr_valid < gaussians.means.shape[1] if has_padding and not self.cfg.update_only_nonzero_grad: raise ValueError( f"update_only_nonzero_grad=False with padded init Gaussians (nr_valid=" f"{gaussians.nr_valid} < N={gaussians.means.shape[1]}): the invisible padding Gaussians " f"would be optimized and turn into real ones. Set update_only_nonzero_grad=True (it skips " f"them, since their gradient is zero), or use an initializer without padding." ) # Handle zero gradient gaussians # We either prune them, or exclude them from the input/output update if self.cfg.update_only_nonzero_grad and gaussian_grads is not None: gaussian_grads, gaussian_grads_raw, gaussians, grad_sign, init_state, input_signal, state = ( self.filter_gaussians_for_optimizer_step(gaussian_grads, gaussian_grads_raw, gaussians, grad_sign, init_state, input_signal, state) ) # During training, skip a scene whose active-Gaussian count is outside the workable # range: too many to fit in memory, or too few for the KNN neighborhood. # NOTE: over-cap scenes are discarded wholesale; subsampling down to # max_active_gaussians instead is a possible future alternative (would change # training dynamics, so it needs validation). active_gaussians_num = state.shape[0] if self.training: if active_gaussians_num > self.cfg.max_active_gaussians: print(f"Skipping batch at iteration {i} with {active_gaussians_num} active gaussians.") raise SkipBatchException() if active_gaussians_num < self.cfg.knn_samples: print( f"Skipping batch at iteration {i} with only {active_gaussians_num} active gaussians (need >= {self.cfg.knn_samples}).") raise SkipBatchException() # Unpack Gaussians means, scales, rotations_unnorm, opacities_raw, shs = unpack_gaussians( gaussians, scales_log=self.cfg.opt_scales_before_act, opacities_logit=True, opacities_unsqueeze=True, detach=True, # stop gradient of last predictions scales_lims=self.cfg.scales_clamp_lims, raw_opacities_lims=(self.cfg.clamp_min_raw_opacities, self.cfg.clamp_max_raw_opacities) ) gaussians_concat = pack_gaussians(means, scales, rotations_unnorm, opacities_raw, shs) # [B, N, C] b, v, c, h, w = context["image"].shape latent_h = h // self.cfg.latent_downsample latent_w = w // self.cfg.latent_downsample # Debugging reprojection error if debug_dict is not None and (not self.training and self.save_every(i, tag="debug")): if "reprojection_error" in debug_dict: self.debug_reprojection_error(means, debug_dict, context, i, latent_h, latent_w) # prepare pt input point_cloud, tmp_batch_size = self.get_point_cloud(latent_h, latent_w, means, v) # Create offset directly on device to avoid CPU-GPU transfer offset = torch.arange(1, b + 1, device=state.device, dtype=torch.long) * tmp_batch_size # reshape gaussians_flat = self.reshape_gaussians_to_nc(latent_h, latent_w, gaussians_concat, v) # [B, N, C] --> [BN, C] # add global attention to exchange info across views if self.cfg.input_error_mv_attn: input_signal = self.apply_global_attn(h, input_signal, latent_h, latent_w, v, w) input_signal_flat = input_signal.reshape(-1, input_signal.shape[ -1]) # [B, N, C] --> [BN, C] - faster than rearrange input_signal_flat = self.append_to_input_signal(b, context, context_render_output, input_signal_flat, v) # Normalize state before input it to the update module if self.cfg.input_normalize_state: state_norm = state.norm(dim=1, keepdim=True) / math.sqrt(state.shape[-1]) # [BG, 1] state = state / (state_norm + 1e-8) # [BG, C] normalized_input = self.input_norm(input_signal_flat) if self.cfg.input_normalize_gaussians: gaussians_flat_mean = gaussians_flat.mean() gaussians_flat_std = gaussians_flat.std() gaussians_flat = (gaussians_flat - gaussians_flat_mean) / (gaussians_flat_std + 1e-8) with torch.amp.autocast(device_type='cuda', enabled=self.cfg.pt_update_amp, dtype=torch.bfloat16): point_cloud, gaussians_flat, state, update_input = self.prepare_update_input(init_state, normalized_input, point_cloud, gaussians_flat, state) updated_state = self._apply_point_transformer(b, latent_h, latent_w, offset, point_cloud, update_input, v, state, i) # Hard coded extract normalized gradients if self.cfg.input_gradient and self.cfg.input_gradient_normalize: normalized_grads = normalized_input else: normalized_grads = None # Recover the state norm if self.cfg.input_normalize_state: # state = state * state_std + state_mean updated_state = updated_state * state_norm # Predict a scale for the updtaed scale for the MLP deltas prediction # The updated state for the next stage remains the same if self.cfg.predict_state_scale: state_scale = self.state_scale_head(update_input.detach()) if self.cfg.predict_state_scale_norm: # Normalize the state vector state_scale = state_scale / (state_scale.norm(p=2, dim=1, keepdim=True) + 1e-8) else: state_scale = torch.tensor([1], device=state.device, dtype=state.dtype) updated_scaled_state = state_scale * updated_state # optionally append time encodiing to normalize input with TimeEncodingWrapper(self.cfg.use_time_encoding, self.time_encoder_fn, optimizer_output.t, self.cfg.time_encoding_max_steps, updated_scaled_state) as embedded_state: if self.cfg.use_time_encoding: assert not self.cfg.concat_init_state assert not self.cfg.replace_init_state # delta gaussian head delta_gaussians = self.apply_delta_gaussian_head(b, context, init_state, embedded_state, v) delta_means, delta_opacities, delta_rotations, delta_scales, delta_shs, init_repeat, delta_gaussians = ( self.postprocess_deltas(b, delta_gaussians, gaussian_grads, gaussians_concat, grad_sign, normalized_grads, state, optimizer_output.t) ) means, opacities_raw, rotations_unnorm, scales, shs = self.repeat_gaussians(means, opacities_raw, rotations_unnorm, scales, shs) covariances, means, scales, rotations, rotations_unnorm, opacities_raw, shs = self._apply_gaussian_deltas( delta_means, delta_scales, delta_rotations, delta_opacities, delta_shs, means, scales, rotations_unnorm, opacities_raw, shs, init_repeat) # Recover the state in non valid gaussians (and grad for logging) if gaussians.sel is not None: sel = gaussians.sel # [B, G] full_state = optimizer_input.prev_output.state.state # Convert full state to the dtype of state full_state = full_state.to(state.dtype) # Use non-in-place index_put to avoid in-place modification of tensors # in the autograd computation graph (fixes version mismatch errors with stability loss) updated_state = full_state.index_put((sel,), updated_state) else: sel = None # update gaussians (only where mask is True) # Use view instead of rearrange for speed shs_reshaped = shs.view(shs.shape[0], shs.shape[1], 3, -1) gaussians = gaussians.update_object_by_curr_mask( means=means, covariances=covariances, harmonics=shs_reshaped, opacities=opacities_raw.squeeze(-1).sigmoid(), scales=scales, rotations=rotations, rotations_unnorm=rotations_unnorm, sel=None, deltas=delta_gaussians if self.training else None, gradients=gaussian_grads_raw if self.training else None, norm_gradients=normalized_grads.unsqueeze(0) if normalized_grads is not None and self.training else None ) updates = { "means": delta_means.detach(), "scales": delta_scales.detach(), "rotations": delta_rotations.detach(), "opacities": delta_opacities.detach(), "shs": delta_shs.detach() } grads_raw = gaussian_grads.detach() if gaussian_grads is not None else None grads_adam = normalized_grads.detach() if normalized_grads is not None else None return gaussians, updated_state, meta_for_adc, updates, grads_raw, grads_adam, updated_scaled_state, sel def postprocess_deltas(self, b, delta_gaussians, gaussian_grads, gaussians_concat, grad_sign, normalized_grads, state, t): # Updates for gradient input (scale, log scale, ) delta_gaussians_raw = delta_gaussians delta_gaussians = self._postprocess_delta_for_gradient_input(delta_gaussians_raw, grad_sign, normalized_grads, ) # Rearrange back to [B, N, C] delta_gaussians, delta_gaussians_raw = self.rearrange_delta_gaussians(b, delta_gaussians, delta_gaussians_raw) # Extra non-weight-sharing heads, each predicting another N Gaussians from the state. # NOTE: these run after rearrange and on the raw state, so their deltas bypass # _postprocess_delta_for_gradient_input (scaling/negation/scalar_scale) that the main # head's delta went through — an asymmetry that only matters if this path is used # (gaussian_head_multiple is always 1 today). See the TODO (release) on the cfg field. with torch.amp.autocast(device_type='cuda', enabled=self.cfg.pt_update_amp, dtype=torch.bfloat16): if self.cfg.gaussian_head_multiple > 1: num_additional_heads = self.cfg.gaussian_head_multiple - 1 delta_gaussian_list = [delta_gaussians] # list of [B, N, C] for j in range(num_additional_heads): curr_delta = self.extra_gaussian_heads[j](state) curr_delta = rearrange(curr_delta, "(b n) c -> b n c", b=b) delta_gaussian_list.append(curr_delta) delta_gaussians = torch.cat(delta_gaussian_list, dim=1) # [B, K*N, C] # TODO (release): remove experimental delta override if self.cfg.experimental_run: self.experimental_update_deltas(delta_gaussians, gaussian_grads, normalized_grads) # Split delta_means, delta_scales, delta_rotations, delta_opacities, delta_shs, init_repeat = ( self.split_delta_gaussians(delta_gaussians) ) # Apply lr delta_means, delta_scales, delta_rotations, delta_opacities, delta_shs = self.scale_deltas_with_lr( t, delta_means, delta_scales, delta_rotations, delta_opacities, delta_shs ) return delta_means, delta_opacities, delta_rotations, delta_scales, delta_shs, init_repeat, delta_gaussians def filter_gaussians_for_optimizer_step(self, gaussian_grads, gaussian_grads_raw, gaussians, grad_sign, init_state, input_signal, state): # Compute a mask for gaussians that did not contribute to any pixel of context views. # Their gradients are strictly zero; skip them for the update but don't prune (they # may be relevant in other views). assert gaussian_grads.shape[0] == 1, "Batch size > 1 not supported with mask" # A gaussian is valid if it has a nonzero gradient in any channel valid_g = gaussian_grads[0].any(dim=-1) # [G] bool sel = None # if everything is valid, skip all slicing work if not valid_g.all(): sel = valid_g.nonzero(as_tuple=True)[0] # [G_valid] input_signal = input_signal[:, sel, :] # [B, G_valid, C] gaussian_grads = gaussian_grads[:, sel, :] # [B, G_valid, D] if gaussian_grads_raw is not None: gaussian_grads_raw = gaussian_grads_raw[:, sel, :] if grad_sign is not None: grad_sign = grad_sign[:, sel, :] state = state[sel, :] # [G_valid, C] if init_state is not None: init_state = init_state[sel, :] # [G_valid, C] gaussians.sel = sel if self.cfg.input_gradient_normalize_type == "adam": self.input_norm.sel = sel return gaussian_grads, gaussian_grads_raw, gaussians, grad_sign, init_state, input_signal, state # TODO (release): remove in public code def experimental_prune_invisible_gaussians(self, context, context_render_output, gaussian_grads, gaussian_grads_raw, gaussians, grad_sign, input_signal, means2d_grads, meta_for_adc, optimizer_input, state): # Get visible gaussians mask, based on the last rendering with torch.no_grad(): visible_mask = self.experimental_get_visible_gaussian_mask(gaussian_grads, gaussians, context_render_output.visibility_filter, context) # [B, N, 1] if visible_mask is None: return gaussian_grads, gaussians, grad_sign, input_signal, state assert visible_mask.shape[0] == 1 visible_mask = visible_mask[0, :, 0] # [N], squeeze batch and last dim # Apply mask gaussians = gaussians[:, visible_mask] state = state[visible_mask] input_signal = input_signal[:, visible_mask] # [B, N, C] if gaussian_grads is not None: gaussian_grads = gaussian_grads[:, visible_mask] # [B, N, C] if grad_sign is not None: grad_sign = grad_sign[:, visible_mask] # [B, N, C] meta_for_adc["visibility_filter"] = context_render_output.visibility_filter[:, :, visible_mask] meta_for_adc["radii"] = context_render_output.radii[:, :, visible_mask] if means2d_grads is not None: meta_for_adc["means_2d_grads"] = means2d_grads[:, :, visible_mask] if self.cfg.input_gradient_normalize and self.cfg.input_gradient_normalize_type == "adam": if not self.input_norm.is_reset(): prune_mask = ~visible_mask self.input_norm.prune(prune_mask) # the prune fn will invert the mask again if self.cfg.any_adc: optimizer_input.prev_output.state.adc_state.external_pruning(visible_mask) return gaussian_grads, gaussians, grad_sign, input_signal, state # TODO (release): remove def experimental_deactivate_updates(self, subset, gaussians, radii_vis_mask, deltas, gaussian_grads): """ Deactivate updates for gaussians that are not visible in any view """ visible_mask = self.experimental_get_visible_gaussian_mask(gaussian_grads, gaussians, radii_vis_mask, subset) deltas = deltas * visible_mask # [B, N, C] return deltas # TODO (release): remove def experimental_get_visible_gaussian_mask(self, gaussian_grads, gaussians, radii_vis_mask, subset): """ Get mask for gaussians that are visible in at least one view. We calculate two criteria: 1. Whether the projected 2d radius is visible in at least one view. 2. Whether the gaussian has a non-zero weight contribution to the rendering. If neither pruning criterion is enabled, returns None. Args: gaussian_grads: [B, N, C] or None gaussians: Gaussians object radii_vis_mask: [B, V, N], bool subset: dict, context or target """ # If no pruning criteria are active, return None if not (self.cfg.local_prune_zero_radii or self.cfg.local_prune_low_weights): return None b, v, n = radii_vis_mask.shape # Criterion 1: Projected radius visibility if self.cfg.local_prune_zero_radii: radii_vis_mask = radii_vis_mask.any(dim=1).unsqueeze(-1) # [B, N, 1] else: radii_vis_mask = torch.ones((b, n, 1), dtype=torch.bool, device=radii_vis_mask.device) # Criterion 2: Weight contribution visibility if self.cfg.local_prune_low_weights: threshold = self.cfg.local_prune_low_weights_thresh weight_vis_contribution, _ = get_visibility_contribution_from_gaussian_obj(subset, gaussians) # [N] weight_cont_mask = (weight_vis_contribution > threshold).view(1, -1, 1) else: weight_cont_mask = torch.ones((b, n, 1), dtype=torch.bool, device=radii_vis_mask.device) visible_mask = radii_vis_mask & weight_cont_mask # [B, N, 1] return visible_mask def experimental_inplace_update_delta(self, deltas, grads, normalized_grads, cfg_attr): # Slicing of the gradients vector per parameter param_num = grads.shape[-1] assert param_num == 11 + self.cfg.sh_d * 3 param_slices = self.param_slices update = getattr(self.cfg.experimental_update, cfg_attr) if update: # Update this parameter use_norm_grad = getattr(self.cfg.experimental_use_norm_grads, cfg_attr) use_grad = self.cfg.experimental_use_grads and not use_norm_grad assert not (use_grad and use_norm_grad) if use_grad: # Use the inverse of the gradients (plain gradient-descent step, for comparison) # TODO (release): hard-coded SGD step size (30) in this experimental gradient-baseline path deltas[..., param_slices[cfg_attr]] = -(grads[..., param_slices[cfg_attr]]).to(deltas.dtype) * 30 elif use_norm_grad: # Use the inverse of the normalized gradients updated_delta = -normalized_grads[..., param_slices[cfg_attr]] * getattr(self.cfg.experimental_lr, cfg_attr) deltas[..., param_slices[cfg_attr]] = updated_delta.to(deltas.dtype) else: # Use the network prediction (already negated before) pass else: # Do not update this parameter deltas[..., param_slices[cfg_attr]] = 0 def experimental_update_deltas(self, deltas, grads, normalized_grads): # Verify that at least one parameter is actually using norm_grads or grads override any_norm_grad = any( getattr(self.cfg.experimental_use_norm_grads, p) for p in self.cfg.experimental_update.param_names) any_grad = self.cfg.experimental_use_grads any_override = any_norm_grad or any_grad assert any_override, ( "experimental_run=true but no parameter has use_norm_grads or use_grads enabled. " "Check that experimental_use_norm_grads._base=true (it gates all other fields via property)." ) if any_norm_grad: assert normalized_grads is not None, ( "experimental_use_norm_grads is enabled but normalized_grads is None. " "Ensure input_gradient=true and input_gradient_normalize=true." ) for p in self.cfg.experimental_update.param_names: self.experimental_inplace_update_delta(deltas, grads, normalized_grads, p) def scale_deltas_with_lr(self, t, delta_means, delta_scales, delta_rotations, delta_opacities, delta_shs): # Scale deltas with learning rates delta_means = delta_means * self.scheduler.get_lr(t, "means") delta_scales = delta_scales * self.scheduler.get_lr(t, "scales") if delta_rotations is not None: delta_rotations = delta_rotations * self.scheduler.get_lr(t, "rotations") delta_opacities = delta_opacities * self.scheduler.get_lr(t, "opacities") # Use view instead of rearrange for speed delta_shs = delta_shs.view(delta_shs.shape[0], delta_shs.shape[1], 3, -1) # [b, g, 3, c] delta_sh0 = delta_shs[..., 0] # [B, N, C] delta_shN = delta_shs[..., 1:] delta_sh0 = delta_sh0 * self.scheduler.get_lr(t, "sh0") delta_shN = delta_shN * self.scheduler.get_lr(t, "shN") delta_shs = torch.cat((delta_sh0.unsqueeze(-1), delta_shN), dim=-1) delta_shs = delta_shs.flatten(-2) # [b, g, d*c] - faster than rearrange return delta_means, delta_scales, delta_rotations, delta_opacities, delta_shs def _unshuffle_to_flat(self, x: torch.Tensor, b: int, v: int) -> torch.Tensor: """Pixel-unshuffle a (b*v, c, h, w) tensor and flatten spatial dims: -> (b*v*h*w, c).""" x = F.pixel_unshuffle(x, downscale_factor=self.cfg.latent_downsample) return rearrange(x, "(b v) c h w -> (b v h w) c", b=b, v=v) def _downsample_error(self, x: torch.Tensor, b: int, v: int) -> torch.Tensor: """Downsample a (b*v, c, h, w) error tensor to latent resolution: -> (b, v*h*w, c).""" if self.cfg.input_error_rgb_no_shuffle: x = F.interpolate(x, scale_factor=1. / self.cfg.latent_downsample, mode='bilinear', align_corners=True) else: # NOTE: pixel_unshuffle assumes per-pixel, grid-aligned Gaussians. A non-grid # subsampling (e.g. farthest-point sampling) would need a different error->Gaussian # mapping to stay aligned — not supported here. x = F.pixel_unshuffle(x, downscale_factor=self.cfg.latent_downsample) return rearrange(x, "(b v) c h w -> b (v h w) c", b=b, v=v) def append_to_input_signal(self, b, context, context_render, input_signal_flat, v): if self.cfg.input_alpha: render_alpha = rearrange(context_render.accumulated_alpha.detach(), "b v h w -> (b v) () h w") render_alpha = self._unshuffle_to_flat(render_alpha, b, v) input_signal_flat = torch.cat((input_signal_flat, render_alpha), dim=-1) if self.cfg.input_depth: render_depth = rearrange(context_render.depth.detach(), "b v h w -> (b v) () h w") render_depth = self._unshuffle_to_flat(render_depth, b, v) input_signal_flat = torch.cat((input_signal_flat, render_depth), dim=-1) if self.cfg.input_depth_smooth_error: disp = 1. / context_render.depth.detach().clamp(min=1e-3, max=1000.) # [B, V, H, W] disp = rearrange(disp, "b v h w -> (b v) () h w") mean_disp = disp.mean(2, True).mean(3, True) norm_disp = disp / (mean_disp + 1e-7) tmp_imgs = rearrange(context["image"], "b v c h w -> (b v) c h w") depth_smooth_error = get_smooth_loss(norm_disp, tmp_imgs, no_mean=True) depth_smooth_error = self._unshuffle_to_flat(depth_smooth_error, b, v) input_signal_flat = torch.cat((input_signal_flat, depth_smooth_error), dim=-1) return input_signal_flat def repeat_gaussians(self, prev_means, prev_opacities_raw, prev_rotations_unnorm, prev_scales, prev_shs): if self.cfg.gaussian_head_multiple > 1: # predict multiple gaussians for each point prev_means = prev_means.repeat(1, self.cfg.gaussian_head_multiple, 1) prev_scales = prev_scales.repeat(1, self.cfg.gaussian_head_multiple, 1) prev_rotations_unnorm = prev_rotations_unnorm.repeat(1, self.cfg.gaussian_head_multiple, 1) prev_opacities_raw = prev_opacities_raw.repeat(1, self.cfg.gaussian_head_multiple, 1) / self.cfg.gaussian_head_multiple # smaller opacities, important prev_shs = prev_shs.repeat(1, self.cfg.gaussian_head_multiple, 1) # NOTE: only repeat at the first iteration refine_repeat = self.cfg.delta_gaussian_multiple if refine_repeat > 1: # predict multiple gaussians for each point prev_means = prev_means.repeat(1, refine_repeat, 1) prev_scales = prev_scales.repeat(1, refine_repeat, 1) prev_rotations_unnorm = prev_rotations_unnorm.repeat(1, refine_repeat, 1) prev_opacities_raw = prev_opacities_raw.repeat(1, refine_repeat, 1) # smaller opacities, important prev_shs = prev_shs.repeat(1, refine_repeat, 1) return prev_means, prev_opacities_raw, prev_rotations_unnorm, prev_scales, prev_shs def split_delta_gaussians(self, delta_gaussians): if self.cfg.init_gaussian_multiple > 1 and not self.cfg.same_num_points: init_repeat = self.cfg.init_gaussian_multiple else: init_repeat = 1 p = get_gaussian_param_sizes(self.cfg.sh_d) frozen = self._frozen_param_groups() # Split the head output among the predicted (non-frozen) groups; frozen groups get # an additive-identity zero delta so downstream apply/rearrange is uniform. b, n = delta_gaussians.shape[0], delta_gaussians.shape[1] active = [(name, key) for name, key in _GAUSSIAN_GROUPS if name not in frozen] active_chunks = delta_gaussians.split([p[key] * init_repeat for _, key in active], dim=-1) deltas = dict(zip((name for name, _ in active), active_chunks)) for name, key in _GAUSSIAN_GROUPS: if name in frozen: deltas[name] = delta_gaussians.new_zeros(b, n, p[key] * init_repeat) delta_means, delta_scales = deltas["means"], deltas["scales"] delta_rotations, delta_opacities, delta_shs = deltas["rotations"], deltas["opacities"], deltas["shs"] if ( self.cfg.delta_gaussian_multiple > 1 or self.cfg.init_gaussian_multiple > 1) and not self.cfg.same_num_points: delta_means = rearrange(delta_means, "b n (c k) -> b (n k) c", k=init_repeat) delta_scales = rearrange(delta_scales, "b n (c k) -> b (n k) c", k=init_repeat) delta_rotations = rearrange(delta_rotations, "b n (c k) -> b (n k) c", k=init_repeat) delta_opacities = rearrange(delta_opacities, "b n (c k) -> b (n k) c", k=init_repeat) delta_shs = rearrange(delta_shs, "b n (c k) -> b (n k) c", k=init_repeat) return delta_means, delta_scales, delta_rotations, delta_opacities, delta_shs, init_repeat def rearrange_delta_gaussians(self, b, delta_gaussians, delta_gaussians_raw): # [BV, C] # update gaussian parameters delta_gaussians = rearrange(delta_gaussians, "(b n) c -> b n c", b=b) delta_gaussians_raw = rearrange(delta_gaussians_raw, "(b n) c -> b n c", b=b) return delta_gaussians, delta_gaussians_raw def _apply_gaussian_deltas(self, delta_means, delta_scales, delta_rotations, delta_opacities, delta_shs, means, scales, rotations_unnorm, opacities_raw, shs, repeat): means = self._apply_mean_delta(delta_means, means) # clamp the scale scales = self._apply_scale_delta(delta_scales, scales, repeat) if self.cfg.opt_scales_before_act: scales = scales.exp() if not self.cfg.freeze_rotation: rotations, rotations_unnorm = self._apply_rotation_delta(delta_rotations, rotations_unnorm) else: rotations = F.normalize(rotations_unnorm, dim=-1) # compute covariance covariances = build_covariance(scales, rotations) # ([1, VHW, 3, 3]) opacities_raw = self._apply_opacity_delta(delta_opacities, opacities_raw, repeat) shs = self._apply_shs_delta(delta_shs, shs) return covariances, means, scales, rotations, rotations_unnorm, opacities_raw, shs def _apply_shs_delta(self, delta_shs, prev_shs): shs = prev_shs + delta_shs # [B, N, 3*sh_d] if self.cfg.clamp_shs_soft: assert self.cfg.clamp_min_shs == -self.cfg.clamp_max_shs, "For soft clamp, min and max should be symmetric around 0" shs = torch.tanh(shs / self.cfg.clamp_max_shs) * self.cfg.clamp_max_shs else: shs = shs.clamp(min=self.cfg.clamp_min_shs, max=self.cfg.clamp_max_shs) return shs def _apply_opacity_delta(self, delta_opacities, prev_opacities_raw, repeat): # update init opacities when predicting multiple gaussians if repeat > 1 and not self.cfg.multi_gaussian_scale_smaller and (self.cfg.init_gaussian_multiple == 1): # Given y = sigmoid(x), to get new x' such that sigmoid(x') = y / K: # The formula is: x' = x + log((1 - y) / (K - y)) # This adjusts x so that the sigmoid output is scaled down by a factor of K tmp_sigmoid = prev_opacities_raw.sigmoid() prev_opacities_raw = prev_opacities_raw + torch.log( (1 - tmp_sigmoid) / (repeat - tmp_sigmoid)) + delta_opacities else: prev_opacities_raw = prev_opacities_raw + delta_opacities # prev_opacities_raw = prev_opacities_raw.clamp(min=-5, max=5) return prev_opacities_raw @staticmethod def _apply_rotation_delta(delta_rotations, prev_rotations_unnorm): assert delta_rotations is not None prev_rotations_unnorm = prev_rotations_unnorm + delta_rotations # normazlie prev_rotations = prev_rotations_unnorm / (prev_rotations_unnorm.norm(dim=-1, keepdim=True) + 1e-8) return prev_rotations, prev_rotations_unnorm def _apply_scale_delta(self, delta_scales, prev_scales, repeat): """Add the predicted scale delta and clamp the result. When one point is represented by `repeat` overlapping gaussians (repeat > 1), the cluster is too dense unless we compensate. `multi_gaussian_scale_smaller` chooses to compensate on size: each gaussian's initial scale is divided by `repeat`, so the cluster covers about the same footprint as the single original gaussian (opacity is left unchanged). When the flag is off, size is kept and the opacity is reduced instead (see _apply_opacity_delta). """ if repeat > 1 and self.cfg.multi_gaussian_scale_smaller: # smaller initial scales (the floor is applied by the shared clamp below). # The divide-by-repeat is in activation space, so this branch assumes scales # are not in raw/log space. assert not self.cfg.opt_scales_before_act, \ "multi_gaussian_scale_smaller assumes opt_scales_before_act=False (activation-space scales)" new_scales = prev_scales / repeat + delta_scales else: new_scales = prev_scales + delta_scales if self.cfg.opt_scales_before_act: min_scale = self.cfg.clamp_min_raw_scales max_scale = self.cfg.clamp_max_raw_scales else: min_scale = self.cfg.clamp_min_scale max_scale = self.cfg.clamp_max_scale new_scales = new_scales.clamp(min=min_scale, max=max_scale) return new_scales @staticmethod def _apply_mean_delta(delta_means, prev_means): prev_means = (prev_means + delta_means) return prev_means def _on_scene_start_impl(self, optimizer_input: OptimizerInput) -> None: # Reset the state if isinstance(optimizer_input.prev_output, InitializerOutput): # New scene from_init = True # Reset the optimizer state for a new scene # We cannot just use super().on_scene_start() because we need to process the InitializerOutput in case it # contain conditioning features self.reset_logs() if self.cfg.input_gradient_normalize_type == "adam": self.input_norm.reset() nr_gaussians = rearrange(optimizer_input.prev_output.gaussians.means, "b n c -> (b n) c").shape[0] param_num = self.gaussian_param_num self.input_norm.initialize(shape=(nr_gaussians, param_num), device=optimizer_input.prev_output.gaussians.means.device) # make sure xyz are contiguous optimizer_input.prev_output.gaussians.means = optimizer_input.prev_output.gaussians.means.contiguous() elif isinstance(optimizer_input.prev_output, OptimizerPreviousOutput): from_init = False if self.cfg.input_gradient_normalize_type == "adam": # Continuing previous optimization from replay buffer self.input_norm.update_state(optimizer_input.prev_output.state.adam_state) # Note: logs are not handled right now for continuing from replay buffer self.reset_logs() else: raise ValueError(f"Unknown prev_output type {type(optimizer_input.prev_output)}") # Preparing the input for a new scene (will handle both new scene and continuing from replay buffer) # Will convert init_output to prev_output internally self.optimizer_preprocessing(optimizer_input, from_init=from_init) # initialize adc state, after converting to prev_output if from_init and self.cfg.any_adc: self.initialize_adc_state(self.cfg, optimizer_input) def reshape_gaussians_to_nc(self, latent_h, latent_w, gaussians_concat, v): if self.cfg.init_gaussian_multiple > 1 and not self.cfg.same_num_points: # gaussians are with more points, reshape factor = self._block_factor() gaussians_flat = rearrange(gaussians_concat, "b (v h x w y) c -> (b v h w) (c x y)", v=v, h=latent_h, w=latent_w, x=factor, y=factor) else: gaussians_flat = rearrange(gaussians_concat, "b n c -> (b n) c") return gaussians_flat def get_point_cloud(self, latent_h, latent_w, prev_means, v): if self.cfg.init_gaussian_multiple > 1 and not self.cfg.same_num_points: # The initializer predicts a (factor x factor) block of gaussians per latent pixel, # so there are more points than the base grid. Subsample back to the base grid so the # point cloud (the KNN structure for the transformer) stays at the base resolution. factor = self._block_factor() point_cloud = rearrange(prev_means, "b (v h w) c -> b v h w c", v=v, h=latent_h * factor, w=latent_w * factor, ) # simply use uniform grid subsample of point cloud to reduce points point_cloud = point_cloud[:, :, ::factor, ::factor] point_cloud = rearrange(point_cloud, "b v h w c -> (b v h w) c") tmp_batch_size = v * latent_h * latent_w else: point_cloud = rearrange(prev_means, "b n c -> (b n) c") tmp_batch_size = prev_means.shape[1] return point_cloud, tmp_batch_size def get_vector_state(self, b, v, n, optimizer_input, from_init): if from_init: # Starting a new scene directly from the initializer # State should not be provided # Create initial state # optimizer_input.prev_output is of type InitializerOutput if optimizer_input.prev_output.features is None or self.cfg.init_state_wo_features: # Creating state without initializer features assert self.cfg.init_state_wo_features with torch.amp.autocast(device_type='cuda', enabled=self.cfg.pt_update_amp, dtype=torch.bfloat16): dtype = torch.get_autocast_dtype('cuda') if self.cfg.init_state_type == "constant": state = torch.ones((b, n, self.cfg.state_channels), device=self.device, dtype=dtype) elif self.cfg.init_state_type == "random": state = torch.randn((b, n, self.cfg.state_channels), device=self.device, dtype=dtype) else: raise ValueError(f"Unknown init_state_type {self.cfg.init_state_type}") state = state * self.cfg.init_state_scale else: # Calculating state from initializer features state = self.get_state_from_condition_features(b, optimizer_input.prev_output.features, v) # [B, N, C] # combine gaussians of all scenes in the batch [B, N, C] -> [B*N, C] state = rearrange(state, "b n c -> (b n) c") else: # Restarting a scene from a replay buffer: the state is already stored flat as [B*N, C]. state = optimizer_input.prev_output.state.state return state @staticmethod def _align_features(features, latent_h: int, latent_w: int) -> list: """Resize each feature map to (latent_h, latent_w) if needed and return as a list.""" out = [] vals = features.values() if isinstance(features, dict) else features for feat in vals: if feat.shape[-2:] != (latent_h, latent_w): feat = F.interpolate(feat, size=(latent_h, latent_w), mode='bilinear', align_corners=True) out.append(feat) return out def _block_factor(self) -> int: """Per-pixel block side length when the initializer predicts multiple gaussians per pixel. init_gaussian_multiple gaussians form a (factor x factor) block, so factor = sqrt(multiple). """ multiple = self.cfg.init_gaussian_multiple factor = math.isqrt(multiple) assert factor * factor == multiple, f"init_gaussian_multiple must be a perfect square, got {multiple}" return factor def _get_latent_size(self, h: int, w: int) -> tuple[int, int]: """Compute latent (H, W) from image (H, W), accounting for init_gaussian_multiple upsampling.""" latent_h = h // self.cfg.latent_downsample latent_w = w // self.cfg.latent_downsample if self.cfg.init_gaussian_multiple > 1 and self.cfg.same_num_points: factor = self._block_factor() latent_h *= factor latent_w *= factor return latent_h, latent_w def render_input_views_for_error_calc(self, context, prev_gaussians, renderer): _, _, _, h, w = context["image"].shape # [B, V, C, H, W] render_res = (h, w) # Use only first N views end = self.cfg.input_error_num_views if self.cfg.input_error_num_views > 0 else None return renderer.forward_batch_subset( prev_gaussians, context, render_res, start=None, end=end, return_radii=False ) def get_state_from_condition_features(self, b, condition_features, v): with torch.amp.autocast(device_type='cuda', enabled=self.cfg.pt_update_amp, dtype=torch.bfloat16): if not self.cfg.pt_update_amp and condition_features.dtype == torch.bfloat16: condition_features = condition_features.float() state = self.state_proj(condition_features.detach()) # [B, C, H, W] if self.cfg.init_gaussian_multiple > 1 and self.cfg.same_num_points: factor = self._block_factor() state = F.interpolate(state, scale_factor=factor, mode='bilinear', align_corners=True) # Convert to vector of Gaussians per batch [B, N, C] state = rearrange(state, "(b v) c h w -> b (v h w) c", b=b, v=v) # N = v * h * w return state def prepare_update_input(self, init_state, input_signal, point_cloud, gaussians_flat, state): if self.cfg.replace_init_state: state = init_state if self.cfg.no_render_error: update_input = torch.cat((gaussians_flat, state), dim=-1) else: update_input = torch.cat((gaussians_flat, state, input_signal), dim=-1) if self.cfg.concat_init_state: update_input = torch.cat((update_input, init_state), dim=-1) return point_cloud, gaussians_flat, state, update_input def _apply_point_transformer(self, b, latent_h, latent_w, offset, point_cloud, update_input, v, state, iter): def recurrent_chunk(update_input, point_cloud, offset): pxo = self.point_transformer[0]([point_cloud, update_input, offset]) state = self.point_transformer[1](pxo, iter=iter, b=b, v=v, h=latent_h, w=latent_w) return state if self.cfg.use_checkpointing or self.cfg.recurrent_use_checkpointing: new_state = torch.utils.checkpoint.checkpoint( recurrent_chunk, update_input, point_cloud, offset, use_reentrant=False, ) else: new_state = recurrent_chunk(update_input, point_cloud, offset) if self.cfg.residual_state: new_state = new_state + state return new_state def apply_delta_gaussian_head(self, b, context, init_state, state, v): if self.cfg.delta_head_concat_img: img_unshuffle = rearrange(context["image"], "b v c h w -> (b v) c h w") img_unshuffle = self._unshuffle_to_flat(img_unshuffle, b, v) head_input = torch.cat((state, img_unshuffle), dim=-1) else: if self.cfg.residual_init_state: head_input = state + init_state else: head_input = state if self.cfg.delta_head_per_param_heads: delta_gaussians = self._apply_per_param_heads(head_input) else: delta_gaussians = self.delta_head(head_input) return delta_gaussians def _apply_per_param_heads(self, head_input): """Run per-parameter-group heads and concatenate results. Each head outputs [N, dim+1] where the last dim is the scalar scale. Per-group normalize + scale is applied independently. """ deltas = [] for name, dim in self._per_param_group_dims.items(): raw = self.delta_head[name](head_input) # [N, dim+1] scale = self.scale_act(raw[:, -1:]) # [N, 1] delta = raw[:, :-1] # [N, dim] if dim > 1: delta = delta / (delta.norm(p=2, dim=-1, keepdim=True) + 1e-8) * scale else: # 1-d (e.g. opacities): no direction to normalize, just scale magnitude delta = delta * scale deltas.append(delta) return torch.cat(deltas, dim=-1) def apply_global_attn(self, h, input_signal, latent_h, latent_w, v, w): assert self.cfg.input_error_resnet_feature assert self.cfg.input_error if self.cfg.input_gradient and self.cfg.input_error: input_render_error = input_signal[..., :self.error_features_channels] else: input_render_error = input_signal with torch.amp.autocast(device_type='cuda', enabled=self.cfg.use_amp, dtype=torch.bfloat16): for blk in self.error_mv_attn: if self.cfg.same_num_points: # no downsample, for re10k 256 input_render_error = blk(input_render_error, v=v, h=h, w=w) else: input_render_error = blk(input_render_error, v=v, h=latent_h, w=latent_w) if self.cfg.input_gradient and self.cfg.input_error: input_signal[..., :self.error_features_channels] = input_render_error else: input_signal = input_render_error return input_signal def prepare_input_signal(self, context, gaussians, renderer): """Build the per-Gaussian signal that drives the update: render gradients and/or feature error. Renders the context views once and turns the errors with the inputs into the network's input signal (gradients from `_calc_input_gradients`, render error, or both concatenated). """ # make sure at least one of the following is True assert self.cfg.input_gradient or self.cfg.input_error input_signal = None input_render_error = None gaussian_grads_raw = None gaussian_grads = None grad_sign = None means2d_grads = None # calculate input gradients if self.cfg.input_gradient: gaussian_grads_raw, gaussian_grads, grad_sign, context_render_output, means2d_grads = ( self._calc_input_gradients(context, gaussians, renderer) ) input_signal = gaussian_grads_raw else: # Render if input_gradient=False. Otherwise, the gradients render can be reused for the error # calculation, all further usage detach before use. context_render_output = self.render_input_views_for_error_calc(context, gaussians, renderer) # calculate input rendering errors if self.cfg.input_error: if means2d_grads is None and self.cfg.need_2d_grads: raise NotImplementedError("Calculating 2dgrad for ADC is not implemented for error input alone") input_render_error = self._calc_input_errors(context, context_render_output) input_signal = input_render_error if self.cfg.input_gradient and self.cfg.input_error: # Concatenate both gradients and errors input_signal = torch.cat((input_render_error, gaussian_grads), dim=-1) return input_signal, gaussian_grads_raw, gaussian_grads, grad_sign, context_render_output, means2d_grads def debug_reprojection_error(self, means, debug_dict, context, i, latent_h, latent_w): # Prepare means (remove singleton dim) means = rearrange(means, "b (v h w) c -> b v (h w) c", h=latent_h, w=latent_w) # [B, V, H*W, 3] # Expand extrinsics/intrinsics for broadcast extrinsics = context["extrinsics"].unsqueeze(2) # [B, V, 1, 4, 4] intrinsics = context["intrinsics"].unsqueeze(2) # [B, V, 1, 3, 3] # Project xy_ray_reconstructed, in_front = project(means, extrinsics, intrinsics) # [B, V, H*W, 2], [B, V, H*W] xy_ray, _ = sample_image_grid((latent_h, latent_w), xy_ray_reconstructed.device) # [B, V, H*W, 1, 2] xy_ray = rearrange(xy_ray, "h w xy -> (h w) () xy") xy_ray = xy_ray.squeeze(-2) # [B, V, H*W, 2] xy_ray_unnorm = xy_ray.clone() xy_ray_unnorm[..., 0] *= latent_w xy_ray_unnorm[..., 1] *= latent_h xy_ray_reconstructed_unnorm = xy_ray_reconstructed.clone() xy_ray_reconstructed_unnorm[..., 0] *= latent_w xy_ray_reconstructed_unnorm[..., 1] *= latent_h reprojection_error = (xy_ray_unnorm - xy_ray_reconstructed_unnorm).abs() if debug_dict["reprojection_error"] is None: # First iteration, first scene debug_dict["reprojection_error"] = [[]] elif i == 0: # New iteration, new scene debug_dict["reprojection_error"].append([]) debug_dict["reprojection_error"][-1].append(reprojection_error.detach().cpu()) # import matplotlib.pyplot as plt # plt.figure(figsize=(12, 6)) # plt.hist(reprojection_error.flatten().detach().cpu(), bins=100, range=(0, 10)) # plt.title(f"Reprojection Error - step {i}") # plt.xlabel("Error (pixels)") # plt.ylabel("Frequency") # plt.show() def _calc_input_errors(self, context, input_render): b, v, _, h, w = context["image"].shape # Detach the last rendered object input_rgb = input_render.color.detach() # compute input view rendering error if self.cfg.input_error_resnet_feature: input0 = rearrange(input_rgb, "b v c h w -> (b v) c h w") if self.cfg.input_error_num_views > 0: gt_input = context["image"][:, :self.cfg.input_error_num_views, :, :, :] else: gt_input = context["image"] input1 = rearrange(gt_input, "b v c h w -> (b v) c h w") transform = _IMAGENET_NORM concat = torch.cat((input0, input1), dim=0) input_tensor = transform(concat) with torch.amp.autocast(device_type='cuda', enabled=self.cfg.pt_update_amp, dtype=torch.bfloat16): with torch.no_grad(): features = self.error_feature_extractor(input_tensor) # align to the latent resolution latent_h, latent_w = self._get_latent_size(h, w) all_features = torch.cat(self._align_features(features, latent_h, latent_w), dim=1) render_view_features = all_features[:input0.shape[0]] gt_view_features = all_features[input0.shape[0]:] feature_error = render_view_features - gt_view_features if self.cfg.input_error_num_views > 0: # pad to V views curr_v = self.cfg.input_error_num_views indices = torch.arange(v) * curr_v // v feature_error = rearrange(feature_error, "(b v) c h w -> b v c h w", b=b) feature_error = feature_error[torch.arange(b).unsqueeze(1), indices, :, :, :] input_render_error = rearrange(feature_error, "b v c h w -> b (v h w) c") else: input_render_error = rearrange(feature_error, "(b v) c h w -> b (v h w) c", b=b) else: input_render_error = (input_render.color - context["image"]).abs() # [B, V, 3, H, W] input_render_error = rearrange(input_render_error, "b v c h w -> (b v) c h w") input_render_error = self._downsample_error(input_render_error, b, v) # [B, N, C] # include both feature error and image error if self.cfg.input_error_add_rgb_feature: rgb_render_error = input_render.color - context["image"] rgb_render_error = rearrange(rgb_render_error, "b v c h w -> (b v) c h w") rgb_render_error = self._downsample_error(rgb_render_error, b, v) # [B, N, C] rgb_render_error = self.rgb_error_proj(rgb_render_error) input_render_error = input_render_error + rgb_render_error return input_render_error def get_input_error_feature_extractor(self): error_feature_extractor = None # resnet feature if self.cfg.input_error_resnet_feature: error_feature_extractor = ResNetFeatureWarpper( shallow_resnet_feature=self.cfg.input_error_shallow_resnet_feature) if self.cfg.input_error_no_freeze_resnet_feature: # remove unused layers # NOTE: layer 3 is also not used error_feature_extractor.layer3 = nn.Identity() error_feature_extractor.train() for params in error_feature_extractor.parameters(): params.requires_grad = True else: error_feature_extractor.eval() for params in error_feature_extractor.parameters(): params.requires_grad = False return error_feature_extractor def _postprocess_delta_for_gradient_input(self, delta_gaussians, grad_sign, normalized_grad): if self.cfg.input_gradient: delta_gaussians = delta_gaussians / self.cfg.input_gradient_scale if self.cfg.input_gradient_log: grad_sign = rearrange(grad_sign, "b n c -> (b n) c") # recover log scale for applying the deltas. # For loss calculation the delta should still be in log scale delta_gaussians = grad_sign * (delta_gaussians.exp() - 1e-8) if self.cfg.input_gradient_log_clip_deltas > 0: # clip the delta to avoid too large updates clip_value = self.cfg.input_gradient_log_clip_deltas clip_mask = delta_gaussians.abs() > clip_value delta_gaussians[clip_mask] = delta_gaussians[clip_mask].sign() * clip_value # NOTE: scalar_scale is a delta-head concern, not a gradient one, but it lives # here because it only runs in the gradient path today. if self.cfg.delta_head_scalar_scale: if self.cfg.delta_head_per_param_heads: # Already handled in _apply_per_param_heads — nothing to do here pass elif self.cfg.delta_head_per_param_scales: # Feature B: per-group scalar scales num_groups = len(self._per_param_group_dims) scales = delta_gaussians[:, -num_groups:] # [G, num_groups] scales = self.scale_act(scales) deltas = delta_gaussians[:, :-num_groups] # [G, D] normalized_deltas = [] offset = 0 for i, (name, dim) in enumerate(self._per_param_group_dims.items()): group_delta = deltas[:, offset:offset + dim] # [G, dim] group_scale = scales[:, i:i + 1] # [G, 1] if dim > 1: group_delta = group_delta / (group_delta.norm(p=2, dim=-1, keepdim=True) + 1e-8) group_delta = group_delta * group_scale normalized_deltas.append(group_delta) offset += dim delta_gaussians = torch.cat(normalized_deltas, dim=-1) else: # Original global scalar scale scale = delta_gaussians[:, -1:] # [G, 1] scale = self.scale_act(scale) # make sure scale is positive deltas_unnorm = delta_gaussians[:, :-1] # [G, D] deltas_norm = deltas_unnorm / (deltas_unnorm.norm(p=2, dim=1, keepdim=True) + 1e-8) # [G, D] delta_gaussians = deltas_norm * scale if self.cfg.scale_residual_grads: delta_gaussians = delta_gaussians * normalized_grad * self.cfg.residual_grad_scale # 1.0 # To match the default behavior of SGD, Adam, and other optimizers, deltas are negated. # SGD applies the gradients as `x = x - lr * grad`, while resaplt applies them as `x = x + lr * deltas`. delta_gaussians = -delta_gaussians return delta_gaussians def _calc_input_gradients(self, context, gaussians, renderer): """Render the context views and backprop the reconstruction loss to per-Gaussian parameter gradients. These gradients are a feature fed to the network (the "what is wrong" signal), not used to update weights — the graph is detached. Chunked over views to bound memory; optionally also returns the 2D means gradients ADC consumes. """ assert not self.cfg.input_gradient_same_loss, "input_gradient_same_loss is not implemented" _, v, _, h, w = context["image"].shape with torch.enable_grad(): # Unpack gaussians means, scales, rotations_unnorm, opacities_raw, shs = unpack_gaussians( gaussians, scales_log=self.cfg.opt_scales_before_act, opacities_logit=True, opacities_unsqueeze=True, detach=True, clone=False, requires_grad=True, scales_lims=self.cfg.scales_clamp_lims, raw_opacities_lims=(self.cfg.clamp_min_raw_opacities, self.cfg.clamp_max_raw_opacities) ) # Create temporary Gaussians with same values but requires_grad=True grad_batch_size = self.cfg.input_gradients_chunk_size if grad_batch_size == -1: grad_batch_size = v gaussian_grads = 0 means2d_grads_chunks = [] nr_chunks = math.ceil(v / grad_batch_size) # Pre-compute shapes and config flags outside the loop shs_shape = (shs.shape[0], shs.shape[1], 3, -1) opt_scales_before_act = self.cfg.opt_scales_before_act # Pre-compute normalized rotations once (not in gradient inputs, so no grad needed) with torch.no_grad(): rotations = rotations_unnorm / (rotations_unnorm.norm(dim=-1, keepdim=True) + 1e-8) for chunk_idx, start, stop in chunk_index_iter(v, grad_batch_size): # zero grads means = means.detach().requires_grad_(True) scales = scales.detach().requires_grad_(True) rotations_unnorm = rotations_unnorm.detach().requires_grad_(True) opacities_raw = opacities_raw.detach().requires_grad_(True) shs = shs.detach().requires_grad_(True) # Apply activation to scales if needed (before calculating covariance) scales_act = scales.exp() if opt_scales_before_act else scales tmp_gaussians = Gaussians( means=means, covariances=None, harmonics=shs.view(shs_shape), opacities=torch.sigmoid(opacities_raw.squeeze(-1)), scales=scales_act, rotations=rotations, rotations_unnorm=rotations_unnorm, ) # render input views, calculate inner loss and backprop to get gradients context_render_output = renderer.forward_batch_subset( tmp_gaussians, context, start=start, end=stop, image_shape=(h, w), ) inputs = [means, scales, rotations_unnorm, opacities_raw, shs] if self.cfg.need_2d_grads: assert context_render_output.means2d is not None, "output_renderer.means2d is None" means2d = context_render_output.means2d # [B, V, N, 2] # means2d.retain_grad() # retain grad for means2d grads computation inputs.append(means2d) inner_loss = inner_loss_for_input_gradients( context["image"][:, start:stop], context_render_output, reduction=self.cfg.input_gradient_loss_reduction, with_ssim=self.cfg.input_gradient_with_ssim_loss, ) if self.cfg.opacity_reg_lambda > 0.0: inner_loss = inner_loss + self.cfg.opacity_reg_lambda * torch.sigmoid(opacities_raw).mean() grads = torch.autograd.grad(outputs=inner_loss, inputs=inputs, create_graph=False, retain_graph=False, ) gaussian_grads = gaussian_grads + torch.cat(grads[:5], dim=-1) # [B, G, D] assert not torch.isnan(gaussian_grads).any(), "NaN detected in gaussian_grads" if self.cfg.need_2d_grads: means2d_grads_chunks.append(grads[5]) # [B, V_chunk, N, 2] gaussian_grads = gaussian_grads / nr_chunks if self.cfg.need_2d_grads: means2d_grads = torch.cat(means2d_grads_chunks, dim=1) # [B, V, N, 2] # Normalize to NDC per the renderer's convention so the ADC stays renderer-agnostic # (gsplat=pixel→×w/2,h/2; inria/fastgs already NDC). means2d_grads = renderer.means2d_grad_to_ndc(means2d_grads, (h, w)) if self.cfg.input_gradient_loss_reduction == "mean_pixels_sum_views": means2d_grads = means2d_grads / v else: means2d_grads = None gaussian_grads_raw = gaussian_grads * self.cfg.input_gradient_scale if self.cfg.input_gradient_log: # log gradients grads_sign = gaussian_grads.sign() gaussian_grads_raw = (gaussian_grads_raw.abs() + 1e-8).log() else: grads_sign = None # Detach gradients to avoid gradient flow through the input gaussian_grads = gaussian_grads.detach() gaussian_grads_raw = gaussian_grads_raw.detach() if grads_sign is not None: grads_sign = grads_sign.detach() # Returning also the render output, but it can only be used for visualization, # as we already backpropagate gradients through it return gaussian_grads_raw, gaussian_grads, grads_sign, context_render_output, means2d_grads