Ref-Nerf Implement And Speeding Up w/ instant-ngp

好耶，复活了。时隔一年，我学会了不少新东西呢。比如，如何使用 instant-ngp 给 RefNerf 加速。

Ref-Nerf Implement And Speeding Up w/ instant-ngp

简介

CV ，炼丹界的内卷之王。 NeRF ，可谓是 CV 界的内卷之王。在今年的 CVPR 上，Ref NeRF 作为 Best Student Paper Honorable Mention 横空出世，它基于 mip-NeRF 框架对反射现象进行了建模，在 NeRF 的基础上增加了 opacity, roughness, tint, normal 的估计。本文将会首先介绍一下 Ref NeRF 论文本身，然后简单聊聊复现思路以及使用 instant-ngp 加速的过程。不过目前，Ref NeRF 的复现结果还没有达到令我满意的程度，只是具备雏形，毕竟融合两篇文章的 idea 可能导致冲突，深度学习这种玄学就更是这样了，模型炸了都不知道从哪一个先开始。

Ref-NeRF 论文简析

反射建模

Ref-NeRF 的一大亮点就在于它对于反射的建模，根据 Phong 模型，被观察到的光线为：

$$I_o=k_aI_a+\sum_ik_d(l_i\cdot n)I_i+\sum_ik_s(r_i\cdot v)^pI_i$$

其中，下标为 a 的部分与环境光（ambient light）有关；下标含有 d 则与漫反射有关，而下标含有 s 则与镜面高光有关。为了更好地解释 Ref NeRF 的反射模型处理以及回顾一下 Phong 模型这个基础模型，这里对模型简单解释：

