| 23 | return sum(p.numel() for p in module.parameters() if p.requires_grad) |
| 24 | |
| 25 | class SplatModel(StaticEncoder): |
| 26 | def __init__(self, opt: Options, **model_kwargs): |
| 27 | super().__init__(opt) |
| 28 | self.opt = opt |
| 29 | self.model = SplatPredictor(opt, **model_kwargs) |
| 30 | if hasattr(opt, 'compile') and opt.compile: |
| 31 | self.model = torch.compile(self.model) |
| 32 | self.gaussian_renderer = gaussian_renderer_dynamic.render |
| 33 | self.background = torch.tensor(opt.background_color, dtype=torch.float32, device="cuda") |
| 34 | |
| 35 | # LPIPS loss |
| 36 | if self.opt.lambda_lpips > 0: |
| 37 | self.lpips_loss = LPIPS(net='vgg') |
| 38 | self.lpips_loss.requires_grad_(False) |
| 39 | self.lpips_loss.eval() |
| 40 | if hasattr(opt, 'compile') and opt.compile: |
| 41 | self.lpips_loss = torch.compile(self.lpips_loss) |
| 42 | |
| 43 | def train(self, mode=True): |
| 44 | super().train(mode) |
| 45 | if 'lpips_loss' in self.__dict__: |
| 46 | self.lpips_loss.eval() |
| 47 | return self |
| 48 | |
| 49 | def state_dict(self, **kwargs): |
| 50 | """Remove non-trainable modules (LPIPS, tracker) from state dict before saving.""" |
| 51 | state_dict = super().state_dict(**kwargs) |
| 52 | for k in list(state_dict.keys()): |
| 53 | if 'lpips_loss' in k or 'tracker_prior' in k: |
| 54 | del state_dict[k] |
| 55 | return state_dict |
| 56 | |
| 57 | def load_state_dict(self, state_dict, strict=True): |
| 58 | """Load state dict, filtering out non-trainable modules if present.""" |
| 59 | filtered_state_dict = {k: v for k, v in state_dict.items() |
| 60 | if not k.startswith('lpips_loss') and not k.startswith('tracker_prior')} |
| 61 | return super().load_state_dict(filtered_state_dict, strict=strict) |
| 62 | |
| 63 | def forward_gaussians(self, frames, depths, cond_times=None): |
| 64 | # frames: [B, V, C, H, W] |
| 65 | # return: gaussians: [B, N, D] |
| 66 | decoder_out = self.model(frames, depths, cond_times=cond_times) |
| 67 | return decoder_out |
| 68 | |
| 69 | def compute_losses(self, input_frames, target_depth, supv_masks, render_pkg): |
| 70 | output_frames = render_pkg["render"] # [B, V, C, H, W] |
| 71 | pred_depths = render_pkg["depth"] |
| 72 | depth_mask = render_pkg["alpha"] > 0.1 # [B, V, 1, H, W] |
| 73 | metrics = {} |
| 74 | with torch.no_grad(): |
| 75 | B, V, C, H, W = input_frames.shape |
| 76 | |
| 77 | # All frames |
| 78 | input_frames_all_256 = F.interpolate(input_frames.reshape(-1, 3, H, W), (256, 256), mode='bilinear', align_corners=False) |
| 79 | output_frames_all_256 = F.interpolate(output_frames.reshape(-1, 3, H, W), (256, 256), mode='bilinear', align_corners=False) |
| 80 | metrics['psnr'] = compute_psnr(input_frames_all_256, output_frames_all_256).mean() |
| 81 | metrics['ssim'] = compute_ssim(input_frames_all_256, output_frames_all_256).mean() |
| 82 | metrics['lpips'] = compute_lpips(input_frames_all_256 * 2 - 1, output_frames_all_256 * 2 - 1).mean() |