NeRF編: 概要
これまでに説明したMLP，エンコーダーに加えて，ボリュームレンダリング(RGB)とその誤差逆伝播，さらにはOccupancy Gridに関する私の実装を説明し，最終的にどのようにNeural Radiance Fieldsが形成されていくかを確認します．
NeRF編: はじめに
Part1, Part2では下準備としてMLPとMultiresolutioin Hash Encoding, Spherical Harmonic Encodingを実装しました．Part3では本格的にNeRFの実装に入っていきます．Part2の最後に画像の近似を行いましたが，あれは2次元平面内部におけるRGB色（RGB場とでも言いましょうか）を近似していくものでした．2次元画像近似は非常に簡単に行えましたが，3次元となった場合はなかなかに複雑化します．それは3次元の状態をそもそも可視化する手法が一筋縄ではいかないためです．色々な手法が提案されてきましたが，その1つにNeRF [Mildenhall et al, 2020]があります．これは煙などの媒質の表現に用いられるボリュームレンダリングを利用したもので，その媒質に関するパラメーターを最適化することにより，空間内部の媒質の状態を近似するというものです．2022年にNVIDIAの発表したInstant NeRFにおいてはOccupancy Gridというものが導入され，それによってボリュームレンダリングのサンプリング効率の向上が実現されました．ではそれぞれの実装を説明していきます．（空間内部の「状態」を近似と書いているのは，空間内部に存在する「物体表面の再構築」と区別するためですが，Occupancy Gridを使っている以上，結局NeRFも後者を行っているのですかね......）
念のため......
内容には注意はしておりますが，記事に誤り等があれば指摘していただけると幸いです．
NeRFの手続き
NeRFと聞くと身構えてしまいそうですが，実はそこまで複雑なことはしていません．まずはざっくりと雰囲気を書くと，
(0): カメラと画像のセットを用意する
(1): カメラから画像のピクセルに対応する方向にレイ（光線）を飛ばす
(2): レイの経路上の点をサンプリング
(3): サンプル点の座標，そしてレイの向きをニューラルネットワーク(NN)に入力
(4): NNを実行し，サンプル点における「色」と「密度」を計算
(5): 得られたサンプル点の情報からボリュームレンダリング方程式を計算
(6): 計算結果と対応するピクセルの色を比較し，誤差を計算
(7): ボリュームレンダリングの処理を逆方向に行う
(8): NNを誤差逆伝播
(9): NN最適化
(10): Occupancy Gridの最適化
となっております．こういう書き方をしているので順番から説明すべきではありますが，実装の説明をする前にボリュームレンダリングについて軽く触れておきます．
ボリュームレンダリング
NeRFではRGBレンダリングとしてボリュームレンダリングを行います．といってもパストレーシングでやるようなランダムウォーク的な経路追跡をするわけではなく， 単純にレイ（半直線）上のランダムにサンプリングされたサンプル点における放射輝度の変化を追うものです（また，距離関数を使用したレイマーチングではありません．）．本来はちゃんと導出するのが良いですし，そもそも光線としての光（幾何光学）の妥当性を書くべきだとは思うのですが，今回はそれを他の記事に任せてフロントエンドだけ書きます．
今回は次の現象を扱います．
・媒質の吸収による放射輝度の減衰
・媒質の散乱による放射輝度の減衰
・媒質の発光(emission)や散乱(in-scattering)による放射輝度の増加
レイ上のある異なる2点を，カメラから近い順にt_near，t_far とします．この時， t_near， t_far における媒質によるカメラに入射するレイの放射輝度への関与は
ここで，はカメラからレイ上のある点までの距離で，
はその点の座標，つまりレイの原点と向きをそれぞれ
として，
です．そして，
はその点における媒質の「密度」で，
はその点における媒質の「色」です．より正確に書くと，「密度」は光学的な消散係数にあたり，「色」は媒質内の粒子における散乱(in-scattering)や発光(emission)による放射輝度への増加に関わるパラメーターにあたります．図に描くとこんな感じのイメージです．
カメラ側から見た時，媒質の「密度」が十分大であれば媒質の「表面」しか見ることが出来ません．これは「表面」より向こう側の放射輝度によるカメラに入射するレイへの寄与が0であるためです．逆に，「密度」が非常に小さい場合，媒質の向こう側が透過して見えます．イイ感じにパラメータを設定してあげると現実世界の空間をある程度近似できそうだというのは何となく想像できます．
計算機でボリュームレンダリングを行う
先ほど示した積分を計算機上で解くためには，理想はレイを無限に分割してやるのが良いのですが，そんなことは出来ないため，有限のサンプル点を使用して近似的に積分を行います．
先ほどの式を離散的に書き直すと図の中に示す式になります．理論的な計算は省略します．["Optical Models for Direct Volume Rendering", Nelson Max, 1995]
このは先ほどの積分の式と同じく「どれだけその点におけるレイの放射輝度がカメラに入るレイの放射輝度に寄与しているか」を表します．さて，このように有限のサンプル点をサンプリングしてあげる必要があります．レンダリングの詳しい実装は後ほど行います．
NeRFの実装: (0): カメラと画像のセットを用意する
さて，やっていきましょう．ただロードすればいいと言われればそうなのですが，学習データと教師データのセットを意識してデータを保持してあげると取り扱いが容易になります．今回の実装においては，最初に画像と対応するカメラの姿勢を一気に全部ホスト側のメモリ上にロードします．そして，それをもとにして，「（画像ID, ピクセル座標）- ピクセルの色」を入力-教師のペアとして構造体にしておきます．次の構造体を書きました．
// 各スレッドではこの構造体1つにしかアクセスしないのでAoSの方が良いと考えた
// [31:21]: PixIndexX
// [20:10]: PixIndexY
// [9:0] : ImageID
struct PixelInfo {
//int ImageID;
//int PixIndexX;
//int PixIndexY;
unsigned int PixInfo = 0;
Vec3h Color;
MNPT_HOST_DEVICE PixelInfo(unsigned int PixInfo = 0) : PixInfo(PixInfo), Color() {}
MFFM_HOST void SetUp(unsigned int ImID, unsigned int IndexX, unsigned int IndexY, Vec3h PixColor) {
if (ImID >= 1024) {
printf("Too large ImageID was set to PixelInfo\n");
exit(1);
}
if (IndexX >= 2048) {
printf("Too large ImageWidth was set to PixelInfo\n");
exit(1);
}
if (IndexY >= 2048) {
printf("Too large ImageHeight was set to PixelInfo\n");
exit(1);
}
PixInfo = 0;
PixInfo = PixInfo | (IndexX << 21) | (IndexY << 10) | ImID;
Color = PixColor;
}
MFFM_HOST void SetUp(int ImID, int IndexX, int IndexY, __half* PixColor) {
if (ImID >= 1024) {
printf("Too large ImageID was set to PixelInfo\n");
exit(1);
}
if (IndexX >= 2048) {
printf("Too large ImageWidth was set to PixelInfo\n");
exit(1);
}
if (IndexY >= 2048) {
printf("Too large ImageHeight was set to PixelInfo\n");
exit(1);
}
PixInfo = 0;
PixInfo = PixInfo | (IndexX << 21) | (IndexY << 10) | ImID;
Color.from_half(PixColor);
}
MFFM_HOST_DEVICE unsigned int PixIndexX() {
unsigned int ret = (PixInfo >> 21); // [31:21]
return ret;
}
MFFM_HOST_DEVICE unsigned int PixIndexY() {
unsigned int ret = ((PixInfo & 0x001FFC00) >> 10); // [20:10]
return ret;
}
MFFM_HOST_DEVICE unsigned int ImageID() {
unsigned int ret = (PixInfo & (0x000003FF)); // [9:0]
return ret;
}
};
ただ単純に入れているだけなのですが，メモリ節約のために32bitの変数を切り分けて情報を保存しておきました．画像IDには10bit，画像上のピクセル座標にはそれぞれ11bitを使用しています．また，Vec3hという，半精度小数点を3個内蔵するベクトル構造体に色の情報を保存しておきます．Vec3やVec3hに関する具体的な説明はのちに行います．画像IDからカメラの姿勢にアクセスすることができ，そしてピクセル座標からカメラから飛ばすレイ（後述します）を計算でき，そして色の情報から教師データにアクセスできる，というビジョンです．ちなみにレイ-ピクセルの色というペアだとダメなのかという問いがあり得ますが，GPU上のメモリ(10GB)に載りませんでした．
このようにしてロードしたデータをGPU上に送り，学習時にはシャッフルを施して多様なレイを学習時に与えることを考えます．このあたりの処理はthrustライブラリの力を借りました．ここが「ほぼ0から実装してみた」になっている要因の一つです．
...
thrust::default_random_engine g;
thrust::shuffle(D_thrust_TargetPixelInfo.begin(), D_thrust_TargetPixelInfo.end(), g);
...
if (IndexOfPixelInfo + nMaxPixelPerBatch > D_thrust_TargetPixelInfo.size()) {
thrust::shuffle(D_thrust_TargetPixelInfo.begin(), D_thrust_TargetPixelInfo.end(), g);
IndexOfPixelInfo = 0;
}
...
全部のピクセルを1イテレーションで全て入力するなんてことはしません．nMaxPixelPerBatchで一回のイテレーションで使用するピクセル数を決めておきます．そして，前シャッフルされてから現在までに使用されているピクセル数をIndexOfPixelInfoで保持しておきます．このif文がおこなっていることは，前回シャッフルしてから初めて，未使用のピクセル数がnMaxPixelPerBatch未満であれば再びシャッフルを行う，ということです．つまりはシャッフルしなおしているだけです．図を置いておきます．
NeRFの実装: (1): カメラから画像のピクセルに対応する方向にレイ（光線）を飛ばす
私たちは光を直線として扱う，すなわちレイ（光線）という概念に馴染みきっております．カメラからレイを飛ばすと言いますが，実際にはカメラに対応する方向から入射してくる光の経路を追っているという方が正しいでしょう．結局何が言いたいかというと，カメラから目的のピクセルの場所に半直線をのばし，その半直線を光路としてカメラに入射してくるレイの放射輝度を求めたいです．なので，その半直線がいかなるものであるかを計算する必要があります．今回は単純にピンホールカメラを考えましょう．
まあCGソフトでよく見る感じのカメラではないでしょうか．基本姿勢としてはカメラの原点（CamOrg），カメラの向いている方向（CamDir），カメラの上向き（CamUp）となります．そして，原点からCamDir方向にCamToScreenDistだけ離れた場所にスクリーンがあると考えてください．スクリーン上の点と，その点を通ってカメラの原点に入り込むレイは一対一に対応します（レイが直線であるため）．このようなレイの持つ放射輝度がカメラの写す画となります（細かい話は置いておきます）．では，そのようなレイを求めます．図のようにカメラの原点から出てスクリーン上の座標(IndexX, IndexY）を通る半直線（橙の点線）は，まずスクリーンの中央に伸ばしたベクトル（赤色の矢印）にスクリーン上でのX方向（緑色の矢印），Y方向のベクトル（青色の矢印）を足し合わせてあげることで求まります．
では実装するうえで便利な構造体を置いておきましょう．
Vec3構造体
名前の通り，浮動小数点の3次元ベクトルを処理するための構造体です．
struct Vec3 {
float x, y, z;
float align;
KGYK_HOST_DEVICE Vec3(const float x = 0.0f, const float y = 0.0f, const float z = 0.0f, const float align = 0.0f) : x(x), y(y), z(z), align(align) {}
KGYK_HOST_DEVICE Vec3 operator+(const Vec3& b) const {
return { x + b.x, y + b.y, z + b.z };
}
KGYK_HOST_DEVICE Vec3 operator-(const Vec3& b) const {
return { x - b.x, y - b.y, z - b.z };
}
KGYK_HOST_DEVICE Vec3 operator*(const float& b) const {
return { x * b, y * b, z * b };
}
KGYK_HOST_DEVICE Vec3 operator/(const float& b) const {
return { x / b, y / b, z / b };
}
KGYK_HOST_DEVICE float length_squared() const {
return x * x + y * y + z * z;
}
KGYK_HOST_DEVICE float length() const {
return sqrtf(length_squared());
}
// Vec3 to float*
KGYK_HOST_DEVICE void to_float(float* Dst) {
Dst[0] = x;
Dst[1] = y;
Dst[2] = z;
}
KGYK_HOST_DEVICE void from_float(float* Src) {
x = Src[0];
y = Src[1];
z = Src[2];
}
KGYK_HOST_DEVICE float operator[](size_t idx) {
return *((float*)this + idx);
}
KGYK_HOST_DEVICE float at(size_t idx) {
if (idx == 3) printf("WARNING: Accessing to Padding Area of Vec3 in Vec3::at()\n");
if (idx > 3) printf("ERROR: Out-of-range access in Vec3::at()\n");
return *((float*)this + idx);
}
};
KGYK_HOST_DEVICE Vec3 operator*(float t, const Vec3 v) {
return v * t;
}
KGYK_HOST_DEVICE bool operator==(const Vec3 a, const Vec3 b) {
return (a.x == b.x && a.y == b.y && a.z == b.z);
}
KGYK_HOST_DEVICE bool operator!=(const Vec3 a, const Vec3 b) {
return (a.x != b.x || a.y != b.y || a.z != b.z);
}
// 長さ1に正規化する
KGYK_HOST_DEVICE inline Vec3 normalize(const Vec3& b, const char* text = "") {
if (b.length() == 0) {
printf("Vec3 Warning: 0-size vector normalization. %s\n", text);
return { 10000, 0, 0 };
}
return Vec3(b / b.length());
}
// AABB内部のベクトルbについて，そのAABBを[0,1]^3と正規化したときのベクトルbを求める．
KGYK_HOST_DEVICE inline Vec3 normalize_by_range(const Vec3& b, Vec3& pmin, Vec3& pmax) {
Vec3 ret;
ret.x = normalize(b.x, pmin.x, pmax.x, 0.0f, 1.0f);
ret.y = normalize(b.y, pmin.y, pmax.y, 0.0f, 1.0f);
ret.z = normalize(b.z, pmin.z, pmax.z, 0.0f, 1.0f);
return ret;
}
KGYK_HOST_DEVICE inline const Vec3 multiply(const Vec3& v1, const Vec3& v2) {
return Vec3(v1.x * v2.x, v1.y * v2.y, v1.z * v2.z);
}
KGYK_HOST_DEVICE inline const Vec3 divide(const Vec3& v1, const Vec3& v2) {
return Vec3(v1.x / v2.x, v1.y / v2.y, v1.z / v2.z);
}
KGYK_HOST_DEVICE inline float dot(const Vec3& v1, const Vec3& v2) {
return v1.x * v2.x + v1.y * v2.y + v1.z * v2.z;
}
KGYK_HOST_DEVICE inline float absdot(const Vec3& v1, const Vec3& v2) {
return std::abs(dot(v1, v2));
}
KGYK_HOST_DEVICE inline Vec3 cross(const Vec3& v1, const Vec3& v2) {
return Vec3(
(v1.y * v2.z) - (v1.z * v2.y),
(v1.z * v2.x) - (v1.x * v2.z),
(v1.x * v2.y) - (v1.y * v2.x));
}
KGYK_HOST_DEVICE inline const Vec3 exponential(const Vec3& v) {
return Vec3(exp(v.x), exp(v.y), exp(v.z));
}
// "TextBefore v.x v.y v.z TextAfter\n"
KGYK_HOST_DEVICE inline void printVec3(const Vec3& v, const char* TextBefore = "", const char* TextAfter = "") {
printf("%s %f %f %f %s\n", TextBefore, v.x, v.y, v.z, TextAfter);
}
KGYK_HOST_DEVICE……等はCUDAのhostやdevice等の修飾子です．特に書くこともないですので説明は省略します．ちなみにメンバ変数のalignは構造体サイズの32byteアラインメントです．Vec3hはこれらを半精度で行う構造体です．
Ray構造体
レイの構造体です．
/*
* レイの情報を持つ構造体
*
* dir: レイの方向ベクトル
* org: レイの原点位置ベクトル
*/
struct NeRFRay {
Vec3 dir;
Vec3 org;
int nSample;
int SampleBeginIdx;
float tmin, tmax, pad1, pad2;
MNPT_HOST_DEVICE NeRFRay(const Vec3& dir = {1,0,0}, const Vec3& org = {0,0,0}) : org(org), dir(dir) {}
};
nSample，SampleBeginIdx, tmin, tmaxはこの後使用する際に説明します．
Camera構造体
さて，ではカメラの実装に取り掛かりましょう．
struct NeRFCamera {
size_t Image_height;
size_t Image_width;
Vec3 CamOrg;
Vec3 CamDir;
float CamToScreenDist;
Vec3 ScreenXDir;
Vec3 ScreenYDir;
float PixelSize = 0.0f;
MNPT_HOST_DEVICE NeRFCamera(const Vec3& CamDir, const Vec3& CamOrg, const float& CamToScreenDist) :
Image_height(0), Image_width(0), CamDir(CamDir), CamOrg(CamOrg), CamToScreenDist(CamToScreenDist) {}
// 指定されたカメラの情報からカメラの姿勢を求める
MNPT_HOST void SetUp(float rotateX, float rotateY, float rotateZ) {
const Vec3 CamUp = rotate(DEG2RAD(rotateX), DEG2RAD(rotateY), DEG2RAD(rotateZ), normalize({ 0.0, 1.0, 0.0 }));
CamDir = rotate(DEG2RAD(rotateX), DEG2RAD(rotateY), DEG2RAD(rotateZ), normalize({ 0.0, 0.0, -1.0 }));
ScreenXDir = cross(CamDir, CamUp);
ScreenYDir = CamUp;
PixelSize = 36.0f / (float)Image_height;
}
MNPT_HOST void SetUp(const size_t InputImage_height, const size_t InputImage_width, std::vector<std::vector<float>>& TransformMatrix, float FOVY) {
// 入力画像のサイズをカメラに記録
Image_height = InputImage_height;
Image_width = InputImage_width;
// 変換行列によりカメラの姿勢を計算
Vec3 CamUp = rotate(TransformMatrix, normalize({ 0.0, 1.0, 0.0 }));
CamDir = rotate(TransformMatrix, normalize({ 0.0, 0.0, -1.0 }));
std::vector<std::vector<float>> Pos(4, std::vector<float>(1));
std::vector<std::vector<float>> Base{{0.0f}, { 0.0f }, { 0.0f }, { 1.0f }};
Pos = multiply_matrix(TransformMatrix, Base);
CamOrg = { Pos[0][0], Pos[1][0], Pos[2][0] };
float top = 18.0f;
ScreenYDir = CamUp;
ScreenXDir = cross(CamDir, CamUp);
PixelSize = 2 * top / (float)Image_height;
CamToScreenDist = top / tanf(FOVY/2.0f);
}
MNPT_HOST void SetUp(Camera* Cam) {
Image_height = ImSizeY;
Image_width = ImSizeX;
CamOrg = Cam->CamOrg;
CamDir = Cam->CamDir;
CamToScreenDist = Cam->CamToScreenDist;
ScreenXDir = Cam->ScreenXDir;
ScreenYDir = Cam->ScreenYDir;
PixelSize = Cam->PixelSize;
}
// 画像上のピクセルインデックス{IndexX, IndexY}に対してのカメラからの1次レイを求める
MNPT_DEVICE inline NeRFRay Generate1stRay(const int IndexX, const int IndexY) {
Vec3 RayDir = CamDir * CamToScreenDist +
ScreenXDir * ((float)IndexX - Image_height / 2.0f) * PixelSize +
ScreenYDir * (Image_width / 2.0f - (float)IndexY) * PixelSize;
RayDir = normalize(RayDir, "FirstRayGen");
return {RayDir, CamOrg};
}
};
これは多少見ておきましょう．
struct NeRFCamera {
size_t Image_height;
size_t Image_width;
Vec3 CamOrg;
Vec3 CamDir;
float CamToScreenDist;
Vec3 ScreenXDir;
Vec3 ScreenYDir;
float PixelSize = 0.0f;
...
Image_heightやImage_widthはこのカメラの映し出す画像のサイズを指します．もっと言えばスクリーンのアスペクトを定義すると言っても良いでしょう．
PixelSizeはスクリーン上のピクセルサイズです．実はスクリーンのY方向のサイズを36に固定しております．つまり，PixelSizeは36/Image_heightとなります．この36という値は特に意味はなく，色々試した結果イイ感じの値というだけです．
後に続くSetUp関数は入力されたカメラの姿勢からメンバ変数を計算してあげる関数です．説明は省略しますが，rotate()は回転行列を作用させる関数です．
では本題のレイの計算に入ります．
// 画像上のピクセルインデックス{IndexX, IndexY}に対してのカメラからの1次レイを求める
MNPT_DEVICE inline NeRFRay Generate1stRay(const int IndexX, const int IndexY) {
Vec3 RayDir = CamDir * CamToScreenDist +
ScreenXDir * ((float)IndexX - Image_height / 2.0f) * PixelSize +
ScreenYDir * (Image_width / 2.0f - (float)IndexY) * PixelSize;
RayDir = normalize(RayDir, "FirstRayGen");
return {RayDir, CamOrg};
}
図中に示した矢印を赤緑青と上から順番に加算していることが分かると思います．細かいことを言えばピクセルの中心に当てるためには多少調整がいるのですが，まあしなくても良いです．計算したレイの長さを1に正規化してあげるのを忘れないようにしましょう．さて，ピクセル座標が分かればレイを計算できるようになりました．ここで先ほど入力データと教師データの対の構造体をつくっていました．ではこの構造体からレイを計算してあげましょう．
// カメラの情報からレイのバッチを作成する．
// これによって（レイ -- ピクセル色）という（入力データ--教師データ）の構造が構築される
__global__ void NeRF_GenerateRay(const int nPixel, NeRFCamera* Cam, PixelInfo *PixInfo, NeRFRay* GeneratedRay) {
int PixelId = blockIdx.x * blockDim.x + threadIdx.x;
if (PixelId >= nPixel) {
return;
}
PixelInfo PixelInfo_this_thread = PixInfo[PixelId];
int ImageId = PixelInfo_this_thread.ImageID();
int PosX = PixelInfo_this_thread.PixIndexX();
int PosY = PixelInfo_this_thread.PixIndexY();
// このピクセルにおけるRay
NeRFRay Ray = Cam[ImageId].Generate1stRay(PosX, PosY);
GeneratedRay[PixelId] = Ray;
}
細かく見ましょう．
__global__ void NeRF_GenerateRay(const int nPixel, NeRFCamera* Cam, PixelInfo *PixInfo, NeRFRay* GeneratedRay) {
...
nPixel: 1イテレーションで使用するピクセル数です．さっき書いたnMaxPixelPerBatchがこれにあたります．
Cam: カメラ配列です．i番目の要素は画像IDがiの画像を撮影したカメラ情報を持つ(NeRF)Camera構造体です．
PixInfo: （画像ID, ピクセル座標）- (教師画像の色)を対として持つ構造体です．
GeneratedRay: 計算されたレイが格納される場所
...
int PixelId = blockIdx.x * blockDim.x + threadIdx.x;
if (PixelId >= nPixel) {
return;
}
...
いつものやつです．スレッド数がnPixelを溢れた分は帰らせる処理です．
... PixelInfo PixelInfo_this_thread = PixInfo[PixelId]; int ImageId = PixelInfo_this_thread.ImageID(); int PosX = PixelInfo_this_thread.PixIndexX(); int PosY = PixelInfo_this_thread.PixIndexY(); // このピクセルにおけるRay NeRFRay Ray = Cam[ImageId].Generate1stRay(PosX, PosY); GeneratedRay[PixelId] = Ray; ...
スレッドIDに対応するPixelInfoにアクセスし，情報を読み出し，それを元にしてレイを計算しています．
なお，先程のシャッフルの際に「すべてのピクセルを1イテレーションでは使用しない」と書いていましたが，イテレーションに関わらずPixelInfoへのアクセスインデックスが0以上nPixel未満となっており，怪しそうですが，カーネル実行時に次のような引数を与えてPixelInfoポインタの先頭をずらしています．
// レイを計算する
NeRF_GenerateRay <<<nMaxPixelPerBatch / 512, 512, 0, stream_train >>>
(
nMaxPixelPerBatch,
D_NeRFCAM,
thrust::raw_pointer_cast(&D_thrust_TargetPixelInfo[IndexOfPixelInfo]),
thrust::raw_pointer_cast(&D_RayBatch[0])
);
gpuErrchk(cudaGetLastError());
gpuErrchk(cudaStreamSynchronize(stream_train));
IndexOfPixelInfoが何たるやは先ほどPixelInfoの説明をした際のシャッフルの話で記述した通りです．また，thrust::raw_pointer_cast()はthrust::device_vectorのデバイスメモリ上でのポインタを返してくれる嬉しい関数です．
NeRFの実装: (2): レイの経路上の点をサンプリング
さあ，下準備も済んでようやくNeRFらしくなってきました．レイの経路上の点をサンプリングするうえではサンプリングする領域を決める必要があります．NeRFを生成する空間をAABB（Axis-Aligned-Bounding-Box），つまりxyz軸に辺がそれぞれ平行である直方体，で定義してあげると計算が楽です．言わずもがなですが，サンプリングする領域はレイの経路の内，このAABB内部にある線分でしょう．ではAABBとレイの交差判定を行いましょう．そのために，AABBの構造体が必要です．
struct NeRFAABB {
Vec3 PosMin;
Vec3 PosMax;
MNPT_HOST_DEVICE NeRFAABB(const Vec3 PosMin = Vec3(SINF, SINF, SINF), const Vec3 PosMax = Vec3(BINF, BINF, BINF)) : PosMin(PosMin), PosMax(PosMax) {}
MNPT_DEVICE bool willIntersectWithAABB(NeRFRay& Ray) {
float t_max = 1e16f;
float t_min = -1e16f;
#pragma unroll
for (int i = 0; i < 3; i++) {
float t1, t2, t_near, t_far;
if (i == 0) {
if (abs(Ray.dir.x) < 1e-5) { // x軸に垂直な平面上のレイを飛ばす場合
if (PosMin.x > Ray.org.x || PosMax.x < Ray.org.x) return false;
else continue;
}
t1 = (PosMin.x - Ray.org.x) / Ray.dir.x;
t2 = (PosMax.x - Ray.org.x) / Ray.dir.x;
}
else if (i == 1) {
if (abs(Ray.dir.y) < 1e-5) { // Y軸に垂直な平面上のレイを飛ばす場合
if (PosMin.y > Ray.org.y || PosMax.y < Ray.org.y) return false;
else continue;
}
t1 = (PosMin.y - Ray.org.y) / Ray.dir.y;
t2 = (PosMax.y - Ray.org.y) / Ray.dir.y;
}
else {
if (abs(Ray.dir.z) < 1e-5) { // Z軸に垂直な平面上のレイを飛ばす場合
if (PosMin.z > Ray.org.z || PosMax.z < Ray.org.z) return false;
else continue;
}
t1 = (PosMin.z - Ray.org.z) / Ray.dir.z;
t2 = (PosMax.z - Ray.org.z) / Ray.dir.z;
}
t_near = min(t1, t2);
t_far = max(t1, t2);
t_max = min(t_max, t_far);
t_min = max(t_min, t_near);
if (t_min > t_max) return false;
}
Ray.tmin = max(t_min, 1e-4f);
Ray.tmax = t_max;
return true;
}
};
Part2でも示しましたが，AABBを定義するにはすべての軸において座標が小さい方の点の座標，そしてその対角点（すべての軸において座標が大きい方の点）の座標があれば十分です．これがPosMinとPosMaxです．さて，肝心の交差判定の具体的な説明は以下の記事に任せて省略します．
さて，先程(NeRF)Ray構造体のところでメンバ変数にtmin, tmaxとありました．これはレイの経路においてAABB内部にある，レイの原点からの距離tの最小値と最大値を保存します．私の実装における
... Ray.tmin = max(t_min, 1e-4f); Ray.tmax = t_max; ...
がそれです．ちなみにt_minが0や負になってほしくはないのでその場合は微小値を入れておきます．図で描くと次の感じです．
CMOSインバーターみたいなのはカメラです．t_min や t_maxが無い場合はそのレイに関してはサンプリングする必要がありません．
さて，サンプリングする領域が求まったのでサンプリングしていきましょう．サンプリングする処理を見ていきましょう．
// サンプル点を求める
/*
* ミニバッチに関する処理
* Rays: 関数内部でインデックス調整（入力には[0]を先頭としたポインタを渡す）
* Pos: 関数内部でインデックス調整（入力には[0]を先頭としたポインタを渡す）
* Dir: 関数内部でインデックス調整（入力には[0]を先頭としたポインタを渡す）
*/
__global__ void NeRF_GenerateSample
(
const uint32_t nMaxBatch,
const uint32_t nMaxPixelBatch,
NeRFInfo* SamplingInfo, NeRFRay* Rays, float* Pos, float* Dir, OccupancyGrid* Grids,
const uint32_t epoch
)
{
uint32_t threadID = blockIdx.x * blockDim.x + threadIdx.x;
if (threadID >= nMaxPixelBatch) {
return;
}
// レイ上の点をサンプリングする
NeRFRay* Ray = &Rays[threadID];
const Vec3 AABBMin = { SamplingInfo->AABB_pMin.x,SamplingInfo->AABB_pMin.y,SamplingInfo->AABB_pMin.z };
const Vec3 AABBMax = { SamplingInfo->AABB_pMax.x,SamplingInfo->AABB_pMax.y,SamplingInfo->AABB_pMax.z };
uint32_t nAcceptedSample = 0;
// AABBとの交差判定
NeRFAABB NeRFBox = { AABBMin, AABBMax};
if (!NeRFBox.willIntersectWithAABB(*Ray)) {
Ray->nSample = 0;
return;
}
float t_min = Ray->tmin;
float t_max = Ray->tmax;
const float dt = (t_max - t_min) / NERF_MAX_SAMPLE_PER_RAY;
// 0になった場合プログラムが停止してしまうため
if (dt < SINF) {
Ray->nSample = 0;
return;
}
// RNG
curandState state, state_old, state_firstAccept;
curand_init(epoch, threadID, 0, &state);
// Uniform sampling
int index = 0;
int index_firstAccept = 0;
constexpr float Throughput_thres = 0.01f;
float Throughput_hat = 1.0f; // これがThroughput_thresを下回れば中断
Vec3 LastSamplePos;
// 前半 ////////////////////////////////////////////////////////////////////////////////////////////////////////
// とりあえずどのレイがどれぐらいサンプル点を取るかを求める
// それによってMLPへの入力データのどの領域にサンプル点の情報が保存されるかを計算する．
while (nAcceptedSample < NERF_MAX_SAMPLE_PER_RAY) {
const float rnd = 0.5f;// curand_uniform(&state);
float t1 = t_min + SINF + dt * index;
float t2 = t_min + SINF + (dt * (index+1));
float t = normalize(rnd, 0.0f, 1.0f, t1, t2);
Vec3 s_Pos = Ray->org + Ray->dir * t;
// AABBの領域に収める
if (t > t_max) {
break;
}
// 念のため
s_Pos.x = clamp(s_Pos.x, AABBMin.x + SINF, AABBMax.x - SINF);
s_Pos.y = clamp(s_Pos.y, AABBMin.y + SINF, AABBMax.y - SINF);
s_Pos.z = clamp(s_Pos.z, AABBMin.z + SINF, AABBMax.z - SINF);
// Occupancy Gridにおいて「物体が無い」と判断された場合はスキップする
if (epoch > 16 && !Grids[0].is_Occupied(s_Pos)) {
index++;
continue;
}
else {
if (nAcceptedSample == 0) {
index_firstAccept = index;
}
nAcceptedSample++;
}
// Throughputの更新
if (epoch > 32 && nAcceptedSample >= 2) {
Vec3 FromLastSample =
{
s_Pos.x - LastSamplePos.x,
s_Pos.y - LastSamplePos.y,
s_Pos.z - LastSamplePos.z
};
Throughput_hat *= expf(-1.0f * Grids[0].Float_at(s_Pos) * FromLastSample.length());
}
// それ以降のサンプル点の寄与の見込みが小さい場合は打ち切る
if (Throughput_hat < Throughput_thres) {
break;
}
//state_old = state;
LastSamplePos = s_Pos;
index++;
}
// サンプリングが完了した
// Dir配列とOrg配列にサンプル情報を記録する（配列上で連続となるようにする）
// 格納先の配列におけるインデックス範囲を記録 : [begin, end]
int begin_idx = atomicAdd(&SamplingInfo->nSamples, nAcceptedSample);
Ray->nSample = nAcceptedSample;
Ray->SampleBeginIdx = begin_idx;
// 後半 ///////////////////////////////////////////////////////////////////////////////////////////////////
// サンプルの保存インデックスがMaxBatchを超えている場合はそのレイを捨てる
if (begin_idx + nAcceptedSample >= nMaxBatch) {
Ray->nSample = 0;
atomicAdd(&SamplingInfo->nSamples, -nAcceptedSample);
return;
}
atomicAdd(&SamplingInfo->RayNum, 1);
// はじめてサンプルを得たところまでスキップ
index = index_firstAccept;
state = state_firstAccept;
nAcceptedSample = 0;
Throughput_hat = 1.0f;
while (nAcceptedSample < NERF_MAX_SAMPLE_PER_RAY) {
const float rnd = 0.5f;//curand_uniform(&state);
float t1 = t_min + SINF + dt * index;
float t2 = t_min + SINF + (dt * (index+1));
float t = normalize(rnd, 0.0f, 1.0f, t1, t2);
Vec3 s_Pos = Ray->org + Ray->dir * t;
// AABBの領域に収める
if (t > t_max) {
break;
}
// 念のため
s_Pos.x = clamp(s_Pos.x, AABBMin.x + SINF, AABBMax.x - SINF);
s_Pos.y = clamp(s_Pos.y, AABBMin.y + SINF, AABBMax.y - SINF);
s_Pos.z = clamp(s_Pos.z, AABBMin.z + SINF, AABBMax.z - SINF);
// Occupancy Gridにおいて「物体が無い」と判断された場合はスキップする
if (epoch > 16 && !Grids[0].is_Occupied(s_Pos)) {
index++;
continue;
}
else {
Ray->dir.to_float(Dir + 3 * (begin_idx + nAcceptedSample));
s_Pos.to_float(Pos + 3 * (begin_idx + nAcceptedSample));
nAcceptedSample++;
}
// Throughputの更新
if (epoch > 32 && nAcceptedSample >= 2) {
Vec3 FromLastSample =
{
s_Pos.x - LastSamplePos.x,
s_Pos.y - LastSamplePos.y,
s_Pos.z - LastSamplePos.z
};
Throughput_hat *= expf(-1.0f * Grids[0].Float_at(s_Pos) * FromLastSample.length());
}
// それ以降のサンプル点の寄与の見込みが小さい場合は打ち切る
if (Throughput_hat < Throughput_thres) {
break;
}
LastSamplePos = s_Pos;
index++;
}
}
さて，とてつもなく長いコードが出てきました．部分で見ていきましょう．
__global__ void NeRF_GenerateSample
(
const uint32_t nMaxBatch,
const uint32_t nMaxPixelBatch,
NeRFInfo* SamplingInfo, NeRFRay* Rays, float* Pos, float* Dir, OccupancyGrid* Grids,
const uint32_t epoch
)
{
...
nMaxBatch: すべてのレイに対してサンプリングする点の総和の上限（NNに入力するバッチサイズの上限）
nMaxPixelBatch: ピクセル数の上限（このイテレーションで考慮するレイの本数）
NeRFInfo* SamplingInfo: サンプリングに関わる情報を記録する構造体
Rays: レイが保存されている配列
Pos, Dir: サンプル点における座標とレイの方向を記録する配列
Grids: Occupancy Gridの情報（最後に説明します）
nMaxPixelBatchとnMaxBatchが両方ある理由ですが，実際の処理においては本当にnMaxBatchだけサンプル点をサンプリングするわけではありません．この話は今回のNeRFのコード設計に大きくかかわっていることです．具体的な理由はこの後で説明します．
さて，ここで出てきた構造体NeRFInfoについて書いておきます．
struct NeRFInfo {
uint32_t nSamples = 0;
uint32_t RayNum = 0;
float3 AABB_pMin = { -1.0f, -1.0f, -1.0f };
float3 AABB_pMax = { 1.0f, 1.0f, 1.0f };
// GUI
uint32_t epoch = 0;
float Loss = 0.0f;
uint32_t BatchSize = 0;
uint32_t nRay = 0;
bool StopTrain = false;
bool StopRender = false;
};
nSamplesはすぐ後で説明します
RayNum: 最終的に使用することとなったレイの本数（詳しいことはすぐ後で説明します）
AABB_pMin, AABB_pMax: NeRFを生成するAABBの設定です．
これ以降のメンバ変数はNeRFの実装に直接かかわってはこないので省略します．（GUI上で情報を得るため等の為の変数です）
さて，さっきから「すべてをサンプリングしない」とか「最終的に使用することとなったレイの本数」とかいう理解を拒ませる変なことを言ってますが，さっさと明らかにしておきます．今回の関数は「各レイにおけるサンプル数を動的に決定する」実装を行っています．具体的な処理の流れを確認しましょう．まず実装の上で次を前提としています．
各レイはNERF_MAX_SAMPLE_PER_RAYという即値よりも多くサンプル点をサンプリングすることを許されません．しかし，すべてのレイがNERF_MAX_SAMPLE_PER_RAYもサンプル点を必要としているわけではありません．例えば，次の図に示すような場合，レイに対するサンプル数が少なくて済みます．
さらに，MLPに入力できるバッチサイズの上限をnMaxBatchとして，具体的には221として固定しております．また，学習の速度を上げたいのでなるべく多くのレイをバッチにはめ込みたいという欲望があります．仮にすべてのレイがNERF_MAX_SAMPLE_PER_RAY個のサンプル点を取ると決めつけてしまうと実装は楽になりますが，かなりレイの数が限られてしまいます．逆に，各レイの持つサンプル点の個数を動的になるべく最小限に抑えてあげることにより，実装は少しめんどくさいですがバッチになるべく多くのレイが収まってくれます．
各レイについて動的にサンプリングを行うとしましたが，じゃあ一体どうやって「扱いやすいデータ」とするのかという疑問が浮かびます．まあ実装してみると分かりますが，次のような問題が発生します．
簡単な実装をするとすれば，各レイに対してNERF_MAX_SAMPLE_PER_RAY要素だけサンプル点情報の領域を確保しておき，余った分は空白としておく，となるでしょう．しかし，これは全くもってメモリ領域の無駄であり，さらに言えばNNに入力するデータとしてはバッチに空白のデータが入り込んでしまい，正しい処理が出来ません．そのため，理想としては下側のように各レイでのサンプル点がメモリ上で連続に配置されている状態となります．
さて，各レイの処理は並列に行っています．ですが，このままでは処理中のスレッドにとって「今自分が処理しているサンプル点の情報をメモリ上のどこに書き込めばいいのか」が不明です．並列処理を諦めることになるのでしょうか．いいえ，実は同じ処理を2回繰り返すことでほとんど並列性を保って実装できます．前半では各スレッドが「自分が何個サンプリングすることになって，サンプル点情報を書き込む配列のどこを始点として書き込めばいいのか」を決定します．
そして後半では「前半と全く同じサンプリングを行い，決めた領域に書き込む」ということを行います．
では具体的な実装を確認していきます．
... uint32_t nAcceptedSample = 0; ...
nAcceptedSample: このレイにおけるサンプリングすると決定されたサンプル点の数を記録します
// AABBとの交差判定
NeRFAABB NeRFBox = { AABBMin, AABBMax};
if (!NeRFBox.willIntersectWithAABB(*Ray)) {
Ray->nSample = 0;
return;
}
float t_min = Ray->tmin;
float t_max = Ray->tmax;
そもそもNeRFを生成する領域（AABB）とレイが交差しなければそのレイについては一切のサンプリングを行いません．なので交差判定を行って処理の必要性，そしてレイの経路上でAABBの内部にある領域を求めておきます．また，(NeRF)Rayのメンバ変数にあったnSampleは，この後も書きますが「そのレイが持つサンプル点の数」です．AABBとの交差が確認されないレイは0となります．
const float dt = (t_max - t_min) / NERF_MAX_SAMPLE_PER_RAY;
// 0になった場合プログラムが停止してしまうため
if (dt < SINF) {
Ray->nSample = 0;
return;
}
さて，サンプリングを行います．今回は単純に，AABB内部にあるレイの経路をNERF_MAX_SAMPLE_PER_RAY等分し，等分された一つの経路上の点をサンプル点としてふさわしいかを考慮することにします．なので「ステップサイズ」としてdtを設定してあげます．ちなみにコーナーケースとして，「レイがAABBの端ぎりぎりに入射した場合」はdtが微小となります．この場合，最悪プログラムがこの後のwhile文内でスタックするので微小の場合はやはりサンプリングせずにreturnしておきます．
// RNG curandState state, state_old, state_firstAccept; curand_init(epoch, threadID, 0, &state);
これは乱数生成のためのCUDAライブラリですが，現状の実装では使わないので無視してください（経路上の点をランダムサンプリングするやり方もありますが，今回は等間隔サンプリングとします．ちなみにInstantNGPの論文では他にも色々サンプリング手法が書かれています）．
// Uniform sampling int index = 0; int index_firstAccept = 0; constexpr float Throughput_thres = 0.01f; float Throughput_hat = 1.0f; // これがThroughput_thresを下回れば中断 Vec3 LastSamplePos;
index: 何個目の等間隔サンプル点を考えているか，つまり考えているサンプル点のインデックスです．
index_firstAccept: 後半の処理では全く同じサンプリングをするので，前半で初めてサンプリングが成立したインデックスを記録しておきます．
Throughput_thres: 先ほどの例の図に示したcase2のように，寄与が非常に小さい場合は無視しても問題ありません．なのでその閾値を設定しておきます．
Throughput_hat: この寄与の大きさはNNを通さないと分からないので，Occupancy Gridに保存されている情報を用いて「ざっくりと」推定しておきます
LastSamplePos: 寄与の大きさを計算するためには前のサンプルからの距離が必要です．
// 前半 ////////////////////////////////////////////////////////////////////////////////////////////////////////
// とりあえずどのレイがどれぐらいサンプル点を取るかを求める
// それによってMLPへの入力データのどの領域にサンプル点の情報が保存されるかを計算する．
while (nAcceptedSample < NERF_MAX_SAMPLE_PER_RAY) {
サンプル数がNERF_MAX_SAMPLE_PER_RAYを超えたらもうそれ以上サンプリングさせないようにします（なお実装上，バグが無ければありえません）
const float rnd = 0.5f;// curand_uniform(&state); float t1 = t_min + SINF + dt * index; float t2 = t_min + SINF + (dt * (index+1)); float t = normalize(rnd, 0.0f, 1.0f, t1, t2);
サンプル点のカメラからの距離を決めてあげます．index番目のサンプル点は[t1, t2]の領域にあります．今回はここの中心を使ってあげます．つまり，サンプル点のカメラからの距離はとなっています．
Vec3 s_Pos = Ray->org + Ray->dir * t;
// AABBの領域に収める
if (t > t_max) {
break;
}
// 念のため
s_Pos.x = clamp(s_Pos.x, AABBMin.x + SINF, AABBMax.x - SINF);
s_Pos.y = clamp(s_Pos.y, AABBMin.y + SINF, AABBMax.y - SINF);
s_Pos.z = clamp(s_Pos.z, AABBMin.z + SINF, AABBMax.z - SINF);
s_Posはサンプル点の座標を記録する変数です．
サンプル点の座標がt_maxを超えた場合は，これ以降のindexにおいてサンプル点がAABBの内部に戻ってくることはあり得ないのでbreakしてあげます．
もしものため，s_Posの座標が絶対にAABBの内部に存在するようにしておきます．ちなみにこれをしなくても外部にあることはあり得ないはずです．
// Occupancy Gridにおいて「物体が無い」と判断された場合はスキップする
if (epoch > 16 && !Grids[0].is_Occupied(s_Pos)) {
index++;
continue;
}
else {
if (nAcceptedSample == 0) {
index_firstAccept = index;
}
nAcceptedSample++;
}
Occupancy Gridの話は最後にしますが，ここでやっているのは「Occupancy Gridに格納されているデータを参照し，そのサンプル点において「物体が存在する」，言い換えると媒質の密度がある程度高いと判定されない場合はサンプル点を無視する」ということをやっています．結局は「カメラに入ってくるレイに関与しないサンプル点を無視する」ということです． 無視しない場合はサンプリングが成立するので，初めてのサンプリング成立であればインデックスを保存し，そしてnAcceptedSampleに1を足してあげます．
// Throughputの更新
if (epoch > 32 && nAcceptedSample >= 2) {
Vec3 FromLastSample =
{
s_Pos.x - LastSamplePos.x,
s_Pos.y - LastSamplePos.y,
s_Pos.z - LastSamplePos.z
};
Throughput_hat *= expf(-1.0f * Grids[0].Float_at(s_Pos) * FromLastSample.length());
}
// それ以降のサンプル点の寄与の見込みが小さい場合は打ち切る
if (Throughput_hat < Throughput_thres) {
break;
}
コメントの通りです．
LastSamplePos = s_Pos; index++;
さて，サンプリングが成立しました．なのでLastSamplePosを更新し，そしてwhile文の終わりなのでindexに1足してあげて次の等間隔領域に移ります．
// サンプリングが完了した // Dir配列とOrg配列にサンプル情報を記録する（配列上で連続となるようにする） // 格納先の配列におけるインデックス範囲を記録 : [begin, end] int begin_idx = atomicAdd(&SamplingInfo->nSamples, nAcceptedSample); Ray->nSample = nAcceptedSample; Ray->SampleBeginIdx = begin_idx;
さて，前半のサンプリングの処理が終わりました．ここで各スレッドは「自分がサンプル点情報を格納する配列に，「どこから」「何個」サンプル点情報を保存するか」を知る必要があります．なので，説明時に出てきたCounterにどれだけサンプリングが成立したかを教えてあげます．このCounterはNeRFInfoのメンバ変数であるSamplingInfo.nSamplesです．SamplingInfo.nSamplesは現在既に何個のサンプリングが他のスレッドで成立してきたか(how many samples "have been" generated)をbegin_idxへと格納すると同時に処理中のスレッドが何個サンプル点を得たかを聞き，加算します．勿論，この処理はatomicで行われる必要があります．こうすることでCounterが更新され，そしてスレッドが「配列のどこから格納すればいいのか」を決定できます．
ここで，Rayのメンバ変数であるnSampleとSampleBeginIdxにそのスレッドが「どこから」「何個」サンプル点情報を保存したかを記憶させておきます．これは後にボリュームレンダリングを行う際に必要となる情報です．
// 後半 ///////////////////////////////////////////////////////////////////////////////////////////////////
// サンプルの保存インデックスがMaxBatchを超えている場合はそのレイを捨てる
if (begin_idx + nAcceptedSample >= nMaxBatch) {
Ray->nSample = 0;
atomicAdd(&SamplingInfo->nSamples, -nAcceptedSample);
return;
}
atomicAdd(&SamplingInfo->RayNum, 1);
さて，後半の開始です．この処理の最初に，サンプル点の総数には上限があり，それがnMaxBatchで定義しておく，とありました．その部分の処理を行います．各スレッドが自分のbegin_idxとnAcceptedSampleから配列上に書き込むインデックスの範囲を知ることが出来ます．それがnMaxBatchを超えていた場合は，サンプル点情報の登録をキャンセルします．そして，キャンセルが発生したのでSamplingInfo->nSamplesからそのスレッドにおけるnAcceptedSampleを引いておきます．ちなみにここでスレッド同期の問題が不安として頭をよぎりましたが，「nMaxBatchから溢れている状態」となっているのが端の問題であるので多分大丈夫なはずです．
// はじめてサンプルを得たところまでスキップ index = index_firstAccept; state = state_firstAccept; nAcceptedSample = 0; Throughput_hat = 1.0f;
後半のサンプリングは前半と全く同じものです．やりましょう．ほとんどが同じなので説明は省略しますが，一か所だけ
// Occupancy Gridにおいて「物体が無い」と判断された場合はスキップする
if (epoch > 16 && !Grids[0].is_Occupied(s_Pos)) {
index++;
continue;
}
else {
Ray->dir.to_float(Dir + 3 * (begin_idx + nAcceptedSample));
s_Pos.to_float(Pos + 3 * (begin_idx + nAcceptedSample));
nAcceptedSample++;
}
前半ではLastSamplePosに格納していましたが，今度はサンプル点の情報を格納する配列，即ちDirとPosに保存します．以上でサンプル点の情報をメモリ上で連続に配置することが出来ました．そして，このままNNに入力することが出来ます．なお，前半の際に静的に確保しているメモリ領域にサンプル点情報を保存すれば後半は不要になりますが，メモリ領域を圧迫しますし非効率なメモリ消費がかなり多くなります．
NeRFの実装: (3)と(4): NNの処理
サンプル点における「座標」と「方向」を入力し，サンプル点における「色」と「密度」を推定するNNを実装する必要があります．InstantNGPの論文では次のネットワークが推されています．
FCというのは全結合層で，活性化関数を下に付記しています．ReLUじゃなくてLeakyReLUを使用しても良いです．図の表記にクセがある（作ったのは半年前の私ですが）ので念のため言葉でも説明しておくと，
図中左半分は「密度」を推定するネットワークで，Density MLPと呼ぶことにし，次の構造を取ります．
入力はサンプル点の座標で，これをMultiresolution Hash Encodingに通しておきます．として32次元の出力にします（厳密ではないので異なっていてもまあ良いと思います）．
MLP: 入力層32次元, 隠れ層64次元，出力層16次元，隠れ層は1層，活性化関数はすべてReLUないしLeakyReLU，そしてexp(出力層の一つ目の要素)を「密度」の推定値とします．
そして後半ではサンプル点の「色」を推定します．Color MLPとします．ここでサンプル点におけるレイの方向をまでのSpherical Harmonic Encoding（実数球面調和関数の基底に変換）を通して16次元にしておきます．そしてdensity networkの出力である16次元（勿論1要素目はexp activationされていない状態のものです）のベクトルと結合し，合わせて32次元としておきます．
MLP: 入力層32次元, 隠れ層64次元，出力層16次元，隠れ層は2層，活性化関数は入力層と隠れ層はReLUないしLeakyReLU，そして出力層はSigmoid（ロジスティック）とします．
なお，私の実装ではColor MLPの隠れ層の数を1にしています．では実装を見ましょう．
template <const uint32_t indim_den, const uint32_t hiddendim_den, const uint32_t outdim_den, const uint32_t nHiddenLayer_den,
const uint32_t indim_col, const uint32_t hiddendim_col, const uint32_t outdim_col, const uint32_t nHiddenLayer_col>
__global__ void MFFM_NeRF_Forward
(
NeRFInfo* SamplingInfo, EncoderInfo* EncInfo_den, EncoderInfo* EncInfo_col,
Encoder Encoder_den, Activation ActHid_den, Activation ActOut_den,
Encoder Encoder_col, Activation ActHid_col, Activation ActOut_col,
float* Pos, float* Dir,
__half* weights_den, __half* buffers_den,
__half* weights_col, __half* buffers_col,
float* Density, float* Color)
{
const int BatchSize = SamplingInfo->nSamples;
// 即値
constexpr uint32_t indim_den_aligned = next_multiple(indim_den, TENSOR_ROW);
constexpr uint32_t hiddendim_den_aligned = next_multiple(hiddendim_den, TENSOR_ROW);
constexpr uint32_t outdim_den_aligned = next_multiple(outdim_den, TENSOR_ROW);
constexpr uint32_t indim_col_aligned = next_multiple(indim_col, TENSOR_ROW);
constexpr uint32_t hiddendim_col_aligned = next_multiple(hiddendim_col, TENSOR_ROW);
constexpr uint32_t outdim_col_aligned = next_multiple(outdim_col, TENSOR_ROW);
const uint32_t maxdim_den = max(indim_den_aligned, max(hiddendim_den_aligned, outdim_den_aligned));
extern __shared__ __half shmem[];
__half* intermediate_col = shmem;
__half* intermediate_den = shmem;
__half* intermediate_dirInput = shmem + (outdim_den_aligned + SKEW) * ONEBATCH_SIZE;
//////////////////////////////////////////////////////// Density MLP //////////////////////////////////////////////////////////////////////////
// 入力のロード
// エンコーダー無しの場合はこの時点でSKEWを与える
// エンコーダーありの場合はSKEWを与えない（ただのコピー）
if (Encoder_den == Encoder::None) {
load_input(indim_den, indim_den_aligned + SKEW, Pos, intermediate_den);
}
else {
load_input(3, 3, Pos, intermediate_den);
}
// Encode
Encode<indim_den_aligned>(Encoder_den, *EncInfo_den,intermediate_den, false, SamplingInfo->AABB_pMin, SamplingInfo->AABB_pMax);
// Density MLP
Kernel_Debug_train_forward<indim_den, hiddendim_den, outdim_den, nHiddenLayer_den>(BatchSize, ActHid_den, ActOut_den, intermediate_den, weights_den, buffers_den);
// Density MLPの出力の1つめの要素をDensityとして保存
store_intermediate<float>(outdim_den_aligned + SKEW, 1, intermediate_den, Density);
//////////////////////////////////////////////////////// END: Density MLP //////////////////////////////////////////////////////////////////////
__syncthreads(); // 必要
//////////////////////////////////////////////////////// Color MLP //////////////////////////////////////////////////////////////////////////
// 入力のロード
// エンコーダー無しの場合はこの時点でSKEWを与える
// エンコーダーありの場合はSKEWを与えない（ただのコピー）
if (Encoder_col == Encoder::None) {
load_input(indim_col - outdim_den, indim_col_aligned - outdim_den_aligned + SKEW, Dir, intermediate_dirInput);
}
else {
load_input(3, 3, Dir, intermediate_dirInput);
}
// Encode
Encode<indim_col_aligned - outdim_den_aligned>(Encoder_col, *EncInfo_col, intermediate_dirInput, false, SamplingInfo->AABB_pMin, SamplingInfo->AABB_pMax);
__syncthreads();
// Density MLPの出力をconcatする．
ConcatShmem<outdim_den, indim_col - outdim_den>(intermediate_den, intermediate_dirInput, intermediate_col);
__syncthreads();
// Color MLP
Kernel_Debug_train_forward<indim_col, hiddendim_col, outdim_col, nHiddenLayer_col>(BatchSize, ActHid_col, ActOut_col, intermediate_col, weights_col, buffers_col);
// 結果を保存
store_intermediate<float>(outdim_col_aligned + SKEW, 3, intermediate_col, Color);
//////////////////////////////////////////////////////// END: Density MLP //////////////////////////////////////////////////////////////////////
}
特に複雑なことはしていないので軽く見ていきましょう．
template <const uint32_t indim_den, const uint32_t hiddendim_den, const uint32_t outdim_den, const uint32_t nHiddenLayer_den,
const uint32_t indim_col, const uint32_t hiddendim_col, const uint32_t outdim_col, const uint32_t nHiddenLayer_col>
この関数ではdenとついているものはDensity MLPに関わる値で，colとついているのはColor MLPに関わる値です．indimは入力層の次元，hiddendimは隠れ層の次元，outdimは出力層の次元，nHiddenLayerは隠れ層の数です．
...
__global__ void MFFM_NeRF_Forward
(
NeRFInfo* SamplingInfo, EncoderInfo* EncInfo_den, EncoderInfo* EncInfo_col,
Encoder Encoder_den, Activation ActHid_den, Activation ActOut_den,
Encoder Encoder_col, Activation ActHid_col, Activation ActOut_col,
float* Pos, float* Dir,
__half* weights_den, __half* buffers_den,
__half* weights_col, __half* buffers_col,
float* Density, float* Color)
{
const int BatchSize = SamplingInfo->nSamples;
...
SamplingInfo: サンプリング処理で出てきたものと同じです．この構造体に「NN実行時のバッチサイズ」と「NeRFを生成するAABBの情報」が保存されているのでそれを読み込むために使用します．最後の行でBatchSizeに読み込ませているのが分かると思います．
EncInfoというのはエンコーダーに関する情報が入っていますが，まだ実装途中（Part2で軽く触れた内容です）なので無視してください．今回の実装では考えなくていいです．
Encoder, Activationというのは次に示すものです
enum class Encoder {
None,
Frequency,
SH, // Spherical Harmonic
HashGrid, // Multiresolution Hash Encoding
UniformGrid
};
enum class Activation {
ReLU,
LeakyReLU,
Sigmoid
};
Pos, Dirはサンプリング処理で得られたサンプル点の情報です．
weightsは全結合層のパラメーターです．
bufferは逆伝播時に使用する順伝播時の各層の入出力を保存しておくための配列です．
Density, ColorはNNの出力（「密度」「色」）を記録するための配列です．
... // 即値 constexpr uint32_t indim_den_aligned = next_multiple(indim_den, TENSOR_ROW); ... const uint32_t maxdim_den = max(indim_den_aligned, max(hiddendim_den_aligned, outdim_den_aligned)); ...
即値を計算していますが詳しい話はPart1を参照してください．
extern __shared__ __half shmem[]; __half* intermediate_col = shmem; __half* intermediate_den = shmem; __half* intermediate_dirInput = shmem + (outdim_den_aligned + SKEW) * ONEBATCH_SIZE;
1行目はshared memoryを動的に確保している部分です．処理中の特徴ベクトルであるintermediateをshared memoryに載せてやるという意図です．実際のところDensity MLPとColor MLPでは同じポインタを使用しますが，実装上区別したいので異なる名前を付けています．なおintermediate_dirInputについてですが，サンプル点の方向ベクトルはDensity MLPの処理を終えた後にshared memoryにロードします．Density MLPの出力がshared memoryに載ったままロードするので，Density MLPの出力を破壊しないように後ろ側の空きスペースにロードします．
//////////////////////////////////////////////////////// Density MLP //////////////////////////////////////////////////////////////////////////
// 入力のロード
// エンコーダー無しの場合はこの時点でSKEWを与える
// エンコーダーありの場合はSKEWを与えない（ただのコピー）
if (Encoder_den == Encoder::None) {
load_input(indim_den, indim_den_aligned + SKEW, Pos, intermediate_den);
}
else {
load_input(3, 3, Pos, intermediate_den);
}
// Encode
Encode<indim_den_aligned>(Encoder_den, *EncInfo_den,intermediate_den, false, SamplingInfo->AABB_pMin, SamplingInfo->AABB_pMax);
// Density MLP
Kernel_Debug_train_forward<indim_den, hiddendim_den, outdim_den, nHiddenLayer_den>(BatchSize, ActHid_den, ActOut_den, intermediate_den, weights_den, buffers_den);
// Density MLPの出力の1つめの要素をDensityとして保存
store_intermediate<float>(outdim_den_aligned + SKEW, 1, intermediate_den, Density);
//////////////////////////////////////////////////////// END: Density MLP //////////////////////////////////////////////////////////////////////
Density MLPの処理です．今回は必ずエンコーダーを通すので最初のif文はelseを通ります．現状実装しているエンコーダーは全て入力次元を3と仮定しているのでここでは3次元の入力としてハードコードしてロードしてますが，変更を加える予定です．load_input()はPart1を参照してください．やっていることとしては，3次元の入力（座標）をintermediate(shared memory)にロードしているだけです．3次元なのでバンクコンフリクトを避けるためのSKEWは入れません．つまりただのコピーです．
Encodeというのは次の関数です．
template <const uint32_t indim_aligned>
MFFM_DEVICE void Encode(Encoder encoder, EncoderInfo &Info, __half* intermediate, bool isInference, const float3 InputRangeMin = { 0,0,0 }, const float3 InputRangeMax = { 1,1,1 }, float* MHEBufferSTU = nullptr, unsigned int* MHEBufferIdxHT = nullptr) {
const int bx = blockIdx.x;
const int tx = threadIdx.x;
const int ty = threadIdx.y;
switch (encoder) {
case Encoder::None:
// NOP
break;
case Encoder::Frequency:
// not implemented...
break;
case Encoder::SH:
SH::Encode_SH_L4(intermediate, intermediate);
break;
case Encoder::HashGrid:
MHE::Encode<indim_aligned>(InputRangeMin, InputRangeMax, intermediate, intermediate);
//MHE2::Encode_inside_kernel<MHE_F>(*Info.MHEInfo, InputRangeMin, InputRangeMax, intermediate, intermediate);
break;
case Encoder::UniformGrid:
UGE::Encode(*Info.UGEInfo, intermediate, intermediate);
break;
default:
printf("Invalid Encoding type\n");
break;
}
__syncthreads();
}
今回はMultiresolution Hash Encodingを使用するのでHashGridの部分に入ります．Part2に示したEncode関数を使用します．SamplingInfoに保存している「NeRFを生成するAABBの情報」をMultiresolution Hash Encodingのエンコード領域に設定しているのが分かると思います．Part2を参照すれば具体的に何をしているかが分かるはずです．
さて，エンコードが終わったのでMLPに入力します．関数名がかなり酷いですがKernel_Debug_train_forward()がこの処理をしています．詳しい処理はPart1にあります．（Debug用の関数が最新版に成り果てました……） MLPの処理が終わったので1要素目をDensityに関わる値として保存しておきます．正確にはここでExp Activationをした方が良いのですが，今回は1要素目だけをActivationする特殊な例として扱い，ボリュームレンダリングを行う際にExp Activationしてあげることにします．store_intermediateを使用して保存してあげます．この関数もPart1に書かれています．
__syncthreads(); // 必要
ブロックレベルでの同期を取りましょう．
Color MLPも基本的に同じことをしますが，先程の図に示した結合処理を行っておく必要があります．
... // Density MLPの出力をconcatする． ConcatShmem<outdim_den, indim_col - outdim_den>(intermediate_den, intermediate_dirInput, intermediate_col); ...
Color MLPの入力次元はDensity MLPの出力次元とエンコードされた方向ベクトルの次元の和であるので，言い換えるとoutdim_den次元のベクトル（Density MLPの出力）と indim_col - outdim_den次元のベクトル（エンコードされた方向ベクトル）の結合を行うことになる，というのがtemplate引数の部分です．具体的にこの関数の中身を見ましょう
/////////////////////////////// CONCAT AND DIVIDE ///////////////////////////////////////////////////
/*
* 1: shmem1
* 2: shmem2
* s: SKEW
* 1s1s1s1s1s...1s1s1s1s2s2s2s2s2s2s...2s2s2s2s → 12s12s12s12s12s12s....12s12s12s12s12s
* 番地的にはshmem = shmem1
*/
template <const uint32_t dim1, const uint32_t dim2>
MFFM_DEVICE void ConcatShmem(__half* shmem1, __half* shmem2, __half* shmem) {
// 即値
constexpr uint32_t dim1_aligned = next_multiple(dim1, TENSOR_ROW);
constexpr uint32_t dim2_aligned = next_multiple(dim2, TENSOR_ROW);
constexpr uint32_t newDim = dim1 + dim2;
constexpr uint32_t newDim_aligned = next_multiple(newDim, TENSOR_ROW);
const int bx = blockIdx.x;
const int tx = threadIdx.x;
const int ty = threadIdx.y;
const int warp_index_required = ONEBATCH_SIZE / 32;
if (ty >= warp_index_required) {
return;
}
float tmp[dim1 + dim2];
__syncthreads();
for (int i = 0; i < dim1; i++) {
tmp[i] = shmem1[(dim1_aligned+SKEW) * 32 * ty + (dim1_aligned+SKEW) * tx + i];
}
for (int i = 0; i < dim2; i++) {
tmp[dim1 + i] = shmem2[(dim2_aligned+SKEW) * 32 * ty + (dim2_aligned+SKEW) * tx + i];
}
__syncthreads();
for (int i = 0; i < dim1 + dim2; i++) {
shmem[(newDim + SKEW) * 32 * ty + (newDim + SKEW) * tx + i] = (__half)tmp[i];
}
for (int i = 0; i < SKEW; i++) {
shmem[(newDim + SKEW) * 32 * ty + (newDim + SKEW) * tx + i + dim1 + dim2] = (__half)0.0f;
}
__syncthreads();
}
Part1を読んだ方であれば何をしているのか分かると思います．複雑そうに見えますが単純にconcatしているだけです．結合前の2つのベクトルは共に[要素][空白][要素][空白]......の構造を取っており，それを[要素1][要素2][空白][要素1][要素2][空白]......と組み直すのがこの関数の処理です．
NeRFの実装: (5)~(7): ボリュームレンダリングの順方向と逆方向
一気に逆方向まで処理してしまいます．最初に示したボリュームレンダリングをやるという話です．NNを通した結果，「色」と「密度」が得られています．説明の時に記した図を再掲します．
この式に代入するだけで計算できます．そして，NNへの逆伝播を行うために，「色」と「密度」への勾配の値を計算する必要があります．詳しい導出はしませんが，勾配は次の式で行えます．
後はこれを実装するだけです．やりましょう．
__global__ void CalculateOutputAndLoss
(
const uint32_t nPixelBatch, NeRFInfo* SamplingInfo,
float* Pos,
NeRFRay* Rays, float* Density, float* Color,
float* Output, PixelInfo* Target, float* Loss, bool NoLossCalc = false
)
{
uint32_t threadID = blockIdx.x * blockDim.x + threadIdx.x;
if (threadID >= nPixelBatch) {
return;
}
// 本スレッドでのポインタ等
NeRFRay* Ray = &Rays[threadID];
const uint32_t IdxBias = Ray->SampleBeginIdx;
const uint32_t nSample = Ray->nSample;
float* pos = Pos + 3 * IdxBias;
float* c = Color + 3 * IdxBias;
float* sigma = Density + IdxBias;
float* out = Output + 3 * threadID;
Vec3h ans = Target[threadID].Color;
float* err = Loss + 3 * threadID;
// サンプル点の厚み（次のサンプル点までの距離）
float dt;
// Throughput
float T = 1.0f;
for (int rgb = 0; rgb < 3; rgb++) {
out[rgb] = 0.0f;
if (!NoLossCalc) {
err[rgb] = 0.0f;
}
}
// Sigma Activation
for (int s = 0; s < nSample; s++) {
NeRF_SigmaExpActivation(sigma[s]); // exp activation
}
// Calc Output
for (int s = 0; s < nSample; s++) {
if (s != nSample - 1) {
Vec3 ToNextSample =
{
pos[3 * (s + 1) + 0] - pos[3 * (s + 0) + 0],
pos[3 * (s + 1) + 1] - pos[3 * (s + 0) + 1],
pos[3 * (s + 1) + 2] - pos[3 * (s + 0) + 2],
};
dt = ToNextSample.length();
}
else {
dt = 0.0f; // 最後のサンプルの厚みを0とする NOTE: これが正しいのかは分からない
}
// Radiance_i = T_i * c_i * (1 - exp(-sigma_i*dt_i))
for (int rgb = 0; rgb < 3; rgb++) {
out[rgb] += T * c[3*s + rgb] * (1.0f - expf(-sigma[s] * dt));
}
T *= expf(-sigma[s] * dt);
}
if (nSample == 0) {
for (int rgb = 0; rgb < 3; rgb++) {
out[rgb] = 0.0f;
err[rgb] = 0.5f * (out[rgb] - (float)ans[rgb]) * (out[rgb] - (float)ans[rgb]);
}
return;
}
if (NoLossCalc) {
return;
}
float dLdout[3] = {0.0f,0.0f,0.0f};
// Calc Loss
for (int rgb = 0; rgb < 3; rgb++) {
Calc_HuborLoss(out[rgb], (float)ans[rgb], dLdout[rgb]);
}
// MSE
for (int rgb = 0; rgb < 3; rgb++) {
err[rgb] = 0.5f * (out[rgb] - (float)ans[rgb]) * (out[rgb] - (float)ans[rgb]);
}
// Calc dLdColor and dLdDensity
float dLdc[3] = {0,0,0};
float dLdsigma = 0.0f;
float suffix[3] = {out[0], out[1], out[2]};
T = 1.0f;
for (int s = 0; s < nSample; s++) {
if (s != nSample - 1) {
Vec3 ToNextSample =
{
pos[3 * (s + 1) + 0] - pos[3 * (s + 0) + 0],
pos[3 * (s + 1) + 1] - pos[3 * (s + 0) + 1],
pos[3 * (s + 1) + 2] - pos[3 * (s + 0) + 2],
};
dt = ToNextSample.length();
}
else {
dt = 0.0f; // 最後のサンプルの厚みを0とする NOTE: これが正しいのかは分からない
}
// dLdc_i = T_i * (1 - exp(-sigma_i*dt_i)) * dLdout
for (int rgb = 0; rgb < 3; rgb++) {
dLdc[rgb] = T * (1.0f - expf(-sigma[s] * dt)) * dLdout[rgb];
}
// Suffix_i := C_hat - C_1 - C_2 - ... - C_i
for (int rgb = 0; rgb < 3; rgb++) {
suffix[rgb] -= T * c[3 * s + rgb] * (1.0f - expf(-sigma[s] * dt));
}
// T_{i+1}
T *= expf(-sigma[s] * dt);
// dLdsigma = dt_i * (T_{i+1} * c_i - Suffix_i) * dLdout
for (int rgb = 0; rgb < 3; rgb++) {
dLdsigma += dt * (T * c[3*s + rgb] - suffix[rgb]) * dLdout[rgb];
}
// 保存
for (int rgb = 0; rgb < 3; rgb++) {
c[3*s + rgb] = dLdc[rgb];
dLdc[rgb] = 0.0f;
}
NeRF_SigmaExpDerivative(dLdsigma, sigma[s]);
sigma[s] = dLdsigma;
dLdsigma = 0.0f;
}
}
部分的にみましょう．
__global__ void CalculateOutputAndLoss
(
const uint32_t nPixelBatch, NeRFInfo* SamplingInfo,
float* Pos,
NeRFRay* Rays, float* Density, float* Color,
float* Output, PixelInfo* Target, float* Loss, bool NoLossCalc = false
)
{
...
nPixelBatch: 使用したレイの本数（ピクセルの数）
SamplingInfo: サンプリング処理で保存した，そのレイに関わるサンプル点情報が「どこから」「何個」あるかを使用します
Pos: サンプル点間の距離がボリュームレンダリングに必要です
Rays, Density, Color: 名前の通りです．ただし，入力時は順方向のデータが入っているDensity, Colorには最終的に勾配のデータを格納して返します．
Output: ボリュームレンダリングの計算結果のRGB色を格納します
Target: PixelInfo構造体には対応する教師データ（RGB色）が格納されています
Loss: 誤差情報を記録する配列です
NoLossCalc : 推論処理では誤差を計算する必要がありません．それを制御するフラグです．
// 本スレッドでのポインタ等 NeRFRay* Ray = &Rays[threadID]; const uint32_t IdxBias = Ray->SampleBeginIdx; const uint32_t nSample = Ray->nSample; float* pos = Pos + 3 * IdxBias; float* c = Color + 3 * IdxBias; float* sigma = Density + IdxBias; float* out = Output + 3 * threadID; Vec3h ans = Target[threadID].Color; float* err = Loss + 3 * threadID;
配列におけるデータの内，処理中のスレッドで扱うレイに関するデータのポインタをあらかじめ計算しておくことで実装ミスが減ります．
// サンプル点の厚み（次のサンプル点までの距離） float dt; // Throughput float T = 1.0f;
コメントの通りです．
for (int rgb = 0; rgb < 3; rgb++) {
out[rgb] = 0.0f;
if (!NoLossCalc) {
err[rgb] = 0.0f;
}
}
0初期化しておきましょう．推論でない場合は誤差も0初期化します．
// Sigma Activation
for (int s = 0; s < nSample; s++) {
NeRF_SigmaExpActivation(sigma[s]); // exp activation
}
さて，NNの実装時には特殊であるため，レンダリング時にexp activationをするとしました．
MFFM_DEVICE void NeRF_SigmaExpActivation(float &sigma) {
sigma = expf(clamp(sigma, -15.0f, 15.0f));
}
単純にexpを取っているだけです．なお，計算時の爆発を避けるため，指数部分を15で抑えています．
// Calc Output
for (int s = 0; s < nSample; s++) {
if (s != nSample - 1) {
Vec3 ToNextSample =
{
pos[3 * (s + 1) + 0] - pos[3 * (s + 0) + 0],
pos[3 * (s + 1) + 1] - pos[3 * (s + 0) + 1],
pos[3 * (s + 1) + 2] - pos[3 * (s + 0) + 2],
};
dt = ToNextSample.length();
}
else {
dt = 0.0f; // 最後のサンプルの厚みを0とする NOTE: これが正しいのかは分からない
}
// Radiance_i = T_i * c_i * (1 - exp(-sigma_i*dt_i))
for (int rgb = 0; rgb < 3; rgb++) {
out[rgb] += T * c[3*s + rgb] * (1.0f - expf(-sigma[s] * dt));
}
T *= expf(-sigma[s] * dt);
}
ボリュームレンダリングの順方向です．
まずはサンプル点間の距離（先ほど示した図中の）を計算しますが，最後のサンプル点については距離を定義できません．今回は最後のサンプル点については
としました．
そして，図に示した式に代入してあげます．RGB独立に代入してあげていいです．
if (nSample == 0) {
for (int rgb = 0; rgb < 3; rgb++) {
out[rgb] = 0.0f;
err[rgb] = 0.5f * (out[rgb] - (float)ans[rgb]) * (out[rgb] - (float)ans[rgb]);
}
return;
}
if (NoLossCalc) {
return;
}
サンプル数が0の場合（そもそもAABBと交差しなかった，もしくはサンプリングをキャンセルした場合）には結果を{0,0,0}として返します．誤差計算はあんまり妥当ではないです．
推論処理である場合は今後の処理は不要なので返します．
float dLdout[3] = {0.0f,0.0f,0.0f};
// Calc Loss
for (int rgb = 0; rgb < 3; rgb++) {
Calc_HuborLoss(out[rgb], (float)ans[rgb], dLdout[rgb]);
}
// MSE
for (int rgb = 0; rgb < 3; rgb++) {
err[rgb] = 0.5f * (out[rgb] - (float)ans[rgb]) * (out[rgb] - (float)ans[rgb]);
}
ボリュームレンダリングの結果と教師画像のデータの誤差を計算します．誤差関数はHubor_Lossで固定してます．
// Calc dLdColor and dLdDensity
float dLdc[3] = {0,0,0};
float dLdsigma = 0.0f;
float suffix[3] = {out[0], out[1], out[2]};
T = 1.0f;
for (int s = 0; s < nSample; s++) {
if (s != nSample - 1) {
Vec3 ToNextSample =
{
pos[3 * (s + 1) + 0] - pos[3 * (s + 0) + 0],
pos[3 * (s + 1) + 1] - pos[3 * (s + 0) + 1],
pos[3 * (s + 1) + 2] - pos[3 * (s + 0) + 2],
};
dt = ToNextSample.length();
}
else {
dt = 0.0f; // 最後のサンプルの厚みを0とする NOTE: これが正しいのかは分からない
}
// dLdc_i = T_i * (1 - exp(-sigma_i*dt_i)) * dLdout
for (int rgb = 0; rgb < 3; rgb++) {
dLdc[rgb] = T * (1.0f - expf(-sigma[s] * dt)) * dLdout[rgb];
}
// Suffix_i := Radiance_hat - Radiance_1 - Radiance_2 - ... - Radiance_i
for (int rgb = 0; rgb < 3; rgb++) {
suffix[rgb] -= T * c[3 * s + rgb] * (1.0f - expf(-sigma[s] * dt));
}
// T_{i+1}
T *= expf(-sigma[s] * dt);
// dLdsigma = dt_i * (T_{i+1} * c_i - Suffix_i) * dLdout
for (int rgb = 0; rgb < 3; rgb++) {
dLdsigma += dt * (T * c[3*s + rgb] - suffix[rgb]) * dLdout[rgb];
}
// 保存
for (int rgb = 0; rgb < 3; rgb++) {
c[3*s + rgb] = dLdc[rgb];
dLdc[rgb] = 0.0f;
}
NeRF_SigmaExpDerivative(dLdsigma, sigma[s]);
sigma[s] = dLdsigma;
dLdsigma = 0.0f;
}
順方向と同じようにして，先程示した計算式に代入しているだけです．式中のは実装中のSuffix_iのことです．NeRF_SigmaExpDerivativeはexp Activationの逆方向です．
MFFM_DEVICE void NeRF_SigmaExpDerivative(float& dLdsigma, float sigma) {
dLdsigma = dLdsigma * sigma;
}
これでボリュームレンダリングの逆方向が終わりました．dLdcとdLdsigmaは処理が終わるごとに0に戻してあげてください．
NeRFの実装: (8): NNを誤差逆伝播
さて，ボリュームレンダリングの逆方向の処理を行い，NNには「色」と「密度」に関する勾配情報が入ってきました．ということで，この勾配情報を元に誤差逆伝播を行います．InstantNGPの論文にて推されているNNの構造を再掲します．
順方向時にはDensity MLP -> Color MLPの順番で処理をしたので逆伝播時には逆に処理します．
template <const uint32_t indim_den, const uint32_t hiddendim_den, const uint32_t outdim_den, const uint32_t nHiddenLayer_den,
const uint32_t indim_col, const uint32_t hiddendim_col, const uint32_t outdim_col, const uint32_t nHiddenLayer_col>
__global__ void MFFM_NeRF_Backward
(
const int epoch,
NeRFInfo* SamplingInfo, EncoderInfo* EncInfo_den, EncoderInfo* EncInfo_col,
Optimize Optim,
Encoder Encoder_den, Activation ActHid_den, Activation ActOut_den,
Encoder Encoder_col, Activation ActHid_col, Activation ActOut_col,
float* Pos, float* Dir,
__half* weights_den, __half* buffers_den,
__half* weights_col, __half* buffers_col,
float* dLdDensity, float* dLdColor,
float* LossDerivativeSumALL_den, float* AdditionalParam_den,
float* LossDerivativeSumALL_col, float* AdditionalParam_col,
float* dLdInput_den = nullptr
)
{
const int BatchSize = SamplingInfo->nSamples;
// 即値
constexpr uint32_t indim_den_aligned = next_multiple(indim_den, TENSOR_ROW);
constexpr uint32_t hiddendim_den_aligned = next_multiple(hiddendim_den, TENSOR_ROW);
constexpr uint32_t outdim_den_aligned = next_multiple(outdim_den, TENSOR_ROW);
constexpr uint32_t indim_col_aligned = next_multiple(indim_col, TENSOR_ROW);
constexpr uint32_t hiddendim_col_aligned = next_multiple(hiddendim_col, TENSOR_ROW);
constexpr uint32_t outdim_col_aligned = next_multiple(outdim_col, TENSOR_ROW);
const uint32_t maxdim_col = max(indim_col_aligned, max(hiddendim_col_aligned, outdim_col_aligned));
extern __shared__ __half shmem[];
__half* intermediate_den = shmem;
__half* intermediate_col = shmem;
__half* intermediate_InDir = shmem + (outdim_den_aligned + SKEW) * ONEBATCH_SIZE;
__half* LossDerivativeSumOfBlock = shmem + (maxdim_col + SKEW) * ONEBATCH_SIZE;
////////////////////////////////////////////// Color MLP ////////////////////////////////////////////////////////////////////////
const int newDim = outdim_col_aligned + SKEW;
load_input(outdim_col, newDim, dLdColor, intermediate_col);
Kernel_Debug_train_backward<indim_col, hiddendim_col, outdim_col, nHiddenLayer_col>
(
BatchSize, Optim, ActHid_col, ActOut_col, epoch,
intermediate_col, LossDerivativeSumOfBlock, weights_col, buffers_col, LossDerivativeSumALL_col, AdditionalParam_col, false
);
// Divide
DivideShmem<outdim_den, indim_col - outdim_den>(intermediate_den, intermediate_InDir, intermediate_col);
//
// 入力をロード
if (Encoder_col == Encoder::HashGrid || Encoder_den == Encoder::UniformGrid) {
load_input(3, 3, Dir, intermediate_InDir + (indim_col_aligned + SKEW) * ONEBATCH_SIZE);
}
BackPropEncoder<indim_col_aligned>(epoch, Encoder_col, *EncInfo_col, intermediate_InDir, SamplingInfo->AABB_pMin, SamplingInfo->AABB_pMax);
////////////////////////////////////////////// END: Color MLP ////////////////////////////////////////////////////////////////////
// この時点で破壊が許容される領域:
// ・LossDerivativeSumOfBlock
// ・intermediate_InDir
// ・
//
// dLdSigmaの誤差伝搬
const int threadID_BlockScope = 32 * threadIdx.y + threadIdx.x;
const int threadID_global = ONEBATCH_SIZE * blockIdx.x + threadID_BlockScope;
if (32 * threadIdx.y + threadIdx.x < ONEBATCH_SIZE) {
intermediate_den[(outdim_den_aligned + SKEW) * threadID_BlockScope] = intermediate_den[(outdim_den_aligned + SKEW) * threadID_BlockScope] + (__half)dLdDensity[threadID_global];
}
__syncthreads();
////////////////////////////////////////////// Density MLP ///////////////////////////////////////////////////////////////////////
Kernel_Debug_train_backward<indim_den, hiddendim_den, outdim_den, nHiddenLayer_den>
(
BatchSize, Optim, ActHid_den, ActOut_den, epoch,
intermediate_den, LossDerivativeSumOfBlock, weights_den, buffers_den, LossDerivativeSumALL_den, AdditionalParam_den, false
);
// 入力をロード
if (Encoder_den == Encoder::HashGrid || Encoder_den == Encoder::UniformGrid) {
load_input(3, 3, Pos, intermediate_den + (indim_den_aligned + SKEW) * ONEBATCH_SIZE);
}
BackPropEncoder<indim_den_aligned>(epoch, Encoder_den, *EncInfo_den, intermediate_den, SamplingInfo->AABB_pMin, SamplingInfo->AABB_pMax);
////////////////////////////////////////////// END: Density MLP //////////////////////////////////////////////////////////////////
// Density MLPの入力側の誤差を保存する
if (dLdInput_den != nullptr) {
store_intermediate<float>(indim_den_aligned + SKEW, indim_den, intermediate_den, dLdInput_den);
}
}
やるだけと言ってしまえばそうなのですが，軽く説明しておきます．
template <const uint32_t indim_den, const uint32_t hiddendim_den, const uint32_t outdim_den, const uint32_t nHiddenLayer_den,
const uint32_t indim_col, const uint32_t hiddendim_col, const uint32_t outdim_col, const uint32_t nHiddenLayer_col>
順伝播時と同じです．（"in", "out"は順伝播時の定義と同じです）
__global__ void MFFM_NeRF_Backward
(
const int epoch,
NeRFInfo* SamplingInfo, EncoderInfo* EncInfo_den, EncoderInfo* EncInfo_col,
Optimize Optim,
Encoder Encoder_den, Activation ActHid_den, Activation ActOut_den,
Encoder Encoder_col, Activation ActHid_col, Activation ActOut_col,
float* Pos, float* Dir,
__half* weights_den, __half* buffers_den,
__half* weights_col, __half* buffers_col,
float* dLdDensity, float* dLdColor,
float* LossDerivativeSumALL_den, float* AdditionalParam_den,
float* LossDerivativeSumALL_col, float* AdditionalParam_col,
float* dLdInput_den = nullptr
)
epoch: イテレーション回数
SamplingInfo: 順伝播時の説明参照
EncInfo: 無視してください
Optimize: 最適化関数の設定です．
enum class Optimize {
GD,
Adam
};
LossDerivativeSumALL, AdditionalParamは詳しい話はPart1, Part2を参照してください（パラメーターの勾配とAdamのmとvです）
dLdInput_den: 無視してください（入力側への勾配伝搬です）
後は順伝播の逆向きをやるだけです．基本的に実装を読めばそのまんまの処理をしていることが分かるはずです．
...
////////////////////////////////////////////// Color MLP ////////////////////////////////////////////////////////////////////////
const int newDim = outdim_col_aligned + SKEW;
load_input(outdim_col, newDim, dLdColor, intermediate_col);
Kernel_Debug_train_backward<indim_col, hiddendim_col, outdim_col, nHiddenLayer_col>
(
BatchSize, Optim, ActHid_col, ActOut_col, epoch,
intermediate_col, LossDerivativeSumOfBlock, weights_col, buffers_col, LossDerivativeSumALL_col, AdditionalParam_col, false
);
// Divide
DivideShmem<outdim_den, indim_col - outdim_den>(intermediate_den, intermediate_InDir, intermediate_col);
...
「色」の勾配をshared memoryにロードし(load_input())，MLPの誤差逆伝播をし（Kernel_Debug_train_backward()），そして順伝播時にはColor MLPの入力を「Density MLPの出力」に「エンコードしたレイの方向ベクトル」を結合したものなので，ここでばらしてあげます
/*
* 1: shmem1
* 2: shmem2
* s: SKEW
* 12s12s12s12s12s12s....12s12s12s12s12s -> 1s1s1s1s1s...1s1s1s1s2s2s2s2s2s2s...2s2s2s2s
* 番地的にはshmem = shmem1
*/
template <const uint32_t dim1, const uint32_t dim2>
MFFM_DEVICE void DivideShmem(__half* shmem1, __half* shmem2, __half* shmem) {
// 即値
constexpr uint32_t dim1_aligned = next_multiple(dim1, TENSOR_ROW);
constexpr uint32_t dim2_aligned = next_multiple(dim2, TENSOR_ROW);
constexpr uint32_t nowDim = dim1 + dim2;
constexpr uint32_t nowDim_aligned = next_multiple(nowDim, TENSOR_ROW);
const int bx = blockIdx.x;
const int tx = threadIdx.x;
const int ty = threadIdx.y;
const int warp_index_required = ONEBATCH_SIZE / 32;
if (ty >= warp_index_required) {
return;
}
float tmp[dim1 + dim2];
__syncthreads();
for (int i = 0; i < dim1 + dim2; i++) {
tmp[i] = shmem[(nowDim_aligned+SKEW) * 32 * ty + (nowDim_aligned + SKEW) * tx + i];
}
__syncthreads();
for (int i = 0; i < dim1; i++) {
shmem1[(dim1_aligned + SKEW) * 32 * ty + (dim1_aligned + SKEW) * tx + i] = (__half)tmp[i];
}
for (int i = 0; i < dim2; i++) {
shmem2[(dim2_aligned + SKEW) * 32 * ty + (dim2_aligned + SKEW) * tx + i] = (__half)tmp[i + dim1];
}
for (int i = 0; i < SKEW; i++) {
shmem1[(dim1_aligned + SKEW) * 32 * ty + (dim1_aligned + SKEW) * tx + dim1 + i] = (__half)0.0f;
shmem2[(dim2_aligned + SKEW) * 32 * ty + (dim2_aligned + SKEW) * tx + dim2 + i] = (__half)0.0f;
}
__syncthreads();
}
Concatの逆の処理をしているだけなので説明は省略します．さて，今回は方向ベクトルのエンコーダーに誤差逆伝播は不要ですが，仮に必要な場合があったとして，その場合は処理を行う必要があります．そこで
...
// 入力をロード
if (Encoder_col == Encoder::HashGrid || Encoder_den == Encoder::UniformGrid) {
load_input(3, 3, Dir, intermediate_InDir + (indim_col_aligned + SKEW) * ONEBATCH_SIZE);
}
BackPropEncoder<indim_col_aligned>(epoch, Encoder_col, *EncInfo_col, intermediate_InDir, SamplingInfo->AABB_pMin, SamplingInfo->AABB_pMax);
...
例えばMultiresolution Hash Encodingの誤差逆伝播時には入力データを必要とするので，ロードしてあげます．BackPropEncoder()においてエンコーダーの誤差逆伝播を行います．
template <const uint32_t indim_aligned>
MFFM_DEVICE void BackPropEncoder(const int epoch, Encoder encoder, EncoderInfo& Info, __half* intermediate, const float3 InputRangeMin = { 0,0,0 }, const float3 InputRangeMax = { 1,1,1 },
float* MHEBufferSTU = nullptr, unsigned int* MHEBufferIdxHT = nullptr) {
const int bx = blockIdx.x;
const int tx = threadIdx.x;
const int ty = threadIdx.y;
switch (encoder) {
case Encoder::None:
// NOP
break;
case Encoder::Frequency:
// NOP
break;
case Encoder::SH:
// NOP
break;
case Encoder::HashGrid:
MHE::Propagate_backward(InputRangeMin, InputRangeMax, intermediate, intermediate + (indim_aligned + SKEW) * ONEBATCH_SIZE, intermediate + (indim_aligned + SKEW) * ONEBATCH_SIZE + 3 * ONEBATCH_SIZE);
//MHE2::Propagate_backward_inside_kernel<MHE_F>(*Info.MHEInfo, InputRangeMin, InputRangeMax, intermediate, intermediate + (indim_aligned + SKEW) * ONEBATCH_SIZE, intermediate + (indim_aligned + SKEW) * ONEBATCH_SIZE + 3 * ONEBATCH_SIZE);
break;
case Encoder::UniformGrid:
UGE::BackPropagate(*Info.UGEInfo, intermediate + (indim_aligned + SKEW) * ONEBATCH_SIZE, intermediate);
break;
default:
printf("Invalid Encoding type\n");
break;
}
__syncthreads();
}
Multiresolution Hash EncodingのPropagate_backward()はPart2に示しました．（詳しくはPart1,2を参照してください）あとはこれの組み合わせです．これでNNの誤差逆伝播が終わり，あとはパラメーターを最適化するだけです．
NeRFの実装: (9): NN最適化
説明は省略します．Part1, Part2を参照してください．
// 最適化のみ
template <const uint32_t indim, const uint32_t hiddendim, const uint32_t outdim, const uint32_t nHiddenLayer>
__global__ void MFFM_Optimize(bool willOptimizeInsideOneKernel, const uint32_t BatchSize, EncoderInfo* EncInfo, Encoder Encoder, Optimize Optim, const int epoch, __half* weights, float* LossDerivative, float* AdditionalParam) {
Optimize_MLP<indim, hiddendim, outdim, nHiddenLayer>(willOptimizeInsideOneKernel, BatchSize, Optim, epoch, weights, LossDerivative, AdditionalParam);
OptimizeEncoder(Encoder, *EncInfo, Optim, BatchSize, epoch);
}
MFFM_DEVICE void OptimizeEncoder(Encoder encoder, EncoderInfo &Info, Optimize Optim, const uint32_t BatchSize, const int epoch) {
switch (encoder) {
case Encoder::None:
// NOP
break;
case Encoder::Frequency:
// NOP
break;
case Encoder::SH:
// NOP
break;
case Encoder::HashGrid:
MHE::Optimization(BatchSize, Optim, epoch);
//MHE2::Optimization(Info.MHEInfo, BatchSize, Optim, epoch);
break;
case Encoder::UniformGrid:
UGE::Optimization(epoch, BatchSize, Optim, *Info.UGEInfo);
break;
default:
printf("Invalid Encoding type\n");
break;
}
__syncthreads();
}
NeRFの実装: (10): Occupancy Gridの最適化
さて，サンプリング処理にも出てきましたが，まずそもそもOccupancy Gridとはなんなんだっていう話です．InstantNGPの論文のAppendix-E.2に書かれているのですが，簡単に言うと「空間に「物体」が存在するかどうかの情報を保持するボクセル」です．2種類のOccupancy Gridがあり，一つは「「物体」が存在するか」という01の情報で，もう一つは「どれくらいの「密度」であるか」というfloatの情報です．前者をBit_OccupancyGridとし，後者をFloat_OccupancyGridとします．どういう設計かの概念を図にしておきます．
グリッド内部のあるランダムに選んだサンプル点における「密度」を計算し，それをFloat_OccupancyGridに反映させた結果，左上のようなFloat_OccupancyGridが得られて，そしてそれを元にして何かしらの判定を行った結果その右に示すBit_OccupancyGridが得られたとします．そして，サンプリング処理で以下の処理がありましたね
...
// Occupancy Gridにおいて「物体が無い」と判断された場合はスキップする
if (epoch > 16 && !Grids[0].is_Occupied(s_Pos)) {
index++;
continue;
}
...
レイ上の点をサンプリングした際にそのサンプル点が果たしてどれだけ不透明であるかをOccupancyGridを用いて判定します．これによって弾かれたサンプル点が図中の青い点で，弾かれずに採用されたサンプル点が橙の点で表されています．このようにすることで，「有効なサンプル点」を増やすことが出来てサンプル効率の向上につながる，というのが簡単なお気持ちです．何もないところサンプリングしても何も得られませんので（屈折率とかの場を作ってそれに応じて屈折させてみるのも面白そうですけどね）．
さて，じゃあOccupancy Gridをどうやって構築してあげたらいいんだって話です．InstantNGPの論文が言うことには
16イテレーションごとにOccupancy Gridを更新する
(1): Float_OccupancyGridの各値を0.95倍する
(2): M個のグリッドを選んで．max(現在の値，グリッド内部のサンプル点をNNに入れた結果得られる「密度」の値)へと更新する
(3): Float_OccupancyGridの各値がより小さければ対応するBit_OccupancyGridを0，そうでければ1にする
なお，(2)のMは最初の256イテレーションは全てのグリッド数だけ一様サンプリング，それ以降は全てのグリッド数の半分を，半分は全てのグリッドから一様サンプリングし，残り半分はBit_OccupancyGridの中身が1になっているグリッドから一様サンプリングする．
ということらしいです．ですが実装コストの面で少しだけ簡略化します．(2)のサンプリングを，サンプル数は同じくして，ただしイテレーション回数に関わらず全てのグリッドから一様サンプリングします（なお，棄却法を使えば論文通りの実装が出来ます）．
また，(3)の通りに実装すると学習が不安定だったので，今回の実装では別の判定方法を取りました．
では，実装を見ていきましょう．
Occupancy Gridの構造体
実装を容易にするためにOccupancy Gridの構造体を作っておきます．
struct OccupancyGrid {
int GridID;
bool* BitGrid;
float* FloatGrid;
Vec3 pMin;
Vec3 pMax;
float ave;
MFFM_HOST_DEVICE OccupancyGrid(const int GridID = 0, const Vec3 pMin = { 0,0,0 }, const Vec3 pMax = { 1,1,1 }) : GridID(GridID), pMin(pMin), pMax(pMax), ave(0.0f) {}
MFFM_DEVICE unsigned int PosToIdx(Vec3& pos) {
return CalcMortonCode7bit(pos, pMin, pMax);
}
MFFM_DEVICE bool& Bit_at(Vec3& pos) {
return BitGrid[CalcMortonCode7bit(pos, pMin, pMax)];
}
MFFM_DEVICE float& Float_at(Vec3& pos) {
return FloatGrid[CalcMortonCode7bit(pos, pMin, pMax)];
}
MFFM_DEVICE bool is_Occupied(Vec3& pos) {
return BitGrid[CalcMortonCode7bit(pos, pMin, pMax)];
}
};
GridID: 論文中では複数のOccupancyGridを作ることを考慮しており，今回の実装でも番号だけ振っておきます．なお今回は1つしか使用しません．
BitGrid, FloatGrid: Bit_OccypancyGrid と Float_OccupancyGridです．
pMin, pMax: Gridの拡がる空間を表します（AABBです）．
ave: FloatGridの各グリッドの中身の値の平均値です．先ほど示した(3)の判定に使用します．
初期化は次のように行います．
__global__ void SetUpOccupancyGrid(int ID, NeRFInfo* Info, OccupancyGrid* Grids, bool* BitGridPtr, float* FloatGridPtr) {
Grids[ID].GridID = ID + 1;
float power = powf(2, ID);
Grids[ID].BitGrid = BitGridPtr;
Grids[ID].FloatGrid = FloatGridPtr;
Grids[ID].pMin = { Info->AABB_pMin.x * power, Info->AABB_pMin.y * power,Info->AABB_pMin.z * power };
Grids[ID].pMax = { Info->AABB_pMax.x * power, Info->AABB_pMax.y * power,Info->AABB_pMax.z * power };
}
BitGridPtr, FloatGridPtrはカーネル外部でthrust::device_vectorで定義している配列のポインタです．
pMinとpMaxは論文通りに実装しましたが，今回は結局1つめのOccupancyGridのみ使用するのでGridの領域は[-1,1]^3となります．
さて，Occupancy Gridの中身は単純な3次元配列を1次元に直したような配列では持たずに，MortonCode（Z-order curve）に載せます．
今回は一辺が128要素のOccupancyGridを使用するので，各方向7bitで処理すればよいです．MortonCodeを計算する関数を置いておきます．
// 1111 -> 1001001001
KGYK_HOST_DEVICE unsigned int ExpandBits(unsigned int v) {
v = (v * 0x00010001u) & 0xFF0000FFu;
v = (v * 0x00000101u) & 0x0F00F00Fu;
v = (v * 0x00000011u) & 0xC30C30C3u;
v = (v * 0x00000005u) & 0x49249249u;
return v;
}
これは元のビット列に2マス隙間を開ける関数です．そして，これを用いて次のようにすることで各方向128分割したグリッドのインデックスを32bitで表現できます．
MFFM_DEVICE unsigned int CalcMortonCode7bit(Vec3 position, Vec3 pMin, Vec3 pMax) {
unsigned long long int MortonCode;
// positionを[0, 1]^3の範囲に圧縮する
float x = normalize(position.x, pMin.x, pMax.x, 0.0f, 1.0f);
float y = normalize(position.y, pMin.y, pMax.y, 0.0f, 1.0f);
float z = normalize(position.z, pMin.z, pMax.z, 0.0f, 1.0f);
x = min(max(x * 128, 0.0f), 127.0f);
y = min(max(y * 128, 0.0f), 127.0f);
z = min(max(z * 128, 0.0f), 127.0f);
unsigned int xx = ExpandBits((unsigned int)x);
unsigned int yy = ExpandBits((unsigned int)y);
unsigned int zz = ExpandBits((unsigned int)z);
MortonCode = xx * 4 + yy * 2 + zz;
return MortonCode;
}
この関数自体は以前にLinear-BVHというデータ構造を実装した際に参考にした次の記事に載っていたものを元にしています．
このようにしてインデックスを定義してあげることで，グリッド内で近いところにある2ブロックはメモリ上で近いところとなる，というわけです．さて，OccupancyGridの要素へのアクセスは，（座標）をMortonCodeに変換し，そのMortonCodeをインデックスとして指定するだけです．それらの処理がPosToIdx(), Bit_at(), Float_at()となっています．そして先に書いておくと，is_Occupied()はBitGridの要素が1であるかを判定しているだけです（図に書いた通りの処理です）．
では更新処理を書いてあげましょう．OccupancyGridの初期値は全て0です．
(1): すべて0.95培
// 全てのブロックの値を0.95倍する
// 全てのブロック数スレッドを立てる
__global__ void NeRF_DecayOccupancyGrid(const int K, OccupancyGrid* Grids) {
uint32_t threadID = blockIdx.x * blockDim.x + threadIdx.x;
if (threadID >= K * 128 * 128 * 128) {
return;
}
const uint32_t GridID = threadID / (128 * 128 * 128);
const uint32_t BlockID = threadID % (128 * 128 * 128);
Grids[GridID].FloatGrid[BlockID] *= 0.95f;
if (BlockID == 0) {
Grids[GridID].ave *= 0.95f;
}
}
128 * 128 * 128 * Kだけスレッドを立てて処理します．(K: OccupancyGridの数．今回は1個としています)
単純に0.95掛けて行っているだけです．また，平均値aveにも0.95を掛けておきます（各Grid内のブロック1つ目を担当するスレッドが行います）
(2): M個サンプリングし，色々やる
この処理は中々めんどくさいです．ランダムに選ばれたM個のブロックにスレッドを割り振り，各スレッドが自分の担当するOccupancyGridにおけるあるブロック内部にあるランダムな点をサンプリングし，NNモデルを通して「密度」の値を得るのですが，M個のサンプリングをまずは行う必要があります．これは簡単のため一様に選ばせます．つまり1~128 * 128 * 128 * Kの順列をシャッフルして最初のM個を選択します．
...
// (2)
// M個のブロックを選ぶ（ここでは前半と後半で分けずに全て一様サンプリングする）
const uint32_t M = (epoch <= 256) ? nBlockALL : nBlockALL / 2;
...
static thrust::device_vector<int> Permutaion(nBlockALL);
if (epoch == 16) {
thrust::sequence(Permutaion.begin(), Permutaion.end());
}
if (epoch > 256) {
thrust::default_random_engine g;
g.seed(epoch + 1);
thrust::shuffle(Permutaion.begin(), Permutaion.end(), g);
}
NeRF_SamplePositionInsideGrid <<<M / 512, 512, 0, Stream >>> (epoch, K, thrust::raw_pointer_cast(&Grids[0]), thrust::raw_pointer_cast(&Permutaion[0]), thrust::raw_pointer_cast(&Position[0]));
gpuErrchk(cudaGetLastError());
gpuErrchk(cudaStreamSynchronize(Stream));
...
Mの決定については論文に従います．また，nBlockALL = 128 * 128 * 128 * Kです．Permutation配列をシャッフルして最初のM個を選択することによりインデックスのサンプリングが完了します．なお，最初の256イテレーションは結局全部サンプリングするのでシャッフルしません
Occupancy Gridのブロックのサンプリングが終わったら今度はブロック内部の点をサンプリングします．ここはCUDAカーネルでM個並列に行います．これは単純に自分の担当しているOccupancyGridにおけるブロックが，NeRFを生成するAABB内部ではどの領域に対応しているのかを計算し，xyz各軸方向に対して乱数を用いて座標を選ぶというやり方で行きます．
// M個グリッドをサンプリングする
// M個のスレッドを立てる
// Indexにはサンプリングするグリッド上の座標が1次元のインデックスで記録されている
// よってグリッド上では(x,y,z)をIndexから計算する
// Indexに入っているインデックスを使用しても実際に(x,y,z)にはアクセスできないことに注意（MortonCodeでOccuoancyGridのインデックスが定義されているためである）
// そのため，サンプル点から得たdensityを元に更新するOccupancyGridのインデックスはサンプル点の座標から計算する
__global__ void NeRF_SamplePositionInsideGrid(const int seed, const int K, OccupancyGrid* Grids, const int* Index, float* Position) {
uint32_t threadID = blockIdx.x * blockDim.x + threadIdx.x;
if (threadID >= K * 128 * 128 * 128) {
return;
}
const int Idx = Index[threadID];
const uint32_t GridID = Idx / (128 * 128 * 128);
const uint32_t BlockID = Idx % (128 * 128 * 128);
// サンプルするグリッドの最小座標
Vec3 pMinGrid = Grids[GridID].pMin;
Vec3 pMaxGrid = Grids[GridID].pMax;
// 今回サンプルする領域
const uint32_t IdxX = BlockID % 128;
const uint32_t IdxY = (BlockID / 128) % 128;
const uint32_t IdxZ = (BlockID / 128 / 128);
// サンプルするブロックの一辺の長さ（Assumption: ブロックは立方体）
const float Range = (pMaxGrid.x - pMinGrid.x) / 128;
// サンプルするブロックの最小座標
const float Xmin = pMinGrid.x + Range * IdxX;
const float Ymin = pMinGrid.y + Range * IdxY;
const float Zmin = pMinGrid.z + Range * IdxZ;
// RNG
curandState state;
curand_init(seed, threadID, 0, &state);
// (0,1]
float t = curand_uniform(&state);
// サンプルする座標
const float s_PosX = Xmin + t * Range;
const float s_PosY = Ymin + t * Range;
const float s_PosZ = Zmin + t * Range;
// MLP入力のthreadID番目にサンプル点の情報を登録
Position[3 * threadID + 0] = s_PosX;
Position[3 * threadID + 1] = s_PosY;
Position[3 * threadID + 2] = s_PosZ;
}
コメントに書いているとおりです．ちなみに本当はこの実装はあまりよろしくなくて，使用するOccupancyGridの拡がっている領域が[-1, 1]^3であり，これはNeRFのAABBと（偶然！）一致しているため，特に正規化を行わずして直接サンプル点の座標の「密度」を評価できます．真似はしない方が良いです．
サンプル点の座標が得られたため，「密度」のみを推定するNNに入れておきます．
template <const uint32_t indim_den, const uint32_t hiddendim_den, const uint32_t outdim_den, const uint32_t nHiddenLayer_den>
__global__ void MFFM_NeRF_InferenceDensityOnly
(
NeRFInfo* SamplingInfo, EncoderInfo* EncInfo_den,
Encoder Encoder_den, Activation ActHid_den, Activation ActOut_den,
float* Pos,
__half* weights_den,
float* Density)
{
// 即値
constexpr uint32_t indim_den_aligned = next_multiple(indim_den, TENSOR_ROW);
constexpr uint32_t hiddendim_den_aligned = next_multiple(hiddendim_den, TENSOR_ROW);
constexpr uint32_t outdim_den_aligned = next_multiple(outdim_den, TENSOR_ROW);
const uint32_t maxdim_den = max(indim_den_aligned, max(hiddendim_den_aligned, outdim_den_aligned));
extern __shared__ __half shmem[];
__half* intermediate_den = shmem;
//////////////////////////////////////////////////////// Density MLP //////////////////////////////////////////////////////////////////////////
// 入力のロード
// エンコーダー無しの場合はこの時点でSKEWを与える
// エンコーダーありの場合はSKEWを与えない（ただのコピー）
if (Encoder_den == Encoder::None) {
load_input(indim_den, indim_den_aligned + SKEW, Pos, intermediate_den);
}
else {
load_input(3, 3, Pos, intermediate_den);
}
// Encode
Encode<indim_den_aligned>(Encoder_den, *EncInfo_den, intermediate_den, true, SamplingInfo->AABB_pMin, SamplingInfo->AABB_pMax);
// Density MLP
Kernel_Debug_inference<indim_den, hiddendim_den, outdim_den, nHiddenLayer_den>(ActHid_den, ActOut_den, intermediate_den, weights_den);
// Density MLPの出力の1つめの要素をDensityとして保存
store_intermediate<float>(outdim_den_aligned + SKEW, 1, intermediate_den, Density);
//////////////////////////////////////////////////////// END: Density MLP //////////////////////////////////////////////////////////////////////
}
やっていることは通常のNNの順伝播の実装と同じです．また，バッチサイズは言うまでもありませんがMです．
これによって「密度」が推定されたのでこれを用いてFloat_OccupancyGridを更新します．
// サンプリングしたM個のブロックを更新する
// M個のスレッドを立てる
// Indexには単純な1次元化した配列のインデックスが格納されている: [x][y][z] -> [idx]
// それゆえ，OccupancyGridへのアクセスはサンプル点の座標から行う（MortonCodeでアクセス）
__global__ void NeRF_UpdateOccupancyGrid_FloatGrids(const int M, OccupancyGrid* Grids, const float* Density, const int* Index, const float* Pos) {
uint32_t threadID = blockIdx.x * blockDim.x + threadIdx.x;
if (threadID >= M) {
return;
}
Vec3 SampledPos = { Pos[3 * threadID], Pos[3 * threadID + 1], Pos[3 * threadID + 2] };
int Idx = Index[threadID];
const uint32_t GridID = Idx / (128 * 128 * 128);
const float old = Grids[GridID].Float_at(SampledPos);
Grids[GridID].Float_at(SampledPos) = max(old, Density[threadID]);
atomicAdd(&Grids[GridID].ave, (Grids[GridID].Float_at(SampledPos) - old) / (128 * 128 * 128));
}
サンプリング時に使用したM個のOccupancyGridのブロックを更新します．処理自体は先ほど示した通り，現在の値と推定値のmaxに更新する，です．今思ったのですが，Exp Activationし忘れてましたね．これが論文通りの(3)が上手く動かない原因かもしれません．あとで試します．ここで更新差分を元にして，aveも更新しておきます．注意点として，Indexではなくて座標（Pos）からMortonCodeを計算してOccupancyGridにアクセスしないとダメです（OccupancyGridの構築の仕方を思い出してください）（1敗）．以上でFloat_OccupancyGridの処理は終わりです．
(3): Bit_OccupancyGridの更新
これは簡単です．閾値を超えているか超えていないかで0と1にするだけです．
// 全てのグリッドを更新する
// 全てのグリッド数だけスレッドを立てる．
__global__ void NeRF_UpdateOccupancyGrid_BitGrids(const int K, OccupancyGrid* Grids) {
uint32_t threadID = blockIdx.x * blockDim.x + threadIdx.x;
if (threadID >= K * 128 * 128 * 128) {
return;
}
const uint32_t GridID = threadID / (128 * 128 * 128);
const uint32_t BlockID = threadID % (128 * 128 * 128);
//constexpr float thres = 0.1f * 1024.0f / 1.732050807f;
const float thres = 2.0f * Grids[GridID].ave;
if (Grids[GridID].FloatGrid[BlockID] > thres) {
Grids[GridID].BitGrid[BlockID] = true;
}
else {
Grids[GridID].BitGrid[BlockID] = false;
}
}
これはFloatGridの値がthres，ここではaveの2倍を超えているかどうかで閾値判定し，対応するBit_OccupancyGridの値を書き換えています．ちなみにこの2倍という値は特に意味はなく，学習が安定した値というだけです．ちなみにこの設定を間違えるとCUDAカーネルが実行時エラー出してプログラムが落ちます（例外処理してないので……）．以上でOccupancy Gridの処理は終わりです．
ちなみに
Occupancy Gridを実装しなくてもある程度絵が出ると思っていましたが意外とそういうことはなく，Occupancy Gridを実装しないと学習精度がかなり低い状態が続きます．なので頑張って最後まで実装しきってください．Occupancy Grid自体は強力な手法です．
NeRFの実行
本当にお疲れ様です．以上でNeRFの実装は終わりました．本当はフロントエンドの実装も見せるべきなのですが，かなり説明が大変な複雑化をしているので（主にGUIのせい）ここでは説明を省略します．自分で実装してみると途方に暮れることは少なくともなく，これまで説明したパーツに相当するものを用いてちゃんと設計できると思います．意外とバグりやすいのでちゃんと実装ノート等に纏めた方が良いです．
さて，NeRFの実装が終わったのでどういう風に学習してくれるかを確認しましょう．まずはお決まりのLegoの空間を近似します．学習データはNVIDIAの出しているInstant NeRFのデモソフトにくっついてきたものを使用します．
この子たちですね．jsonファイルで画像とカメラの姿勢等を指定してファイルに読み込ませます．まずは改めて実行の設定を確認しましょう（図中に示したColor MLPの構成と少し異なります）．
ボリュームレンダリングに関わる設定
NERF_MAX_SAMPLE_PER_RAY: 512
学習データ
画像: 100枚
サイズ: 800x800
NNの設定
DensityMLP:
InDim: 32
HiddenDim: 64
OutDim: 16
nHiddenLayer: 1
活性化関数: ReLU
エンコーダー: HashGrid
HashGrid: L = 16, F = 2, T = 220, Nmin = 2, b = 2.2
ColorMLP:
InDim: 32
HiddenDim: 64
OutDim: 3
nHiddenLayer: 1
活性化関数: 入力層と隠れ層: ReLU / 出力層: Sigmoid(Logistic)
エンコーダー: Spherical Harmnic ()
最適化関数: Adam
誤差関数; Hubor_Loss
学習率: 0.07
NeRFの設定
AABB: [-1,1]^3
OccupancyGridの数: 1個 / 領域: [-1,1]^3
さて，では結果を見ていきましょう．とりあえず10000イテレーション回してみて出力画像の変化を見ました．
最初は何も見えませんが，モヤモヤし始めて目的の状態に近似されていくのが分かります．本当は動画で載せたいのですが，容量の問題で載せられないので最後にTwitterにあげておいた動画へのリンクを纏めて貼っておきます．今回の実装ではGUI上で動かしており，純粋な学習のみの処理時間を計測することは現状では不可能なのですが，GUI上のカメラにほとんどNeRFを生成するAABBを映していない状態で，約27秒で1000イテレーション回ります．学習処理だけであればもっと速いはずです．
せっかくなので，レイ1本あたりのサンプル点数の平均値でもグラフにプロットしておきます．
最初は設定した通り512サンプルきっかり全部サンプリングしてくれています．そしてiter = 16とiter = 32において激しく減少していることが分かります．これは16イテレーションごとにOccupancy Gridを更新し，サンプリング中にサンプルが弾かれやすくなっているためです．もっと縦軸の範囲を狭く，横軸の範囲をのばして見てみましょう．
もちろんサンプリングに使用するバッチによって多少の変化はあるのですが，全体として減少する傾向にあることが分かります．
他のデータセットも試してみましょう．
これはマイクですね．原著NeRFのプロジェクトページで公開されているデータセットにBlenderで開ける3Dモデルファイルがあるのですが，それをBlender上で多少編集し，100視点ランダムに選んでレンダリングしました．この作業は次の記事を参考にしました．
qiita.com
得られた画像群はこんな感じです．
さて，設定は先ほどと全く同じです．なお，画像のサイズが1024x1024になっております．
同様の変化を取っています．マイクの金網の部分も最終的には再現できていることが見えると思います．
最後に，同じデータセットにあるDrumを同様にしてレンダリングして入力データを作成しました．
さて，マイクと同じくして学習させました．
他の2つと比べると少し精度は低めですね．細い物体（高周波成分）の学習にはなかなか手こずります．ちなみにこれ床の色を黒にすると椅子等が上手く学習されません（色としての知覚上，床も椅子も黒に見えると区別がつかないためだと思います）．
最後に動画への埋め込みリンクを載せておきます．録画したのは少し前となるので全く同じ実装ではないですが，今でも同じような（改善された）結果が得られます．
Lego pic.twitter.com/bHjFGaqgQu
— Yoshi’s Dreann (@TTRS_Yoshi_CG) 2024年2月23日
Mic pic.twitter.com/pVpr4CQhW6
— Yoshi’s Dreann (@TTRS_Yoshi_CG) 2024年2月23日
Drum pic.twitter.com/1EWtaT5Xgz
— Yoshi’s Dreann (@TTRS_Yoshi_CG) 2024年2月23日
今よりかは多少遅いですが，まあでも数分でイイ感じに学習できています．現状の私の実装ではかなり荒いところもあるので改善はまだまだ出来そうです．
さいごに
Part1ではMLPの，Part2ではエンコーダーの，Part3ではNeRF全体の実装を示し，説明しました．こういう記事を書くのは初めてなのですが，記事を書いて初めて気づくバグや自分の理解不足な箇所とかを洗い出せてなかなか良い作業ではありました（記事書くために全部で数十時間かかったのでコスパといえばどうなのかとなりますが）．さて，NeRF自体がどういうものであるか，というものは論文を読めば大体わかりますが，いざ実装をしてみると意外と手が動きません．そもそも現在では実装しなくても他人が実装したプログラムやライブラリがあります．しかし，何かを実装するということは実装対象をある意味で言語化的にアウトプットする作業とも言え，ぼんやりとした理解の存在を前面に押し出してくれます．一度は泥臭く実装してみるのも良いと思います．
実装に際してUshioさんには色々と助けていただきました．ありがとうございました．
参考文献
・NeRF: Representing Scenes as Neural Radiance Fields for View Synthesis [Ben Mildenhall et al. ECCV2020]
・Instant Neural Graphics Primitives with a Multiresolution Hash Encoding [Thomas Müller et al. SIGGRAPH2022]
・Optical Models for Direct Volume Rendering [Nelson Max. 1995]
・マルペケつくろーどっとコム その18 直線とAABB
・Bounding Volume Hierarchy (BVH) の実装 - 交差判定編
・memoRANDOM ボリュームレンダリング方程式 (Volume Rendering Equation) 1
・NVIDIA Developer Blog, Thinking Parallel, Part III: Tree Construction on the GPU
・live2d_dev NeRF検討用データセットをBlender Pythonスクリプトで作成する方法
・Ushioさんのredpillの実装と解説スライド