• 第一项代表环境光对观察结果的影响：环境光只对结果产生常数 offset
• 第二项代表了所有光源在物体上产生的漫反射。漫反射与观察方向无关，只要能观察到，就是 “恒定” 的。光源 i ($l_i$ 为其发射的某一光线的方向）发射光越是能垂直表面（与法向平行），漫反射越强。
• 第三项代表了镜面高光：不一定是完全的镜面反射（所以称之为镜面高光），可以稍显模糊，但是其强度是与视角有关的。$r_i$ 代表了反射光线的方向，$v$ 则代表了观察方向。两者重合（点乘结果接近 1）时，反射较强。指数 p 用于加强衰减，可知 p 越大，$(r_i\cdot v)^p$ 变化越快。也即视线与反射方向重合发生变化时，镜面高光的变化越明显。

而我们的 Ref-NeRF 对后面两部分进行了建模：

$$L_{out}(\hat{\omega_o})\propto\int L_{in}(\hat{\omega_i})f(\hat{\omega_i},\hat{\omega_r})\mathrm d\hat{\omega_i}$$

其中，$L_{in}$ 为入射光，$L_{out}$ 为出射光，$f$ 为 BRDF（Bidirectional Reflectance Distribution Function）。很明显，这就是一个关于 $\omega_r$ 的函数，而这个 BRDF 我们试图通过神经网络直接学习出来。不过这也让我有些思考，如果把 Phong 模型换成 Blinn-Phong 模型会不会有所加速？毕竟 Blinn-Phong 模型不需要计算反射方向，而只需要通过平行四边形定理得到观察方向和入射方向的 “法线” 即可。

directional MLP 的修改

在这里作者进行了一个类似于 mip NeRF 光锥采样的操作。在 mip NeRF 中，我们为了达到 mipmap 那样的感觉，进行了个一个光锥上的采样。而在 Ref NeRF 中，作者也认为 “反射” 也不能仅仅考虑单个反射方向。由于物体实际是凹凸不平的，并非完美的镜面，表面法向量并非完全一致。可以认为，物体表面的凹凸起伏（噪声），使得反射发生了一些改变（distortion），物体表面噪声的概率分布，经过反射的数学操作后被映射成了新的分布。

作者在这里使用了一个被称为 vMF 的概率分布，看上去就和高斯分布差不多，甚至有个别称叫做 normalized spherical gaussian 。个人推测作者选这个函数是为了方便进行后续的数学推导以及近似，就 像 mip NeRF 中，将光锥用混合高斯模型进行近似一样。有了此分布，自然需要使用积分将所有可能的反射方向考虑进去。在这里，作者使用了球谐函数进行方向编码，然后对其进行积分。

$$E_{\hat\omega\sim vMF(\hat\omega_r,\kappa)}\left[Y^m_l(\hat\omega)\right]=\exp(\frac{-l(l+1)}{2\kappa})Y^m_l(\hat\omega)$$

其中，$Y^m_l$ 为球谐函数，$\kappa$ 为 vMF 分布的参数，等于粗糙度的倒数。这里的 $l$ 代表了球谐函数的阶数，$l$ 越大，球谐函数的分布越均匀，而 $l$ 越小，球谐函数的分布越集中。$\kappa$ 越小，对于高阶（频）球谐影响越大。这也就反映了这样一个事实：光滑程度减小，高低频信息均有所衰减，但高频信息衰减更严重。

这里作者还给 directional MLP 喂了 $l_i\cdot n$ ，当然不是要建模漫反射，而是考虑到一些稍微复杂一些的效应，如菲涅尔效应（并非全反射，部分反射部分折射，反射所占的比例需要另外计算）等等。

Normal 估计

作者有言：

normal vectors estimated from its volume density gradient as in Equation (上面那个) are often extremely noisy

NeRF tends to “fake” specular highlights by embedding emitters inside the object and partially occluding them with a “foggy” diffuse surface

为此，作者设计了两个正则项，

• 针对第一个问题，作者通过 spatial MLP 预测法向量 $\hat n$ ，将其和梯度法向量 $n^\prime$ 求一个 MSE。

$$L_n=\sum_i\omega_i||{\hat n_i-n_i^\prime}||^2$$

由于 $\hat n$ 比较光滑，而 $n^\prime$ 是通过梯度求一阶导数求出来的，不是很光滑，所以我们可以通过光滑的数据给不光滑的梯度学一学，让它更加光滑。至于为何，小编也不知道（bushi），我认为大概是网络比较菜，所以只能学习低频部分，所以比较光滑。

• 针对第二个问题，作者设计了一个正则项，用于惩罚反向的法向量。

$$R_o=\sum_i\omega_i\max\left(0, \hat n_i\cdot\hat{\mathrm d}\right)^2$$

其实就是惩罚法向量和光线方向相同的部分，而且和 weight 有关。这就说明：被遮挡的有效背面不会被影响（weight 低），而 fake surface（在半透明 surface 后的 embedded emitter 表面）将会被惩罚（其 density 衰减方向与光线方向一致）。

然后作者说，把这两个正则项放在一起就可以得到一个比较好的效果。

• 对于 $\hat n$ 和 $n^\prime$ 相差较大时，前者可以让网络使得两者相等，从而平滑梯度。
• 对于两者相差不大时，第二个正则项可以达到惩罚的目的得到更加合理的法向量。

复现中的细节

proposal network

这里我们使用了 proposal network distillation ，而没有选择 stratified sampling 。stratified sampling非常慢（由于coarse network需要forward所有点，fine network也要forward），proposal network则是用浅MLP（5层），forward点后直接输出预测的density，再利用fine network输出的density（weight）进行监督。这样就能完成从fine network到proposal network的蒸馏。

但是雀氏也有一些问题，

如图，这个球上有很多的噪点，怎么辉石呢？我们推断这是 density 有一些问题，于是进行了一些改动。

​原始prop net输入coarse points之后，输出density，此后coarse points将会被弃用。根据density计算的weight，将指导inverse sampling，fine network的输入只为inverse sampling的结果（也就是说，集中在weight高的地方）。假设，prop net计算的density有缺陷，也即weight有缺陷，在实际的表面附近weight很小，在空域中weight大，那么inverse sampling可能无法在此条光线上采到有效的点。

很简单，就是复用coarse points，将coarse depths（采样的长度）与inverse sampling的采样长度（fine depths）进行拼接，排序。但其实在实现中，要考虑proposal network的weight bound计算。啊，这一步很复杂。可以这么说：

• proposal network需要预测每一个采样点（fine采样点）的weight上界（weight bound）。上界如何计算？两个fine采样点之间会存在一个采样区间，此区间将会与coarse采样的区间重合，则此fine采样区间的weight上界应该是所有与之有交集的coarse区间weight之和。
• coarse points合并到fine points相当于修改了fine采样区间。那么就需要计算更多区间交集。

​这样修改，也就使得每条光线上，既有均匀采样的部分（保证了coverage），又使得density大的部分可以有更多采样点。可能有人觉得，不就是增加了一些采样点吗？这样为什么能保证proposal network的学习是正确的呢？很简单，网络不仅有更加充足的输入，fine network提供给proposal network的监督也更加充足了。

instant ngp

instant ngp ，一种使用哈希编码用于加速神经网络训练的方法。简单来说，就是通过一个多分辨率的哈希映射到一个向量中，然后再进行插值得到一个训练出来的编码值。

虽然这个东西看上去很好很fancy，但是会有哈希碰撞的问题，作者说这玩意儿可以通过神经网络来学习减少问题，不过我们觉得不行，这会让我们训练出来的玩意儿非常的“粉”。

这个就很难受了，两个相距很远的点获得了类似的性质，导致一些没有意义的点也被赋予了奇怪的颜色，有意义的点也变得有很多空洞。为此我们尝试了增大哈希表的大小，但是这样就炸显存了，非常神秘。所以目前正在着手于实现一种更好的采样方法，这样或许可以减少冲突的问题。