File size: 10,728 Bytes
ea66a99 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 | import torch
from utils.graphics_utils import apply_rotation, batch_quaternion_multiply
def _compute_smooth_regularization_71(self,
xyz,
scales_emb,
rotations_emb,
opacity_emb,
shs_emb,
time,
out_pts,
out_scales,
out_rot,
out_opacity,
out_shs,
smooth_cfg,
orig_cache=None):
"""
Compute smooth regularization loss on the deformation outputs.
The idea is: for a small perturbation in the input (both spatial and temporal),
the deformation network should produce similar outputs.
This function re-computes the hidden representation and each branch's deformation
for the original and perturbed inputs, then calculates an L2 loss for each branch,
weighted by cfg, and returns the averaged loss.
"""
# Retrieve perturbation magnitudes from config (or use defaults)
spatial_eps = smooth_cfg.get("spatial_eps", 5e-3)
temporal_eps = smooth_cfg.get("temporal_eps", 5e-3)
# temporal_eps = 0.0
choice = None
if smooth_cfg['downsample_ratio'] < 1.0:
# == Choose Random Points ==
device = xyz.device
N_all = xyz.shape[0]
subN = int(max(1, N_all * smooth_cfg['downsample_ratio']))
choice = torch.randperm(N_all, device=device)[:subN]
# (a) Subset + detach
xyz = xyz[choice,:3].clone().detach()
scales_emb = scales_emb[choice, :3].clone().detach()
rotations_emb = rotations_emb[choice, :4].clone().detach()
time = time[choice,:1].clone().detach() # [subN,1]
# === Compute original deformation outputs ===
if orig_cache is None:
hidden_orig = self.query_time(xyz, scales_emb, rotations_emb, None, time)
dx_orig = self.pos_deform(hidden_orig) if not self.args.no_dx else None
ds_orig = self.scales_deform(hidden_orig) if not self.args.no_ds else None
dr_orig = self.rotations_deform(hidden_orig) if not self.args.no_dr else None
do_orig = self.opacity_deform(hidden_orig) if not self.args.no_do else None
dshs_orig = (
self.shs_deform(hidden_orig).reshape(xyz.shape[0], 16, 3)
if not self.args.no_dshs else None
)
else:
dx_orig = orig_cache.get("dx") if orig_cache else None
ds_orig = orig_cache.get("ds") if orig_cache else None
dr_orig = orig_cache.get("dr") if orig_cache else None
do_orig = orig_cache.get("do") if orig_cache else None
dshs_orig = orig_cache.get("dshs") if orig_cache else None
if not self.args.no_dx and dx_orig is None:
raise ValueError("orig_cache missing 'dx' for smooth regularization.")
if not self.args.no_ds and ds_orig is None:
raise ValueError("orig_cache missing 'ds' for smooth regularization.")
if not self.args.no_dr and dr_orig is None:
raise ValueError("orig_cache missing 'dr' for smooth regularization.")
if not self.args.no_do and do_orig is None:
raise ValueError("orig_cache missing 'do' for smooth regularization.")
if not self.args.no_dshs and dshs_orig is None:
raise ValueError("orig_cache missing 'dshs' for smooth regularization.")
if choice is not None:
if dx_orig is not None:
dx_orig = dx_orig.index_select(0, choice)
if ds_orig is not None:
ds_orig = ds_orig.index_select(0, choice)
if dr_orig is not None:
dr_orig = dr_orig.index_select(0, choice)
if do_orig is not None:
do_orig = do_orig.index_select(0, choice)
if dshs_orig is not None:
dshs_orig = dshs_orig.index_select(0, choice)
# === Generate perturbed inputs ===
# Perturb the spatial (rays_pts_emb) and temporal (time_emb) inputs.
xyz = self.normalize_aabb(xyz, self.grid.get_aabb)
perturbed_rays = xyz + (torch.rand_like(xyz) * 2 - 1) * spatial_eps
perturbed_time = time + (torch.rand_like(time) * 2 - 1) * temporal_eps
# Compute the hidden representation for the perturbed inputs.
hidden_pert = self.query_time(perturbed_rays, scales_emb, rotations_emb, None, perturbed_time, need_normalize=False)
# hidden_pert = self.feature_out(hidden_pert)
dx_pert = self.pos_deform(hidden_pert) if dx_orig is not None else None
ds_pert = self.scales_deform(hidden_pert) if ds_orig is not None else None
dr_pert = self.rotations_deform(hidden_pert) if dr_orig is not None else None
do_pert = self.opacity_deform(hidden_pert) if do_orig is not None else None
dshs_pert = (self.shs_deform(hidden_pert)
.reshape(xyz.shape[0], 16, 3)) if dshs_orig is not None else None
# === Compute L2 differences per branch ===
loss = 0.0
count = 0
if dx_orig is not None:
loss_dx = torch.mean((dx_orig - dx_pert)**2)
loss += smooth_cfg.get("w_dx", 1.0) * loss_dx
count += 1
if ds_orig is not None:
loss_ds = torch.mean((ds_orig - ds_pert)**2)
loss += smooth_cfg.get("w_ds", 1.0) * loss_ds
count += 1
if dr_orig is not None:
loss_dr = torch.mean((dr_orig - dr_pert)**2)
loss += smooth_cfg.get("w_dr", 1.0) * loss_dr
count += 1
if do_orig is not None:
loss_do = torch.mean((do_orig - do_pert)**2)
loss += smooth_cfg.get("w_do", 1.0) * loss_do
count += 1
if dshs_orig is not None:
loss_dshs = torch.mean((dshs_orig - dshs_pert)**2)
loss += smooth_cfg.get("w_dshs", 1.0) * loss_dshs
count += 1
assert count > 0, "No deformation branches enabled for smooth regularization."
# Return the averaged loss across the branches used.
return loss / (count if count > 0 else 1.0)
def _compute_smooth_regularization_72(self,
xyz,
scales_emb,
rotations_emb,
opacity_emb,
shs_emb,
time,
out_pts,
out_scales,
out_rot,
out_opacity,
out_shs,
smooth_cfg):
"""
Compute smooth regularization loss on the deformation outputs.
The idea is: for a small perturbation in the input (both spatial and temporal),
the deformation network should produce similar outputs.
This function re-computes the hidden representation and each branch's deformation
for the original and perturbed inputs, then calculates an L2 loss for each branch,
weighted by cfg, and returns the averaged loss.
"""
# Retrieve perturbation magnitudes from config (or use defaults)
spatial_eps = smooth_cfg.get("spatial_eps", 5e-3)
temporal_eps = smooth_cfg.get("temporal_eps", 5e-3)
# temporal_eps = 0.0
if smooth_cfg['downsample_ratio'] < 1.0:
# == Choose Random Points ==
device = xyz.device
N_all = xyz.shape[0]
subN = int(max(1, N_all * smooth_cfg['downsample_ratio']))
choice = torch.randperm(N_all, device=device)[:subN]
# (a) Subset + detach
xyz = xyz[choice,:3].clone().detach()
scales_emb = scales_emb[choice, :3].clone().detach()
rotations_emb = rotations_emb[choice, :4].clone().detach()
time = time[choice,:1].clone().detach() # [subN,1]
# === Compute original deformation outputs ===
# For consistency, re-run query_time and feature_out on the original input.
hidden_orig = self.query_time(xyz, scales_emb, rotations_emb, None, time)
dx_orig = self.pos_deform(hidden_orig) if not self.args.no_dx else None
ds_orig = self.scales_deform(hidden_orig) if not self.args.no_ds else None
dr_orig = self.rotations_deform(hidden_orig) if not self.args.no_dr else None
do_orig = self.opacity_deform(hidden_orig) if not self.args.no_do else None
dshs_orig = (self.shs_deform(hidden_orig)
.reshape(xyz.shape[0], 16, 3)) if not self.args.no_dshs else None
# === Generate perturbed inputs ===
# Perturb the spatial (rays_pts_emb) and temporal (time_emb) inputs.
xyz = self.normalize_aabb(xyz, self.grid.get_aabb)
perturbed_rays = xyz + (torch.rand_like(xyz) * 2 - 1) * spatial_eps
perturbed_time = time + (torch.rand_like(time) * 2 - 1) * temporal_eps
# Compute the hidden representation for the perturbed inputs.
hidden_pert = self.query_time(perturbed_rays, scales_emb, rotations_emb, None, perturbed_time, need_normalize=False)
# hidden_pert = self.feature_out(hidden_pert)
dx_pert = self.pos_deform(hidden_pert) if dx_orig is not None else None
ds_pert = self.scales_deform(hidden_pert) if ds_orig is not None else None
dr_pert = self.rotations_deform(hidden_pert) if dr_orig is not None else None
do_pert = self.opacity_deform(hidden_pert) if do_orig is not None else None
dshs_pert = (self.shs_deform(hidden_pert)
.reshape(xyz.shape[0], 16, 3)) if dshs_orig is not None else None
# === Compute L2 differences per branch ===
loss = 0.0
count = 0
if dx_orig is not None:
loss_dx = torch.mean((dx_orig - dx_pert)**2)
loss += smooth_cfg.get("w_dx", 1.0) * loss_dx
count += 1
if ds_orig is not None:
loss_ds = torch.mean((ds_orig - ds_pert)**2)
loss += smooth_cfg.get("w_ds", 0.3) * loss_ds
count += 1
if dr_orig is not None:
loss_dr = torch.mean((dr_orig - dr_pert)**2)
loss += smooth_cfg.get("w_dr", 0.3) * loss_dr
count += 1
if do_orig is not None:
loss_do = torch.mean((do_orig - do_pert)**2)
loss += smooth_cfg.get("w_do", 0.1) * loss_do
count += 1
if dshs_orig is not None:
loss_dshs = torch.mean((dshs_orig - dshs_pert)**2)
loss += smooth_cfg.get("w_dshs", 0.1) * loss_dshs
count += 1
assert count > 0, "No deformation branches enabled for smooth regularization."
# Return the averaged loss across the branches used.
return loss / (count if count > 0 else 1.0)
|