Gaussian Splatting 학습이 빨라졌다고요?

Hwan Heo·2024년 12월 18일

3D GS Gaussian Splatting Neural Rendering

Neural Rendering

목록 보기

17/17

1. Introduction

2024년 하반기엔 공적으로는 사내 데모와 연구 논문 작업 등, 사적으로는 기초군사훈련 등 바쁜 일이 많아 블로그를 잠시 방치하고 있었다. 훈련 수료 후 블로그 복귀를 위해 적절한 아이템이 없나 찾아보던 중, original 3D GS repo 에서 diff-gaussian-rasterizer 가 업데이트 된 것을 알게 되었다.

특히 training 속도가 이전보다 꽤나 빨라졌는데, Taming-3DGS 에서 제시된

fused SSIM
sparseAdam

을 적용했다고 한다. 각 요소를 3D GS 에 적용하면 다음과 같은 성능 이점이 있다.

1.6× speedup using -optimizer_type default
2.7× speedup using -optimizer_type sparse_adam

이 중에서 사실 ‘sparse_adam’ 요 친구는, SH coefficient update 를 original 보다 1/16 주기로 적게 하는 것에 불과하다.

따라서 이번 글에서는, 그보다는 살짝 복잡한

왜 ‘fused SSIM’ 을 사용할 때 1.6배의 속도 향상이 생기는지

에 대해서 분석해보자.

2. FusedSSIM vs SSIM

2.1. Original SSIM

Gaussian smoothing은 일반적으로 2D 필터를 사용하여 수행된다. 이는 이미지의 모든 픽셀에 대해 2D 필터를 convolution 하는 것이다.

SSIM 계산은 이 Gaussian Filter 를 이용해 총 3단계의 연산을 거치는데,

Gaussian filter 을 통해 평균 $(\mu_x, \mu_y)$ 계산
분산 $(\sigma_x, \sigma_y)$ 및 공분산 계산
최종 SSIM 식 적용:

SSIM(x, y) = \frac{(2\mu_x \mu_y + C_1 ) (2\sigma_{xy}+C_2)}{(\mu_x^2 + \mu_y^2 + C_1)(\sigma_x^2+\sigma_y^2 + C_2)}

이 과정에서 Gaussian filter 은 여러 번 호출되고, 중간 결과가 별도로 저장되면서 메모리 접근 비용과 연산 중복이 발생한다.

각 픽셀마다 평균과 분산을 계산할 때, GPU에서 메모리 접근 (read/write) 비용이 증가.
중간 결과를 반복적으로 GPU 메모리에 저장하고 다시 읽어오는 과정에서 latency 가 발생.

2.2. Separable Filters

Gaussian kernel $G(x,y)$ 의 정의를 상기하자.

G(x,y) = \frac{1}{2 \pi \sigma^2 } \exp \left ( - \frac{x^2 + y^2}{2 \sigma^2} \right )

여기서, 2D Gaussian kernel 은 두 개의 1D Gaussian kernel $G_x(x)$ 와 $G_y (y)$ 로 분리할 수 있음을 알 수 있다.

G(x,y) = G_x(x) \cdot G_y(y)

where

G_x (x) = \frac{1}{ \sqrt{2 \pi} \sigma } \exp \left ( - \frac{x^2}{2 \sigma^2} \right )

G_y (y) = \frac{1}{ \sqrt{2 \pi} \sigma } \exp \left ( - \frac{y^2}{2 \sigma^2} \right )

이를 인지한 상태로, convolution 계산 시의 각각의 방법의 computational cost 를 비교해보자.

2.3. Comparison of Computational Cost

2.3.1. Orignal 2D Convolution

I_\text{out}(i,j)= \sum_{m=-k}^{k} \sum_{n=-k}^{k} G(m,n) \cdot I(i+m,j+n)

total computation cost 는 $O(N \times K^2 )$ 이 된다.

2.3.2. Separable Convolution

Seperable filters 를 이용해 2D convolution을 다음과 같이 두 단계로 나눌 수 있으며,

가로 방향 1D Convolution (Horizontal pass):

I_\text{horizontal}(i,j)= \sum_{m=-k}^{k} G_x(m) \cdot I(i+m,j)

세로 방향 1D Convolution (Vertical pass):

I_\text{vertical}(i,j)= \sum_{n=-k}^{k} G_y(n) \cdot I(i,j+n)

이 때 total computation cost 가 $O(N \times K + N \times K )= O(2N \times K)$ 이다.

이 뿐만 아니라 Fused SSIM 방식은

SSIM 계산의 Gaussian smoothing, 평균, 분산, 공분산 계산을 하나의 CUDA 커널로 병합
중간 결과를 GPU 메모리에 저장하는 대신, 필요한 모든 계산을 단일 CUDA kernel에서 수행하여 메모리 병목 현상을 제거

하여, GPU의 shared memory를 적극 활용해 속도를 크게 높이고, redundant한 memory access를 최소화한다고 한다.

3. 정리

Separable Filters로 Gaussian Convolution 최적화
- 기존 2D convolution 대신, 두 번의 1D convolution을 사용하여 연산량을 $O(N\times K^2)$ 에서 $O(2N \times K)$ 로 대폭 감소.
- CUDA 병렬 처리를 통해 추가적인 성능 향상을 제공.
SSIM 계산의 통합 (Fused)
- Gaussian smoothing, 평균, 분산, 공분산 계산을 단일 CUDA 커널에서 처리하여 메모리 접근 횟수를 최소화.
- GPU의 shared memory 활용으로 병목 현상을 제거, redundant memory access를 제거.
Training budget 과 Gaussian primitves 수가 적은 경우
- Gaussian 개수가 적고 이미지 해상도가 낮을수록 전체 계산 중 SSIM 계산의 비중이 높아지며, 이 경우 Fused SSIM이 병렬화의 강점을 더욱 발휘한다.

엔지니어링적인 발견 같으면서도, 간단한 수학적 직관을 통해서 동일 결과에 성능 향상만 이끌어내는 좋은 아이디어인 것 같다. NeRF 이후 GS 쪽은 사실 과연 ML 보단 Graphics 쪽 스키마가 중요해져 버려서 renderer / rasterizer 최적화에 대한 공부를 소홀히 할 수 없는데, 역시 기본이 가장 중요하다.

여담으로 오랜만에 작업한 연구도 곧 project page 가 공개 될 것 같다. 결과가 꽤 좋으니 기대하시라. Stay Tuned!

Hwan Heo

기타치는AI Researcher

이전 포스트