X-KANeRF [KANeRF-benchmarking]: KAN-based NeRF with Various Basis Functions

June 30, 2024 · View on GitHub

Is there any basis function can explain the NeRF formula?!

c,σ=FΘ(x,d),\mathbf{c}, \sigma = F_{\Theta}(\mathbf{x}, \mathbf{d}), where c=(r,g,b)\mathbf{c}=(r,g,b) is RGB color, σ\sigma is density, x\mathbf{x} is 3D position, d\mathbf{d} is the direction.

To explore this issue, I used Kolmogorov-Arnold Networks (KAN) with various basis functions to fit the NeRF equation based on nerfstudio.

The code might be a bit COARSE, any suggestions and criticisms are welcome!

X-KAN Models (Here are various KANs!)

TODOBasis FunctionsMathtypeAcknowledgement
1B-SplineSi(x)=ai+bi(xxi)+ci(xxi)2+di(xxi)3S_i(x) = a_i + b_i(x - x_i) + c_i(x - x_i)^2 + d_i(x - x_i)^3Efficient-Kan
2Fourierϕk(x)=sin(2πkx),ϕk(x)=cos(2πkx)\phi_k(x) = \sin(2\pi kx), \phi_k(x) = \cos(2\pi kx)FourierKAN
3Gaussian RBFϕ(x,c)=exc22σ2\phi(x, c) = e^{-\frac{\|x - c\|^2}{2\sigma^2}}FastKAN
4Radial Basis Functionϕ(x,c)=f(xc)\phi(x, c) = f(\|x - c\|)RBFKAN
5FCN-FCN-KAN
6FCN-Interpolation-FCN-KAN
71st Chebyshev PolynomialsTn(x)=cos(ncos1(x))T_n(x) = \cos(n \cos^{-1}(x))ChebyKAN
82nd-Chebyshev PolynomialsUn(x)=sin((n+1)cos1(x))sin(cos1(x))U_n(x) = \frac{\sin((n+1)\cos^{-1}(x))}{\sin(\cos^{-1}(x))}OrthogPolyKANs
9Jacobi polynomialsPn(α,β)(x)=12nn!dndxn[(1x)α+n(1+x)β+n]P_n^{(\alpha, \beta)}(x) = \frac{1}{2^n n!} \frac{d^n}{dx^n} \left[ (1-x)^{\alpha+n} (1+x)^{\beta+n} \right]JacobiKAN
10Hermite polynomialsHn(x)=(1)nex2dndxn(ex2)H_n(x) = (-1)^n e^{x^2} \frac{d^n}{dx^n}(e^{-x^2})OrthogPolyKANs
11Gegenbauer polynomialsCn+1(λ)(x)=2(n+λ)n+1xCn(λ)(x)(n+2λ1)n+1Cn1(λ)(x)C_{n+1}^{(\lambda)}(x) = \frac{2(n+\lambda)}{n+1}x C_n^{(\lambda)}(x) - \frac{(n+2\lambda-1)}{n+1}C_{n-1}^{(\lambda)}(x)OrthogPolyKANs
12Legendre polynomialsPn(x)=12nn!dndxn(x21)nP_n(x) = \frac{1}{2^n n!} \frac{d^n}{dx^n} \left( x^2 - 1 \right)^nOrthogPolyKANs
13Laguerre polynomialsLn(x)=exn!dndxn(xnex)L_n(x) = \frac{e^x}{n!} \frac{d^n}{dx^n} \left( x^n e^{-x} \right)OrthogPolyKANs
14Bessel polynomialsJn(x)=k=0(1)kk!(n+k)!(x2)2k+nJ_n(x) = \sum_{k=0}^{\infty} \frac{(-1)^k}{k!(n+k)!} \left( \frac{x}{2} \right)^{2k+n}OrthogPolyKANs
15Fibonacci polynomialsFn(x)=xFn1(x)+Fn2(x),for n2.F_n(x) = xF_{n-1}(x) + F_{n-2}(x), \quad \text{for } n \geq 2.OrthogPolyKANs
16Lucas polynomialsLn(x)=xLn1(x)+Ln2(x)L_n(x) = xL_{n-1}(x) + L_{n-2}(x)OrthogPolyKANs
17Mexican hat waveletψ(x)=23aπ14(1x2a2)ex22a2\psi(x) = \frac{2}{\sqrt{3a}\pi^{\frac{1}{4}}} \left(1 - \frac{x^2}{a^2}\right) e^{-\frac{x^2}{2a^2}}Wav-KAN
18Morlet wavelet (Gabor wavelet)ψ(t)=π14eiω0tet22\psi(t) = \pi^{-\frac{1}{4}} e^{i\omega_0 t} e^{-\frac{t^2}{2}}Wav-KAN
19Difference of Gaussians(DoG)DoG(x,y)=12πσ1ex2+y22σ1212πσ2ex2+y22σ22\text{DoG}(x, y) = \frac{1}{\sqrt{2\pi}\sigma_1} e^{-\frac{x^2 + y^2}{2\sigma_1^2}} - \frac{1}{\sqrt{2\pi}\sigma_2} e^{-\frac{x^2 + y^2}{2\sigma_2^2}}Wav-KAN
20Meyer waveletψ(x)=2Tk=1N(1(kN)2)[cos(2πxkT)sin(πxk/T)πxk/T]\psi(x) = \sqrt{\frac{2}{T}} \sum_{k=1}^{N} \left(1 - \left(\frac{k}{N}\right)^2\right) \left[ \cos\left(\frac{2\pi x k}{T}\right) - \frac{\sin(\pi x k / T)}{\pi x k / T}\right]Wav-KAN
21Shannon waveletψ(t)=sin(πt)sin(πt2)πt\psi(t) = \frac{\sin(\pi t) - \sin\left(\frac{\pi t}{2}\right)}{\pi t}Wav-KAN
22Bump waveletψ(t)=e11t2\psi(t) = e^{-\frac{1}{1 - t^2}}Wav-KAN
More and More!!!---

