Облегченные модели для многозадачного вывода в реальном времени

Введение

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

В этом посте мы узнаем, как обучить одну модель для одновременного выполнения задач классификации и регрессии. Код для этого поста можно найти на GitHub. Вот обзор:

Мотивация

Зачем нам использовать облегченную модель? Не снизит ли это производительность? Если мы не развертываем на периферии, не должны ли мы использовать как можно большую модель?

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

Легкие модели могут помочь снизить затраты и снизить выбросы углерода за счет меньшего энергопотребления.

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

Подход

Чтобы построить нашу многозадачную архитектуру, мы в общих чертах рассмотрим подход из этой статьи, где одна модель была обучена для одновременной сегментации и оценки глубины (и мы сделаем то же самое здесь). Основная цель состояла в том, чтобы выполнять эти задачи быстро и эффективно, с приемлемой потерей производительности. В многозадачном обучении мы обычно группируем похожие задачи вместе. Во время обучения мы также можем добавить вспомогательную задачу, которая может помочь обучению нашей модели, но мы можем не использовать ее во время вывода [1, 2]. Для простоты мы не будем использовать какие-либо вспомогательные задания при обучении.

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

Набор данных

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

Подготовка карты глубины — несоответствие по умолчанию

Для данных о глубине мы рассмотрели использование предоставленных карт несоответствия SteroSGBM, но они зашумлены множеством отверстий, соответствующих бесконечной глубине, и частью, где всегда отображается эго-транспортное средство. Общий подход к шумоподавлению этих карт несоответствий включает:

  1. Обрежьте нижние 20% вместе с частями левого и верхнего краев.
  2. Изменить размер до исходного масштаба
  3. Применение сглаживающего фильтра
  4. Выполнить рисование

Что приводит к:

Эти тонкие детали этого подхода выходят за рамки этого поста, но если вам интересно, вот видео объяснение.

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

Подготовка карты глубины — несоответствие CreStereo

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

С любой картой несоответствия CityScapes мы можем вычислить карту глубины с помощью:

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

Теперь, когда у нас есть данные, давайте посмотрим на архитектуру.

Архитектура модели

После [1] архитектура будет состоять из магистрали/кодировщика MobileNet, декодера LightWeight RefineNet и головок для каждой отдельной задачи. Общая архитектура показана ниже на рисунке 3.

Для кодера/магистрали мы будем использовать MobileNetV3 и пропускать соединения с разрешениями 1/4, 1/8, 1/16 и 1/32 в облегченную сеть уточнения. Наконец, выходные данные передаются каждой головке, отвечающей за разные задачи. Обратите внимание, как мы можем даже добавить больше задач в эту архитектуру, если захотим.

Для реализации кодировщика мы используем предварительно обученный кодировщик MobileNetV3, где мы передадим кодировщик MobileNetV3 в пользовательский модуль PyTorch. Результатом его прямой функции является ParameterDict пропусков соединений для ввода в LightWeight Refine Net. Фрагмент кода ниже показывает, как это сделать.

class MobileNetV3Backbone(nn.Module):
    def __init__(self, backbone):
        super().__init__()
        self.backbone = backbone
    
    def forward(self, x):
        """ Passes input theough MobileNetV3 backbone feature extraction layers
            layers to add connections to
                - 1:  1/4 res
                - 3:  1/8 res
                - 7, 8:  1/16 res
                - 10, 11: 1/32 res
           """
        skips = nn.ParameterDict()
        for i in range(len(self.backbone) - 1):
            x = self.backbone[i](x)
            # add skip connection outputs
            if i in [1, 3, 7, 8, 10, 11]:
                skips.update({f"l{i}_out" : x})

        return skips

Декодер LightWeight RefineNet очень похож на реализованный в [1], за исключением нескольких модификаций, чтобы сделать его совместимым с MobileNetV3, а не с MobileNetV2. Мы также отмечаем, что часть декодера состоит из головок Segmentation и Depth. Полный код модели доступен здесь. Мы можем собрать модель следующим образом:

from torchvision.models import mobilenet_v3_small
    
mobilenet = mobilenet_v3_small(weights='IMAGENET1K_V1')

encoder = MobileNetV3Backbone(mobilenet.features)
decoder = LightWeightRefineNet(num_seg_classes)
model = MultiTaskNetwork(encoder, freeze_encoder=False).to(device)

Подход к обучению

Мы разделяем обучение на три этапа: первый с разрешением 1/4, второй с разрешением 1/2 и последний с полным разрешением. Все веса были обновлены, так как замораживание весов энкодера, похоже, не дало хороших результатов.

Преобразования

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

Преобразование глубины

