Поля нейронного излучения с нуля (Часть II) — от теории к практике

В части I мы рассмотрели теорию полей нейронного излучения (NeRF) и объяснили, как обучается модель. Как говорится, в теории между практикой и теорией нет разницы, но на практике она есть… Так что присоединяйтесь ко мне и давайте изучим детали реализации этих умных моделей.

Эта статья на 100% написана человеком!

1. Введение

Я представил Neural Radiance Fields (NeRFs) [1] в Части I, поэтому я просто подытожу здесь. Пожалуйста, обратитесь к предыдущей истории для более подробной информации.

NeRF — это класс генеративных моделей, которые могут генерировать новые виды 3D-сцены (естественные или синтетические). Для их обучения требуется набор данных со многими (обычно несколькими сотнями) изображениями сцены, снятыми с разных ракурсов с одной и той же камеры. Интересно, что они действуют как форма алгоритма сжатия: вся сцена хранится в весах нейронной сети, и любой вид может быть сгенерирован очень быстро. С другой стороны, обучение может быть довольно медленным.

Обучение NeRF включает следующие этапы:

  • Извлеките координаты начала координат и вектора направления всех лучей. В наборе данных представлений каждый пиксель соответствует тому месту, где луч, исходящий из фокуса камеры и пересекающий сцену, пересекает плоскость изображения.

  • Переведите координаты начала координат и направления в общую систему координат сцены. Это необходимо, потому что система отсчета камеры не является фиксированной.
  • Проведите вдоль каждого луча и аккумулируйте (интегрируйте) прогнозы модели для плотности (~ непрозрачности) и значений красного, зеленого и синего цвета.
  • Используя уравнение рендеринга, вычислите результирующий цветовой вектор пикселя, соответствующего каждому лучу.
  • Вычислите ошибку между фактическими цветами пикселей, найденными в изображениях истинной земли, и визуализированными. Обратно распространите ошибку, чтобы постепенно научить модель лучшему внутреннему представлению профиля излучения сцены в разных точках и под разными углами.
  • Модель должна знать только положение камеры, чтобы создать новый вид. Он применяет все вышеперечисленные шаги, кроме последнего. Пиксели, полученные таким образом, составляют конечное изображение.

Это теория в двух словах. На практике различия начинаются с уравнения рендеринга.

2. Дискретное уравнение рендеринга

Непрерывная версия уравнения рендеринга, которую мы объяснили в предыдущем посте, выглядит так:

Если оно кажется вам пугающим, прочитайте мои пояснения в части I. В этом уравнении мы идем по лучу r и непрерывно суммируем значения под интегралом: накопленный коэффициент пропускания T (функция положения t вдоль луча), локальная плотность σ(функция луча, за которым мы следуем, и положения на луче), и цветовой вектор R, G, B c (зависит от положения на луче, но также и от его направления, поскольку воспринимаемый цвет может зависеть от угла обзора).

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

На практике линейный участок пространства между ближней границей tₙ и дальней границей t_f делится на N равных ячеек, и берется выборка в случайном месте внутри каждого бина (см. раздел «Уточнения», чтобы узнать о более сложном подходе к выборке).

Затем мы пишем дискретное уравнение рендеринга:

В уравнении (2), мы распознаем коэффициент пропускания T (теперь Tᵢв позиции tᵢ) и испускаемый цвет c (теперь cна позиции tᵢ). Что касается δᵢ, этоэторасстояние между точкой в ​​положении i+1 и точкой в ​​я. Вы можете видеть, что оно играет роль бесконечно малого расстояния dsв выражении Tᵢ и dt в общем выражении С.

Менее очевидно, почему терм σ(r(t)).dt стал (1 — exp( -σᵢ δᵢ)). Поскольку σᵢδᵢ≥ 0, размещение этого количества внутри функции (1 -exp(-x)) ограничивает результат областью [0, 1). Поскольку у нас 0 ‹ Tᵢ≤ 1, веса
wᵢ= Tᵢ(1 — exp(-σᵢδᵢ)) находятся в области [0, 1), так что вклады всех положений вдоль луча положительны, и ни один из них не может в одиночку вносить вклад в результирующий цвет.

3. Прогнозирование цвета и плотности

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

Его входными данными являются координаты местоположения и направления всех точек вдоль пакета лучей с одной особенностью: вместо того, чтобы передавать значения напрямую в виде двух трехмерных векторов (или одного шестимерного вектора), они сначала проходят через кодирование. процесс, который я подробно описываю в разделе Уточнения ниже. Для ясности: каждый входной вектор 6D в пакете соответствует одной точке на одном луче. Вы можете вернуться к Части I, чтобы понять, как получаются эти координаты.

Что выходит из нейронной сети? Две вещи: пакет скалярных значений σᵢ (исправленных функцией ReLU, чтобы убедиться, что они положительны) и набор векторов (R, G, B) c (проходит через сигмовидную функцию, чтобы ограничить их интервалом (0, 1)). Они соответствуют предсказаниям сети о плотности и излучаемом цвете в каждой позиции пакета лучей. Это значения, которые используются в уравнении дискретного рендеринга для рендеринга отдельных пикселей.

