Иерархический NeRF на JAX3D: туториал по объёмному рендерингу и синтезу новых ракурсов

В свежем туториале автор собирает полный пайплайн иерархического Neural Radiance Field (NeRF) на стеке JAX, Flax и Optax, используя примитивы объёмного рендеринга из jax3d. Сначала создаётся синтетический мультивью-датасет из аналитической сцены с объёмной геометрией и зависящим от направления излучением — через sample_along_rays и volume_rendering настраивается прямой рендеринг.

Затем реализуется сама модель: позиционное кодирование, skip-соединения, отдельные coarse и fine сети, кондиционирование по направлению взгляда. Иерархический сэмплинг делается через sample_piecewise_constant_pdf. Обучение идёт с JIT-компиляцией JAX, оптимизатором Adam, экспоненциальным затуханием learning rate и клиппингом градиентов.

В финале — оценка синтеза новых ракурсов по PSNR, визуализация глубины и непрозрачности, диагностика сэмплирования, рендеринг на 360 градусов и извлечение геометрии через marching cubes. Для CPU предусмотрен облегчённый конфиг, на GPU — полноценный. Это практичный рецепт для тех, кто хочет разобраться в NeRF на современном JAX-стеке.