В Simli мы больше всего заботимся о задержке. В конце концов, это то, чем мы занимаемся: видео с низкой задержкой. С другой стороны, некоторые из наиболее часто используемых алгоритмов машинного обучения аудио имеют очень медленную реализацию. Чтобы внести ясность, эти реализации обычно подходят для создания самих моделей или для пакетного вывода. Но для нас в Simli пара миллисекунд может означать разницу между заиканием и плавностью видео.
К счастью для меня (и, по доверенности, для вас, читателя), это руководство не требует больших знаний в математике, гораздо более умные люди уже выяснили, как получить правильный ответ, мы просто делаем вычисления более эффективными. Если вам нужна дополнительная информация, чтобы понять, что такое MelSpectrogram, вы можете прочитать эту статью. Существует несколько способов расчета спектрограммы, это сильно зависит от вашего приложения. Итак, ради удобства автора мы концентрируемся на мелах, необходимых для запуска наших внутренних моделей.

Общее решение: Librosa

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

import librosa

# load in any audio to test
sampleAudio, sr = librosa.load("sample.mp3", sr=None) # sr=None means the original sampling rate
spectrogram = librosa.feature.melspectrogram(
    y=sampleAudio,
    sr=sr,
    n_fft=int(0.05 * sr),  # 50ms
    hop_length=int(0.0125 * sr),  # 12.5ms
    win_length=int(0.05 * sr),
)

Проще говоря, на виртуальной машине GCP g2 это занимает в среднем около 2 мс. Итак, есть две основные проблемы:

1. Обычно при работе с моделями DL вам необходимо запустить модель на графическом процессоре. Это означает, что часть вашей цепочки выполняется на процессоре, а затем вы копируете результаты обратно на графический процессор. Для пакетного вывода это в основном нормально, поскольку вам следует собрать столько данных, сколько вы можете поместить на графический процессор/передачу. Однако в нашем случае мы часто работаем с одним кадром за раз, чтобы сократить время ожидания и обработки.
2. Наш общий временной бюджет составляет примерно 33 мс/кадр. Сюда входит задержка передачи с сервера API на сервер вывода ML, копирование ЦП на ГП, предварительная и постобработка для моделей, включая мелспектрограмму. Каждая миллисекунда имеет значение, когда вы работаете с таким ограниченным бюджетом. Эти две миллисекунды фактически способствовали созданию работающего живого видеопотока для Simli (ну, это было много оптимизаций, каждая из которых стоила одной или двух миллисекунд).

Ищем решения в Интернете

Пытаясь посмотреть, как это сделали другие люди (к счастью, это не уникальная проблема для нас), я нашел эту статью, в которой объяснялось, как работают мелспектрограммы, и предоставлялась эталонная реализация, которая по какой-то причине заняла всего 1 мс (50 % улучшение). Это хорошее начало, но есть еще первая проблема: не все было на графическом процессоре. Мы используем PyTorch и полагаемся на torch.compile с mode=reduce-overhead для максимального повышения скорости. Однако подобная передача данных может снизить производительность, поскольку компилятор PyTorch также не сможет оптимизировать функцию. Решение немного утомительное, но относительно простое: перепишите его в torch. Команда PyTorch позаботилась о том, чтобы большая часть их синтаксиса и функциональности была максимально приближена к NumPy (с некоторыми крайними случаями, которые обычно хорошо документированы, за исключением одного, из-за которого я потерял пару дней, но это история для другого блога). .

Переписывание PyTorch

Итак, нам нужно сделать пару шагов, чтобы успешно переписать все в Pytorch. Мелспектрограммы можно разделить на три этапа:

• Вычисление кратковременного преобразования Фурье
• Создание частотных банков мел-шкалы
• Создание спектрограммы.

Есть хорошие и плохие новости. Хорошей новостью является то, что все необходимые функции легко доступны в pytorch или torchaudio. Плохая новость заключается в том, что поведение по умолчанию сильно отличается от поведения librosa, поэтому требуется много настроек, проб и ошибок, чтобы сделать все правильно. Я прошел через это и делюсь информацией, потому что даже не могу пожелать такого ада своему злейшему врагу. Нам нужно понять одну вещь: этот код в значительной степени зависит от кэширования некоторых результатов, которые будут использоваться позже. Это делается с помощью функции инициализации, которая предварительно генерирует все статические массивы (например, банки частот мелов зависят от частоты дискретизации и количества нужных мелов). Вот наша оптимизированная версия функции мелспектрограммы с использованием PyTorch

import torch

if torch.cuda.is_available
    @torch.compile(mode="reduce-overhead")
else:
    @torch.compile
def melspecrogram_torch(wav:torch.Tensor,sample_rate:int, hann_window: torch.Tensor, mel_basis: torch.Tensor):
    stftWav = torch.stft(
            wav,
            n_fft=int(sample_rate*0.05),
            win_length=int(sample_rate*0.05),
            hop_length=int(sample_rate*0.0125),
            window=hann_window,
            pad_mode="constant",
            return_complex=True,
        ).abs()
    stftWav = stftWav.squeeze()
    mel_stftWav = torch.mm(mel_basis, stftWav)
    return mel_stftWav

device = "cuda" if torch.cuda.is_available() else "cpu"

melspectrogram_torch(
    sampleAudio,
    sr,
    torch.hann_window(int(sample_rate*0.05), device=device, dtype=torch.float32),
    torchaudio.functional.melscale_fbanks(
        sample_rate=sr,
        n_freqs=(int(sample_rate*0.05) // 2   1),
        norm="slaney", # this is the normalization algorithm used by librosa
        # this is an example that's related to our own pipeline, check what you need for yours
        n_mels=80,
        f_min=55,
        f_max=7600,
    )
    .T.to(device)
)

После первоначального запуска компиляции мы измерили, что эта функция занимает 350 микросекунд с использованием графического процессора Nvidia L4 (с кэшированием hann_window и melscale_fbanks). Скорректированный вызов будет выглядеть так:

hann=torch.hann_window(int(sample_rate*0.05), device=device, dtype=torch.float32),
melscale=torchaudio.functional.melscale_fbanks(
        sample_rate=sr,
        n_freqs=(int(sample_rate*0.05) // 2   1),
        norm="slaney", # this is the normalization algorithm used by librosa
        # this is an example that's related to our own pipeline, check what you need for yours
        n_mels=80,
        f_min=55,
        f_max=7600,
    )
    .T.to(device)
melspectrogram_torch(
    sampleAudio,
    sr,
    hann,
    melscale,
)

Это одна часть серии статей о том, как мы оптимизировали наши развернутые предварительно обученные модели, оптимизируя этапы предварительной и постобработки. Вы можете проверить https://www.simli.com/demo, чтобы увидеть развернутые модели и аватары с наименьшей задержкой, которые мы предоставляем