U-Net стал популярным методом сегментации изображений. Но как это произошло?

Оглавление
1. Поставленная задача
2. Кодировщик-декодер
3. Пропустить соединения
4. Детали реализации
а. Функция потерь
б. Методы повышения частоты дискретизации
c. Прокладывать или не прокладывать?
5. U-Net в действии

Задача под рукой

U-Net разработан для задачи семантической сегментации. Когда нейронная сеть получает изображения в качестве входных данных, мы можем выбрать классификацию объектов либо в целом, либо по экземплярам. Мы можем предсказать, какой объект находится на изображении (классификация изображения), где расположены все объекты (локализация изображения/семантическая сегментация) или где расположены отдельные объекты (обнаружение объекта/сегментация экземпляра). На рисунке ниже показаны различия между этими задачами компьютерного зрения. Чтобы упростить дело, мы рассматриваем классификацию только для одного класса и одной метки (бинарная классификация).

В задаче классификации мы выводим вектор размера k, где k — количество классов. В задачах обнаружения нам нужно вывести вектор x, y, height, width, class, который определяет ограничивающие рамки. Но в задачах сегментации нам нужно вывести изображение того же размера, что и исходный ввод. Это представляет собой довольно сложную инженерную задачу: как нейронная сеть может извлечь соответствующие функции из входного изображения, а затем спроецировать их в маски сегментации?

Кодер-декодер

Если вы не знакомы с кодировщиком-декодером, рекомендую прочитать эту статью:



Понимание скрытого пространства в машинном обучении
Изучите фундаментальную, но часто «скрытую концепцию глубокого обученияtowardsdatascience.com»



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

В части кодировщика я использовал сверточные слои, а затем ReLU и MaxPool в качестве экстракторов признаков. В части декодера я транспонировал свертки, чтобы увеличить размер карты признаков и уменьшить количество каналов. Я использовал заполнение, чтобы сохранить размер карт объектов одинаковым после операций свертки.

Одна вещь, которую вы можете заметить, это то, что в отличие от сетей классификации, эта сеть не имеет полностью связанного/линейного слоя. Это пример полностью сверточной сети (FCN). Было показано, что FCN хорошо справляется с задачами сегментации, начиная с Shelhamer et al. статья «Полностью сверточные сети для семантической сегментации» [1].

Однако у этой сети есть проблема. По мере того, как мы увеличиваем количество слоев в слоях кодировщика и декодера, мы эффективно «сжимаем» карту объектов все больше и больше. Таким образом, кодировщик может отказаться от более подробных функций в пользу более общих функций. Если мы имеем дело с сегментацией медицинских изображений, каждый пиксель, классифицированный как больной/нормальный, может быть важен. Как мы можем убедиться, что эта сеть кодировщик-декодер использует как общие, так и подробные функции?

Пропустить соединения



Введение в ResNets
Эта статья основана на глубоком остаточном обучении распознаванию изображений от He et al. [2] (Microsoft Research)…towardsdatascience.com



Поскольку глубокие нейронные сети могут «забывать» определенные функции при передаче информации через последовательные слои, пропуск соединений может повторно вводить их, чтобы сделать обучение более сильным. Пропустить соединение было введено в Residual Network (ResNet) и показало улучшения классификации, а также более плавные градиенты обучения. Вдохновленные этим механизмом, мы можем добавить пропущенные подключения к U-Net, чтобы каждый декодер включал карту функций из соответствующего кодировщика. Это определяющая особенность U-Net.

U-Net имеет два определяющих качества:

  1. Сеть кодер-декодер, которая извлекает больше общих признаков, чем глубже она идет.
  2. Пропускное соединение, которое повторно вводит подробные функции в декодер.

Эти два качества означают, что U-Net может сегментировать, используя детализированные и общие характеристики. Первоначально U-Net была введена для обработки биомедицинских изображений, где очень важна точность сегментации [2].

Детали реализации

В предыдущих разделах был дан очень общий обзор U-Net и того, почему он работает. Однако детали стоят между общим пониманием и фактической реализацией. Здесь я дал обзор некоторых вариантов реализации U-Net.

Функции потерь

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

