[NoPe-NeRF, CUDA 11.8] 코드 빌드 및 Tanks and Temples & NeRF LLFF 전체 데이터 셋 학습 Python 코드

아래 github 페이지를 참고하여 진행

https://github.com/ActiveVisionLab/nope-nerf/

Installation

기본적인 설치 방법은 아래와 같다.

git clone https://github.com/ActiveVisionLab/nope-nerf.git
cd nope-nerf
conda env create -f environment.yaml
conda activate nope-nerf

 

그러나 environment.yaml 파일을 까보면 pytorch 1.7, cudatoolkit 10.1을 쓰고 있는데, 현재 cuda 11.8을 쓰고 있어서 여기서 돌아가도록 python, pytorch, torchvision, torchaudio, cudatoolkit 부분을 주석 처리 후에 새롭게 인스톨했다.

 

이는 수정된 environment.yaml

name: nope-nerf
channels:
  - pytorch
  - conda-forge
  - anaconda
  - defaults
dependencies:
  # - python=3.9
  # - pytorch=1.7
  # - torchvision=0.8.2 
  # - torchaudio 
  # - cudatoolkit=10.1
  - cffi
  - cython
  - imageio
  - numpy
  - scipy
  - matplotlib
  - pandas
  - tensorboard
  - yaml
  - pillow
  - wheel
  - pip
  - tqdm
  - pip:
    - ipdb
    - ipython
    - ipython-genutils
    - jedi
    - opencv-python
    - scikit-image
    - pyyaml
    - trimesh
    - imageio
    - matplotlib
    - plyfile
    - timm
    - lpips
    - setuptools
    - kornia==0.5.0
    - imageio-ffmpeg

 

그리고 새롭게 설치. 사실 아래 옵션에 맞게 environment.yaml 파일을 수정하는 게 제일 간편하다(만약 위의 Installation 코드를 있는 그대로 돌렸다면, 파일을 수정한 다음에 conda env update --file environment.yaml 하면 됨)

conda install python=3.10
conda install pytorch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2 pytorch-cuda=11.8 -c pytorch -c nvidia
conda install cudatoolkit=11.8

 

이렇게 환경 구성은 끝.

 

Training with Tanks and Temples Dataset 

위 Github 레포지토리에서 제공하는 데이터 셋을 그대로 다운로드(https://www.robots.ox.ac.uk/~wenjing/Tanks.zip)하고, nope-nerf 루트 디렉토리에 data 디렉토리를 생성한 다음 해당 위치에 unzip하였다.

이미 DPT 결과 데이터도 포함되어 있어 따로 전처리를 진행할 필요는 없다. 아래는 Ignatius 데이터 셋을 돌리는 예시.

# Ignatius.yaml
depth:
  type: None
pose:
  learn_pose: True
dataloading:
  path: data/Tanks
  scene: ['Ignatius']
  customized_focal: False
  random_ref: 1
training:
  out_dir:  out/Tanks/Ignatius
  auto_scheduler: True
extract_images:
  resolution: [540, 960]

python train.py configs/Tanks/Ignatius.yaml

 

Training with LLFF Dataset 

LLFF 데이터 셋을 다운로드 받은 후, 이또한 data 디렉토리에 압축 해제하였다. 그리고 LLFF는 DPT 결과물이 필요하다. 따라서 전처리를 진행해줘야 한다. 전처리는 configs/preprocess.yaml에 적혀진 정보를 기반으로 진행하는 것이 기본 세팅. 당연 다른 yaml 파일을 만들어도 되지만, 귀찮아서 얘만 썼다.

python preprocess/dpt_depth.py configs/preprocess.yaml

 

그런데 주의할 점이 있다. LLFF 데이터 셋에 대한 학습 config 파일을 보면 1/4로 rescale해서 쓰도록 되어있다(resize_factor: 4).

# fern.yaml
depth:
  type: None
pose:
  learn_pose: True
rendering:
  depth_range: [0.0, 1.0]
  dist_alpha: True
  sample_option: ndc
dataloading:
  path: data/nerf_llff_data
  scene: ['fern']
  random_ref: 1
  resize_factor: 4
training:
  out_dir:  out/llff/fern
  vis_resolution: [75, 100]
extract_images:
  resolution: [756, 1008]

 

그래서 전처리를 할 때도 이에 맞춰서 설정을 진행해줘야 한다. 아니면 data/nerf_llff_data/fern/images_4가 아니라 data/nerf_llff_data/fern/images를 기준으로 DPT를 진행하고, 이러면 DPT 결과물의 파일 명과 학습할 때의 파일 명이 일치하지 않아 학습이 진행되지 않는다(에러 뿜뿜).

 

기본 preprocess.yaml 파일

depth:
  type: DPT
dataloading:
  path: data/nerf_llff_data
  scene: ['room']
  resize_factor:
  load_colmap_poses: False
training:
  mode: 'all'

 

수정 파일

depth:
  type: DPT
dataloading:
  path: data/nerf_llff_data
  scene: ['room']
  resize_factor: 4
  load_colmap_poses: False
training:
  mode: 'all'

 

resize_factor 뒤에 resize한 scale만 적어주면 된다.

 

전체 학습 코드

import subprocess
import time
import os
import torch

datasets = ['LLFF/fern.yaml', 'LLFF/flower.yaml', 'LLFF/fortress.yaml', 'LLFF/horns.yaml', 'LLFF/leaves.yaml', 'LLFF/orchids.yaml', 'LLFF/room.yaml', 'LLFF/trex.yaml',
            'Tanks/Ballroom.yaml', 'Tanks/Barn.yaml', 'Tanks/Church.yaml', 'Tanks/Family.yaml', 'Tanks/Horse.yaml', 'Tanks/Ignatius.yaml', 'Tanks/Museum.yaml']  # 여러 데이터셋의 이름

# 설정하고 싶지 않으면 ""로 만들면 됨
set_target_gpu = "CUDA_VISIBLE_DEVICES=0"

# [개선 코드] 파일 입출력에 대한 오버헤드를 줄인 코드
def gs_process_with_fileio_nrt2():
    for dataset_name in datasets:
        print(f'target dataset: {dataset_name}')
        data_name = dataset_name.split('/')[1].split('.')[0]

        # 디렉토리 생성
        output_directory = f"output/{data_name}"
        os.makedirs(output_directory, exist_ok=True)

        # CUDA_VISIBLE_DEVICES를 설정하고 Python 스크립트를 실행하는 단계
        train_cmd = f"{set_target_gpu} python train.py configs/{dataset_name}"

        train_process = subprocess.Popen(train_cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        train_output, train_error = train_process.communicate()
        with open(f"{output_directory}/train_{data_name}_{int(time.time())}.txt", "w") as f:
            f.write(train_output.decode())
            f.write(train_error.decode())



if __name__ == "__main__":
    gs_process_with_fileio_nrt2()

각각의 데이터 셋에 대한 학습 결과를 file로 저장한다.

 

학습 시간 진짜 무지막지하게 걸린다