Performance Comparision on RTX-3090

Modelhidden_dimhidden_dim_colornum_layersnum_layers_colorgeo_feat_dimappearance_embed_dim
Nefacto-MLP-A32322278
Nefacto-MLP-B888878
Others881178
  • nerf_synthetic: lego / 30k

Note that the current Train Rays/Sec and Train Time(ETA Time) are not accurate, they are the values ​​when the number of iterations reaches 100.

ModelLayer Params \downarrowTrain Rays/Sec \uparrowTrain Time \downarrowFPS \uparrowPSNR \uparrowSSIM \uparrowLPIPS \downarrow
Nerfacto-MLP-A9902~170K~14m0.7132.530.9680.0167
Nerfacto-MLP-B3382~165K~14m0.7527.110.9150.0621
Nerfacto-MLP1118~190K~13m0.9928.600.9520.0346
BSplines-KAN8092~37K~54 m0.1932.330.9650.0174
GRBF-KAN3748~115K~19 m0.5032.390.9670.0172
RBF-KAN3512~140K~15m0.7132.570.9660.0177
Fourier-KAN5222~80K~25 m0.4231.720.9560.0241
FCN-KAN(Iters: 4k)5184~4K~90m0.0229.670.9380.0401
FCN-Interpolation-KAN6912~52K~40m0.2132.670.9650.0187
1st Chebyshev-KAN4396~53K~40m0.3428.560.9240.0523
1st Chebyshev-KAN(faster)4426~86K~23m0.4528.540.9240.0514
Jacobi-KAN3532~72K~30m0.3727.880.9150.0553
Bessel-KAN3532~76K~28m0.3325.790.8780.1156
2nd Chebyshev-KAN4396~55K~39m0.3328.530.9240.0500
Fibonacci-KAN4396~65K~32m0.3428.300.9220.0521
Gegenbauer-KAN4396~53K~40m0.3228.390.9220.0514
Hermite-KAN4396~55K~38m0.3727.580.9130.0591
Legendre-KAN4396~55K~38m0.3326.640.8930.0986
Lucas-KAN3532~75K~28m0.4227.950.9160.0550
Laguerre-KAN3532~74K~28m0.3927.390.9120.0593
MexicanHat-KAN3614~66K~32m0.3531.230.9610.0221
Morlet-KAN3614~67K~31m0.3813.060.6860.2583
DoG-KAN3614~75K~28m0.4132.590.9660.0174
Meyer-KAN3614~36K~55m0.1711.910.7280.2991
Shannon-KAN3614~73K~28m0.499.150.7380.4434
Bump-KAN3614~70K~28m0.331.790.1470.9072

Nerfstudio Installation

# create python env
conda create --name nerfstudio -y python=3.8
conda activate nerfstudio
python -m pip install --upgrade pip

# install torch
pip install torch==2.0.1+cu117 torchvision==0.15.2+cu117 --extra-index-url https://download.pytorch.org/whl/cu117
conda install -c "nvidia/label/cuda-11.7.1" cuda-toolkit

# install tinycudann
pip install ninja git+https://github.com/NVlabs/tiny-cuda-nn/#subdirectory=bindings/torch

# install nerfstudio
pip install nerfstudio==0.3.4

# pip install torchmetrics==0.11.4

# Tab command
ns-install-cli

# !!! If you use `ns-process-data`, please install this version opencv
pip install opencv-python==4.3.0.36

Run

# Train
############# kan_basis_type #############
# mlp, bspline, grbf, rbf, fourier,
# fcn, fcn_inter, chebyshev, jacobi
# bessel, chebyshev2, finonacci, hermite
# legendre, gegenbauer, lucas, laguerre
# mexican_hat, morlet, dog, meyer, shannon, bump
bash train_blender.sh [kan_basis_type]

# eval
bash run_eval.sh [exp_path]

# render RGB & Depth
bash run_render.sh [exp_path]

Docs

Citation

If you use this benchmark in your research, please cite this project.

@misc{xkanerf,
	title={X-KANeRF: KAN-based NeRF with Various Basis Functions},
	author={Linfei Li},
	howpublished = {\url{https://github.com/lif314/X-KANeRF}},
	year={2024}
}

Acknowledgement

  • KANeRF, A big thank you for this awesome work!
    @Manual{kanerf,
    	title = {Hands-On NeRF with KAN},
    	author = {Delin Qu, Qizhi Chen},
    	year = {2024},
    	url = {https://github.com/Tavish9/KANeRF},
    }
    
  • nerfstudio
    @inproceedings{nerfstudio,
    	title = {Nerfstudio: A Modular Framework for Neural Radiance Field Development},
    	author = {
    		Tancik, Matthew and Weber, Ethan and Ng, Evonne and Li, Ruilong and Yi, Brent
    		and Kerr, Justin and Wang, Terrance and Kristoffersen, Alexander and Austin,
    		Jake and Salahi, Kamyar and Ahuja, Abhik and McAllister, David and Kanazawa,
    		Angjoo
    	},
    	year = 2023,
    	booktitle = {ACM SIGGRAPH 2023 Conference Proceedings},
    	series = {SIGGRAPH '23}
    }