Привет, хабровчане! Меня зовут Дмитрий, я студент первого курса Вышки магистерской программы «Системный анализ и математические технологии».

В октябре прошлого года я прошёл отбор в Инженерно-математическую школу (ИМШ). Это совместная образовательная программа VK Education и НИУ ВШЭ. Я стал участником мастерской по прикладному искусственному интеллекту. В одном из проектов мы командой из двух человек (да, нас было мало, но мы были сильны!) решили реальную технологическую задачу под руководством эксперта VK и научных сотрудников университета. Нам нужно было придумать технологию разделения аудиосигнала на источники. В команде я отвечал за подготовку данных, а мой коллега — за проведение экспериментов. Хочу рассказать вам подробно, как мы решали эту задачу — надеюсь, что это будет для вас полезно.

Кстати, в ИМШ есть и другие практикоориентированные мастерские: по аппаратному обеспечению искусственного интеллекта, по компиляторам и высоконагруженным приложениям, по виртуальным платформенным решениям «Робби» (в СПб). Если вас что-то из этого интересует и вы являетесь студентом Вышки, узнать об условиях поступления можно здесь. Следующий набор в мастерские ИМШ откроется осенью этого года, а в магистратуры школы поступить можно уже сейчас. Обучение в мастерских бесплатное. А еще: те, кто хорошо покажут себя в работе мастерских, получают преимущества при рассмотрении на вакансии VK.

Разделение аудиосигнала на источники — одна из многих задач ИИ в области обработки звука. К ней близка задача по созданию технологии шумоподавления, в рамках которой необходимо выделить один источник сигнала и максимально увеличить качество выделенного материала. Но фундаментально две эти задачи отличаются друг от друга.  В отличие от шумоподавления в результате разделения аудиосигнала на источники на выходе получается множество очищенных сигналов.

Из всех моделей искусственного интеллекта к решению этой задачи максимально приблизилась модель из семейства DEMUCS — HT Demucs (Hybrid Transformer Demucs). Она заняла первое место в соревновании “Music Demixing Challenge 2021”, которое провела компания Sony. 

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

Расскажем подробнее, что мы делали с моделькой. 

Представление звука

Для нашей задачи звук можно представить как зависимость интенсивности от времени для композиции волновых колебаний. Этот временной сигнал можно перевести в частотное представление для каждого (в непрерывном случае) временного окна. Такое преобразование называют оконным преобразованием Фурье (STFT).

Поскольку STFT представляет входной сигнал в комплексной форме, мы можем рассмотреть по отдельности амплитудное и фазовое представление преобразованного сигнала. Следующие представления мы можем представлять как тепловые карты, на которых по оси абсцисс будет отложено время, по оси ординат — частота, а амплитуда или фаза будет спроецирована на плоскость, кодируя значение цветом точки на карте. Полученные тепловые карты — это изображения, которые можно обрабатывать нейросетевыми методами компьютерного зрения. Ради них мы здесь сегодня и собрались.

Про модели DEMUCS

DEMUCS — это глубокая модель для выделения аудиосигналов, берущая за основу подход, используемый при сегментировании изображений, а именно модель U-Net. Архитектура U-Net состоит из кодировщика, декодировщика и промежуточного слоя между ними, который обычно именуется bottleneck. 

Основное свойство U-Net — наличие skip-соединений между слоями кодировщика и декодировщика, что позволяет на этапе декодирования не терять изначальную информацию. В итоге получается, что DEMUCS берёт подход для сегментирования изображений и применяет его для обработки волны аудиосигнала. В дальнейшем на основе базовой модели DEMUCS были разработаны Hybrid DEMUCS, Hybrid Transform DEMUCS и другие:

• Hybrid DEMUCS использует bi-U-Net (две параллельные U-net сети): одну для волны во временном представлении, а другую для амплитудной спектрограммы в частотном представлении.

• Hybrid Transform DEMUCS — развитие Hybrid DEMUCS. В ней слои кодировщика в U-net сетях заменяют на трансформеры “cross-domain”, которые в bi-U-Net связывают механизмом внимания две части: временную и частотную.

Измеряем качество

Возникает логичный вопрос: «Как можно судить о качестве модели, если единственное, что можно сделать, — это послушать выход?» Слух — хороший показатель качества, но, к сожалению, субъективный. По этой причине для оценки качества моделей по разделению звука существуют числовые метрики:

• SNR — Signal-to-Noise Ratio.

• SAR — Source-to-Artifact Ratio.

• SIR — Source-to-Interference Ratio.

• SDR — Source-to-Distortion Ratio.

SDR= 10log10‖starget‖2‖einterf+enoise+eartif‖2, dB

SDR — метрика, которая обобщает SNR, SAR и SIR, характеризуя общую зашумленность сигнала. Чем SDR больше, тем менее искажен сигнал. Значение SDR  ≥ 10 — удовлетворительно для человеческого уха.

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

Обучение модели

