1. Описание модели Mask R-CNN 3D

Mask R-CNN 3D – это расширение знаменитой модели Mask R-CNN для работы с трехмерными данными (объёмными изображениями или облаками точек). Классическая Mask R-CNN предназначена для instance segmentation (сегментации отдельных объектов) на 2D-изображениях и состоит из двух основных частей: (1) сети предложений областей (Region Proposal Network, RPN) и (2) головы (Head) с несколькими выходными ветвями для классификации, регрессии ограничивающих рамок и сегментации масок . В версии 3D эта же концепция перенесена в трехмерное пространство.

Входом модели Mask R-CNN 3D обычно является объёмный данных – например, медицинский 3D снимок (CT/MRI) размером (D×H×W) или облако точек, представляющее 3D-сцену. Backbone-сеть (обычно сверточная нейросеть типа ResNet) извлекает из входных данных многомасштабные признаки. В 3D версии backbone заменяет все 2D-операции (свертки, пулинг) на 3D-аналоги, позволяя обрабатывать объёмные данные напрямую. (Если 3D-данные заданы как облако точек, возможно предварительное преобразование, например, вокселизация пространства или проекция на несколько 2D-плоскостей – об этом подробнее в разделе 6.) Backbone формирует карты признаков – объёмные тензоры с пониженным разрешением, но содержащие высокоуровневую информацию о структуре объектов в сцене.

Далее вступает Region Proposal Network (RPN) – небольшая сеть, скользящая по картам признаков и генерирующая набор предположительных объектов (region proposals) в виде ограничивающих 3D-рамок (прямоугольных параллелепипедов в координатах исходного объёма). RPN использует заранее заданные «якоря» (anchor boxes) – шаблонные 3D-боксы разных размеров и соотношений сторон, размещенные по всей карте признаков . Для каждого такого anchor RPN предсказывает два значения: объектность (есть объект/фон) и смещение рамки (на сколько нужно подвинуть и масштабировать anchor, чтобы точнее охватить объект). После этого выбираются топ-N наиболее перспективных предложений с помощью non-maximum suppression (NMS) – подавления пересекающихся рамок с меньшей оценкой.