В оригинальной статье U-Net к функции потерь добавляется дополнительный вес. Этот параметр веса выполняет две функции: компенсирует дисбаланс классов и придает большее значение границам сегментации. Во многих реализациях U-Net, которые я нашел в Интернете, этот дополнительный весовой коэффициент используется нечасто.

Другой часто встречающейся функцией потерь является проигрыш в кубиках. Потеря в кости измеряет, насколько похожи два набора изображений, сравнивая площадь их пересечения с их общей площадью. Обратите внимание, что потеря кубиков — это не то же самое, что Intersection-over-Union (IOU). Они измеряют схожие вещи, но имеют разный знаменатель. Чем выше коэффициент кубика, тем меньше потеря кубиков.

Здесь добавляется член эпсилон, чтобы избежать деления на 0 (обычно это эпсилон 1). В некоторых реализациях, например в Milletari et al., значения пикселей в знаменателе возводятся в квадрат перед их суммированием [3]. По сравнению с кросс-энтропийной потерей, потеря кубиков очень устойчива к несбалансированной маске сегментации, что типично для задач сегментации биомедицинских изображений.

Методы повышения частоты дискретизации

Еще одной деталью является выбор метода повышения частоты дискретизации для декодера. Вот несколько распространенных методов:

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

Максимальное удаление. Этот метод противоположен Max-pooling. Он использует индексы операции maxpool и заполняет эти индексы максимальным значением. Все остальные значения установлены на 0. Как правило, слой свертки следует за max-unpooling, чтобы «сгладить» все пропущенные значения.

Деконволюция/транспонированная свертка. О деконволюции написано много сообщений в блогах. Рекомендую эту статью как хороший наглядный путеводитель.



Введение в различные типы сверток в глубоком обучении
Позвольте мне дать вам краткий обзор различных типов сверток и их преимуществ. Ради…towardsdatascience.com



Деконволюция состоит из двух этапов: добавьте отступы к каждому пикселю исходного изображения, затем примените свертку. В оригинальной U-Net транспонированная свертка 2x2 с шагом 2 используется для изменения как пространственного разрешения, так и глубины канала.

Перемешивание пикселей. Этот метод использовался в сетях сверхвысокого разрешения, таких как SRGAN. Для начала мы используем свертку, чтобы перейти от C x H x W карты объектов к (Cr^2) x H x W. Затем перетасовка пикселей возьмет это и «перемешает» пиксели мозаичным образом, чтобы получить результат размером C x (Hr) x (Wr).

Прокладывать или не прокладывать?

Слой свертки с ядром больше 1x1 и без заполнения будет давать выходные данные, которые меньше, чем входные данные. Это проблема для U-Net. Напомним, что на рисунке U-Net в предыдущем разделе мы объединяем часть изображения с его декодированным аналогом. Если мы не используем заполнение, то декодированное изображение будет иметь меньший пространственный размер по сравнению с закодированным изображением.

Однако в оригинальной статье U-Net отступы не использовались. Несмотря на то, что никаких обоснований не было дано, я полагаю, что это произошло потому, что авторы не хотели вводить ошибки сегментации на полях изображения. Вместо этого они обрезали закодированное изображение по центру перед конкатенацией. Для изображения с входным размером 572 x 572 на выходе будет 388 x 388, потеря ~50%. Если вы хотите запустить U-Net без заполнения, вам нужно запустить его несколько раз на перекрывающихся плитках, чтобы получить полное изображение сегментации.

U-Net в действии

Здесь я реализовал очень простую сеть, подобную U-Net, для сегментации только эллипсов. U-Net имеет глубину всего 3 слоя, использует такое же заполнение и бинарную перекрестную энтропийную потерю. Более сложные сети могут использовать больше слоев свертки при каждом разрешении или увеличивать глубину по своему усмотрению.

import torch
import numpy as np
import torch.nn as nn

class EncoderBlock(nn.Module):        
    # Consists of Conv -> ReLU -> MaxPool
    def __init__(self, in_chans, out_chans, layers=2, sampling_factor=2, padding="same"):
        super().__init__()
        self.encoder = nn.ModuleList()
        self.encoder.append(nn.Conv2d(in_chans, out_chans, 3, 1, padding=padding))
        self.encoder.append(nn.ReLU())
        for _ in range(layers-1):
            self.encoder.append(nn.Conv2d(out_chans, out_chans, 3, 1, padding=padding))
            self.encoder.append(nn.ReLU())
        self.mp = nn.MaxPool2d(sampling_factor)
    def forward(self, x):
        for enc in self.encoder:
            x = enc(x)
        mp_out = self.mp(x)
        return mp_out, x