Как правило, карты глубины содержат в основном меньшие значения, поскольку большая часть информации, содержащейся в карте глубины, состоит из объектов и поверхностей, близких к камере. Поскольку большая часть карты глубины сосредоточена вокруг более низких значений (см. слева на рис. 4 ниже), ее необходимо преобразовать, чтобы нейронная сеть могла эффективно изучить ее. Карта глубины обрезается между 0 и 250, потому что данные о стереодиспаритете/глубине на больших расстояниях обычно ненадежны, и в этом случае нам нужен способ отбросить их. Затем мы берем натуральный логарифм и делим его на 5, чтобы сжать распределение вокруг меньшего диапазона чисел. Подробности смотрите в этой блокноте.

Честно говоря, я не был уверен, как лучше преобразовать данные о глубине. Если есть лучший способ или если бы вы сделали это по-другому, мне было бы интересно узнать больше в комментариях :).

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

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

Скорость обучения

Мы используем скорость обучения с косинусным отжигом за один цикл с максимальным значением 5e-4 и тренируемся в течение 150 эпох с разрешением 1/4. Блокнот, используемый для обучения, находится здесь.

Затем мы настраиваемся с разрешением 1/2 на 25 эпох и снова на полное разрешение еще на 25 эпох со скоростью обучения 5e-6. Обратите внимание, что нам нужно было уменьшать размер пакета каждый раз, когда мы настраивали повышенное разрешение. Помимо настройки скорости обучения для нескольких пробных прогонов, никакие другие гиперпараметры обучения не настраивались.

Вывод

Для вывода мы нормализовали входное изображение и выполнили прямой проход по модели. На рис. 6 показаны результаты обучения как по проверочным, так и по тестовым данным.

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

Интересный провал

В нижней части рисунка 6 показан интересный случай отказа полностью сегментировать фонарный столб в левой части изображения. Сегментация охватывает только нижнюю половину фонарного столба, в то время как глубина показывает, что нижняя половина фонарного столба намного ближе, чем верхняя половина. Ошибка глубины может быть вызвана смещением нижних пикселей, обычно соответствующих меньшей глубине; обратите внимание на линию горизонта вокруг пикселя 500, есть четкое разделение между более близкими пикселями и дальними пикселями. Похоже, что эта предвзятость могла просочиться в задачу сегментации модели. Этот тип утечки задач следует учитывать при обучении многозадачных моделей.

Учебные данные одной задачи могут повлиять на производительность другой задачи.

Распределение глубины

Давайте проверим, как распределена предсказанная глубина по сравнению с правдой. Для простоты мы просто будем использовать выборку из 94 пар карт глубины полного разрешения истинного/прогнозированного.

Похоже, что модель изучила два распределения: распределение с пиком около 4 и распределение с пиком около 30. Обратите внимание, что артефакт отсечения, похоже, не имеет значения. Общее распределение содержит длинный хвост, который характерен тем фактом, что только небольшая часть изображения будет содержать данные о далекой глубине.

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

Бонус: скорость вывода

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

# find optimal backend for performing convolutions 
torch.backends.cudnn.benchmark = True 

# rescale to half size
rescaled_sample = Rescale(400, 1024)(sample)
rescaled_left = rescaled_sample['left'].to(DEVICE)

# INIT LOGGERS
starter, ender = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
repetitions = 300
timings=np.zeros((repetitions,1))
#GPU-WARM-UP
for _ in range(10):
    _, _ = model(rescaled_left.unsqueeze(0))
# MEASURE PERFORMANCE
with torch.no_grad():
    for rep in range(repetitions):
        starter.record()
        _, _ = model(rescaled_left.unsqueeze(0))
        ender.record()
        # WAIT FOR GPU SYNC
        torch.cuda.synchronize()
        curr_time = starter.elapsed_time(ender)
        timings[rep] = curr_time

mean_syn = np.sum(timings) / repetitions
std_syn = np.std(timings)
print(mean_syn, std_syn)

Тест на вывод показывает, что эта модель может работать со скоростью 18,69+/-0,44 мс или около 55 Гц. Важно отметить, что это всего лишь прототип Python, запущенный на ноутбуке с графическим процессором NVIDIA RTX 3060, скорость вывода будет меняться на другом оборудовании. Следует также отметить, что такой SDK, как Torch-TensorRt, может обеспечить значительное ускорение при развертывании на графическом процессоре NVIDIA.

Заключение

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

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

Рекомендации

[1] Некрасов, Владимир, и др. Совместная семантическая сегментация в реальном времени и оценка глубины с использованием асимметричных аннотаций. CoRR, том. abs/1809.04766, 2018, http://arxiv.org/abs/1809.04766.

[2] Стэндли, Тревор и др. Какие задачи следует изучать вместе при многозадачном обучении? CoRR, vol. abs/1905.07553, 2019, http://arxiv.org/abs/1905.07553.