Ref NeRF复现

Posted on 2022-08-13 Edited on 2022-08-20 In learning Views: 904 Symbols count in article: 15k Reading time ≈ 14 mins.

Ref NeRF复现

I. Intros

科研恢复性训练。之前把CUDA加速的全并行shadow caster写完了，算是复习了CUDA、Rust以及FFI的使用。深度学习方面就以复现Ref NeRF为一个小任务。论文CVPR 2022 Best Student Honorable Mention: Ref NeRF - Structured View-Dependent Appearance for Neural Radiance Fields对反射现象进行了良好的建模。之前我一直在研究折射现象的NeRF建模，就折射建模思路而言，此文对我有很大的启发。并且个人认为，基于Ref NeRF以及mip NeRF作为框架是比较好的选择（除此之外就是要考虑如何让训练变快了，Instant NGP当然是不二选择）。当然，由于在复现过程中也遇到了许多问题，本文在阐述复现思路以及论文理解的同时，也会探讨踩过的坑。目前，Ref NeRF的复现结果还没有达到令我满意的程度，只是具备雏形，毕竟融合两篇文章的idea可能导致冲突，深度学习这种玄学就更是这样了，模型炸了都不知道从哪一个先开始。复现见 Enigmatisms/NeRF

Figure 1. Ref NeRF半成品(Shinny Blender - Helmet)。从左至右：RGB、depth与“奇怪的法向量”

II. 论文梳理

2.1 Ref NeRF的建模

Ref NeRF建模了表面反射，此工作比之前所看过的一篇处理反射的NeRF（NeRFReN）更有价值。个人认为，Ref NeRF可以被认为是一个不完整的光线传播模型，较好地解决了反射问题，只需要加上折射部分就可以使得NeRF处理基本的光线传播现象。本文对Ref NeRF的主要思路进行梳理，并且基于笔者对本论文的理解，在之前实现的NeRF（mip NeRF 360）基础上，增加Ref NeRF模块。

Ref NeRF处理的首要问题是光线的反射问题，反射一般分为两种：镜面反射（specularity）以及漫反射（diffusion）。那么根据Phong反射模型（以及之前学过半吊子图形学基础），被观察到的光线（强度）将为： $\begin{matrix} (1) & I_{o} = k_{a} I_{a} + \sum_{i} k_{d} (l_{i} \cdot n) I_{i} + \sum_{i} k_{s} (r_{i} \cdot v)^{p} I_{i} \end{matrix}$ 其中，下标为a的部分与环境光（ambient light）有关；下标含有d则与漫反射有关，而下标含有s则与镜面高光有关。为了更好地解释Ref NeRF的反射模型处理以及回顾一下Phong模型这个基础模型，这里对公式 $(1)$ 模型简单解释：

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

Ref NeRF对于后两个部分都进行了建模： $\begin{matrix} (2) & L_{o u t} (\hat{ω_{o}}) \propto \int L_{i n} (\hat{ω_{i}}) p (\hat{ω_{i}} \cdot \hat{ω_{r}}) d \hat{ω_{i}} = F (\hat{ω_{r}}) \end{matrix}$ 其中 $\hat{ω_{o}}$ 是观察方向（相机到物体）， $\hat{ω_{i}}$ 是光线的入射方向（光源到物体）， $\hat{ω_{r}}$ 是观察方向在物体上的反射（可以认为是相机按照 $\hat{ω_{o}}$ 方向发射光后反射的方向）。那么显然， $\hat{ω_{i}} \cdot \hat{ω_{r}}$ 就是公式 $(1)$ 中的 $r_{i} \cdot v$ 。在此工作中，函数p等均是神经网络（学出来的反射，强），并且作者强调，此函数就是反射方向 $\hat{ω_{r}}$ 的函数（内部包含了观察视角方向信息）。反射方向的求解非常简单，这里不赘述。不过我不禁思考，如果把作者的反射模型从Phong模型换成Bling Phong模型会怎么样？在复现中笔者进行了尝试，结果见实验部分。

