Модель U-net и её варианты

Базовая модель U-net

Архитектура U-net

Модель U-net [1] осуществляет семантическую сегментацию, используя свёрточный кодировщик и декодировщик.

Кодировщик постепенно сжимает пространственное разрешение, применяя свёртки 3x3 без расширения (padding), а также используя пулинги 2x2 с шагом 2. Сжатие разрешения компенсируется увеличением числа слоёв после каждого пулинга.

Декодировщик постепенно увеличивает пространственное разрешение, используя операции повышения пространственного разрешения (upsampling) с одновременным уменьшением числа каналов. В декодировщике также применяются свёртки 3x3.

Проблема недостаточно точного восстановления границ из низкоразмерного промежуточного представления кодировщика решается тем, что внутренние представления кодировщика с более высоким пространственным разрешением передаются на соответствующие слои декодировщика, как показано серыми линиями на схеме [1]:

Объединение информации из декодировщика и кодировщика происходит путём конкатенации (объединения) внутренних представлений вдоль каналов.

Выходом U-net является тензор $\hat{Y}\in\mathbb{R}^{C\times H\times W}$, где

• $H,W$ - высота и ширина сегментируемого изображения;

• $C$ - число классов, включая фоновый (на схеме $C=2$).

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

Поскольку в сети многократно применяются свёртки 3x3 без расширения, то пространственная размерность постепенно снижается на границах. Поэтому при переносе внутренних представлений с кодировщика, они обрезаются по краям, чтобы конкатенировались тензоры одинакового пространственного размера.

Для того, чтобы после всех свёрток на выходе получить ту же пространственную размерность, которой обладало сегментируемое изображение, входное изображение перед обработкой расширяется с помощью отражения пикселей по краям (mirror padding), как показано на рисунке [1]:

Настройка сети U-net

Сеть настраивается, используя принцип максимума правдоподобия, т.е. максимизируя взвешенное правдоподобие отнесения каждого пикселя $x_{ij}$ к верному классу:

$$\sum_i\sum_j w_{ij}\log p(\hat{c}_{ij})\to \max_w,$$

Вес учёта каждого пикселя $(i,j)$ считался по формуле

$$w_{ij} = w(c_{ij})+\alpha e^{-\frac{d_1(i,j)+d_2(i,j)}{2\sigma^2}},$$

где

• $\alpha>0, \sigma>0$ - гиперпараметры,

• $w(c_{ij})$ - вес класса (выше для более редких, чтобы классы вносили сопоставимый вклад в оптимизацию),

• $d_1(i,j),d_2(i,j)$ - минимальные расстояние до ближайшей границы с областью другого класса и ближайшее расстояние до границы ещё одного другого класса.

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

Ниже на графике (d) показана пространственная карта весов для сегментируемого изображения (a), корректной сегментации (b) и сегментационной карты (с) [1]:

Архитектура U-net задала стандарт переноса промежуточных представлений кодировщика в декодировщик и сейчас этот принцип используется во многих задачах, где по входному изображению нужно сгенерировать выходное (image-to-image).

Варианты архитектуры U-net

Рассмотрим варианты изменения архитектуры U-net для семантической сегментации.

LinkNet

Модель LinkNet [2] построена из принципа упрощения вычислений и настройки модели.

Архитектура LinkNet приведена ниже [2]:

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

Детализация каждого блока кодировщика и декодировщика приведена ниже [2]:

Здесь

• conv [(AxA),(B,C)] обозначает свёрточный слой с ядром свёртки AxA, переводящий внутреннее представление из B в С каналов.

• /2 означает уменьшение разрешения в 2 раза за счёт свёртки с шагом 2,

• *2 означает увеличение разрешения в 2 раза.

Вычислительная эффективность LinkNet обеспечивается тем, что

• промежуточные представления кодировщика не конкатенируются, а суммируются с промежуточными представлениями кодировщика;

