NVP(2) - Code Execution

구명규·2023년 3월 15일

INR

'23 Individual Research

목록 보기

7/19

이제 NVP 논문의 코드를 하나씩 뜯어보자.

Code 구성

[Utilities]

/config/config_nvp_l.json $\rarr$ 모델 전반에 대한 hyperparameter 명시.
dataio.py $\rarr$ training/testing data 로드.
loss_functions.py $\rarr$ loss function 정의.
modulation.py modules.py $\rarr$ NVP model structure 정의.
sparsegrid.py $\rarr$ Latent grid 생성.
utils.py $\rarr$ Parameter 수 계산, training summary 등의 기타 기능 수행.

[Training]

training.py train_video.py $\rarr$ 모델 학습(video encoding).

[Evaluation]

1. w/o compression
eval.py $\rarr$ 학습된 latent code quantization & de-quantization, evaluation 수행.
2. w/ compression using codecs
compression.py $\rarr$ 학습된 latent code quantization 수행.
compression.ipynb $\rarr$ Quantized latent code에 대한 HEVC/JPEG compression.
eval_compression.py $\rarr$ Latent code de-quantization, video reconstruction 수행.

Code 실행

학습 데이터의 화질과 모델 크기를 축소하여 Colab GPU 리소스로 실행시켜 보았으며, 그 과정은 아래와 같다.
(https://colab.research.google.com/drive/1Iw2b71uJJxOBnfY9Mpw55_uIsHEdrsGh#scrollTo=OoTtJ4LHtmHn)

[Environment Settings]

처음엔 environment.yaml 상의 모든 requirements를 condacolab을 통해 설치하려 했으나(약 24분 소요), 그럴 필요가 없다는 것을 깨달았다. (웬만한 package들은 Colab environment 상에 이미 설치되어 있더라.)
다만 tiny-cuda-nn 및 기타 package들은 아래와 같이 설치해주어야 한다.

!pip install git+https://github.com/subin-kim-cv/tiny-cuda-nn/#subdirectory=bindings/torch
!pip install sk-video
!pip install pytorch-msssim

논문의 깃헙 코드들 역시 clone 해준다.

!git clone https://github.com/subin-kim-cv/NVP.git
%cd NVP

[Dataset Preparation]

논문에서 사용한 UVG-HD dataset 중 Jockey video를 사용하였으며, 아래와 같이 1920x1080 YUV RAW 파일을 RGB png 파일로 extract하였다.

ffmpeg -f rawvideo -vcodec rawvideo -s 1920X1080 -r 60 -pix_fmt yuv420p -t 10 -i Jockey_1920x1080_120fps_420_8bit_YUV.yuv rgb_converted/f%05d.png

이후 별도로 640x360의 jpg로 downsampling 하였으며, /NVP/rgb_reduced 파일에 저장하였다. (이로 인해 dataio.py 파일의 'png'를 'jpg'로 수정해주어야 한다.)

[Training]

Colab의 제한된 자원을 사용해야 하기에 모델의 크기 역시 아래와 같이 축소하였다.

$L$ (layers of 2D latent grid): 16 $\rarr$ 8
$H, W, S$ (size of 3D latent grid): 300, 300, 600 $\rarr$ 100, 100, 300
$n$ _ $neurons$ (modulator network): 128 $\rarr$ 64
위 정보가 담긴 config_mini_nvp.json 파일을 config 폴더에 추가해준 뒤, 아래 커맨드로 모델을 학습시킨다. (약 1시간 45분 소요)

!CUDA_VISIBLE_DEVICES=0 python experiment_scripts/train_video.py --logging_root ./logs_nvp --experiment_name mini_nvp --num_epochs 10000 --epochs_til_ckpt 2500 --steps_til_summary 100 --dataset ./rgb_reduced --num_frames 300 --config ./config/config_mini_nvp.json

[Evaluation]

Codec을 통한 latent code의 compression 수행 없이 모델을 evaluate 하기 위해서는 아래 커맨드로 eval.py를 실행시키면 된다.

!CUDA_VISIBLE_DEVICES=0 python experiment_scripts/eval.py --logging_root ./logs_nvp --experiment_name mini_nvp --dataset ./rgb_reduced --num_frames 300 --config ./logs_nvp/mini_nvp/config.json

만약, codec compression을 수행하고자 할 경우, 아래 두 커맨드(compression.py, eval_compression.py) 사이에 compression.ipynb 파일의 코드를 별도로 실행시켜 codec의 압축과정을 거쳐야 한다. 이 때, compression.ipynb 파일의 direction을 적절히 수정해주어야 한다.

!CUDA_VISIBLE_DEVICES=0 python experiment_scripts/compression.py --logging_root ./logs_nvp --experiment_name mini_nvp --config ./logs_nvp/mini_nvp/config.json

!CUDA_VISIBLE_DEVICES=0 python experiment_scripts/eval_compression.py --logging_root ./logs_nvp --experiment_name mini_nvp --dataset rgb_reduced --num_frames 300 --config ./logs_nvp/mini_nvp/config.json --qscale 2 3 3 --crf 21 --framerate 25

실행 결과

기존의 모델을 그대로 돌리고자 하였지만 Colab 상에서 자꾸 CUDA OUT OF MEMORY 에러가 발생하여 모델과 데이터셋의 크기를 줄여 학습시켜 보았다.
결과는 BPP score 0.735에 PSNR score 30.633으로, 논문 내 성능과 비슷한 수치를 보이는 것으로 보아 모델과 데이터셋의 규모가 비슷한 수준으로 줄어든 것을 확인할 수 있었다.
앞으로의 모델 개선 등에 대한 실험은 위 실험 결과를 기준으로 진행해볼 계획이다.