作者花大力气改造原始NeRF的directional MLP，spatial MLP则只是输出了一些新的辅助信息。directional MLP必须要能够输出真正的view dependent颜色。作者认为：spatial MLP也需要直接预测颜色，但是此颜色应该是受环境光、漫反射等影响产生的（由公式 $(1)$ 可知，环境光与漫反射均与视角无关），视角不变的颜色。严格来说，此部分颜色并非albedo（由于漫反射存在）。而directional MLP应当产生镜面高光部分。作者对directional MLP作了如下修改：

2.2 directional MLP修改

首先，作者做了一个类似于mip NeRF中光锥处理的操作。mip NeRF中，为了达到mip map的效果，作者将点采样变为了光锥采样，以考虑整个光锥中的所有信息。Ref NeRF中，作者认为“反射”也不能仅仅考虑单个反射方向。由于物体实际是凹凸不平的，并非完美的镜面，表面法向量并非完全一致。可以认为，物体表面的凹凸起伏（噪声），使得反射发生了一些改变（distortion），物体表面噪声的概率分布，经过反射的数学操作后被映射成了新的分布。并且：

新的分布应该是各方向对称的。由于我们并不知道物体表面能有什么样的各向异性特征。
法向量分布的期望应该是 $\hat{ω_{r}}$ ，也即根据求得的表面法向量以及观察方向求出的反射方向。其背后的intuition（总觉得，不应该翻译成直觉，故我这里用了一个英文单词，以表示：我觉得这里用中文表达不了那种意思）很好理解。
法向量的方差应该受到物体表面粗糙程度进行控制。表面越粗糙，方差越大（法向量分布较广、分散）。这很符合现实。

作者使用了一种叫做 von Mises-Fisher（简写为vMF）的概率分布。此分布的pdf非常像高斯分布pdf，估计也有差不多的性质。个人推测作者选这个函数是为了方便进行后续的数学推导以及近似，就像mip NeRF中，将光锥用混合高斯模型进行近似一样。有了此分布，自然需要使用积分将所有可能的反射方向考虑进去。当然不是直接求期望（求出来就是 $\hat{ω_{r}}$ ，就直接trivial了）。考虑到，在将 $\hat{ω_{r}}$ 输入到directional MLP之前，需要进行positional encoding。在mip NeRF中，positional encoding参与了积分，在此处实际上也类似：我们需要在高维空间中求每一维度的期望。但作者并没有用sinusoidal positional encoding，而是使用球谐函数（spherical harmonics）对方向进行表示。最后对每一维球谐函数进行积分： $\begin{matrix} (3) & E_{\hat{ω} \sim vMF (\hat{ω_{r}}, κ)} [Y_{l}^{m} (\hat{ω})] \approx \exp (\frac{- l (l + 1)}{κ}) Y_{l}^{m} (\hat{ω}) \end{matrix}$ 其中， $κ$ 为表面粗糙度的倒数（此值越大，越光滑）， $Y_{l}^{m} (\hat{ω})$ 表示球谐函数（m，l都表示了其阶）。不难从公式 $(3)$ 看出，积分近似后的球谐encoding实际上多了一个系数，此系数受到球谐阶（频率的指数）以及表面光滑程度（ $κ$ ）的影响。显然， $κ$ 越小，对于高阶（频）球谐影响越大。这也就反映了这样一个事实：光滑程度减小，高低频信息均有所衰减，但高频信息衰减更严重。通过 $κ$ 就能直接影响输出结果中镜面高光的强度（模糊性以及保留细节多少）。

当然作者不仅仅给directional MLP提供了 $\hat{ω_{r}}$ ，如公式 $(1)$ 所示第二项的 $l_{i} \cdot n$ 也被传入网络（ $d \cdot n$ ，d为观察方向，n为法向量），这并不是要建模漫反射，而是考虑到一些稍微复杂一些的效应，如菲涅尔效应（并非全反射，部分反射部分折射，反射所占的比例需要另外计算）等等。

