[3DGS 구현] 8. Densification 및 기타 util코드

utils.py에 함수 추가

학습을 위한 코드를 짜보자.

먼저 손실함수를 보면 SSIM이라는 지표가 있다.

이를 구하는 식은 아래와 같다.

두 이미지의 평균,분산을 골고루 이용해서 유사성을 평가하는 지표이다.

따라서 utils.py에 이를 코드로 구현해 추가하자.

def ssim(img1, img2, window_size=11, size_average=True):
    # 0~1 사이의 텐서
    channel = img1.size(-3)
    window = torch.ones((channel, 1, window_size, window_size), device=img1.device) / (window_size**2)
    
    mu1 = F.conv2d(img1, window, padding=window_size//2, groups=channel)
    mu2 = F.conv2d(img2, window, padding=window_size//2, groups=channel)

    mu1_sq = mu1.pow(2)
    mu2_sq = mu2.pow(2)
    mu1_mu2 = mu1 * mu2

    sigma1_sq = F.conv2d(img1 * img1, window, padding=window_size//2, groups=channel) - mu1_sq
    sigma2_sq = F.conv2d(img2 * img2, window, padding=window_size//2, groups=channel) - mu2_sq
    sigma12 = F.conv2d(img1 * img2, window, padding=window_size//2, groups=channel) - mu1_mu2

    C1 = 0.01**2
    C2 = 0.03**2

    ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2))

    return ssim_map.mean() if size_average else ssim_map

 

추가적으로 이미 qvec2rotmat가 있지만, 이는 CPU(numpy)에서 돌아가는 코드이다.

나중을 위해 GPU버전인 quat_to_rotmat도 만들었다.

def quat_to_rotmat(q):

    norm = torch.nn.functional.normalize(q)
    r, x, y, z = norm[:, 0], norm[:, 1], norm[:, 2], norm[:, 3]

    res = torch.zeros((q.shape[0], 3, 3), device=q.device)

    res[:, 0, 0] = 1 - 2 * (y**2 + z**2)
    res[:, 0, 1] = 2 * (x * y - r * z)
    res[:, 0, 2] = 2 * (x * z + r * y)
    res[:, 1, 0] = 2 * (x * y + r * z)
    res[:, 1, 1] = 1 - 2 * (x**2 + z**2)
    res[:, 1, 2] = 2 * (y * z - r * x)
    res[:, 2, 0] = 2 * (x * z - r * y)
    res[:, 2, 1] = 2 * (y * z + r * x)
    res[:, 2, 2] = 1 - 2 * (x**2 + y**2)

    return res

 

 

Densification

학습과정 중에는 가우시안을 쪼개거나 복사하는 과정인 densification이 필요하다.

 

gaussianModel.py 의 GaussianModel클래스의

__init__부분에 아래와 같이 코드를 추가한다.

densification을 위한 변수를 선언하는 과정이다.

 

기본적으로 여러번의 iter에서 누적된 gradient의 평균을 보고 densify를 할지말지 정할 것이기 때문에 

누적된 gradient합을 저장할 변수와 평균을 위한 분모에 들어갈 개수를 저장할 변수를 만들어 주었다.

또 화면을 과도하게 가리는 가우시안인지를 판단하기 위해 반지름을 저장하는 변수도 만들었다.

        #densification을 위해서
        self.xyz_gradient_accum = torch.empty(0) # 누적된 그래디언트 합
        self.denom = torch.empty(0)              # 가우시안이 화면에 나타난 횟수
        self.max_radii2D = torch.empty(0)       # 회면에 맻힌 최대 반지름

 

create_from_pcd에서 모든 값 0으로 초기화한다.

	num_points = self._xyz.shape[0]
        self.xyz_gradient_accum = torch.zeros((num_points, 1), device="cuda")
        self.denom = torch.zeros((num_points, 1), device="cuda")
        self.max_radii2D = torch.zeros((num_points), device="cuda")

 

 

add_densification_stats함수,  reset_densification_stats함수

각 가우시안들이 화면에 나올 때마다 gradient들을 더하는 함수이다.

이 때 같이 반지름의 최댓값도 저장한다.

 

밑의 reset함수에서는 이 모든 값을 0으로 만든다.

이 값을 기준으로 densify를 했다면 다시 새로운 누적값을 저장하기 전에 초기화를 해야하기 때문이다.

 

    def add_densification_stats(self, viewspace_grad, visibility_filter, radii):

        # 현재 뷰의 gradient계산
        grad_norm = torch.norm(viewspace_grad[visibility_filter, :2], dim=-1, keepdim=True)
        
        # 2. 보이는 가우시안한테만 누적
        self.xyz_gradient_accum[visibility_filter] += grad_norm
        
        # 3. 가우시안들이 화면에 나타난 횟수 1씩 증가
        self.denom[visibility_filter] += 1

        self.max_radii2D[visibility_filter] = torch.max(
            self.max_radii2D[visibility_filter], 
            radii[visibility_filter]
        )
    
    def reset_densification_stats(self):
        # 누적된 합과 카운터를 다시 0으로 초기화
        self.xyz_gradient_accum.fill_(0.)
        self.denom.fill_(0.)
        self.max_radii2D.fill_(0.)

 

densify_and_clone함수

핵심 함수 중 하나이다.

위에서 구한 평균 gradient가 임계값 이상이고, 크기가 작으면 가우시안을 복제한다.

이렇게 새로 만들어진 가우시안을 기존의 가우시안 리스트에 추가해줘야한다.

이 때 densification_postfix라는 함수를 사용할 것인데 세부 동작은 나중에 알아보자.

    def densify_and_clone(self, grad_threshold, scene_extent):
        # 평균 그래디언트 계산
        grads = self.xyz_gradient_accum / self.denom
        grads[torch.isnan(grads)] = 0
        
        # 평균 그래디언트 > 임계 , AND 크기가 작아야함
        selected_pts_mask = torch.where(grads >= grad_threshold, True, False).squeeze()
        selected_pts_mask &= torch.where(torch.max(self.get_scaling, dim=1).values <= 0.01 * scene_extent, True, False)
        
        # 복제
        new_xyz = self._xyz[selected_pts_mask]
        new_features_dc = self._features_dc[selected_pts_mask]
        new_features_rest = self._features_rest[selected_pts_mask]
        new_opacity = self._opacity[selected_pts_mask]
        new_scaling = self._scaling[selected_pts_mask]
        new_rotation = self._rotation[selected_pts_mask]

        # 기존 파라미터 리스트 뒤에 붙이기
        self.densification_postfix(new_xyz, new_features_dc, new_features_rest, new_opacity, new_scaling, new_rotation)

 

densify_and_split함수

이번에는 gradient가 임계값 이상이고, 크기가 크면 가우시안을 쪼개는 함수를 살펴보자.

한 개의 큰 가우시안 안에서 쪼개질 두 개의 가우시안의 중심점을 찾아야 한다.

코드 상에서 이는 부모 가우시안을 기준으로 계산되기 때문에 두 점을 추출한 후에는 부모의 회전/이동 행렬을 적용해 월드좌표계로 변환한다.

논문에 따라 쪼개진 가우시안들은 부모 가우시안 보다 1.6배 작아야 한다.

    def densify_and_split(self, grad_threshold, scene_extent):
        grads = self.xyz_gradient_accum / self.denom
        grads[torch.isnan(grads)] = 0
        
        # 평균 그래디언트 > 임계, AND 크기가 커야함
        selected_pts_mask = torch.where(grads >= grad_threshold, True, False).squeeze()
        selected_pts_mask &= torch.where(torch.max(self.get_scaling, dim=1).values > 0.01 * scene_extent, True, False)

        #분포 내 위치 2개 뽑기
        stds = self.get_scaling[selected_pts_mask].repeat(2, 1)
        means = torch.zeros((stds.size(0), 3), device="cuda")
        samples = torch.normal(means, stds) #(0,0,0)중심이고 편차가 std인 분포에서 뽑기
        rots = self.get_rotation[selected_pts_mask].repeat(2, 1)
        
        # 변환
        new_xyz = torch.bmm(quat_to_rotmat(rots), samples.unsqueeze(-1)).squeeze(-1) + self._xyz[selected_pts_mask].repeat(2, 1)
        
        # 스케일 감소 
        new_scaling = torch.log(self.get_scaling[selected_pts_mask].repeat(2, 1) / 1.6)
        
        # 나머지 속성들은 그대로 복사
        new_features_dc = self._features_dc[selected_pts_mask].repeat(2, 1, 1)
        new_features_rest = self._features_rest[selected_pts_mask].repeat(2, 1, 1)
        new_opacity = self._opacity[selected_pts_mask].repeat(2, 1)
        new_rotation = self._rotation[selected_pts_mask].repeat(2, 1)

        # 새로운 점들 추가, 부모 점들은 삭제
        self.densification_postfix(new_xyz, new_features_dc, new_features_rest, new_opacity, new_scaling, new_rotation)
        padded_prune_mask = torch.zeros(self._xyz.shape[0], device="cuda", dtype=bool)
        
        # 앞부분(부모들 영역)에만 원래의 selected_pts_mask
        # selected_pts_mask의 길이는 딱 부모들이 있던 구간까지만 해당함
        padded_prune_mask[:selected_pts_mask.shape[0]] = selected_pts_mask
        
        
        self.prune_points(padded_prune_mask)

padded_prune_mask는 기존 가우시안 개수 + 추가된 가우시안 개수 만큼의 길이를 가지고 모든 값이 false로 초기화된 리스트라고 생각하면 된다.

원래 있던 가우시안들 중 분열이 될(부모) 가우시안의 정보가 써져있는 selected_pts_mask를 사용해 padded_prune_mask 중 해당되는 값을 true로 바꾼다.

prune_points를 통해 해당 true 가우시안들을 삭제한다.

 

densification_postfix함수

위에서 계속 사용된 함수이다.

한마디로 기존 가우시안 리스트에 새 가우시안들을 붙이는 함수라고 볼 수 있다.

    def densification_postfix(self, new_xyz, new_features_dc, new_features_rest, new_opac, new_scaling, new_rot):

        old_params = [
            self._xyz, self._features_dc, self._features_rest, 
            self._opacity, self._scaling, self._rotation
        ]

        # 기존 + 신규 합치기
        # 각 텐서의 차원을 맞춰서 cat 
        d = {
            "_xyz": torch.cat([self._xyz, new_xyz], dim=0),
            "_features_dc": torch.cat([self._features_dc, new_features_dc], dim=0),
            "_features_rest": torch.cat([self._features_rest, new_features_rest], dim=0),
            "_opacity": torch.cat([self._opacity, new_opac], dim=0),
            "_scaling": torch.cat([self._scaling, new_scaling], dim=0),
            "_rotation": torch.cat([self._rotation, new_rot], dim=0)
        }

        #새 텐서로 교체 
        for name, tensor in d.items():
            setattr(self, name, nn.Parameter(tensor.cuda()))

        # Optimizer 업데이트
        self.update_optimizer_after_densify(new_xyz.shape[0], old_params)

        # 4. 통계 변수도 새 개수(N + M)에 맞춰 확장
        num_new = new_xyz.shape[0]
        self.xyz_gradient_accum = torch.cat([self.xyz_gradient_accum, torch.zeros((num_new, 1), device="cuda")], dim=0)
        self.denom = torch.cat([self.denom, torch.zeros((num_new, 1), device="cuda")], dim=0)
        self.max_radii2D = torch.cat([self.max_radii2D, torch.zeros(num_new, device="cuda")], dim=0)

 

update_optimizer_after_densify함수

denfisication_postfix함수에서 개수가 늘어난 새 텐서가 생성되었다.

이는 원래의 텐서와 다른 주소를 가진다.

이 함수는 옵티마이저에게 이 새 주소를 넣어주는 역할이다.

    def update_optimizer_after_densify(self, num_new, old_params):

        for i, group in enumerate(self.optimizer.param_groups):
            #xyz찾아 새 파라미터로 교체
            name = group["name"]

            old_p = old_params[i]
            stored_state = self.optimizer.state.get(old_p, None)
            
            new_p = getattr(self, f"_{name}")
            group['params'][0] = new_p

            if stored_state is not None:
                # Adam 상태 확장
                stored_state["exp_avg"] = torch.cat([stored_state["exp_avg"], torch.zeros((num_new, *stored_state["exp_avg"].shape[1:]), device="cuda")], dim=0)
                stored_state["exp_avg_sq"] = torch.cat([stored_state["exp_avg_sq"], torch.zeros((num_new, *stored_state["exp_avg_sq"].shape[1:]), device="cuda")], dim=0)

                # 옛날 주소 데이터 삭제 후 새 주소에 이식
                del self.optimizer.state[old_p]
                self.optimizer.state[new_p] = stored_state

 

prune_points함수

위에서 봤던 가우시안을 삭제하는 함수이다.

전체 가우시안 리스트에서 mask에 표시된 가우시안들을 지우고 이 결과를 새로 메모리에 할당한다.

mask는 지울 것들이 true라서 '~'을 통해 비트 반전시켜 살릴 것들을 true로 바꾼다.

	def prune_points(self, mask):
        valid_points_mask = ~mask
        

        optimizable_tensors = self.update_optimizer_after_pruning(valid_points_mask)


        self._xyz = optimizable_tensors["_xyz"]
        self._features_dc = optimizable_tensors["_features_dc"]
        self._features_rest = optimizable_tensors["_features_rest"]
        self._opacity = optimizable_tensors["_opacity"]
        self._scaling = optimizable_tensors["_scaling"]
        self._rotation = optimizable_tensors["_rotation"]


        self.xyz_gradient_accum = self.xyz_gradient_accum[valid_points_mask]
        self.denom = self.denom[valid_points_mask]
        self.max_radii2D = self.max_radii2D[valid_points_mask]

 

update_optimizer_after_pruning함수

옵티마이저가 보는 주소를 바꿔준다.

가우시안이 늘어났을 경우 gradient의 모멘텀에 대한 정보가 없기 때문에 이를 0으로 넣어놔 에러가 나지 않게 한다.

	def update_optimizer_after_pruning(self, valid_points_mask):
        optimizable_tensors = {} 
        
        for group in self.optimizer.param_groups:
            old_p = group['params'][0]
            stored_state = self.optimizer.state.get(old_p, None)

            # 새 파라미터 생성
            new_p = nn.Parameter(old_p[valid_points_mask])
            group['params'][0] = new_p

            if stored_state is not None:
                stored_state["exp_avg"] = stored_state["exp_avg"][valid_points_mask]
                stored_state["exp_avg_sq"] = stored_state["exp_avg_sq"][valid_points_mask]

                del self.optimizer.state[old_p]
                self.optimizer.state[new_p] = stored_state
            

            optimizable_tensors[group["name"]] = new_p 
            
        return optimizable_tensors

 

3DGS는 SH함수를 사용하는데 우리가 아는 0~1사이의 RGB값을 SH함수의 계수로 변환해야한다.

C0 = 0.28209479177387814

def RGB2SH(rgb):
    return (rgb - 0.5) / C0

    def create_from_pcd(self, points3d_dict, knn_distances):
        print("가우시안 파라미터 초기화 및 GPU로 로딩 시작")
        
        xyz = np.array([p['xyz'] for p in points3d_dict.values()])
        rgb = np.array([p['rgb'] for p in points3d_dict.values()]) / 255.0
        sh_dc = RGB2SH(rgb)
        
        # 가우시안 좌표, 색 관련
        self._xyz = nn.Parameter(torch.tensor(xyz, dtype=torch.float32).cuda())
        self._features_dc = nn.Parameter(torch.tensor(sh_dc, dtype=torch.float32).unsqueeze(1).cuda())

 

가우시안을 늘리고 줄이는 과정을 하나에 담은 함수이다.

또 opacity값을 통채로 0으로 바꿀 때 옵티마이저와 연결이 끊기지 않도록 하는 replace_tensor_to_optimizer함수도 만들었다.

    def densify_and_prune(self, grad_threshold, min_opacity, scene_extent, max_screen_size):
        # 늘리기 
        self.densify_and_clone(grad_threshold, scene_extent)
        self.densify_and_split(grad_threshold, scene_extent)

       
        # 투명도가 너무 낮은(0.005 미만) 애들
        prune_mask = (self.get_opacity < min_opacity).squeeze()

        # 화면에 너무 크게 맺히거나 공간상에서 너무 큰 애들
        if max_screen_size:
            big_points_vs = self.max_radii2D > max_screen_size
            #big_points_vs = torch.zeros_like(self.get_opacity, dtype=torch.bool).squeeze()
            big_points_ws = self.get_scaling.max(dim=1).values > 0.1 * scene_extent
            prune_mask = torch.logical_or(torch.logical_or(prune_mask, big_points_vs), big_points_ws)

        
        self.prune_points(prune_mask)
        torch.cuda.empty_cache()

	
    def replace_tensor_to_optimizer(self, tensor, name):
        optimizable_tensors = {}
        for group in self.optimizer.param_groups:
            if group["name"] == name:
                stored_state = self.optimizer.state.get(group['params'][0], None)
                if stored_state is not None:
                    stored_state["exp_avg"] = torch.zeros_like(tensor)
                    stored_state["exp_avg_sq"] = torch.zeros_like(tensor)

                    del self.optimizer.state[group['params'][0]]
                    group["params"][0] = nn.Parameter(tensor.requires_grad_(True))
                    self.optimizer.state[group['params'][0]] = stored_state

                    optimizable_tensors[group["name"]] = group["params"][0]
        return optimizable_tensors

 

논문에서 3000번마다 투명도를 리셋하는 과정이 있기 때문에 이를 해주는 함수를 만들었다.

    def reset_opacity(self):
        opacities_new = torch.min(self.get_opacity, torch.ones_like(self.get_opacity) * 0.01)
        eps = 1e-7
        opacities_new = torch.clamp(opacities_new, min=eps, max=1-eps)
        opacities_logit = torch.log(opacities_new / (1 - opacities_new))
        
        
        optimizable_tensors = self.replace_tensor_to_optimizer(opacities_logit, "_opacity")
        self._opacity = optimizable_tensors["_opacity"]