NeRF 3D синтез на JAX: пошаговый гайд от установки до готовой сцены

Оригинальный туториал, на котором построен этот гайд, описывает сборку иерархического NeRF с использованием JAX, Flax, Optax и примитивов объёмного рендеринга из библиотеки jax3d. Подход включает позиционное кодирование, skip-соединения (пропускные связи, когда сигнал перепрыгивает через несколько слоёв сети), раздельные «грубую» и «тонкую» сети, а также иерархическую выборку по значимости. Ниже разложено всё по шагам так, чтобы повторить результат мог и разработчик без опыта в объёмном рендеринге.

Что понадобится?

• Python 3.9+ и менеджер пакетов pip
• JAX с поддержкой GPU (на CPU тоже работает, но медленнее и с урезанными параметрами)
• Flax (библиотека для описания нейросетей поверх JAX), Optax (оптимизаторы), etils (утилиты массивов)
• Клон репозитория google-research/jax3d (клонируется автоматически скриптом)
• Google Colab с GPU T4 или локальная машина с видеокартой NVIDIA
• Около 30 минут на GPU (на CPU с уменьшенными параметрами хватит примерно 15 минут, но качество будет ниже)

Пошаговая инструкция

1. Установите зависимости и загрузите jax3d

Скрипт сам ставит пакеты и клонирует репозиторий. Если вы работаете в Colab, просто запустите ячейку:

import os, sys, subprocess

def _sh(cmd):
    subprocess.run(cmd, shell=True, check=False,
                   stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)

_sh(f'{sys.executable} -m pip install -q '
    f'"etils[array-types,epy,etree,enp]" chex flax optax scikit-image')

REPO_DIR = "/content/jax3d"
if not os.path.isdir(REPO_DIR):
    _sh(f"git clone -q --depth 1 "
        f"https://github.com/google-research/jax3d.git {REPO_DIR}")

Модуль volume_rendering загружается напрямую из клона без полной установки пакета, это позволяет обойти тяжёлые зависимости вроде gin и tfds.

2. Настройте конфигурацию под своё железо

Все ключевые параметры собраны в датакласс Config. На GPU используются полные значения, на CPU скрипт автоматически уменьшает разрешение, число обучающих ракурсов и ширину сети:

import dataclasses

@dataclasses.dataclass
class Config:
    H: int = 64          # высота изображения в пикселях
    W: int = 64          # ширина
    n_train_views: int = 24
    n_coarse: int = 64   # точки выборки для грубой сети
    n_fine: int = 64     # точки для тонкой сети
    deg_pos: int = 10    # степень позиционного кодирования
    deg_dir: int = 4
    width: int = 128     # ширина скрытых слоёв
    depth: int = 6       # глубина сети
    skip: int = 3        # слой skip-соединения
    steps: int = 2500    # шагов обучения
    lr_init: float = 5e-4
    lr_final: float = 5e-6

Если GPU нет, скрипт сам переключится на облегчённый режим: разрешение 40 на 40, 400 шагов, ширина сети 64.

3. Постройте камеры и сгенерируйте лучи

Камеры расставляются по куполу вокруг сцены методом «золотого угла» (golden-angle), это даёт равномерное покрытие ракурсов. Для каждого ракурса строятся лучи через пинхол-модель (модель камеры-обскуры, когда все лучи проходят через одну точку):

import math, numpy as np

FOCAL = 0.5 * cfg.W / math.tan(0.5 * math.radians(cfg.fov_deg))
poses = orbit_poses(cfg.n_train_views, cfg.cam_radius)
origins, directions = rays_from_pose(poses[0], cfg.H, cfg.W, FOCAL)

Функция rays_from_pose возвращает начала и направления лучей для каждого пикселя. Направления нормализованы, поэтому глубины соответствуют реальным расстояниям в мировых координатах.

4. Создайте синтетическую сцену

В туториале используется аналитическая сцена из трёх цветных сфер с объёмной геометрией и освещением, зависящим от направления взгляда. Сцена нужна, чтобы иметь идеальный ground truth (эталон), с которым потом сравнивается результат NeRF 3D синтеза.

5. Соберите архитектуру NeRF

Сеть состоит из двух частей:

• Грубая сеть (coarse): получает позиционно закодированные координаты точки и предсказывает плотность и цвет
• Тонкая сеть (fine): работает с точками, выбранными по иерархической выборке из sample_piecewise_constant_pdf (функция, которая сгущает точки там, где грубая сеть нашла что-то непрозрачное)

Обе сети используют skip-соединения (сигнал со входа подаётся напрямую в средний слой) и отдельный вход для направления взгляда, чтобы моделировать блики и зеркальные отражения.

6. Обучите модель

Обучение запускается с JIT-компиляцией (just-in-time, когда JAX компилирует функцию при первом вызове и дальше выполняет её на GPU без накладных расходов):

• Оптимизатор: Adam
• Расписание скорости обучения: экспоненциальное затухание от 5e-4 до 5e-6
• Обрезка градиентов (gradient clipping), чтобы предотвратить «взрывы» весов
• Размер пакета: 2048 лучей на шаг (1024 на CPU)

7. Оцените результат

После обучения скрипт считает PSNR (Peak Signal-to-Noise Ratio, отношение сигнала к шуму в децибелах: чем выше, тем ближе к оригиналу), визуализирует карту глубин и непрозрачности, строит 360-градусный облёт камеры и извлекает геометрию методом marching cubes (алгоритм, который из поля плотности строит полигональную поверхность).

Как оценить, что всё получилось?

Три индикатора успеха:

• PSNR выше 25 дБ на тестовых ракурсах, значит, сеть научилась воспроизводить сцену
• Карта глубин показывает чёткие контуры объектов, а не шум
• 360-градусный рендер выглядит плавным, без «дыр» и артефактов

Если вы дошли до 360-градусного рендера и видите чистую картинку без артефактов, вы освоили полный цикл NeRF 3D синтеза: от лучей до полигональной модели. Следующий шаг, подставить вместо синтетических сфер фотографии реального объекта и проверить, как сеть справится с настоящими данными.