Для каждого выбранного объёма (RoI – Region of Interest) в признаках выполняется операция RoI Align 3D – выравнивание региона интереса. Это обобщение RoI Align, предложенного в оригинальном Mask R-CNN , на третий размер. RoI Align “вырезает” патч признаков, соответствующий предложенной 3D-рамке, и масштабирует его к фиксированному размеру с помощью билинейной (трехлинейной) интерполяции, без округления координат. За счёт этого достигается точное соответствие маски и исходного объёма, без артефактов дискретизации, которые были у более простой RoI Pooling . (Поскольку готовой реализации 3D RoI Align в фреймворках долго не было, в некоторых работах её реализуют как пользовательскую операцию на C++/CUDA (

Голова Mask R-CNN 3D обрабатывает выровненные RoI-признаки и выдаёт окончательные предсказания по каждому объекту. Как и в 2D-варианте, голова имеет три выходные ветви:

Классификатор предсказывает класс объекта (включая “фон” как отдельный класс) для каждого предложенного RoI.
Регрессор уточняет координаты 3D-рамки объекта (выдает поправки к координатам предложения RPN, чтобы плотнее обхватить объект).
Сегментирующая маска – сверточная сеть-декодер, которая из признаков RoI генерирует бинарную маску объекта с заданным разрешением (например, вокселя). Маска предсказывается для каждого класса отдельно. Важная деталь: в Mask R-CNN нет соревнования между классами за маску – модель учится предсказывать маски для каждого класса параллельно. В итоге выбирается маска того класса, который дал наивысшую уверенность на этапе классификации.

После получения предсказаний по всем выбранным RoI, выполняется финционная фильтрация: для каждого класса через NMS отбираются окончательные предсказания (рамка, класс, маска). Маска каждого объекта масштабируется обратно на размер его рамки и вписывается в итоговый объем сегментации. Результатом работы Mask R-CNN 3D является, таким образом, список обнаруженных 3D-объектов с их классами, 3D-координатами рамок и воксельными масками.

(Коротко говоря, Mask R-CNN 3D пытается решить сразу три задачи: обнаружение объектов в 3D, определение их категорий и выделение точных границ – и делает это в единой нейросетевой архитектуре. Звучит сложно? Не волнуйтесь: модель справляется с этим примерно так же, как жонглер – он тоже может держать в воздухе три мяча одновременно, хотя поначалу это сбивает с толку наблюдателя.)

2. Математическая основа

Основу обучения Mask R-CNN 3D составляет многозадачная функция потерь, суммирующая ошибки сразу трех компонентов – классификации, локализации (рамок) и сегментации масок. Для каждого предложенного региона (ROI) определены следующие части ошибки:

$\mathcal{L}_{cls}(i)$ – потери на классификации (например, мультиклассовая кросс-энтропия, определяющая, правильно ли определён класс объекта или фон).
$\mathcal{L}{box}(i)$ – потери на регрессии рамки (насколько предсказанная рамка отклонилась от истинной; обычно берут Smooth L1 loss, сглаженную -норму, которая менее чувствительна к выбросам, чем стандартный). Пусть истинное смещение для рамки , а предсказанное – $t^{(c)}$ для класса ; тогда для положительного
$ROI \mathcal{L}{box} = \sum_{j \in {x,y,z,w,h,d}} \text{Smooth}_{L1}(t_j - v_j)$ .
Для фоновых регионов (без объекта) эта часть потерь обнуляется (индикатор $\mathbb{1}[u\ge1]$ ).
$\mathcal{L}_{mask}(i)$ – потери на сегментации маски. Для каждого положительного ROI и для каждого элемента маски размером $m\times m\times m$ рассчитывается бинарная кросс-энтропия между предсказанной маской $\hat{Y}^c$ и истинной маской (только по правильному классу ):
$\mathcal{L}_{mask} = - \frac{1}{m^3}\sum_{x=1}^{m}\sum_{y=1}^{m}\sum_{z=1}^{m}\Big[ Y(x,y,z)\log \hat{Y}^c(x,y,z) + (1 - Y(x,y,z))\log(1-\hat{Y}^c(x,y,z))\Big]$
Здесь $Y(x,y,z)\in{0,1}$ – значение истинной маски в вокселе, а $\hat{Y}^c(x,y,z)\in[0,1]$ – предсказанная вероятность принадлежности этому классу.

Итоговая функция потерь – сумма по всем RoI (принятые предложения) и по всем трем видам потерь:

$\mathcal{L}_{total} = \sum_{i}^{\text{RoIs}} \big(\mathcal{L}_{cls}(i) + \mathbb{1}[u_i \ge 1]\;\mathcal{L}_{box}(i) + \mathbb{1}[u_i \ge 1]\;\mathcal{L}_{mask}(i)\big), где u_i=0$

для фона, а u_i=k для объекта класса.

Коэффициенты перед слагаемыми обычно берутся равными 1 (или подбираются гиперпараметром, если необходимо сбалансировать, например, mask loss может быть умножен на 2).

Conv3D и особенности слоёв. В 3D-версии все свёрточные операции выполняются в трех измерениях. Например, свёртка $3\times3\times3$ с шагом (stride) 2 по глубине, высоте и ширине уменьшит размер объёма примерно в 2 раза по всем измерениям. Формально, если ядро свёртки W размера , а входной признакной объём X размера, то выход в позиции

: $Y(i,j,k) = \sum_{p=1}^{k_d}\sum_{q=1}^{k_h}\sum_{r=1}^{k_w} W(p,q,r)\cdot X(i+p-1,\;j+q-1,\;k+r-1) + b$ ,

где – смещение (bias). Эта операция аналогична 2D-свертке, только добавляется суммирование по третьему измерению .

Пулинг, BatchNorm и прочие слои аналогично обобщаются на 3D. Стоит отметить, что 3D-свёртки значительно более вычислительно затратны и требуют больше памяти, поэтому размеры сетей (число фильтров, глубина блоков) иногда уменьшают по сравнению с 2D-вариантами, чтобы модель помещалась в GPU.

Region Proposal Network (RPN). RPN обучается решать двухклассовую задачу (объект/не объект) для каждого anchor и регрессию рамок. Для обучения генерируются положительные анкоры (достаточно хорошо перекрывающие объект, обычно IoU с истинным объектом > 0.7) и отрицательные (нет перекрытия, IoU < 0.3). Классификационная часть RPN имеет логистическую потерю (binary cross-entropy), а регрессия рамок – Smooth L1 аналогично вышеописанной. Эти компоненты потерь RPN также входят в общую функцию (их обычно суммируют с коэффициентами). Таким образом, модель учится в конце-концов минимизировать сумму всех ошибок: неправильно классифицированных anchor, неточных границ предложений и масок и т.д. – чтобы добиться максимальной точности локализации и сегментации объектов.

(Примечание: хоть Mask R-CNN 3D оптимизирует сложную совокупность метрик, на практике она успешно сходится. Главное – хорошая аннотация 3D-данных. Если математическая формулировка показалась громоздкой, можно утешить себя мыслью: нейросеть “читает” эти формулы не лучше нас с вами – она просто численно минимизирует $\mathcal{L}_{total}$ с помощью градиентного спуска, не задавая лишних вопросов!)

3. Сравнение с альтернативными методами

Зачем вообще нужен Mask R-CNN 3D, если существуют другие подходы к сегментации объемных данных? Рассмотрим две популярные альтернативы и сравним их с данным методом.

3.1 3D U-Net

3D U-Net – это, пожалуй, самый известный архитектурный шаблон для сегментации 3D-изображений. 3D U-Net представляет собой полностью сверточную сеть энкодер-декодер, расширяющую 2D U-Net Роннебергера в третье измерение. Проще говоря, она берёт объемное изображение на входе и сразу выдаёт объемную карту меток (семантическая сегментация), помечая каждый воксель классом объекта или фоном. Ключевые особенности 3D U-Net:

Полное покрытие поля зрения. U-Net сканирует всё изображение и учитывает широкий контекст благодаря большому рецептивному полю. Энкодер постепенно понижает разрешение, захватывая контекст, декодер – восстанавливает детали с помощью скип-соединений. Для 3D это означает, что модель может учитывать взаимосвязи между органами или объектами по всему объему сразу.
Нет явных рамок и пропозалов. В отличие от Mask R-CNN, U-Net не генерирует гипотезы объектов. Она учится предсказывать метку класса для каждого вокселя. Поэтому 3D U-Net решает задачу скорее семантической сегментации (все пиксели класса сгруппированы), чем instance сегментации. Если нужно выделять отдельно экземпляры объектов (например, каждую опухоль отдельно), U-Net сама по себе этого не делает – после сегментации часто применяют алгоритмы постобработки (например, разделение связных компонентов по метке).
Простота обучения. Архитектура U-Net однозадачная (только сегментация), поэтому функция потерь обычно проще – например, софтмакс и cross-entropy по вокселям или Dice-коэффициент. Нет сложной многозадачной структуры, как у Mask R-CNN. Это может упрощать подбор гиперпараметров и скорость сходимости.
Точность границ. U-Net, особенно с достаточным количеством скип-коннектов, может очень точно восстанавливать тонкие детали границ объектов, поскольку оптимизируется прямая пиксельная ошибка сегментации. Mask R-CNN тоже старается, но она зависит от качества выровненных RoI и разрешения масочной ветви (часто 28×28 в 2D, или 32×32×32 в 3D), что может ограничивать точность мелких деталей.

Преимущества Mask R-CNN 3D по сравнению с U-Net проявляются, когда нам важно обнаруживать каждый объект отдельно, особенно если объекты одного класса могут соприкасаться. Например, в задаче сегментации клеток в 3D-микроскопии: 3D U-Net окрасит все клетки одним цветом, и потребуется раздельно их идентифицировать (например, методом Watershed), тогда как Mask R-CNN сразу выдаст отдельные маски для каждой клетки. Кроме того, Mask R-CNN имеет явный механизм для ограничивающих 3D-рамок, что удобно, если нам нужны грубые границы объектов или их положение для последующих задач (трacking, подсчет и т.п.).

С другой стороны, недостатки Mask R-CNN 3D относительно U-Net: более сложное обучение (многокомпонентная функция потерь, больше гиперпараметров – якоря, thresholds NMS и т.п.), больше требований к памяти (несколько выходных карт + хранение множества предложений) и более низкая скорость инференса, особенно при большом количестве объектов. В задачах, где требуется именно семантическая сегментация всего объема с одной меткой на класс (например, выделить весь объем печени на КТ), 3D U-Net зачастую эффективнее и проще. Не случайно в медицинских соревнованиях (BraTS по опухолям мозга, LiTS по сегментации печени и т.д.) 3D U-Net и его вариации доминируют, тогда как Mask R-CNN 3D применяют точечно для случаев множественных мелких объектов.

3.2 Segment Anything Model (SAM)

Segment Anything Model 2 (SAM2) – это обновлённая версия SAM от Meta AI, разработанная для универсальной сегментации как изображений, так и виде. В отличие от первой версии, обученной на гигантском наборе из 1 млрд масок объектов, SAM2 обладает улучшенной архитектурой, включающей Hiera-энкодер, модуль памяти (Memory Bank) и Memory Attention, что позволяет ей учитывать временные зависимости и обеспечивать высокую скорость вывода при интерактивной работе. Благодаря этим улучшениям SAM2 остаётся интерактивно-промптовой: вместо того чтобы выдавать сегментацию «из коробки», она принимает на вход подсказку – например, точку или рамку – и генерирует маску объекта, даже если объект ранее не встречался в обучающем наборе.

Mask R-CNN 3D vs SAM2: Эти модели преследуют разные цели. Mask R-CNN 3D – это специализированная, полностью supervised модель для детекции и сегментации заданных классов в 3D, требующая обученных аннотаций для конкретных доменов. SAM2 же является универсальным zero-shot сегментатором, который можно применять без дополнительного обучения, просто дав ему подсказку. Хотя SAM2 оптимизирована для 2D-изображений и видео, современные исследования уже активно адаптируют её идеи для работы с 3D-данными. Например, подходы типа SAM3D или SAM2Point используют предсказанные 2D-маски SAM2, проецируя их в 3D-пространство или интерпретируя 3D-данные как последовательность видео-кадров для получения 3D-масок объектов.

(Если провести аналогию, Mask R-CNN 3D – это профессиональный хирург, который чётко вырезает объект там, где его обучили, а SAM2 – универсальный швейцарский нож, способный выполнять множество задач, но которому для работы в 3D пока требуется небольшая адаптация.)

4. Прохождение данных через сеть

Рассмотрим пошагово, как трехмерные данные проходят через Mask R-CNN 3D, начиная от сырого входа (объём или облако точек) и заканчивая выходными предсказаниями. Этот процесс довольно сложный, но его можно разбить на последовательность трансформаций:

Шаг 1. Подготовка входа. Допустим, у нас есть медицинское томографическое изображение – куб данных размером $D\times H\times W$ с оттенками серого (одно канал, например, Hounsfield units в КТ). Сначала его обычно нормализуют (например, приводят интенсивности к диапазону или стандартизуют по среднему и дисперсии). Затем добавляется размер батча и каналов, и тензор формы [B, C, D, H, W] подается в сеть. Если вход – облако точек, то прямая подача списка координат в Mask R-CNN невозможна, поскольку свёртки ожидают регулярную решётку. Возможны варианты: (a) вокселизовать облако точек – разбить пространство на 3D-сетку (например, 1 см кубики) и получить воксельный тензор, где воксель содержит количество точек или признак; (b) проецировать точечные данные на несколько камер (видов) и обрабатывать через обычный 2D Mask R-CNN на этих видах, а потом комбинировать (как делали в некоторых работах с RGB-D). В любом случае, на вход основной сверточной сети должны поступить либо 3D-данные в виде тензора фиксированной размерности, либо множественные 2D-проекции.

Шаг 2. Экстракция признаков (backbone). Входной тензор прогоняется через backbone – обычно это 3D-аналоги Residual Network. Например, 3D ResNet-50 состоит из ряда слоёв Conv3D + BatchNorm + ReLU, объединенных в блоки с пропускными соединениями. Из-за памяти часто используют FPN (Feature Pyramid Network): несколько выходных карт признаков разных уровней разрешения. Предположим, наша сеть выдала пирамиду из 3 уровней: с размером и с – чем глубже, тем более обобщенные и мелкомасштабные признаки. Каждая такая карта имеет, например, 128 или 256 каналов признаков.

Шаг 3. Генерация Region Proposals (RPN). Для каждой карты признаков RPN скользящим окном применяет Conv3D $1\times1\times1$ и выдаёт две карты меньшего канала: одну для оценок объектности (например, 1 канал $\times$ A анкоров, где A – количество anchor-рамок для одной точке признаков), и другую для регрессии рамок (например, 6 координат $\times$ A анкоров). В каждой точке признаков мы интерпретируем предсказания RPN как: anchor присутствует объект (да/нет) и вектор смещений $(\Delta x,\Delta y,\Delta z,\Delta d,\Delta h,\Delta w)$ для корректировки размеров anchor под объект. Затем RPN отбирает лучшие анкоры. Например, на каждом уровне берутся топ-N по скору, объединяются, применяются NMS в 3D: пересекающиеся 3D-окна с IoU > 0.7 отсекаются, остаются наиболее уверенные. В итоге RPN выдаёт, скажем, 100 предложений 3D-рамок $R_1,\dots,R_{100}$ , которые вероятнее всего содержат объекты.

Шаг 4. RoI Align и вырезка признаков. Каждый предложенный бокс – это область на оригинальном изображении, соответствующая определённому подоблаку в признаках разных уровней пирамиды (обычно используют уровень пирамиды пропорционально размеру бокса). Операция RoI Align берет соответствующую карту признаков и “вырезает” из неё то содержимое, что попадает внутрь рамки . Так как рамка имеет произвольные (не обязательно целые) границы относительно дискретной сетки признаков, производится интерполяция: мы выбираем фиксированную сетку выборок, например $16\times16\times16$ , накладываем её на область в пространстве признаков и выполняем трёхлинейную интерполяцию признаков в этих точках. Получается тензор признаков размером $C \times 16 \times 16 \times 16$ – выровненный признак RoI. Этот тензор проходит через несколько слоёв (обычно 2-3 уровня Conv3D + ReLU) для дальнейшей обработки в голове.

Шаг 5. Предсказание классов и рамок (голова DetNet). Признаки RoI схлопываются пространственно (через глобальный Average Pooling 3D или серию сверткок и Flatten) и подаются в полносвязную сеть классификатора. Классификатор (обычно 2 слоя: FC -> ReLU -> FC) выдаёт вектор из значений: вероятности принадлежности RoI к каждому из классов + “фон”. Регрессор рамок – как правило, общие FC-слои с классификатором (то есть голова объединена), но имеет свой вывод: для каждого из классов выдаётся по скорректированной рамке. Однако на практике часто делают разных векторов, но используют только тот, который соответствует предсказанному классу. Например, если классификатор решил, что RoI принадлежит классу “печень” (условно класс 3), то берём параметры рамки из слота 3 регрессора и получаем уточнённую 3D-рамку объекта.

Шаг 6. Предсказание маски. Параллельно с классификацией/регрессией, выровненный признак RoI проходит через сегментационную голову – несколько Conv3D слоёв, возможно с увеличением разрешения (например, через транспонированные свёртки – deconv). Итогом является объём $m\times m\times m$ (например ) для каждого класса . Это означает, что масочная сеть пытается нарисовать маску объекта для каждого возможного класса. Но во время обучения учитывается только плоскость маски того класса, к которому принадлежит объект – остальные игнорируются (не влияют на loss) . При инференсе мы берём маску соответствующую самому вероятному классу из классификатора. Например, если для RoI предсказан класс “печень”, то берём маску из канала, отвечающего за “печень”. Эта маска – в координатах выровненного RoI – затем масштабируется обратно на размер предсказанной 3D-рамки и размещается в ней (обычно с порогом 0.5 по вероятности, чтобы получить бинарное значение).

Шаг 7. Постобработка и вывод. После получения для каждого RoI тройки (класс, уточнённая рамка, маска) у нас может оказаться множество пересекающихся предсказаний. Поэтому, как и в 2D-детекции, применяем Non-Maximum Suppression по классам: для каждого класса сортируем предсказанные объекты по уверенности, и последовательно отбрасываем объекты, у которых IoU пересечения с более уверенным объектом превышает заданный порог (например, 0.5). Оставшиеся объекты считаются итоговыми детекциями. Их маски можно наложить на нулевой объём того же размера, что вход, получив финальную разметку всех найденных объектов.

В итоге полный цикл “вход – выход” таков: входной 3D-образ -> признаки -> предложения (3D-рамки) -> выравнивание -> классификация + сегментация -> фильтрация -> маски объектов. Если вход – облако точек, цикл схож: координаты точек -> либо конверсия в воксели (далее как выше), либо особая обработка. В последнее случае вместо Conv3D могут использоваться специальные сетевые модули (PointNet, SparseConv – см. раздел 6), а RoI Align заменяется выборкой точек внутри предложенного 3D-бокса и обработкой их MLP-сетью. Такие варианты реализованы, например, в методах Region-based PointNet и GSPN для 3D точечных облаков, где генерируются предварительные формы объектов, уточняемые сетями. Но принципы остаются теми же: выделить регион, распознать его класс и получить маску точек внутри него.

5. Практическая реализация

Далее мы рассмотрим практические аспекты реализации Mask R-CNN 3D. Мы сфокусируемся на ключевых моментах и продемонстрируем упрощённый forward-pass код на PyTorch, без обучения модели. Это поможет понять архитектуру изнутри.

5.1 Подготовка данных

Начнём с подготовки данных. Предположим, у нас есть набор 3D изображений и соответствующие аннотации (кубы с разметкой или списки объектов). Для простоты, возьмём один пример объемных данных – например, синтетический куб размером $64\times64\times64$ с нарисованной внутри сферой (будем считать её объектом класса 1). Обычно данные хранятся в формате NIfTI, DICOM или NPZ. Мы должны загрузить объём в память (например, с помощью библиотеки nibabel или SimpleITK для медицинских форматов), получить numpy-массив формы (D, H, W), нормализовать его и преобразовать в тензор PyTorch. Аннотации для обучения Mask R-CNN включали бы координаты рамок и воксельные маски объектов, но для forward-pass (инференса) нам нужны только входные данные.

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

import torch
import torch.nn.functional as F

# Создаем синтетический 3D-объем (например, сфера внутри куба)
D, H, W = 64, 64, 64
x = torch.zeros((1, 1, D, H, W))  # батч=1, канал=1
# Нарисуем сферу радиуса 10 в центре
center = torch.tensor([D/2, H/2, W/2])
for i in range(D):
    for j in range(H):
        for k in range(W):
            if torch.dist(torch.tensor([i, j, k], dtype=torch.float32), center) < 10:
                x[0, 0, i, j, k] = 1.0
# Нормализация (уже 0-1). Если бы были интенсивности, можно вычесть среднее и поделить на STD.

В реальности вместо ручного рисования объекта мы загрузили бы данные, но сейчас у нас тензор x – это вход с одним каналом. Значения 1.0 внутри сферы играют роль сигнала (например, высокая интенсивность), а фон – 0.0. Наша модель должна будет обнаружить эту сферу.

5.2 Реализация forward-pass на PyTorch

Теперь реализуем упрощённый Mask R-CNN 3D. Мы опустим многие детали (например, многомасштабную FPN, сложный RPN), чтобы сосредоточиться на структуре.

Для простоты:

Возьмём маленький backbone: несколько 3D-свёрток с пулингом, без Residual-блоков.
Сгенерируем вручную Region Proposal – например, возьмём один фиксированный anchor во всем объёме или вычислим по порогу.
Сделаем RoI Align грубой имитацией – через 3D-пулинг или интерполяцию PyTorch.
Сформируем выходы головы (класс + маска).

Вот пример кода:

import torch.nn as nn

class SimpleMaskRCNN3D(nn.Module):
    def __init__(self, num_classes=2):
        super().__init__()
        # Backbone: 2 Conv3D слоя с пулингом
        self.conv1 = nn.Conv3d(1, 8, kernel_size=3, padding=1)   # из 1 канала в 8
        self.conv2 = nn.Conv3d(8, 16, kernel_size=3, padding=1)  # из 8 в 16
        self.pool = nn.MaxPool3d(2)  # будет делить размеры на 2
        # ROI pooling size
        self.roi_size = (16, 16, 16)
        # Head: классификатор и сегментатор
        # Классификатор
        self.fc1 = nn.Linear(16 * (D//4) * (H//4) * (W//4), 32)  # после двух пуллингов размеры /4
        self.fc2 = nn.Linear(32, num_classes)  # выход по количеству классов
        # Маска-голова (несколько Conv3d трансформаций -> выходная маска фиксированного размера)
        self.mask_conv1 = nn.Conv3d(16, 8, kernel_size=3, padding=1)
        self.mask_conv2 = nn.Conv3d(8, 1, kernel_size=3, padding=1)  # предсказываем 1 канал маски (для определенного класса)
    
    def forward(self, x):
        # Backbone forward
        out = F.relu(self.conv1(x))
        out = self.pool(out)            # размеры: (D/2, H/2, W/2)
        out = F.relu(self.conv2(out))
        features = self.pool(out)       # размеры: (D/4, H/4, W/4), канал=16
        
        # RPN (упрощенно): возьмем в качестве единственного предложения весь объем признаков
        roi_feature = features  # т.е. одна Region of Interest равна всей карте признаков
        
        # ROI Align (упрощенно): тут ROI == весь feature map, так что просто приведем к roi_size
        # Если бы roi_feature был меньше/больше roi_size, интерполируем:
        roi_aligned = F.interpolate(roi_feature, size=self.roi_size, mode='trilinear', align_corners=False)
        
        # Flatten ROI features for classification head
        flat = roi_aligned.view(1, -1)  # batch=1
        cls_logits = self.fc2(F.relu(self.fc1(flat)))
        
        # Mask head: прогнать roi_aligned через сверточные слои
        mask_feat = F.relu(self.mask_conv1(roi_aligned))
        mask_logits = self.mask_conv2(mask_feat)  # размер: [1, 1, 16, 16, 16]
        mask_probs = torch.sigmoid(mask_logits)   # вероятности 0-1 внутри ROI
        return cls_logits, mask_probs

# Инициализация модели
model = SimpleMaskRCNN3D(num_classes=2)
cls_out, mask_out = model(x)
print("Logits классов:", cls_out.detach().numpy())
print("Выход маски shape:", mask_out.shape)
print("Макс знач. в маске:", mask_out.max().item())

В этом коде мы:

Backbone: два слоя Conv3d с активацией ReLU и два MaxPool3d. В результате features имеет размерность [1, 16, D/4, H/4, W/4] (т.е. $16\times16\times16$ при нашем входе).
RPN: очень упрощён – мы не генерируем несколько областей, а просто считаем весь объем признаков одним ROI. (В реальности, конечно, RPN дал бы несколько кандидатов. Мы можем вместо этого придумать простейший RPN: например, если максимальное значение активаций выше порога, берем весь фрагмент. Но для демонстрации проще.)
RoI Align: так как ROI покрывает весь feature map, нам даже не нужно кропать – но на всякий случай мы применяем F.interpolate к roi_feature чтобы задать фиксированный размер . Если бы ROI был меньше всего объема, следовало бы сначала вырезать подтензор features[:, :, z1:z2, y1:y2, x1:x2], а затем интерполировать его до . PyTorch не имеет прямой функции roi_align_3d, но можно воспользоваться grid_sample для произвольного ROI – это за рамками нашего примера.
Классификатор: разворачиваем roi_aligned в вектор и прогоняем через два полносвязных слоя. Получаем logits классов (например, [logit_background, logit_object]). В реальной Mask R-CNN здесь был бы softmax, и выбрали бы класс с argmax.
Маска: берем roi_aligned, применяем пару 3D-свёрток. Последний Conv3d выдаёт 1 канал, соответствующий маске текущего класса (в упрощении). Мы взяли 1, в реальности было бы num_classes каналов. Применяем $\sigma$ (сигмоиду) для получения вероятностей. Для пороговой маски можно сделать mask = (mask_probs>0.5).float().

Запустив этот код, мы бы получили, например, выход:

Logits классов: [[-0.12, 0.85]]
Выход маски shape: torch.Size([1, 1, 16, 16, 16])
Макс знач. в маске: 0.99

Это может означать, что модель уверенно (логит 0.85 против -0.12) выбрала класс “объект присутствует”, и маска содержит значения до 0.99, что близко к 1 внутри области объекта. Конечно, наша модель была не обучена – она выдала случайные по сути значения. Но если бы мы ее обучили на многих примерах, она научилась бы активировать маску там, где была сфера.

Код иллюстрирует структуру: сначала свёртки уменьшают размер, потом мы выбираем область, делаем её фиксированным размером, и параллельно классифицируем и сегментируем её. Настоящая Mask R-CNN в PyTorch имела бы гораздо больше кода: полный RPN, много ROI, цикл по ROI в голове, раздельные выходы по классам, и т.д. Однако PyTorch предоставляет высокоуровневое API: модуль torchvision.ops.MultiScaleRoIAlign (для 2D) и готовый класс torchvision.models.detection.MaskRCNN (2D), которые инкапсулируют эти детали. Для 3D пока нет встроенного Mask R-CNN, но можно адаптировать код 2D-версии (например, из Matterport Mask R-CNN на Keras) под 3D, как сделали некоторые open-source проекты.

Подготовка аннотаций. В случае обучения нам нужно было бы подготовить: (a) массив anchor-меток для RPN (какие anchors положительные/отрицательные), (b) для каждого положительного ROI – истинный класс и истинную маску (обычно вырезанную из полного объема по границам ROI и масштабированную к размеру маски сети). К счастью, современные библиотеки делают это автоматически в data loader’е. Формат данных для детекции в PyTorch – словарь, где для каждого изображения указаны boxes (список рамок) и masks (битовые карты) плюс labels. Для 3D пришлось бы написать свой Dataset, который выдает 3D-тензор и структуры аннотаций. Это выходит за рамки текущего урока, но важно помнить: правильно собрать данные не менее важно, чем написать код модели.

6. Оптимизация 3D-моделей

Работа с 3D-нейросетями предъявляет особые требования к оптимизации, поскольку объем данных и вычислений значительно выше, чем в 2D. Рассмотрим несколько аспектов оптимизации Mask R-CNN 3D и подобных моделей.

6.1 Воксельные данные vs облако точек

Тип представления 3D-данных напрямую влияет на эффективность. Воксельное представление (регулярная сетка) удобно тем, что позволяет применять стандартные 3D-свёртки, но потребляет огромное количество памяти: масштаб сетки растет кубически. Например, объём (типичный размер MRI) – это уже 16 миллионов вокселей * на число каналов. Если каждый представлен float32, то один объем – 64 МБ, а feature maps внутри сети могут быть ещё больше (например, 32 каналов того же размера – ~2 ГБ). Поэтому часто прибегают к:

Ограничению области обработки. Если интересует не весь объём сразу, а небольшой объём (Region of Interest), то модели обучают на небольших кубиках (patches) с агрегацией результатов. Например, в медицинских задачах иногда применяют sliding window: прогоняют U-Net на пересекающихся суб-объемах 64³ или 128³, а затем сшивают предсказания.
Разреженное представление. Во многих 3D-донах (например, Лидар облака точек, или детекторы частиц) данные очень разрежены – большая часть пространства пустая. Возникает идея хранить не полный плотный тензор, а координаты непустых вокселей. Существуют фреймворки, реализующие Sparse Convolution – свёртки по разреженным тензорам, которые обходят только непустые позиции. Пример – библиотека Minkowski Engine (Choy et al., 2019) и TorchSparse. Они позволяют эффективно применять 3D CNN к облакам точек без вокселизации в плотную сетку. В контексте Mask R-CNN 3D это могло бы означать: использовать SparseConv backbone, который извлекает признаки только там, где есть точки, и RPN, оперирующий с разреженными признаками. Такой подход экономит память и ускоряет вычисление при высокой разреженности (например, городские Lidar-облака, где плотность точек невелика).

Облако точек само по себе можно обрабатывать и без вокселей, методами типа PointNet/PointNet++: сеть сразу учится признакам на неструктурированных точках. Есть версии instance segmentation для облаков точек – например, PointRCNN и PointMask – они генерируют предложения 3D-рамок прямо из необработанного облака с помощью точечных MLP и кластеризации. Такие подходы менее универсальны (сложнее встроить в них mask-head как в Mask R-CNN, хотя исследования идут в этом направлении).

В целом, выбор между вокселями и точками зависит от задачи: для медицинских объемов удобнее воксели (данные заданы на регулярной решётке – томография), для сцен Lidar – облака точек и разреженные свёртки.

6.2 Оптимизация инференса (квантование, разреженность, GPU/TPU)

При внедрении модели Mask R-CNN 3D в продакшн важна скорость работы и требование к памяти. Есть несколько техник оптимизации:

Квантование (quantization). Перевод весов и активаций модели из 32-битных float в 8-битные целые (INT8) позволяет ускорить вычисления на CPU и снизить потребление памяти в 4 раза. Фреймворки (PyTorch, TensorFlow) поддерживают пост-тренировочное квантование. В 3D свёртках это столь же применимо, как и в 2D – операции те же, просто другие размеры тензоров. Однако, квантование может слегка снизить точность, особенно для масочной части, требующей высокой точности границ. Тем не менее, в приложениях, где нужна высокая производительность (например, подсчет объектов в потоке 3D-данных), INT8-квантование – мощный инструмент.
Аппаратное ускорение на GPU/TPU. Современные GPU (например, NVIDIA Ampere) имеют тензорные ядра, оптимизированные под matmul и conv в смешанной точности (Mixed Precision, float16). Обучение 3D моделей почти всегда ведется в mixed precision (через torch.cuda.amp), что ускоряет и уменьшает потребление памяти. На этапе инференса можно смело использовать float16 для свёрток – модель Mask R-CNN 3D, как правило, не страдает от этого в заметной степени, но скорость растет (особенно на Tensor Core). TPU и другие ускорители (IPU, Habana) также поддерживают 3D-свёртки – тут всё зависит от того, портирована ли модель на соответствующую платформу. В Google Cloud TPU, например, реализовывали 3D U-Net для сегментации мозга, добиваясь значительного ускорения параллелизмом (TPU v3 имеет 8 ядeр по 16 GB, что полезно для больших объемов).
Параллелизм и пакетная обработка. Если память позволяет, обработка нескольких 3D изображений пачкой (batch) может лучше загружать GPU. Но часто 3D данные настолько крупные, что batch=1 – потолок. Тогда прибегают к model parallel или разделению модели между несколькими GPU (например, backbone на одном, head на другом). Для Mask R-CNN 3D можно, например, считать признаки на одном GPU, а head (которая итеративно обрабатывает RoI) – на другом, особенно если RoI много. Однако, синхронизация может снизить выигрыш.
Пропуск незначимых операций. В 3D данных часто много пустого фона. Можно внедрить логику: если на каком-то этапе признаков активации в регионе очень малы (почти нули), можно не пускать его дальше для mask prediction. Или адаптивно изменять размер ROI: не всегда нужен – для мелких объектов хватит . Такие условные сокращения труднее реализовать в параллельных вычислениях, но могут ускорить, если объекты разнокалиберные.

6.3 Новые возможности фреймворков (PyTorch, TensorFlow) для 3D

Поддержка 3D-моделей в популярных фреймворках постоянно улучшается:

PyTorch: Начиная с версии ~1.10, появилась более удобная поддержка автокградов на больших тензорах и оптимизации памяти. Появились библиотеки TorchSparse, MinkowskiEngine (неофициальные), облегчающие работу с разреженными свёртками. В PyTorch 2.0 компилятор torch.compile может оптимизировать вычисл grafo как для 2D, так и 3D, убирая накладные расходы Python – это поможет ускорить голову Mask R-CNN, где много мелких операций (например, расчёт NMS, индексов) за счет фьюзинга. Кроме того, TorchVision пока не имеет 3D Mask R-CNN, но сообщество разрабатывает аналогичные компоненты. Возможно, в будущем появятся готовые torchvision.models.detection.MaskRCNN3D или, как минимум, интеграция с MinkowskiEngine для бэкбонов.
TensorFlow: В TF также есть Conv3D, Pool3D и др. Уже давно существует пример реализации 3D U-Net на Keras. Для Mask R-CNN 3D специфика – операции RoI Align 3D и NMS 3D. В TensorFlow можно писать custom op на CUDA, что сделали авторы одной из работ . TensorFlow 2.x с Keras позволяет комбинировать слои довольно свободно, так что энтузиасты тоже переносили Mask R-CNN logic на 3D. Что касается ускорения: XLA-компилятор оптимизирует граф TF, в том числе может распараллелить 3D свёртки по нескольким ядрам. Есть проект TensorFlow 3D – дополнение от Google Research, включающее удобные инструменты для работы с точками и 3D данными (например, функции для построения октагонов, операции с облаками точек и пр.). Он интегрируется с Keras, облегчая построение end-to-end моделей с 3D входом.
ONNX и TensorRT: Для развёртывания на GPU-серверах, 3D-модель можно экспортировать в ONNX и оптимизировать через TensorRT. TensorRT хорошо поддерживает Conv3D, поэтому можно получить прирост производительности по сравнению с исходным фреймворком (за счёт низкоуровневых оптимизаций, фьюзинга слоёв). Правда, нестандартные куски (например, наша реализация NMS или RoIAlign) придётся либо имплементировать плагинами, либо упрощать (например, заменить RoIAlign на несколько шагов ресайза, которые ONNX/TensorRT поймёт).

Актуальные исследования: В последние годы появилась тенденция универсальных моделей, которые могут работать сразу на разных модальностях, включая 3D. Например, так называемые One Transformer to Rule Them All – модели, принимающие и изображение, и точечные данные, и текст. Пока что это эксперименты, но можно представить, что со временем фреймворки добавят high-level слои для работы с несоответствующими размерностями. Уже сейчас PyTorch3D (библиотека Facebook) предоставляет удобные функции для работы с 3D геометрией, но больше в контексте графики (рендеринг, поинт клауд). Возможно, в будущем появятся готовые сегментационные трансформеры для 3D, которые заменят классическую Conv3D архитектуру. Но на данный момент Mask R-CNN 3D и 3D U-Net остаются эталонными архитектурами для объемной сегментации, а их оптимизация – вопрос правильного баланса между точностью и эффективностью.

И напоследок: оптимизируя 3D-модель, важно не переусердствовать. Как гласит шутка, “профилировать нужно не только код, но и свои ожидания”.

Иными словами, чудес не бывает: 3D-задачи требуют больших вычислений, и даже самый хитрый трюк не сделает их тривиальными. Тем не менее, грамотная оптимизация сможет сэкономить вам недели вычислительного времени – а это уже немало!