Aquí en Simli, lo que más nos importa es la latencia. Después de todo, de eso se trata: vídeo de baja latencia. Por otro lado, algunos de los algoritmos más utilizados en Audio Machine Learning tienen implementaciones realmente lentas. Para ser claros, estas implementaciones suelen ser adecuadas para crear los propios modelos o para realizar inferencias por lotes. Pero para nosotros en Simli, un par de milisegundos podrían significar la diferencia entre un vídeo entrecortado o un vídeo fluido.
Afortunadamente para mí (y para usted, el lector), esta guía no requiere muchos conocimientos de matemáticas; personas mucho más inteligentes ya han descubierto cómo obtener la respuesta correcta; simplemente estamos haciendo que el cálculo sea más eficiente. Si necesita más información para comprender qué es el MelSpectrogram, puede leer este artículo. Hay varias formas de calcular el espectrograma, depende en gran medida de su aplicación. Por lo tanto, nos estamos centrando en las combinaciones necesarias para ejecutar nuestros modelos internos por conveniencia para el autor.

La solución común: Librosa

Lo más probable es que estés aquí después de encontrar un repositorio que utiliza Librosa. Para ser honesto, es una biblioteca bastante útil. Hay un montón de utilidades, formas sencillas de leer el audio en el disco y acceso rápido a muchas funciones comúnmente requeridas, como remuestreo de audio, mezcla de canales y otras. En nuestro caso, nos interesa una funcionalidad en particular: el cálculo del melspectrograma. En librosa, obtener el melspectrograma es sencillo.

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),
)

Sencillo y tarda en promedio alrededor de 2 ms en una máquina virtual GCP g2. Bueno, hay dos problemas principales:

1. Por lo general, cuando se trabaja con modelos DL, es necesario ejecutar el modelo en una GPU. Esto significa que parte de su cadena se ejecuta en la CPU y luego copia los resultados nuevamente a la GPU. Para la inferencia por lotes, esto está bien, ya que debe recopilar tantos datos como pueda en la GPU/transferencia. Sin embargo, en nuestro caso, solemos trabajar con un cuadro a la vez para reducir el tiempo de espera y procesamiento.
2. Nuestro presupuesto de tiempo total es de aproximadamente 33 ms/cuadro. Esto incluye la latencia de transferencia desde el servidor API al servidor de inferencia ML, copia de CPU a GPU, preprocesamiento y posprocesamiento para los modelos, incluido el melspectrograma. Cada milisegundo importa cuando trabajas con un presupuesto tan ajustado. Estos dos milisegundos en realidad contribuyeron a tener una transmisión de video renderizada en vivo funcional para Simli (bueno, fueron muchas optimizaciones, cada una de las cuales valió uno o dos milisegundos).

Buscando soluciones en línea

Mientras intentaba ver cómo lo han hecho otras personas (afortunadamente, este no es un problema exclusivo para nosotros), encontré este artículo que explica cómo funcionan los melspectrogramas y proporciona una implementación de referencia que por alguna razón tomó solo 1 ms (50 % mejora). Es un buen comienzo, pero aún queda el primer problema: no todo estaba en la GPU. Estamos usando PyTorch y hemos confiado en torch.compile con mode=reduce-overhead para obtener mejoras máximas de velocidad. Sin embargo, una transferencia de datos como esta puede afectar el rendimiento, ya que el compilador de PyTorch tampoco podrá optimizar la función. La solución es un poco tediosa pero relativamente fácil, reescríbela en antorcha. El equipo de PyTorch se ha asegurado de que gran parte de su sintaxis y funcionalidad sea lo más cercana posible a NumPy (con algunos casos extremos que generalmente están bien documentados, aparte de uno que me hizo perder un par de días, pero esa es una historia para otro blog) .

La reescritura de PyTorch

Entonces, hay un par de pasos que debemos seguir para reescribir todo con éxito en Pytorch. Los melspectrogramas se pueden dividir en tres pasos:

• Cálculo de la transformada de Fourier de corto tiempo
• Generando los bancos de frecuencias de la escala mel
• Generando el espectrograma.

Hay buenas y malas noticias. La buena noticia es que todas las funciones necesarias están disponibles en pytorch o torchaudio. La mala noticia es que el comportamiento predeterminado es muy diferente al de librosa, por lo que hay mucha configuración y prueba y error para hacerlo bien. He pasado por eso y estoy compartiendo la información porque ni siquiera puedo desearle ese infierno a mi peor enemigo. Una cosa que debemos entender es que este código depende en gran medida del almacenamiento en caché de algunos de nuestros resultados para usarlos más adelante. Esto se hace en una función de inicialización que pregenera todas las matrices estáticas (los bancos de frecuencias mel, por ejemplo, dependen de la frecuencia de muestreo y la cantidad de mels que necesita). Aquí está nuestra versión optimizada de la función melspectrograma usando 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)
)

Después de la ejecución de compilación inicial, medimos que esta función tomara 350 microsegundos usando una GPU Nvidia L4 (con almacenamiento en caché de hann_window y melscale_fbanks). La llamada ajustada se verá así:

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,
)

Esta es una parte de una serie de artículos sobre cómo optimizamos nuestros modelos previamente entrenados implementados, optimizando los pasos de preprocesamiento y posprocesamiento. Puede consultar https://www.simli.com/demo para ver los modelos implementados y los avatares de latencia más baja que ofrecemos