4. Оптимизация NeRF

Теперь, когда мы можем реконструировать цвет каждого пикселя, предсказывая плотности и испускаемые цвета вдоль его луча, у нас есть почти все, что нужно для обучения модели.

Процесс выглядит следующим образом:

  • Возьмите кучу пикселей/лучей из набора данных и получите их координаты направления и начала координат из соответствующих матриц позы,
  • Возьмите контрольные точки вдоль каждого из этих лучей,
  • Используйте нейронную сеть, чтобы предсказать плотность и испускаемый цвет в каждой точке в направлении луча,
  • Используйте эти значения для вычисления результирующих цветов в выбранных пикселях.
  • Вычислите среднеквадратичную ошибку между визуализированным и истинным цветами (из исходных изображений), суммированную по пакету.

Процесс оптимизации, используемый в статье, относительно прост и не требует сложной аугментации или регуляризации данных; они используют оптимизатор Адама с уменьшением веса. Основным источником регуляризации является стохастичность процесса выборки точек вдоль лучей.

5. Вывод

После того, как модель обучена, новые изображения могут быть созданы для любой позы путем построения матрицы позы (которая зависит только от положения и ориентации камеры) и использования ее для получения ориентации лучей и начала координат для каждого пикселя в будущем изображении.

Затем по каждому лучу производится выборка точек, и нейросеть предсказывает соответствующие цвета и плотности. Используя дискретное уравнение рендеринга, мы вычисляем цвет каждого пикселя и вуаля: у вас есть готовое изображение!

6. Уточнения

Авторы улучшают качество получаемых изображений, применяя два приема.

6.1. Кодирование положения

Я упоминал об этом в разделе 3. Вспомним, что (трехмерное) положение каждой точки выборки и (трехмерная) ориентация соответствующего луча образуют единый шестимерный вектор (xᵢ, yᵢ, zᵢ, u, в, ш). Предыдущие исследования, особенно в моделях на основе трансформаторов, показали, что нейронные сети могут значительно выиграть от более богатого кодирования положения. Вместо того, чтобы дать нашей модели один скаляр для каждой координаты, мы расширяем каждую координату в вектор, отображая функции синуса и косинуса с различной периодичностью. В документе используется десять факторов для координат положения и четыре для направлений. Для позиций мы создаем десять функций синуса и десять функций косинуса для каждой координаты, что приводит к 60-мерному кодированию. Точно так же направления кодируются в 24 измерениях. Авторы показывают, что такое кодирование значительно улучшает восстановление мелких деталей.

6.2. Иерархическая выборка

Вы помните, когда я сказал вам, что выборка вдоль каждого луча выполняется путем выбора случайных позиций в бинах постоянного размера? Ну, это была не вся правда… Авторы используют два распределения выборки — грубое и тонкое — и, соответственно, обучают две нейронные сети. Выборка, которую я описал выше, является грубой.

Как определяется тонкое распределение? В уравнении (2) мы видели, что цвет пикселя представляет собой взвешенную сумму цветов в каждой точке луча, где веса имеют вид wᵢ= Tᵢ(1 — exp(-σᵢδᵢ)) и, таким образом, зависят от плотности в каждой точке. Поэтому мы используем грубую модель для вычисления этих плотностей, что дает нам набор весов. Эти веса нормализуются, чтобы в сумме дать единицу для каждого луча, а затем используются для определения кусочно-постоянной функции плотности вероятности.

Точная выборка выполняется путем выборки из этого нового распределения (и мы также берем больше выборок, чем при грубой выборке, 128 в исходной статье). Веса по своей сути помогают определить важные части изображения, т. е. те места, которые могут блокировать отображаемый цвет или влиять на него. Следовательно, это места, где требуется более высокая плотность выборки.

7. Заключение

Вот и все! Мы рассмотрели методологию от позы камеры до лучей, точечных выборок, прогнозируемых локальных плотностей и цветов и визуализированных цветов пикселей. NeRF — это пример умного исследования, которое не требует целой фермы серверов TPU и при этом дает выдающиеся результаты. Самое примечательное в них то, что вся 3D-сцена захватывается моделью, которая весит всего несколько МБ.

С другой стороны, обучение может занять несколько часов GPU. В документе сообщается об 1 GPU-день для большого изображения, а TinyNeRF на моем GPU заняло около 8 часов. Тем не менее, результат узнаваем, но зашумлен (без сомнения, гиперпараметры и код могут быть дополнительно оптимизированы). Я также заметил, что с набором данных TinyNeRF точная выборка не приносит огромных преимуществ и значительно замедляет обучение, поэтому я бы посоветовал изначально придерживаться грубой выборки. Вы можете увидеть мою реализацию с нуля здесь.

Если вам понравился этот пост, пожалуйста, спамьте значок хлопка и поделитесь им с друзьями!

Использованная литература:

[1] Милденхолл, Б., Сринивасан, П.П., Танчик, М., Бэррон, Дж.Т., Рамамурти, Р. и Нг, Р., 2021. «NeRF: представление сцен как полей нейронного излучения для синтеза представлений». Сообщения ACM, 65(1), стр. 99–106.