Привет, хабровчане! Меня зовут Дмитрий, я студент первого курса Вышки магистерской программы «Системный анализ и математические технологии». В октябре прошлого года я прошёл отбор в Инженерно-математическую школу (ИМШ). Это совместная образовательная программа VK Education и НИУ ВШЭ. Я стал участником мастерской по прикладному искусственному интеллекту. В одном из проектов мы командой из двух человек (да, нас было мало, но мы были сильны!) решили реальную технологическую задачу под руководством эксперта VK и научных сотрудников университета. Нам нужно было придумать технологию разделения аудиосигнала на источники. В команде я отвечал за подготовку данных, а мой коллега — за проведение экспериментов. Хочу рассказать вам подробно, как мы решали эту задачу — надеюсь, что это будет для вас полезно. Кстати, в ИМШ есть и другие практикоориентированные мастерские: по аппаратному обеспечению искусственного интеллекта, по компиляторам и высоконагруженным приложениям, по виртуальным платформенным решениям «Робби» (в СПб). Если вас что-то из этого интересует и вы являетесь студентом Вышки, узнать об условиях поступления можно здесь. Следующий набор в мастерские ИМШ откроется осенью этого года, а в магистратуры школы поступить можно уже сейчас. Обучение в мастерских бесплатное. А еще: те, кто хорошо покажут себя в работе мастерских, получают преимущества при рассмотрении на вакансии VK. Разделение аудиосигнала на источники — одна из многих задач ИИ в области обработки звука. К ней близка задача по созданию технологии шумоподавления, в рамках которой необходимо выделить один источник сигнала и максимально увеличить качество выделенного материала. Но фундаментально две эти задачи отличаются друг от друга.  В отличие от шумоподавления в результате разделения аудиосигнала на источники на выходе получается множество очищенных сигналов. Из всех моделей искусственного интеллекта к решению этой задачи максимально приблизилась модель из семейства DEMUCS — HT Demucs (Hybrid Transformer Demucs). Она заняла первое место в соревновании “Music Demixing Challenge 2021”, которое провела компания Sony.  Несмотря на это, в своей работе мы имели дело с моделью HDemucs (Hybrid Demucs). Казалось бы, логичнее работать с более успешным подходом, но есть нюанс — HT Demucs тяжело обучить: а все потому, что данных, на которых обучалась эта модель, нет в открытом доступе.  Расскажем подробнее, что мы делали с моделькой.  Для нашей задачи звук можно представить как зависимость интенсивности от времени для композиции волновых колебаний. Этот временной сигнал можно перевести в частотное представление для каждого (в непрерывном случае) временного окна. Такое преобразование называют оконным преобразованием Фурье (STFT). Поскольку STFT представляет входной сигнал в комплексной форме, мы можем рассмотреть по отдельности амплитудное и фазовое представление преобразованного сигнала. Следующие представления мы можем представлять как тепловые карты, на которых по оси абсцисс будет отложено время, по оси ординат — частота, а амплитуда или фаза будет спроецирована на плоскость, кодируя значение цветом точки на карте. Полученные тепловые карты — это изображения, которые можно обрабатывать нейросетевыми методами компьютерного зрения. Ради них мы здесь сегодня и собрались. 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 и НИУ ВШЭ и постараемся улучшить наше решение, сделав разделение на «произвольное» количество источников. Для решения этой задачи попробуем интегрировать векторные представления разделенных источников. Полученные векторы источников будут передавать информацию в выходные слои модели, позволяя ей настроиться на сегментирование определенного источника.  **edited**