[3DGS 구현] 9. train 코드 작성, 학습 결과

 

'프로젝트, 연구/3DGS 구현' 카테고리의 글 목록

 

이제 마지막 코드인 train.py를 만들어보자.

os.environ["CUDA_VISIBLE_DEVICES"] = "2"

이 코드는 개인별로 맞게 수정해야한다.

나의 경우 그래픽카드가 3개가 있어 그중 어떤 것을 사용할지 지정했다.

 

resol_scale, num_iterations등의 하이퍼파라미터를 원하는대로 조정하면 된다.

나는 resol_scale은 메모리를 고려하여 4로, num_iterations는 원본 논문을 따라 30000번으로 정했다.

30001로 코드가 되어있는 것은 30000일 때의 png, ply파일을 얻기 위함이다.

 

논문에서처럼 처음 몇백번의 iter를 제외하고 15000 iter까지는 densification을 수행했다.

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "2"

import torch
import torchvision
import random
from tqdm import tqdm
from renderer import render
from gaussianModel import GaussianModel
from camera_loader import load_cameras
import config


from utils import ssim 

class Pipe:
    def __init__(self):
        self.sh_degree = 2
        self.debug = False

def train():

    from points3d_bin_parser import Point3DLoader
    import knn
    
    if not os.path.exists("output"):
        os.makedirs("output")

    loader = Point3DLoader(config.POINTS3D_BIN_PATH)
    points = loader.load()
    initial_scales = knn.compute_initial_scaling(points)

    model = GaussianModel().cuda()
    model.create_from_pcd(points, initial_scales)
    

    resol_scale = 4.0

    print("camera 클래스 로딩")
    cameras = load_cameras(resolution_scale=resol_scale)
    num_iterations = 30001

    #cam_centers = torch.stack([cam.camera_center for cam in cameras])
    #scene_extent = (cam_centers.max(dim=0).values - cam_centers.min(dim=0).values).norm().item()


    cam_centers = torch.stack([cam.camera_center for cam in cameras])
    avg_cam_center = cam_centers.mean(dim=0)
    distances = torch.norm(cam_centers - avg_cam_center, dim=1)
    scene_extent = distances.max().item() * 1.1
    

    # 초기값과 최종값 모두 scene_extent를 곱한다
    def get_lr_decay(iteration, max_iters=30000, lr_init=0.00016 * scene_extent, lr_final=0.0000016 * scene_extent):
            if iteration > max_iters: return lr_final
            return lr_init * ((lr_final / lr_init) ** (iteration / max_iters))
    
    

    # 1. 파라미터별 학습률
    lrs = {
        "xyz": 0.00016 * scene_extent,
        "f_dc": 0.0025,
        "f_rest": 0.0025 / 20,
        "opacity": 0.05,
        "scaling": 0.005,
        "rotation": 0.001
    }
    
    params = [
        {'params': [model._xyz], 'lr': lrs["xyz"], "name": "_xyz"},
        {'params': [model._features_dc], 'lr': lrs["f_dc"], "name": "_features_dc"},
        {'params': [model._features_rest], 'lr': lrs["f_rest"], "name": "_features_rest"},
        {'params': [model._opacity], 'lr': lrs["opacity"], "name": "_opacity"},
        {'params': [model._scaling], 'lr': lrs["scaling"], "name": "_scaling"},
        {'params': [model._rotation], 'lr': lrs["rotation"], "name": "_rotation"}
    ]
    optimizer = torch.optim.Adam(params, lr=0.0, eps=1e-15)
    model.optimizer = optimizer
    
    progress_bar = tqdm(range(0, num_iterations), desc="Training progress")
    
    pipe = Pipe()
    bg_color = torch.tensor([1,1,1], dtype=torch.float32, device="cuda")

    
    


    #threshold = 0.0002 / resol_scale
    threshold = 0.0002
    test_cam = cameras[0]

    for iteration in progress_bar:
        if iteration % 1000 == 0:
            pipe.sh_degree = min(iteration // 1000, 2)
                
        current_xyz_lr = get_lr_decay(iteration)


        
        for param_group in optimizer.param_groups:
            if param_group["name"] == "_xyz":
                param_group["lr"] = current_xyz_lr



        viewpoint_cam = random.choice(cameras)
        
        render_pkg = render(viewpoint_cam, model, pipe, bg_color)
        image = render_pkg["render"]
        viewspace_point_tensor = render_pkg["viewspace_points"] # (N, 3) 2D 좌표+depth
        visibility_filter = render_pkg["visibility_filter"]     # 화면 내 가우시안 표시
        
        # Loss (L1 + SSIM)
        gt_image = viewpoint_cam.original_image
        
        l1_loss = torch.abs(image - gt_image).mean()
        ssim_loss = 1.0 - ssim(image, gt_image)
        loss = (1.0 - 0.2) * l1_loss + 0.2 * ssim_loss
        
        loss.backward()

        model.add_densification_stats(viewspace_point_tensor.grad, visibility_filter, render_pkg["radii"])


        # densification
        if iteration > 600 and iteration < 15000:
            if iteration % 100 == 0:
                size_threshold = 20 if iteration > 3000 else None
                model.densify_and_prune(threshold, 0.01, scene_extent, size_threshold)

                model.reset_densification_stats()
            
            if iteration % 3000 == 0 or iteration == 3000:
                model.reset_opacity()

        optimizer.step() #실제로 가우시안 값 업데이트
        optimizer.zero_grad(set_to_none=True)

        
        if iteration % 1000 == 0:
            count = model.get_xyz.shape[0]
            progress_bar.write(f"Checkpoint saved at iteration {iteration}, gaussian Num : {count}")
            with torch.no_grad(): # 저장할 때는 그래디언트 계산이 필요 없으니 메모리 절약
                test_render = render(test_cam, model, pipe, bg_color)["render"]
                torchvision.utils.save_image(test_render, f"output1/iter_{iteration}.png")

            if iteration % 5000 == 0:
                model.save_ply(f"output1/points_iter_{iteration}.ply")

        
        if iteration % 10 == 0:
            progress_bar.set_description(f"Training ")
            progress_bar.set_postfix({"Loss": f"{loss.item():.4f}"})

if __name__ == "__main__":
    train()

 

실행 시키면 1000 iter마다 png파일이, 5000 iter마다 ply파일이 저장된다.

ply파일은 용량이 크기 때문에 5000번 마다 하게끔 했다.

이 수치를 바꾸고 싶다면 바꿔도 된다.

 

또 결과가 저장되는 폴더도 지정되어 있는데, 폴더가 없을 때 처음 코드를 돌리면 폴더가 자동생성되지 않을 수도 있으니 이러한 에러가 나온다면 폴더를 만들고 실행하면 된다.

전체적인 프로젝트 폴더 구조는 아래와 같으니 이를 참고해서 폴더를 만들면 수월하다.

 

실행시켜 나온 결과를 약 3초 간격으로 연결하여 변화를 볼 수 있는 동영상을 만들었다.

 

10000번 대의 iter만 되어도 물체는 높은 퀄리티로 렌더링된 것을 볼 수 있다.