• в блоке декодировщика перед применением свёртки 3x3 действует слой свёрток 1x1, снижающий число каналов в 4 раза. После действия свёрток 3x3 число каналов возвращается к исходному также свёртками 1x1.

Для упрощения настройки сети в каждом блоке кодировщике используются ResNet блоки. Тождественные связи этих блоков (с понижением разрешения) упрощают

• перетекание информации об исходном изображении через кодировщик,

• распространение градиента на ранние слои при настройке сети.

One Hundred Layers Tiramisu

One Hundred Layers Tiramisu [3] также основана на идее U-net, но использует dense-блоки как базовые элементы.

Её архитектура представлена ниже [3]:

• Блок Transition Down снижает разрешение в 2 раза, используя максимизирующий пулинг.

• Блок Transition Up повышает разрешение в 2 раза, используя транспонированную свёртку.

Серые связи обозначают перенос промежуточных представлений кодировщика (слева) в декодировщик (справа). Объединение представлений производится не суммированием, как в LinkNet, а конкатенацией вдоль каналов, как в U-net.

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

За счёт конкатенаций внутренних представлений в dense-блоках и при переносе информации из кодировщика в декодировщик, сети приходится обрабатывать повышенное число каналов.

Поэтому сеть One Hundred Layers Tiramisu более требовательна к вычислениям, чем LinkNet, зато теоретически может обеспечить более высокую точность сегментации.

U-net++

Встроенная проблема методологии U-net заключена в том, что более низкоуровневые и простые признаки объединяются с более высокоуровневыми и сложными.

Для решения этой проблемы в архитектуре U-net++ [4] предложено вместо непосредственной конкатенацией внутренних представлений кодировщика и декодировщика конкатенировать преобразованные представления кодировщика через dense-блок, как показано ниже [4]:

Dense-блок призван приводить в семантическое соответствие более простые представления кодировщика с более сложными представлениями декодировщика.

Разница в семантической сложности выше для более высоких ярусов сети на схеме, поэтому там dense-блоки глубже. В этих блоках действует свёртка с нелинейностью, принимающая на вход конкатенацию всех предыдущих слоёв блока.

Архитектура U-net++ показала более высокое качество по сравнению с U-net [1].

Настройка сети велась, минимизируя сумму кросс-энтропийных потерь и коэффициента dice, причём потери считались по выходам всех верхних ярусов сети $X^{0,1},X^{0,2},X^{0,3},X^{0,4}$ (принцип deep supervision). На выходе это позволяет считать прогнозы не только по $X^{0,4}$, но и по урезанным версиям U-net++, выдающим прогнозы на слоях $X^{0,3},X^{0,2}$ и $X^{0,1}$. Урезанные (prunned) версии сети показаны ниже [4]:

Более короткие версии U-net++ позволяют производить семантическую сегментацию менее точно, зато более быстро.

Литература

1. Ronneberger O., Fischer P., Brox T. U-net: Convolutional networks for biomedical image segmentation //Medical image computing and computer-assisted intervention–MICCAI 2015: 18th international conference, Munich, Germany, October 5-9, 2015, proceedings, part III 18. – Springer International Publishing, 2015. – С. 234-241.
2. Chaurasia A., Culurciello E. Linknet: Exploiting encoder representations for efficient semantic segmentation //2017 IEEE visual communications and image processing (VCIP). – IEEE, 2017. – С. 1-4.
3. Jégou S. et al. The one hundred layers tiramisu: Fully convolutional densenets for semantic segmentation //Proceedings of the IEEE conference on computer vision and pattern recognition workshops. – 2017. – С. 11-19.
4. Zhou Z. et al. Unet++: A nested u-net architecture for medical image segmentation //Deep Learning in Medical Image Analysis and Multimodal Learning for Clinical Decision Support: 4th International Workshop, DLMIA 2018, and 8th International Workshop, ML-CDS 2018, Held in Conjunction with MICCAI 2018, Granada, Spain, September 20, 2018, Proceedings 4. – Springer International Publishing, 2018. – С. 3-11.

Литература