"""Test-time 3DGS post-processing applied after the learned optimizer. `PostProcessing3DGS` runs extra vanilla-3DGS refinement on the optimized Gaussians at test time: a torch Adam/SGD optimizer descends the parameters directly (built as a `GaussiansModule`), optionally with adaptive density control (clone/split/prune). This is the `meta_trainer.test.postprocessing` step and also provides the standalone 3DGS baseline numbers; it is off by default. """ from dataclasses import dataclass from typing import List import tqdm as tqdm import numpy as np import torch from torch import Tensor import math from pytorch_optimizer import load_optimizer from torch.optim.lr_scheduler import LambdaLR from einops import rearrange from optgs.evaluation.metrics import compute_rgb_metrics from optgs.misc.benchmarker import Benchmarker from optgs.misc.io import FrequencyScheduler from optgs.scene_trainer.gaussian_module import GaussiansModule, gaussians2module, module2gaussians from optgs.model.types import Gaussians from optgs.scene_trainer.optimizer.optimizer import OptimizerOutput from optgs.scene_trainer.optimizer.optimizer_utils import Number3DGSCfg from optgs.misc.detaching_cpu_list import DetachingCPUList from optgs.dataset.camera_datasets.camera import get_scene_scale from optgs.paths import DEBUG from optgs.misc.general_utils import get_expon_lr_func from fused_ssim import fused_ssim from optgs.model.decoder.decoder import DecoderOutput @dataclass class PostProcessADCCfg: """ADC (Adaptive Density Control) config for postprocessing. Defaults match vanilla 3DGS (config/scene_trainer/scene_optimizer/refiner/default.yaml). """ # Densification strategy: "vanilla" (3DGS) or "fastgs" (multi-view metric + Abs-GS split). # "fastgs" requires the FastGS decoder (it needs the abs-gradient and metric-counts render). strategy: str = "vanilla" do_densify: bool = True do_prune: bool = True do_opacity_reset: bool = True # Scheduling pause_refine_after_reset: int = 0 refine_every: int = 100 reset_every: int = 3000 refine_start_iter: int = 500 refine_stop_iter: int = 15000 refine_scale2d_stop_iter: int = 0 # Densification thresholds grow_grad2d: float = 0.0002 grow_scale3d: float = 0.01 # aka percent_dense grow_scale2d: float = 0.05 # Pruning thresholds prune_scale3d: float = 0.1 prune_scale2d: float = 0.15 min_opacity: float = 0.005 # FastGS strategy knobs (strategy == "fastgs"); see adc/fastgs.py. grad_abs_thresh: float = 0.0012 # Abs-GS split threshold (split signal) importance_thresh: int = 5 # min averaged high-error view count to densify (integer count) prune_budget_frac: float = 0.5 # fraction of eligible Gaussians pruned per step opacity_clamp: float = 0.8 # opacity ceiling applied after each densification loss_thresh: float = 0.1 # normalized-L1 threshold for the high-error metric map num_score_cams: int = 10 # context views sampled to score each densification @dataclass class PostProcessCfg: name: str steps: int compute_metrics_every: int lr_data: Number3DGSCfg scheduler: str | None scheduler_warm_up_ratio: float # SGD-specific momentum: float = 0.0 nesterov: bool = False # Adam-specific betas: List[float] | None = None eps: float = 1e-8 amsgrad: bool = False # Shared weight_decay: float = 0.0 # LR scheduling: steps already done by scene trainer (offsets the schedule) prior_steps: int = 0 # Means LR scheduling (defaults match vanilla optimizer behavior) means_lr_final_ratio: float = 0.0625 # ratio of final/initial means LR (vanilla: 1e-5 / 1.6e-4) means_lr_delay_mult: float = 0.01 # ramp-up delay multiplier (vanilla default: 0.01) means_lr_scale_by_scene_extent: bool = True # scale means LR by scene extent (vanilla default) # View chunking for gradient accumulation chunk_size: int = -1 # -1 = all views at once # ADC (Adaptive Density Control) adc: PostProcessADCCfg | None = None @property def is_active(self) -> bool: return self.name != "none" and self.steps > 0 def get_dir_name(self, with_name=True): dir_str = self._get_dir_name() return f"{self.name}_{dir_str}" if with_name else dir_str def _get_dir_name(self): if self.name == "sgd": return f"lr{self.lr_data.base}_mom{self.momentum}" elif self.name == "adam": return f"lr{self.lr_data.base}_betas{'-'.join(map(str, self.betas or []))}_eps{self.eps}" return "" def _module_to_deactivated_gaussians(gm: GaussiansModule) -> Gaussians: """Convert GaussiansModule to Gaussians with deactivated (raw) values for ADC.""" return Gaussians( means=gm.means.detach().unsqueeze(0), scales=gm.scales_raw.detach().unsqueeze(0), # log space opacities=gm.opacities_raw.detach().unsqueeze(0), # logit space rotations=gm.rotations.detach().unsqueeze(0), rotations_unnorm=gm.rotations_unnorm.detach().unsqueeze(0), harmonics=gm.harmonics.detach().unsqueeze(0), stores_activated=False, ) def _deactivated_gaussians_to_module(gaussians: Gaussians, device: torch.device) -> GaussiansModule: """Convert deactivated Gaussians back to GaussiansModule.""" assert not gaussians.stores_activated return GaussiansModule( means=gaussians.means[0].to(device), harmonics=gaussians.harmonics[0].to(device), opacities=torch.sigmoid(gaussians.opacities[0]).to(device), scales=torch.exp(gaussians.scales[0]).to(device), rotations_unnorm=gaussians.rotations_unnorm[0].to(device), ) def _to_strategy_cfg(adc_cfg: PostProcessADCCfg): """Build the dispatch's strategy config (a BaseStrategyCfg subclass) from the postprocessing ADC config, so postprocessing runs the shared apply_adc_strategy instead of its own copy of the clone/split/prune logic. Postprocessing imposes no cap and no MCMC/reduce, so those base knobs take inert defaults; loss_thresh / num_score_cams stay on PostProcessADCCfg (render params).""" from optgs.scene_trainer.adc.vanilla import VanillaStrategyCfg from optgs.scene_trainer.adc.fastgs import FastGSStrategyCfg common = dict( do_densify=adc_cfg.do_densify, do_prune=adc_cfg.do_prune, do_opacity_reset=adc_cfg.do_opacity_reset, cap_max=-1, pause_refine_after_reset=adc_cfg.pause_refine_after_reset, refine_every=adc_cfg.refine_every, reset_every=adc_cfg.reset_every, refine_start_iter=adc_cfg.refine_start_iter, refine_stop_iter=adc_cfg.refine_stop_iter, refine_scale2d_stop_iter=adc_cfg.refine_scale2d_stop_iter, grow_grad2d=adc_cfg.grow_grad2d, grow_scale3d=adc_cfg.grow_scale3d, prune_scale3d=adc_cfg.prune_scale3d, prune_scale2d=adc_cfg.prune_scale2d, grow_scale2d=adc_cfg.grow_scale2d, min_opacity=adc_cfg.min_opacity, prune_zero_radii=False, reduce_opacity=False, reduce_factor=0.0, reduce_every=0, fallback_means_lr=0.0, ) if adc_cfg.strategy == "fastgs": return FastGSStrategyCfg( name="fastgs", **common, grad_abs_thresh=adc_cfg.grad_abs_thresh, importance_thresh=adc_cfg.importance_thresh, prune_budget_frac=adc_cfg.prune_budget_frac, opacity_clamp=adc_cfg.opacity_clamp, ) return VanillaStrategyCfg(name="default", **common) class PostProcessing3DGS: def __init__(self, cfg: PostProcessCfg, save_every: FrequencyScheduler): self.cfg = cfg self.save_every = save_every # Per-iteration timing + stats. benchmarker.time("iter") records CUDA events with a deferred # sync, so the loop is not stalled every step; the *_log properties read back from it. self.benchmarker = Benchmarker() self.reset_logs() def reset_logs(self): # All per-iteration logs (timings + counts) live in self.benchmarker; one clear resets them. self.benchmarker.clear_history() @property def iter_time_log(self) -> list[float]: """Total ms per post-processing iteration (from benchmarker.time("iter")).""" return self.benchmarker.execution_times["iter"] # Per-iteration stats, stored in self.benchmarker next to the timings (one entry per step). # Counts are ints, so the element type is float | int (see Optimizer for the same pattern). @property def nr_gaussians_log(self) -> list[float | int]: return self.benchmarker.execution_times["gaussians"] @property def nr_nonzero_grad_log(self) -> list[float | int]: return self.benchmarker.execution_times["nonzero_grads"] @property def nr_cloned_log(self) -> list[float | int]: return self.benchmarker.execution_times["cloned"] @property def nr_splitted_log(self) -> list[float | int]: return self.benchmarker.execution_times["splitted"] @property def nr_pruned_log(self) -> list[float | int]: return self.benchmarker.execution_times["pruned"] def _calc_loss( self, context, output_renderer: DecoderOutput ) -> Tensor: # compute scalar loss # assume batch size 1 assert context["image"].shape[0] == 1 assert context["image"].shape == output_renderer.color.shape l1_render_error = (output_renderer.color - context["image"]).abs().mean() ssim_score = fused_ssim( rearrange(output_renderer.color, "b v c h w -> (b v) c h w"), rearrange(context["image"], "b v c h w -> (b v) c h w"), padding="valid" ) loss = 0.8 * l1_render_error + 0.2 * (1 - ssim_score) return loss def _chunked_forward_backward(self, gaussian_module, iter_context, decoder, render_res, adc_state): """Render views in chunks, accumulate gradients, and collect ADC metadata. Matches the gradient accumulation approach of calc_input_gradients in the vanilla optimizer: each chunk computes a mean loss, gradients accumulate, then are averaged by nr_chunks. """ v = iter_context["image"].shape[1] chunk_size = self.cfg.chunk_size if self.cfg.chunk_size > 0 else v nr_chunks = math.ceil(v / chunk_size) # Accumulate means2d grads and radii for ADC across chunks need_adc = adc_state is not None h, w = render_res if need_adc: N = gaussian_module.means.shape[0] means2d_grads_all = torch.zeros((1, v, N, 2), device=gaussian_module.means.device) # FastGS Abs-GS split signal (only populated when the decoder exposes means2d_abs). means2d_abs_grads_all = torch.zeros((1, v, N, 2), device=gaussian_module.means.device) has_abs_grad = False radii_all = torch.zeros((1, v, N, 2), device=gaussian_module.means.device) visibility_all = torch.zeros((1, v, N), dtype=torch.bool, device=gaussian_module.means.device) for chunk_start in range(0, v, chunk_size): chunk_end = min(chunk_start + chunk_size, v) # Slice views for this chunk chunk_context = { "image": iter_context["image"][:, chunk_start:chunk_end], "extrinsics": iter_context["extrinsics"][:, chunk_start:chunk_end], "intrinsics": iter_context["intrinsics"][:, chunk_start:chunk_end], "near": iter_context["near"][:, chunk_start:chunk_end], "far": iter_context["far"][:, chunk_start:chunk_end], } # Render chunk_output = decoder.forward_batch_subset(gaussian_module, chunk_context, render_res) # Retain means2d grad(s) for ADC if need_adc and chunk_output.means2d is not None: chunk_output.means2d.retain_grad() if need_adc and chunk_output.means2d_abs is not None: chunk_output.means2d_abs.retain_grad() # Loss and backward (gradients accumulate across chunks) chunk_loss = self._calc_loss(chunk_context, chunk_output) chunk_loss.backward() # Collect ADC metadata from this chunk if need_adc: if chunk_output.radii is not None: radii_all[:, chunk_start:chunk_end] = chunk_output.radii.detach() if chunk_output.visibility_filter is not None: visibility_all[:, chunk_start:chunk_end] = chunk_output.visibility_filter.detach() if chunk_output.means2d is not None and chunk_output.means2d.grad is not None: # 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_all[:, chunk_start:chunk_end] = decoder.means2d_grad_to_ndc( chunk_output.means2d.grad.detach(), render_res) if chunk_output.means2d_abs is not None and chunk_output.means2d_abs.grad is not None: means2d_abs_grads_all[:, chunk_start:chunk_end] = decoder.means2d_grad_to_ndc( chunk_output.means2d_abs.grad.detach(), render_res) has_abs_grad = True # Average gradients across chunks (matches vanilla behavior) if nr_chunks > 1: for param in gaussian_module.parameters(): if param.grad is not None: param.grad /= nr_chunks # Return ADC metadata if need_adc: return { "radii": radii_all, "visibility_filter": visibility_all, "means_2d_grads": means2d_grads_all, "means_2d_abs_grads": means2d_abs_grads_all if has_abs_grad else None, } return None def _apply_adc(self, step, gaussian_module, adc_state, strategy_cfg, batch, decoder, render_res, device): """Run the configured ADC strategy through the shared dispatch (apply_adc_strategy) — no reimplementation of clone/split/prune. Returns (gaussian_module, optimizer_needs_rebuild).""" from optgs.scene_trainer.adc import apply_adc_strategy # FastGS multi-view scores: only at densification steps; vanilla ignores these (None). scores = {} if strategy_cfg.name == "fastgs" and self._is_densify_step(step, strategy_cfg): imp, pr = self._compute_fastgs_scores(gaussian_module, batch, decoder, render_res, self.cfg.adc) scores = dict(importance_score=imp, pruning_score=pr) # ADC mutates deactivated Gaussians in place; postprocessing rebuilds the optimizer on # change, so it threads no smoothers (empty dict -> the strategy's object hooks are no-ops). gaussians = _module_to_deactivated_gaussians(gaussian_module) nr_cloned, nr_splitted, nr_pruned, _, _ = apply_adc_strategy( strategy_cfg, step=step, gaussians=gaussians, adc_state=adc_state, smoothers={}, **scores ) self.benchmarker.record("cloned", nr_cloned) self.benchmarker.record("splitted", nr_splitted) self.benchmarker.record("pruned", nr_pruned) changed = self._strategy_modifies_module(step, strategy_cfg, nr_cloned, nr_splitted, nr_pruned) if changed: gaussian_module = _deactivated_gaussians_to_module(gaussians, device) return gaussian_module, changed def _is_densify_step(self, step: int, cfg) -> bool: """The densification gate shared by apply_vanilla / apply_fastgs — used to time the FastGS score render and to detect module-modifying steps.""" return ( step < cfg.refine_stop_iter and step > cfg.refine_start_iter and step % cfg.refine_every == 0 and step % cfg.reset_every >= cfg.pause_refine_after_reset ) def _strategy_modifies_module(self, step, cfg, nr_cloned, nr_splitted, nr_pruned) -> bool: """Whether this ADC step changed the Gaussians (so the module + optimizer must be rebuilt): any clone/split/prune, a periodic opacity reset, or — for FastGS — the per-densify opacity clamp. Mirrors the gates inside apply_vanilla_strategy / apply_fastgs_strategy.""" if nr_cloned or nr_splitted or nr_pruned: return True # Periodic opacity reset: vanilla resets unconditionally; fastgs only while still densifying. if cfg.do_opacity_reset and step > 0 and step % cfg.reset_every == 0: if cfg.name != "fastgs" or step < cfg.refine_stop_iter: return True # FastGS clamps opacity to a ceiling on every densification step. if cfg.name == "fastgs" and self._is_densify_step(step, cfg): return True return False @torch.no_grad() def _compute_fastgs_scores(self, gaussian_module, batch, decoder, render_res, adc_cfg): """FastGS multi-view importance/pruning scores: sample context views, flag high-error pixels, and ask the FastGS rasterizer which Gaussians contributed to them. Mirrors ``compute_gaussian_score_fastgs`` in submodules/FastGS/utils/fast_utils.py.""" from optgs.scene_trainer.adc.fastgs import compute_fastgs_scores ctx = batch["context"] n_views = ctx["extrinsics"].shape[1] n = min(adc_cfg.num_score_cams, n_views) idx = torch.randperm(n_views, device=ctx["extrinsics"].device)[:n] ext, intr = ctx["extrinsics"][:, idx], ctx["intrinsics"][:, idx] near, far = ctx["near"][:, idx], ctx["far"][:, idx] gt = ctx["image"][:, idx] # [1, n, 3, H, W] sub = {"extrinsics": ext, "intrinsics": intr, "near": near, "far": far} render = decoder.forward_batch_subset(gaussian_module, sub, render_res).color # [1, n, 3, H, W] eps = torch.finfo(render.dtype).eps metric_maps, losses = [], [] for j in range(n): r, g = render[0, j], gt[0, j] # [3, H, W] l1 = (r - g).abs().mean(0) # [H, W] l1n = (l1 - l1.min()) / (l1.max() - l1.min()).clamp_min(eps) metric_maps.append((l1n > adc_cfg.loss_thresh).int().reshape(-1)) # [H*W] ssim = fused_ssim(r.unsqueeze(0), g.unsqueeze(0), padding="valid") losses.append((0.8 * (r - g).abs().mean() + 0.2 * (1.0 - ssim)).item()) counts = decoder.render_metric_counts( gaussian_module, ext, intr, near, far, render_res, torch.stack(metric_maps) ) # [n, N] return compute_fastgs_scores(list(counts), losses, densify=True) @torch.no_grad() def apply( self, batch, gaussians: Gaussians, decoder, metrics=["psnr", "ssim"], iter_batch_size: int = -1, batchify_fn=None, visualization_dump=None ) -> OptimizerOutput | None: # This instance is reused across scenes (created once in scene_trainer), so clear the # per-iteration logs (benchmarker: iter_time_log, nr_gaussians_log, ...) for a per-scene reset. self.reset_logs() target_render_list = DetachingCPUList() context_render_list = DetachingCPUList() if self.cfg.steps == 0: return None # Calculate scene_scale from both context + target (matches vanilla optimizer) camtoworlds_context = batch['context']['extrinsics'][0].cpu().numpy() # [Vc, 4, 4] camtoworlds_target = batch['target']['extrinsics'][0].cpu().numpy() # [Vt, 4, 4] camtoworlds = np.concatenate([camtoworlds_context, camtoworlds_target], axis=0) scene_scale = get_scene_scale(camtoworlds) device = batch['context']['image'].device # convert Gaussians to GaussiansModule gaussian_module = gaussians2module(gaussians, device=device) optimizer = self.get_optimizer(gaussian_module, scene_scale) scheduler = self.get_scheduler(optimizer, scene_scale=scene_scale, prior_steps=self.cfg.prior_steps) assert batch["context"]["extrinsics"].shape[0] == batch["context"]["extrinsics"].shape[0] == 1, \ "Batch size > 1 not supported for post-processing" nr_context_views, _, h, w = batch["context"]["image"][0].shape # controlling number of context views seen at each iteration (for rendering chunk size) _iter_batch_size = iter_batch_size if iter_batch_size > 0 else nr_context_views render_res = (h, w) # Initialize ADC state via the shared dispatch, from a real strategy config. adc_state = None strategy_cfg = None if self.cfg.adc is not None: from optgs.scene_trainer.adc import init_strategy_state if self.cfg.adc.strategy == "fastgs" and not hasattr(decoder, "render_metric_counts"): raise ValueError( "postprocessing adc.strategy='fastgs' requires the FastGS decoder " "(decoder=fastgs); it needs the abs-gradient and metric-counts render." ) strategy_cfg = _to_strategy_cfg(self.cfg.adc) nr_points = gaussian_module.means.shape[0] adc_state = init_strategy_state( strategy_cfg, nr_points=nr_points, device=device, scene_extent=scene_scale ) # render before first step context_render_output = decoder.forward_batch_subset(gaussian_module, batch["context"], render_res, iter_batch_size=_iter_batch_size) context_render_list.append(context_render_output, detach_and_cpu=True) # initial rendering target_render_output = decoder.forward_batch_subset(gaussian_module, batch["target"], render_res, iter_batch_size=_iter_batch_size) target_render_list.append(target_render_output, detach_and_cpu=True) # initial rendering # Reset viewpoint stack for fresh sampling in postprocessing batch["context"].viewpoint_stack = None pbar = tqdm.tqdm(range(self.cfg.steps), desc=f"PP {self.cfg.name}", ncols=120) pbar_postfix = {} for i in pbar: with self.benchmarker.time("iter"): with torch.enable_grad(): # Log number of gaussians self.benchmarker.record("gaussians", gaussian_module.means.shape[0]) # reset gradients optimizer.zero_grad() # Sample context views using the same strategy as the optimizer iter_context, _ = batchify_fn(batch, "context") # Render in chunks, accumulate gradients, collect ADC metadata meta_for_adc = self._chunked_forward_backward( gaussian_module, iter_context, decoder, render_res, adc_state ) # step optimizer.step() # update scheduler if scheduler is not None: scheduler.step() # ADC: update state and apply densification/pruning via the dispatch. if adc_state is not None and meta_for_adc is not None: from optgs.scene_trainer.adc import update_strategy_state v_chunk = iter_context["image"].shape[1] # means_2d_abs_grads is the FastGS split signal; the dispatch ignores it for the # vanilla state and consumes it for the FastGS state. update_strategy_state( adc_state, radii_2d=meta_for_adc["radii"], means2d_grads=meta_for_adc["means_2d_grads"], means2d_abs_grads=meta_for_adc.get("means_2d_abs_grads"), visibility_mask=meta_for_adc["visibility_filter"], v=v_chunk, w=w, h=h, ) gaussian_module, adc_changed = self._apply_adc( i, gaussian_module, adc_state, strategy_cfg, batch, decoder, render_res, device ) if adc_changed: # Rebuild optimizer and scheduler after ADC changed the Gaussian count. # Caveat: rebuilding drops the Adam moments (exp_avg/exp_avg_sq) for every # Gaussian, so they restart from zero after each densification. Reference 3DGS # instead resizes the optimizer state in place to keep the moments across # clone/split/prune (as the in-loop optimizers do via AdamInputSmoothing); the # nn.Parameter + torch.optim path here cannot do that without rebuilding. optimizer = self.get_optimizer(gaussian_module, scene_scale) scheduler = self.get_scheduler( optimizer, scene_scale=scene_scale, prior_steps=self.cfg.prior_steps ) # Fast-forward scheduler to current step for _ in range(i + 1): scheduler.step() if scheduler is not None else None if self.save_every(i + 1, tag="context"): with torch.no_grad(): context_render_output = decoder.forward_context(gaussian_module, batch, (h, w)) context_render_list.append(context_render_output, detach_and_cpu=True) context_rgb = context_render_output.color[0] # [Vc, 3, Hc, Wc] ctx_scores: dict = compute_rgb_metrics( rgb=context_rgb, rgb_gt=batch["context"]["image"][0], metrics=metrics, iter_batch_size=iter_batch_size if "lpips" in metrics else -1 ) for k, v in ctx_scores.items(): pbar_postfix[f"ctx_{k}"] = f"{v.item():.2f}" if self.save_every(i + 1, tag="target"): with torch.no_grad(): target_render_output = decoder.forward_target(gaussian_module, batch, (h, w)) target_render_list.append(target_render_output, detach_and_cpu=True) target_rgb = target_render_output.color[0] # [Vt, 3, Ht, Wt] tgt_scores: dict = compute_rgb_metrics( rgb=target_rgb, rgb_gt=batch["target"]["image"][0], metrics=metrics, iter_batch_size=iter_batch_size if "lpips" in metrics else -1 ) for k, v in tgt_scores.items(): pbar_postfix[f"tgt_{k}"] = f"{v.item():.2f}" pbar_postfix["gs"] = gaussian_module.means.shape[0] pbar.set_postfix(pbar_postfix) if visualization_dump is not None and "grads" in visualization_dump: self.debug_grads(gaussian_module, visualization_dump, i) # convert back to Gaussians postprocessed_gaussians = module2gaussians(gaussian_module) postprocessed_gaussians_list = DetachingCPUList() postprocessed_gaussians_list.append(postprocessed_gaussians, detach_and_cpu=True) output = OptimizerOutput( target_render_list=target_render_list, context_render_list=context_render_list, gaussian_list=postprocessed_gaussians_list, info = {} ) return output def debug_grads(self, gaussians: GaussiansModule, debug_dict, step): if debug_dict["grads"] is None: # First iteration, first scene debug_dict["grads"] = [[]] elif step == 0: # New iteration, new scene debug_dict["grads"].append([]) grads = [param.grad for name, param in gaussians.named_parameters() if param.grad is not None] gaussian_num = gaussians.means.shape[0] grads = [g.view(gaussian_num, -1) for g in grads] grads = [g.detach().cpu() for g in grads] grads = torch.cat(grads, dim=-1) # [num_gaussians, total_param_dim] debug_dict["grads"][-1].append(grads) def get_optimizer(self, gaussians: GaussiansModule, scene_scale: float): # NOTE: post-processing runs one scene at a time (batch size 1) batch_size: int = 1 # Build params list (name, parameter, lr) named_parameters = dict(gaussians.named_parameters()) params = [] for key in named_parameters.keys(): lr_data_attr = key lr_data_attr = lr_data_attr.replace("_raw", "") lr_data_attr = lr_data_attr.replace("_unnorm", "") params.append((key, named_parameters[key], getattr(self.cfg.lr_data, lr_data_attr))) world_size = torch.distributed.get_world_size() if torch.distributed.is_initialized() else 1 BS = batch_size * world_size # Build parameter groups for a single optimizer param_groups = [ { "name": name, "params": param, "lr": lr * math.sqrt(BS), # individual learning rate } for name, param, lr in params ] # Get other optimizer parameters opt_params = self.extract_opt_params() # Manipulate opt_params with BS if needed if "weight_decay" in opt_params: opt_params["weight_decay"] *= BS if "eps" in opt_params: opt_params["eps"] /= math.sqrt(BS) if "betas" in opt_params: beta1, beta2 = opt_params["betas"] opt_params["betas"] = (1 - BS * (1 - beta1), 1 - BS * (1 - beta2)) # Instantiate a single optimizer with all parameter groups optimizer_class = load_optimizer(self.cfg.name) optimizer = optimizer_class( param_groups, **opt_params ) return optimizer _OPT_PARAMS = { "sgd": ("momentum", "weight_decay", "nesterov"), "adam": ("betas", "eps", "weight_decay", "amsgrad"), } def extract_opt_params(self): allowed = self._OPT_PARAMS.get(self.cfg.name, ()) return {k: getattr(self.cfg, k) for k in allowed if getattr(self.cfg, k, None) is not None} def get_scheduler(self, optimizer, scene_scale: float = 1.0, prior_steps: int = 0): if self.cfg.scheduler is None: return None total_steps = prior_steps + self.cfg.steps if self.cfg.scheduler == "exponential": # Per-param-group scheduling: # - Means: exponential decay optionally scaled by scene_extent (matching vanilla optimizer) # - Other params: constant LR lambdas = [] for group in optimizer.param_groups: if group["name"] == "means" and self.cfg.means_lr_scale_by_scene_extent: # Vanilla-style means LR: exponential decay with scene_extent scaling base_lr = group["lr"] # initial means LR from param group means_lr_func = get_expon_lr_func( lr_init=base_lr * scene_scale, lr_final=base_lr * scene_scale * self.cfg.means_lr_final_ratio, lr_delay_mult=self.cfg.means_lr_delay_mult, max_steps=total_steps, ) # LambdaLR computes: effective_lr = base_lr * lambda(step) # We want: effective_lr = means_lr_func(step) # So: lambda(step) = means_lr_func(step) / base_lr _base_lr = base_lr # capture for closure _func = means_lr_func lambdas.append(lambda step, f=_func, b=_base_lr: f(step) / b) else: # Constant LR for all other param groups lambdas.append(lambda step: 1.0) scheduler = LambdaLR(optimizer, lr_lambda=lambdas) # Fast-forward to prior_steps so LR continues from where scene trainer left off for _ in range(prior_steps): scheduler.step() return scheduler else: raise ValueError(f"Unknown scheduler: {self.cfg.scheduler}")