2.3 法向量估计

作者使用volume density的梯度可以估计法向量，但作者说：

normal vectors estimated from its volume density gradient as in Equation 3 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

确实，（1）很显然，已经有很多论文做这方面的工作了，不是添加正则化项就是采用一些表面重建技术获得好的表面。（2）的话通常也是由“集中性”正则化项解决的（只不过作者的想法不太一样）。

问题（1）的解决（prediction penalty）：通过spatial MLP预测某一点的法向量 $\hat{n}$ ，与梯度法向量 $n^{'}$ 求MSE： $\begin{matrix} (4) & L_{n} = \sum_{i} w_{i} ‖ {\hat{n}}_{i} - n_{i}^{'} ‖ \end{matrix}$ 由于 $\hat{n}$ 相对比较光滑，而 $n^{'}$ 是对不光滑的density求的一阶导，就更不光滑了。用光滑数据帮助不光滑的梯度进行学习，可以起到对梯度的光滑作用。但为什么 $\hat{n}$ 比较光滑？这是作者说的，其实我比较纳闷，毕竟这玩意是学出来的。我的理解是：可以认为，此处需要网络也同时将 $n^{'}$ 学出来（利用 $n^{'}$ 作为监督）。但是网络的表示能力有限，特别是参数较少时，只能达到对 $n^{'}$ 低频部分的学习。
问题（2）的解决（orientation loss）：作者设计了一个正则化项：此正则化惩罚反向法向量，也即“背面”。但也不是盲目对背面进行惩罚，毕竟volume sampling过程中是可以在有效的背面进行采样的，故用下式： $\begin{matrix} (5) & R_{o} = \sum_{i} w_{i} max (0, {\hat{n}}_{i} \cdot d) \end{matrix}$ 其中 $w_{i}$ 为点weight值， $d$ 为光线方向。此处也即惩罚法向量与光线方向相同部分。并且与weight有关，这就说明：被遮挡的有效背面不会被影响（weight低），而fake surface（在半透明surface后的embedded emitter表面）将会被惩罚（其density衰减方向与光线方向一致）