Хоть у нас и есть архитектура HDemucs, но без определенных дополнений никак не обойтись. Во-первых, можно составить нетривиальный функционал качества, который поможет модели лучше оценивать качество разделения, выделяя больше полезной информации из спектрограммы сигнала. Во-вторых, для обучения модели необходимо немало примеров для сходимости, поэтому даже крупные наборы данных желательно увеличить ещё больше. Для этого был бы полезен целый пайплайн, который позволит максимально накормить модель. В-третьих, это вычислительные мощности и оптимизация, а главное — те жертвы, на которые приходится идти, чтобы модель сходилась не вечность. Каждый из этих этапов далее я опишу подробнее.

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

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

• L1 Loss — классическая функция потерь, определяемая модулем разности предсказанного сигнала и эталонного. Для текущей задачи определена во временном домене.

• MultiSpec Loss — квадрат разницы между пикселями предсказанной и эталонной спектрограмм.

• SiSDR Loss — функция потерь, направленная на максимизацию метрики SDR (Source to Distortion Ratio). Считается во временном представлении сигнала. 

• Perceptual loss — используется для сравнения высокоуровневых различий таких, как несоответствие содержания и стиля между изображениями, в рамках задачи это относится к спектрограммам. Мы применили обученную VGG на слоях, которой считаем различия в паттернах.

Каждая из этих функций постепенно добавлялась при проведении экспериментов. В итоге получившуюся функцию потерь для обучения можно записать следующим образом: L(y, y) = L1(y, y) +MultiSpec(y, y) +SiSDR(y, y) +Perceptual(y, y).

Набор данных

В качестве набора данных было решено использовать свободно распространяемый, размеченный набор данных MUSDB18. Это классический набор для подобных задач. Давайте рассмотрим его подробнее. Структура набора: 100 композиций на обучение и 50 композиций на валидацию.

Этот набор данных состоит из 150 полноформатных музыкальных треков суммарной длительностью примерно 10 часов.

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

Аугментации

MUSDB18 — это хорошо, а увеличенный и расширенный — еще лучше. Для такой модели, как HDemucs, необходимо не просто много, а ошеломительно много данных. Чтобы обеспечить модель материалом, давайте, словно диджеи, миксовать и всячески менять сэмплы. К тому же ремиксы, особенно всем набившее оскомину “slow+reverb remix”, можно считать по сути отдельными треками.

А теперь рассмотрим те приемы, которые мы использовали для расширения:

• FlipChannels. Поскольку в MUSDB аудиосигнал стерео, то есть у него два канала на каждое, то можно поменять местами левый и правый.

• Shift — можно просто сдвинуть трек во времени, то есть начать его.

• Remix — источники трека перемешиваются в рамках батча.

• FlipSign — изменение знака источников.

• Scale — увеличение или уменьшение амплитуды во сколько-то  раз, то есть повышение или понижение громкости.

• TimeChange — глобальное увеличение или уменьшение действующих частот, то есть ускорение или замедление трека.

• PitchShift — увеличение или уменьшение частоты только вокальной части, то есть повышение или понижения тональности вокалиста.

• Fade — затухание случайного сегмента источника.

• Reverse segment — реверс случайного сегмента соответствующего источника вокала (позаимствовано из фильма «Малыш на драйве»).

• Дублирование канала — оба канала становятся либо левыми, либо правыми.

Подробности реализации

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

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

Во-вторых, поскольку генерировать аугментации нужно много и быстро, необходимо использовать специальные числовые типы, такие как bf16 (Brain Floating Point) — это формат чисел с плавающей точкой, жертвующий некоторой точностью, а взамен дающий прирост по обучению.

К тому же нехватка железа в рамках учебного проекта не позволяет просто закинуть на 8 видеокарт и ждать чуда; нужно выжимать и танцевать с бубном на скромном количестве железа. Это ограничение мотивировало нас уменьшить глубину модели. В отличие от оригинального HDemucs, который имел глубину 6, мы уменьшили модель до глубины 4, что пропорционально уменьшило количество параметров аж в 9 раз! И знаете, что самое интересное? То, что получилось в результате.

Результаты

«Магия запрещена вне Хогвартса», но что вы скажете на то, что, несмотря на уменьшение параметров в 9 раз, удалось не просто сравниться с HDemucs в качестве разделения вокала, но и побить его по качеству? Это точно волшебство, хотя, конечно, другие источники похуже, но в целом также достаточно адекватны при прослушивании.

Вот примеры работы нашей модели на одном легендарном треке:

• Queen — We Will Rock You (Live in Montreal 1981)

• Оригинал: original.wav

• Вокал: vocals.wav 

• Ударные: drums.wav 

• Бас: bass.wav 

• Другое: other.wav 

Сравнительная таблица результатов: 

HDemucs d4 — уменьшенная версия оригинального HDemucs с четырьмя слоями для кодировщика и декодировщика в bi-U-Net.

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

Мы продолжим работу над проектом в рамках Инженерно-математической школы VK и НИУ ВШЭ и постараемся улучшить наше решение, сделав разделение на «произвольное» количество источников. Для решения этой задачи попробуем интегрировать векторные представления разделенных источников. Полученные векторы источников будут передавать информацию в выходные слои модели, позволяя ей настроиться на сегментирование определенного источника.