TopoGPU ускорил расчёт персистентных гомологий в среднем в 53 раза

Персистентные гомологии всё активнее используют в глубоком обучении, но считать их внутри цикла тренировки мучительно долго. Авторы новой работы на arXiv предлагают TopoGPU, потоковый GPU-конвейер, который строит диаграммы персистентности для кубических комплексов, порождённых 2D- и 3D-изображениями. Конкурент, Cubical Ripser, считается одной из лучших существующих реализаций. TopoGPU опередил его на каждом протестированном бенчмарке: в среднем в 53,24 раза по полному времени расчёта (end-to-end) и максимум в 198,01 раза. Код открыт, есть готовые бинарники, ноутбуки для Google Colab и образы Docker.

Где был затык

Персистентные гомологии отслеживают, как при изменении порога яркости в изображении возникают и исчезают компоненты связности, петли и полости. На выходе получается диаграмма персистентности, компактное описание формы. В задачах вроде сегментации её используют как часть функции потерь: сеть штрафуют за разорванные сосуды или лишние дырки в маске. Но тогда диаграмму приходится пересчитывать на каждом шаге обучения. Cubical Ripser работает на CPU, и на крупных 3D-объёмах расчёт съедает значительную долю времени тренировки.

Как устроен конвейер

TopoGPU режет входное изображение на куски и подаёт их на GPU по очереди. Каждый кусок обрабатывают массово-параллельные ядра на сетке GPU-блоков. Найденные граничные соотношения копятся в памяти хоста, и уже там CPU выполняет редукцию граничной матрицы. Самая параллелизуемая часть работы уходит на видеокарту, последовательный финальный шаг остаётся на процессоре. Потоковая схема, судя по описанию, не требует держать весь объём на GPU одновременно, хотя точные ограничения по памяти в аннотации не приведены.

Главная техническая идея связана с дискретной теорией Морса. Она позволяет «схлопнуть» комплекс, убрав пары клеток и не изменив гомологии, так что граничная матрица становится заметно меньше. Проблема в том, что при потоковой обработке куски видят только часть данных, и наивное сопоставление Морса может исказить результат. Авторы вводят stratification-aware discrete Morse matching и доказывают, что он сохраняет персистентные гомологии при потоковой схеме. Для работы на GPU добавлены параллельная топологическая сортировка и параллельный алгоритм чётности V-путей, который выводит границы Морса.

• Среднее ускорение end-to-end против Cubical Ripser: 53,24x
• Максимальное ускорение: 198,01x
• Входные данные: 2D- и 3D-изображения, кубические комплексы
• Разделение труда: GPU строит граничные соотношения, CPU редуцирует матрицу
• Лицензия и поставка: открытый код, бинарники, Colab, Docker

Кому это пригодится

Авторы проверили метод в топологически осознанной сети и показали, что стоимость расчёта персистентных гомологий во время обучения заметно падает. Это главный практический вывод. Если ваша модель использует топологические потери, например для сегментации тонких структур в медицинских 3D-сканах, TopoGPU может убрать узкое место без смены подхода к обучению. Благодаря готовым образам Docker и ноутбукам Colab проверить его можно за вечер.

Есть и оговорки. Сравнение идёт с Cubical Ripser, и цифры 53x и 198x зависят от набора бенчмарков и железа, которое в аннотации не описано. Редукция матрицы остаётся на CPU, так что на сложных данных именно она может стать новым ограничением. Прежде чем менять пайплайн, стоит прогнать TopoGPU на собственных изображениях и сравнить время расчёта на ваших размерах объёмов, а не только средний выигрыш из статьи.