methods/wrapper_compgs.py

import os
import sys
import random
import torch
import torch.nn.functional as F
from argparse import ArgumentParser

from core.registry import register_method
from core.base_method import BaseMethod

sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../../CompGS_offy')))
from Modules.TrainerCompGS import TrainerCompGS
from Modules.Common import RenderSettings
from pytorch_msssim import ssim

@register_method("compgs")
class CompGSWrapper(BaseMethod):
    def __init__(self, dataset_config, hyperparams):
        self.source_path = dataset_config["source_path"]
        self.model_path = dataset_config["model_path"]
        self.resolution = dataset_config.get("resolution", 1)
        self.track_decoupling = hyperparams.get("track_decoupling", False)

        dummy_config = os.path.join(os.path.dirname(__file__), '../../CompGS_offy/Configs/TanksAndTemplates.yaml')
        import yaml
        base_cfg_path = "/root/autodl-tmp/CompGS_offy/Configs/TanksAndTemplates.yaml"
        with open(base_cfg_path, "r") as f:
            cfg = yaml.safe_load(f)
            
        # 强制接管核心路径与参数
        cfg["dataset"]["root"] = dataset_config["source_path"]
        cfg["training"]["save_directory"] = dataset_config["model_path"]
        
        # 动态对齐迭代步数
        iters = hyperparams.get("iterations", 30000)
        cfg["training"]["max_iterations"] = iters
        
        os.makedirs(dataset_config["model_path"], exist_ok=True)
        runtime_cfg_path = os.path.join(dataset_config["model_path"], "runtime_config.yaml")
        with open(runtime_cfg_path, "w") as f:
            yaml.dump(cfg, f)
            
        self.trainer = TrainerCompGS(config_path=runtime_cfg_path, override_cfgs={})
        self.gaussian_model = self.trainer.gaussian_model
        self.dataset = self.trainer.dataset
        self.optimizer = self.trainer.gaussian_optimizer
        self.aux_optimizer = self.trainer.aux_optimizer
        
        self.last_n_gaussians = self.gaussian_model.num_coupled_primitive

    def train_iteration(self, step):
        self.trainer.gaussian_lr_scheduler(iteration=step)
        
        sample = self.dataset[random.randint(0, len(self.dataset) - 1)]
        
        render_settings = RenderSettings(
            cam_idx=sample.cam_idx, image_height=sample.image_height, image_width=sample.image_width,
            tanfovx=sample.tan_half_fov_x, tanfovy=sample.tan_half_fov_y, campos=sample.camera_center,
            viewmatrix=sample.world_to_view_proj_mat, projmatrix=sample.world_to_image_proj_mat)

        retain_grad = step < self.trainer.configs['adaptive_control']['stop_iteration']
        render_results = self.gaussian_model.render(render_settings=render_settings, retain_grad=retain_grad)

        ssim_weight = self.trainer.configs['training']['ssim_weight']
        l1_loss = F.l1_loss(render_results.rendered_img, sample.img)
        ssim_loss = 1 - ssim(render_results.rendered_img.unsqueeze(dim=0), sample.img.unsqueeze(dim=0), data_range=1., size_average=True)
        rendering_loss = (1 - ssim_weight) * l1_loss + ssim_weight * ssim_loss

        reg_loss = 0.01 * render_results.scales.prod(dim=1).mean()
        bpp = render_results.bpp
        rate_loss = self.trainer.configs['training']['lambda_weight'] * sum(v for v in bpp.values()) if step > self.trainer.configs['training']['rate_loss_start_iteration'] else torch.tensor(0.0, device=self.trainer.device)

        loss_target = rendering_loss
        loss_parasitic = reg_loss + rate_loss
        loss = loss_target + loss_parasitic

        grad_cos_sim = 0.0
        parasitic_ratio = 0.0
        stats = {}

        if self.track_decoupling and step % 100 == 0:
            params = self.gaussian_model.gaussian_params
            param_groups_map = {
                "spatial": [params.means],
                "geometry": [params.scales_before_exp, params.rotations_before_norm],
                "opacity": [params.ref_feats],
                "appearance": [params.res_feats],
            }

            self.optimizer.zero_grad(set_to_none=True)
            self.aux_optimizer.zero_grad(set_to_none=True)
            loss_target.backward(retain_graph=True)
            
            grad_target = params.means.grad.clone() if params.means.grad is not None else torch.zeros_like(params.means)
            grads_target = {}
            for group_name, plist in param_groups_map.items():
                grads_target[group_name] = torch.cat([p.grad.clone().reshape(-1) for p in plist if p.grad is not None])

            self.optimizer.zero_grad(set_to_none=True)
            self.aux_optimizer.zero_grad(set_to_none=True)
            loss_parasitic.backward(retain_graph=True)
            
            grad_parasitic = params.means.grad.clone() if params.means.grad is not None else torch.zeros_like(params.means)
            grads_parasitic = {}
            for group_name, plist in param_groups_map.items():
                grads_parasitic[group_name] = torch.cat([p.grad.clone().reshape(-1) for p in plist if p.grad is not None])

            valid_mask = (torch.norm(grad_target, dim=1) > 0) & (torch.norm(grad_parasitic, dim=1) > 0)
            if valid_mask.any():
                grad_cos_sim = float(F.cosine_similarity(grad_target[valid_mask], grad_parasitic[valid_mask], dim=1).mean())
            parasitic_ratio = float(torch.norm(grad_parasitic, dim=1).mean() / (torch.norm(grad_target, dim=1).mean() + 1e-7))

            for group_name in param_groups_map:
                gt, gp = grads_target.get(group_name), grads_parasitic.get(group_name)
                if gt is not None and gp is not None and gt.norm() > 0 and gp.norm() > 0:
                    cos = float(F.cosine_similarity(gt.unsqueeze(0), gp.unsqueeze(0)))
                    r = float(gp.norm() / (gt.norm() + gp.norm() + 1e-7))
                    ti = r * max(0.0, -cos)
                else:
                    ti = 0.0
                stats[f"sti_{group_name}"] = ti

            self.optimizer.zero_grad(set_to_none=True)
            self.aux_optimizer.zero_grad(set_to_none=True)
            loss.backward()
            aux_loss = self.gaussian_model.aux_loss
            aux_loss.backward()
        else:
            loss.backward()
            aux_loss = self.gaussian_model.aux_loss
            aux_loss.backward()

        self.trainer.optimize(iteration=step, render_results=render_results)

        num_gaussians = self.gaussian_model.num_coupled_primitive
        metrics = {
            "loss": float(loss), "loss_l1": float(l1_loss), "loss_ssim": float(ssim_loss),
            "num_gaussians": int(num_gaussians), "delta_N": int(num_gaussians - self.last_n_gaussians),
            "peak_vram_GB": float(torch.cuda.max_memory_allocated() / (1024 ** 3)),
            "grad_cos_sim": float(grad_cos_sim), "parasitic_ratio": float(parasitic_ratio)
        }
        metrics.update(stats)
        self.last_n_gaussians = num_gaussians
        
        histograms = {}
        if step % 1000 == 0:
            histograms["scaling"] = torch.exp(self.gaussian_model.gaussian_params.scales_before_exp).clone().detach()
            scales_2d = histograms["scaling"][:, :2]
            histograms["anisotropy"] = scales_2d.max(dim=-1)[0] / (scales_2d.min(dim=-1)[0] + 1e-7)
            histograms["sh_dc_mag"] = self.gaussian_model.gaussian_params.ref_feats.detach().norm(dim=-1)

        return metrics, histograms
        
    def render(self, camera):
        with torch.no_grad():
            render_settings = RenderSettings(
                cam_idx=camera.cam_idx, image_height=camera.image_height, image_width=camera.image_width,
                tanfovx=camera.tan_half_fov_x, tanfovy=camera.tan_half_fov_y, campos=camera.camera_center,
                viewmatrix=camera.world_to_view_proj_mat, projmatrix=camera.world_to_image_proj_mat)
            rendered_img, _, _ = self.gaussian_model.render_inference(render_settings=render_settings)
        return {"image": rendered_img, "depth": None}

    def save(self, save_dir, step):
        ckpt_folder = os.path.join(save_dir, f'iteration_{step}')
        os.makedirs(ckpt_folder, exist_ok=True)
        self.gaussian_model.save_uncompressed_params(os.path.join(ckpt_folder, 'point_cloud.ply'))
        self.gaussian_model.save_weights(os.path.join(ckpt_folder, 'weights.pth'))

    def load(self, model_path, iteration):
        ckpt_folder = os.path.join(model_path, f'iteration_{iteration}')
        self.gaussian_model.load_uncompressed_params(os.path.join(ckpt_folder, 'point_cloud.ply'))
        self.gaussian_model.load_weights(os.path.join(ckpt_folder, 'weights.pth'))

    def get_spatial_centers(self):
        return self.gaussian_model.means

    def compute_physical_metrics(self, cameras=None):
        metrics = {}
        with torch.no_grad():
            scales = torch.exp(self.gaussian_model.gaussian_params.scales_before_exp)
            scales_2d = scales[:, :2] if scales.dim() > 1 and scales.shape[1] >= 2 else scales.unsqueeze(-1).expand(-1, 2)
            
            max_S, _ = torch.max(scales_2d, dim=1)
            min_S, _ = torch.min(scales_2d, dim=1)
            gamma = max_S / (min_S + 1e-7)
            
            metrics["gamma_median"] = float(torch.median(gamma))
            metrics["gamma_90th_percentile"] = float(torch.quantile(gamma, 0.90))
            metrics["scale_mean"] = float(torch.mean(scales_2d))
            
            ref_f = self.gaussian_model.gaussian_params.ref_feats
            res_f = self.gaussian_model.gaussian_params.res_feats
            if res_f is not None and res_f.shape[1] > 0:
                metrics["sh_energy_ratio"] = float(res_f.norm(dim=-1).mean() / (ref_f.norm(dim=-1).mean() + 1e-7))

        return metrics

    def evaluate_spatial_field(self, query_points: torch.Tensor, cameras=None) -> torch.Tensor:
        with torch.no_grad():
            V = query_points.shape[0]
            densities = torch.zeros(V, device=self.trainer.device)
            xyz = self.gaussian_model.means
            scales = torch.exp(self.gaussian_model.gaussian_params.scales_before_exp)
            sigma_sq = scales[:, :2].max(dim=1)[0].pow(2) if scales.shape[1] >= 2 else scales.squeeze().pow(2)
            
            N_gaussians = xyz.shape[0]
            chunk_size = max(1, 30000000 // (N_gaussians + 1))
            for i in range(0, V, chunk_size):
                end = min(i + chunk_size, V)
                dist_sq = torch.cdist(query_points[i:end], xyz, p=2).pow(2)
                weights = torch.exp(-0.5 * dist_sq / (sigma_sq.unsqueeze(0) + 1e-7))
                densities[i:end] = torch.sum(weights, dim=1)
            return densities