Привет, хабровчане! Меня зовут Дмитрий, я студент первого курса Вышки магистерской программы «Системный анализ и математические технологии».
В октябре прошлого года я прошёл отбор в Инженерно-математическую школу (ИМШ). Это совместная образовательная программа 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
Сравнительная таблица результатов:
Модель | Ударные SDR, db | Бас SDR, db | Другое SDR, db | Вокал SDR, db |
Our Demucs | 5,1 | 2,1 | 3,3 | 8,6 |
HDemucs d4 | 5 | 1,1 | 3,6 | 2,5 |
HDemucs | 8,24 | 8,76 | 5,59 | 8,13 |
HDemucs d4 — уменьшенная версия оригинального HDemucs с четырьмя слоями для кодировщика и декодировщика в bi-U-Net.
Результаты экспериментов основаны на метриках SDR и свидетельствуют об эффективности применения составной функции потерь, а также приемов, использованных при обучении.
Мы продолжим работу над проектом в рамках Инженерно-математической школы VK и НИУ ВШЭ и постараемся улучшить наше решение, сделав разделение на «произвольное» количество источников. Для решения этой задачи попробуем интегрировать векторные представления разделенных источников. Полученные векторы источников будут передавать информацию в выходные слои модели, позволяя ей настроиться на сегментирование определенного источника.