作者自己也解释了一下，为什么这两个loss结合在一起就work了：网络对于法向量的预测（ $\hat{n}$ ），与density梯度对应的法向量 $n^{'}$ 之间的关系大致有两种：

$\hat{n}$ 偏离 $n^{'}$ 。那么此时，prediction penalty较大，网络偏向于使得两者相等，则可以达到平滑梯度的目的。
$\hat{n}$ 很接近 $n^{'}$ 。那么此时，由于 $\hat{n} \approx n^{'}$ ，我们可以认为公式 $(5)$ 作用在 $n^{'}$ 上。这也就达到了惩罚的目的。

2.4 扩展知识

2.4.1 一些物理概念

（1）能量，功率与辐射通量

能量的单位一般采用（J），而功率则是单位时间内的能量（J/s，定义为W），此概念与通量（flux）是同一个概念。与接收能量的面积没有关系。

（2）辐照度与辐出度

辐照度（irradiance）与辐射度（radiosity）是两个类似的概念，与面积有关。单位面积的通量（或者说，功率），单位是： $W / m^{2}$ 。注意，此面积是光照的“正对面积”，也即需要进行投影处理。当表面法向量与光线入射（出射）方向垂直时，辐照度以及辐射度最大。当然，以上讨论仅限于面平行光源。对于点光源而言，其辐照（射）度与距离的关系遵循平方反比定律（距离越远，辐照（射）度越低）。

（3）立体角与辐射强度（intensity）

立体角表达的物理意义是：1立体角对应在单位球上截取的大小为 $r^{2}$ ，其中 $r$ 为球的半径。也即： $\begin{matrix} (6) & Ω = \frac{A_{s r}}{r^{2}} sr \end{matrix}$ $sr$ 是单位，而不是变量或常量。其中 $A_{s r}$ 则是对应大小的球面切片面积。有了立体角之后，可以使用立体角 + 3D点 + 3D方向来描述任意光线。光线从给定的3D点出射，方向是给定的3D方向，光锥的形状由立体角确定。通过单位立体角的辐射功率（通量）称为辐射强度，单位是 $W / sr$ 。辐射强度描述的是光源的属性，而非空间中某一点受光照的特性。故给定某一光源，不考虑空间中的其他因素的影响，空间中任意一点的辐射照度都是一样的。

（4）辐射率（radiance）

搞NeRF也有一段时间了，还没有搞清楚radiance的意义。此前只是将radiance当作是观察到的物体的颜色，而现在这里将对辐射率的物理含义进行介绍。上文已经介绍了辐射强度的概念，并且可知，辐射强度与空间特性无关，只与光源有关，而我们又希望获得某一光源在空间中某一位置可产生的影响。则定义辐射率：单位时间、单位面积、单位立体角通过的光能。也即单位面积、单位立体角的功率。单位是： $W (m^{2} sr)^{- 1}$ 。光源辐照能力改变（如灯光的亮暗），反映于功率变化（ $W$ ）。而光源的“能量聚集度”（可以想象一个可以调焦的手电筒）反映在辐射强度上（影响单位为 $sr$ 的项）。物体距离光源的远近则反映在面积项上（平方反比）。这也恰好与成像结果直接相关，故可以将radiance理解为空间中一点成像的结果（如RGB值）。

2.4.2 光场

四维光场是否只是一个近似？四维光场认为，在empty space内部，radiance（也就是产生的RGB值）是恒定的。也就是说，无论我从此光线上的哪一点开始积分，最后得到的RGB值都是一样的，此RGB值与光线上的积分距离（或是深度）无关。那么这样可以从五维的光场建模中，去掉一个维度，实际则为四维的。个人更愿意将“四维”理解为一个五维空间中的四维流形，因为此深度维度并非五维度基中之一，也不能用其线性组合描述。理解了四维光场的由来，我们回头检查一下假设：空域内部radiance为常数。正常情况下：无半透明物体或者介质，这是成立的。引入半透明物体与介质后（比如混浊溶液，胶体，雾），个人认为此假设不成立，那么表四光场还是应该使用五维函数。

2.5 球谐函数的使用

打开维基百科，查之。woc，我看到了什么。总之很复杂了，本人本科所学的数学知识不够用（或者说是遗忘了）。关于球谐函数本身的理解，个人的看法是：类似于傅里叶变换，展开成基函数的线性组合（这些基函数独立分析时其实形式较为简单，但组合起来就能表示非常丰富的内容），并且不像泰勒展开（泰勒展开的物理意义不甚好理解），球谐函数有频率概念，并且三维空间中的球谐函数仍然很好理解，非常具有物理直观性。如果想要知道如何初步理解，这里推荐【知乎：球谐函数介绍（Spherical Harmonics）】，我就不赘述了。下面只给出结论，实球谐函数的形式如下： $\begin{matrix} (7) & Y_{l m} (θ, ϕ) = {\begin{cases} (- 1)^{m} \sqrt{2 \frac{(2 l + 1)}{4 π} \frac{(l - | m |)!}{(l + | m |)!}} P_{l}^{| m |} (\cos θ) \sin (| m | ϕ), if m < 0 \\ \sqrt{\frac{2 l + 1}{4 π}} P_{l}^{m} (\cos θ), if m = 0 \\ (- 1)^{m} \sqrt{2 \frac{(2 l + 1)}{4 π} \frac{(l - m)!}{(l + m)!}} P_{l}^{m} (\cos θ) \cos (m ϕ), if m > 0 \end{cases} \end{matrix}$ 其中， $ϕ, θ$ 是方位角（azimuthal angles）， $l, m$ 为此球谐函数的阶？（此处实际上涉及到一个微分方程求解问题，此二变量与解的个数有关）。另外，比较令人疑惑的是 $P_{l}^{| m |}$ ，这实际上代表的是伴随勒让德多项式（associated Legendre polynomials）。这里，Wikipedia会告诉你勒让德多项式的奇妙性质，但是我不关心（说实在的，一开始看到这样的公式我还在想应该如何高效实现这么多微分）。这里直接给出维基上写的闭式公式： $\begin{matrix} (8) & P_{l}^{m} (x) = (- 1)^{m} \cdot 2^{l} \cdot (1 - x^{2})^{m / 2} \sum_{k = m}^{l} \frac{k!}{(k - m)!} \cdot x^{k - m} \cdot (\begin{matrix} l \\ k \end{matrix}) (\begin{matrix} \frac{l + k - 1}{2} \\ l \end{matrix}) \end{matrix}$ 此公式的代码实现就简单多了，我们只需要关心最后的求和项。求和项中涉及到求组合数，可以选择使用lookup table，预先计算一个大小为 $2^{L} \cdot 2^{L}$ 的矩阵（L应该不会很大）。那么： $\begin{matrix} (9) & \frac{k!}{(k - m)!} \cdot x^{k - m} \end{matrix}$ 应该如何计算？能不用for循环就不用for循环（多用GPU操作！！）。观察容易知道：求和项中的首项是 $m! x^{0}$ ，此后每一项与前一项的关系是： $\begin{matrix} (10) & T_{k + 1} = T_{k} \cdot x \cdot \frac{k + 1}{k + 1 - m} \end{matrix}$ 故我们应该有这样的一个张量（一行， $l - m$ 列）：第一列是 $m!$ ，此后所有列是： $[\frac{m + 1}{1} x, \frac{m + 2}{2} x, \frac{m + 3}{3} x, . . ., \frac{l}{l - m} x]$ 。那么实现就是：

# device 某个传入函数的张量t的t.device
rare = (torch.arange(m+1,l+1, device = device) / torch.arange(1,l-m+1, device = device)).unsqueeze(0) * x
whole = torch.cat([torch.ones(1, 1, device = device) * m_factorial, rare], dim = -1)
whole_prod = torch.cumprod(whole, dim = -1)

此后，根据lookup或者在线求解公式 $(8)$ 中的组合数，就能求得求和式中的每一项，进而也就求出了整个多项式。但这只是针对 $m, l$ 求出了结果，而 $m, l$ 如何进行并行计算我就不懂了。不过其实也没必要懂：

NeRF的positional encoding，不同的相位（ $\sin, \cos$ ）以及不同的频率是串行计算最后concatenate的
大可以自定义CUDA算子，写一个forward再写一个backward函数。forward好写而backward不好写罢了（这玩意显然是有闭式解的，推起来比较麻烦），故CUDA自定义算子：全并行但费脑子（哈，而且还不一定有pytorch快，例如之前我用thrust的sort库对比pytorch的张量sort，前者就比后者慢），pytorch实现弱并行，但不费脑子。

球谐函数的使用也说明了，我们在处理方向时，需要将三维的方向转化为二维的方向角（才能喂给球谐函数，当然，球谐函数确实有笛卡尔坐标系表示，但这不是我们的讨论范畴）。下面，笔者简单讨论一下文中球谐函数在vMF分布下的期望求解（论文附录部分）。作者要证明： $\begin{matrix} (11) & E_{\hat{ω} \sim vMF ({\hat{ω}}_{r}, κ)} [Y_{l}^{m} (\hat{ω})] = A_{l} (κ) Y_{l}^{m} ({\hat{ω}}_{r}) \end{matrix}$ 我们只讨论上式中 $A_{l} (κ)$ 怎么来的，而不关心（所以在这不推导） $A_{l} (κ)$ 如何近似为那个只与 $l, κ$ 有关的指数形式（近似的推导部分，但凡有点耐心。。。emmm）。首先，作者把积分空间进行了线性变换：进行了旋转。将 ${\hat{ω}}_{r}$ 旋转至z轴，那么根据vMF分布的形式以及随机变量函数的期望的计算： $\begin{aligned} (12) & P_{vMF ({\hat{ω}}_{r}, κ)} (ω) = c (κ) \exp (κ {\hat{ω}}_{r}^{T} ω) \\ (13) & E_{\hat{ω} \sim vMF ({\hat{ω}}_{r}, κ)} [Y_{l}^{m} (\hat{ω})] = c (κ) \int_{S} Y_{l}^{m} (ω) \exp (κ {\hat{ω}}_{r}^{T} ω) d ω \end{aligned}$ 则将 $ω^{'} = R ω$ 中，可以得到： $\begin{aligned} (14) & E_{\hat{ω} \sim vMF ({\hat{ω}}_{r}, κ)} [Y_{l}^{m} (\hat{ω})] & = c (κ) \int_{S} Y_{l}^{m} (R^{- 1} ω^{'}) \exp (κ (R^{- 1} z)^{T} R^{- 1} ω^{'}) d R^{- 1} ω^{'} \\ (15) & = c (κ) \int_{S} Y_{l}^{m} (R^{- 1} ω^{'}) \exp (κ z^{T} ω^{'}) d R^{- 1} ω^{'} \\ (16) & = c (κ) \int_{S} Y_{l}^{m} (R^{- 1} ω^{'}) \exp (κ \cos θ) d ω^{'} \end{aligned}$ 从公式 $(14)$ 到 $(15)$ 很好理解，旋转向量的的正交性可以立刻抵消exp中的R。而公式 $(15)$ 到 $(16)$ 则稍微有点绕：

$z^{T} ω^{'}$ 恰好是 $ω^{'}$ 的方位角（中的俯仰角 $θ$ ）
$d R^{- 1} ω^{'}$ 中的 $R^{- 1}$ 原先需要被提出来，但由于其行列式为1（旋转矩阵嘛），不会影响最后函数输出的单个值的结果（注意，对于每一个 $m, l$ 而言，公式 $(14)$ 的输出都是一个值），故可以省略

此后，根据球谐函数的性质（比如展开为Wigner D Matrix表示以及其正交性），我们最后可以将公式 $(14)$ 左边写为： $\begin{matrix} (17) & E_{\hat{ω} \sim vMF ({\hat{ω}}_{r}, κ)} [Y_{l}^{m} (\hat{ω})] = c (κ) D_{m 0}^{(l)} ({\hat{ω}}_{r}) \int_{S^{2}} Y_{l}^{0} (\hat{ω}) e^{κ \cos θ} d \hat{ω} \end{matrix}$ 其中： $D_{m 0}^{(l)} ({\hat{ω}}_{r})$ 以及 $Y_{l}^{0} (\hat{ω})$ 的形式都是已知的，这里不在赘述了（确实很麻烦）。我们只将上式的积分部分进行拆解。由于积分实际上是在球面上完成的，这里实际进行了一次二重积分（球面积分），将笛卡尔坐标系转化为球面的天顶角： $(x, y, z) \to (θ, ϕ)$ ，这里给出一张图：

Figure 2. Ref NeRF半球积分示意图

图中的黑色圆圈就是“积分对象”：可以看做球体的上半球表面是由无数个圆环组成的，对应了角度为 $θ$ （确定了位置），半径为 $\sin θ$ （确定了形状）的圆圈（可知 $ϕ$ 的积分范围是 $- π \to π$ ）：故积分可以拆成（首先知道 $Y_{l}^{0} (\hat{ω}) = P_{l} (\cos θ)$ ）: $\begin{matrix} (18) & \int_{S^{2}} Y_{l}^{0} (\hat{ω}) e^{κ \cos θ} d \hat{ω} = \int_{- π}^{π} \int_{0}^{π} P_{l} (\cos θ) e^{κ \cos θ} \sin θ d θ d ϕ \end{matrix}$ $θ$ 的范围是 $[0, π]$ ，这是因为我们只需要积分上半球。由于反射不可能发生在下半球，可知下半球 $[- π, 0]$ 的所有输出都是0。这样，根据论文中所说的步骤，进一步替换积分变量，展开并简化可以推出IDE公式。

但是，选择实函数形式的球谐并不是一个很好的选择。首先，公式 $(7)$ 已经展示了实函数形式的复杂性（主要是有分支）；此外，我们需要考虑到，角度形式的SH需要把方向向量转化为天顶角，这种额外的操作看着就很不优雅。建议使用笛卡尔坐标系形式的球谐函数。参见Wikipedia/Shperical Harmonics - Separated Cartesian form，可刺激了，虚函数运算与二项式公式的结合哦。不过只要使用这种方法实现，利用Pytorch就很简单：

首先Pytorch支持虚数运算。这里吐槽一下，Pytorch支持的虚数运算基于torch.complex64类型，一看到64就知道怎么回事了，运算量可能等同于float32运算（两个float32存在一块），在使用auto mixed precision时，complex32将可能产生一系列的问题，很多算子都没有对complex32有支持。而假设我们使用APEX库的O2等级优化，所有的float32都会变成float16，那么这里参与运算的复数就会产生一系列未定义算子的问题。（什么ComplexHalf未定义这些算子啊之类的报错）
SH在这就是用作encoding，所以其实只需要将结果的实部、虚部取出来作为encoding的成分就好了。可以认为，原来的positional encoding是用了正交相位的两种基底进行表示，而SH其实也是正交的两种基底（实与虚在虚平面上是正交的）进行表示。

III. 复现细节

3.1 值得一提的点

Figure 3. Ref-NeRF网络结构

其实，没什么好说的。根据Ref NeRF论文图，复现的内容如下：

首先修改spatial network，在spatia network之后追加三个head。三个head的输入就是spatial network的输出：

opacity & roughness head: 由于τ（density）与ρ（roughness）都是没有输出范围限制的（但是都需要大于0，故ρ之后与τ类似，单独处理即可）。注意roughness与opacity都需要经过bias之后再通过softplus激活（软ReLU）
color & tint head: $c_{d}, s$ 都是[0, 1] 之间的值，故输出直接变6维。color与tint都使用sigmoid进行激活，不过bias不同。
normal head：预测法向量，三维，需要进行normalize，不需要任何激活

需要绕一下的部分是法向量。法向量的来源有两个：（1）spatial network计算的density，可以对位置求导。Pytorch中，要实现对位置求导需要使得输入的位置向量可导：

fine_pos.requires_grad = True										    # 此步必要	
fine_rgbo, pred_normal = mip_net.forward(fine_pos, fine_dir)
r_num, p_num, _ = fine_rgbo.shape
density_grad, = torch.autograd.grad(fine_rgbo[..., -1], fine_pos, 			  # fine_pos 是可优化的
       torch.ones(r_num, p_num, device = fine_rgbo.device), retain_graph = True # retain graph使得计算图可以继续backward
)

另一个来源就是spatial network之后的normal head，预测。预测的结果在经过正则化（orientation loss）以及预测误差（与density grad的差别）backward之后，应当可以学出较好的结果。

此外注意，directional network变深了。可以认为directional network有spatial network一样的结构（包括8层256 hidden units的MLP，并且有skip connection）。

3.2 一些坑

由于笔者复现Ref NeRF基于之前实现的Enigmatisms/NeRF，这个repo中实现的NeRF并非原始NeRF，而是mip NeRF 360（部分），主要包含这样的特性：

本repo实现了mip NeRF 360中的proposal network distillation。也即：coarse-to-fine的stratified sampling并没有使用。stratified sampling非常慢（由于coarse network需要forward所有点，fine network也要forward），proposal network则是用浅MLP（5层），forward点后直接输出预测的density，再利用fine network输出的density（weight）进行监督。这样就能完成从fine network到proposal network的蒸馏。此蒸馏的实现较为复杂（weight bound的计算），这里不赘述。我也不想单开一篇博客讨论。
使用了自动混合精度（amp），可以选用apex或者torch.native_scaler

proposal network着实坑了我一把。复现完成（第一版）之后，跑了一下shinny blender的ball数据集：

Figure 4. 为什么那么多噪点？

如上图第一列：很多噪点。这些噪点是什么？一开始我也不敢下结论，毕竟不知道是不是自己的ref nerf复现有问题。于是我将深度以及normal都可视化了出来（后两列）。可见，深度图中也存在很多空洞，normal就不说了吧（事实上，normal学习一直有问题，normal本身就比较难学）。空洞产生的可能原因有两个：

density学习有问题，有些位置density非常低
采样（不管是训练还是测试时）位置不对，采样在了一些空的位置

于是我开了一个新的分支，在此分支上我将proposal network删掉了，使用原始NeRF的coarse-to-fine框架进行测试。对比图如下图所示：左边为原始proposal network存在时，训练60轮之后的结果，而右边则为原始NeRF训练40轮之后的结果。可以看出，原始NeRF框架下的采样是（至少相对上是）没问题的。但笔者并不想放弃proposal network（个人觉得这个方法比较优雅），故笔者将proposal network进行了小的改动：

原始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采样区间。那么就需要计算更多区间交集。这里我不赘述方法，详见：NeRF/coarseFineMerge以及NeRF/getBounds

这样修改，也就使得每条光线上，既有均匀采样的部分（保证了coverage），又使得density大的部分可以有更多采样点。可能有人觉得，不就是增加了一些采样点吗？这样为什么能保证proposal network的学习是正确的呢？很简单，网络不仅有更加充足的输入，fine network提供给proposal network的监督也更加充足了（见【修改后的Proposal Network】部分，其中我们说到“更多的区间交集”）。

Figure 5. shinny blender helmet数据集，训练20轮

无proposal net 训练40轮（4k次）	有proposal net 训练60轮（6k次）

修改后的方法，只训练20轮输出后的结果就让我觉得肯定是改对了，如上图所示。接下来将展示【未完成训练】时的训练结果。由于原始Ref NeRF使用 $2^{14}$ 大小的batch，需要训练250k轮。我设备能力有限，只batch大小为 $2^{9}$ （差了32倍），也没训练到250k就先暂停了（笑死，我的3060 for laptop，有个风扇坏了，训练的时候能飙到83度，所以我得开一台电风扇对着我电脑吹）。下面是简单的结果展示：

Figure 6. 训练不是一次完成的，保存checkpoint短点续训。收敛性并不好

从上图看出，虽然训练应该远没有结束（这才不到7小时），PSNR曲线的上升能力属实堪忧。笔者感觉PSNR无法上升至论文中的29，个人认为可能的问题仍然是出在proposal network上。目前正在尝试一种改进的方法。

Figure 7. 有趣的手动退火

效果大概就是这样（如下图所示）。从左到右分别是：RGB，深度与法向量。emm，法向量大概是可视化错了。不过学的应该也有一些问题，在某些视角能看到头盔上的条纹对应的法向量与其他位置不一致。PSNR为19，还是太模糊了一点。

Figure 8. 旋转头盔

Reference

https://en.wikipedia.org/wiki/Spherical_harmonics
https://en.wikipedia.org/wiki/Associated_Legendre_polynomials
扩展阅读: https://dellaert.github.io/NeRF22/