class DecoderBlock(nn.Module):
    # Consists of 2x2 transposed convolution -> Conv -> relu
    def __init__(self, in_chans, out_chans, layers=2, skip_connection=True, sampling_factor=2, padding="same"):
        super().__init__()
        skip_factor = 1 if skip_connection else 2
        self.decoder = nn.ModuleList()
        self.tconv = nn.ConvTranspose2d(in_chans, in_chans//2, sampling_factor, sampling_factor)

        self.decoder.append(nn.Conv2d(in_chans//skip_factor, out_chans, 3, 1, padding=padding))
        self.decoder.append(nn.ReLU())

        for _ in range(layers-1):
            self.decoder.append(nn.Conv2d(out_chans, out_chans, 3, 1, padding=padding))
            self.decoder.append(nn.ReLU())

        self.skip_connection = skip_connection
        self.padding = padding
    def forward(self, x, enc_features=None):
        x = self.tconv(x)
        if self.skip_connection:
            if self.padding != "same":
                # Crop the enc_features to the same size as input
                w = x.size(-1)
                c = (enc_features.size(-1) - w) // 2
                enc_features = enc_features[:,:,c:c+w,c:c+w]
            x = torch.cat((enc_features, x), dim=1)
        for dec in self.decoder:
            x = dec(x)
        return x

class UNet(nn.Module):
    def __init__(self, nclass=1, in_chans=1, depth=5, layers=2, sampling_factor=2, skip_connection=True, padding="same"):
        super().__init__()
        self.encoder = nn.ModuleList()
        self.decoder = nn.ModuleList()

        out_chans = 64
        for _ in range(depth):
            self.encoder.append(EncoderBlock(in_chans, out_chans, layers, sampling_factor, padding))
            in_chans, out_chans = out_chans, out_chans*2

        out_chans = in_chans // 2
        for _ in range(depth-1):
            self.decoder.append(DecoderBlock(in_chans, out_chans, layers, skip_connection, sampling_factor, padding))
            in_chans, out_chans = out_chans, out_chans//2
        # Add a 1x1 convolution to produce final classes
        self.logits = nn.Conv2d(in_chans, nclass, 1, 1)

    def forward(self, x):
        encoded = []
        for enc in self.encoder:
            x, enc_output = enc(x)
            encoded.append(enc_output)
        x = encoded.pop()
        for dec in self.decoder:
            enc_output = encoded.pop()
            x = dec(x, enc_output)

        # Return the logits
        return self.logits(x)

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

Заключение

Если бы мне нужно было объяснить U-Net в одном предложении, я бы сказал, что U-Net похож на кодер-декодер для изображений, но с пропуском соединений, чтобы не потерять мелкие детали. U-Net часто используются во многих задачах сегментации, а в последние годы также стали применяться в задачах создания изображений.

Если вы хотите увидеть код, который я использовал для создания рисунков и обучения своего U-Net, вот ссылка на Github. Удачного кодирования!

Если вам нравится читать эту статью и вы хотели бы прочитать больше подобных в будущем, рассмотрите возможность подписаться на меня в Medium или Linkedin.

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

[1] Лонг, Джонатан, Эван Шелхамер и Тревор Даррелл. «Полностью сверточные сети для семантической сегментации». Материалы конференции IEEE по компьютерному зрению и распознаванию образов. 2015.

[2] Роннебергер, Олаф, Филипп Фишер и Томас Брокс. «U-net: сверточные сети для сегментации биомедицинских изображений». Международная конференция по обработке медицинских изображений и компьютерным вмешательствам. Спрингер, Чам, 2015.

[3] Миллетари, Фаусто, Насир Наваб и Сейед-Ахмад Ахмади. «V-net: полностью сверточные нейронные сети для объемной сегментации медицинских изображений». Четвертая международная конференция 2016 года по 3D-зрению (3DV). ИИЭР, 2016.