Chez Simli, nous nous soucions le plus de la latence. Après tout, c'est ce qui compte pour nous : une vidéo à faible latence. D’un autre côté, certains des algorithmes les plus utilisés dans l’Audio Machine Learning ont des implémentations vraiment très lentes. Pour être clair, ces implémentations conviennent généralement à la création des modèles eux-mêmes ou à l'inférence par lots. Mais pour nous chez Simli, quelques millisecondes pourraient faire la différence entre un désordre saccadé ou une vidéo fluide.
Heureusement pour moi (et par procuration pour vous le lecteur), ce guide ne nécessite pas beaucoup de connaissances en mathématiques, des personnes beaucoup plus intelligentes ont déjà compris comment obtenir la bonne réponse, nous rendons simplement le calcul plus efficace. Si vous avez besoin de plus d'informations pour comprendre ce qu'est le MelSpectrogram, vous pouvez lire cet article. Il existe plusieurs façons de calculer le spectrogramme, cela dépend fortement de votre application. Nous nous concentrons donc sur les éléments requis pour exécuter nos modèles internes pour des raisons de commodité pour l'auteur.

La solution commune : Librosa

Vous êtes probablement ici après avoir rencontré un dépôt utilisant Librosa. C'est une bibliothèque assez pratique, pour être honnête. Il existe une tonne d'utilitaires, des moyens simples de lire l'audio sur le disque et un accès rapide à de nombreuses fonctionnalités couramment requises telles que le rééchantillonnage audio, le sous-mixage des canaux, etc. Dans notre cas, nous nous intéressons à une fonctionnalité particulière : le calcul de melspectrogramme. Dans librosa, obtenir le melspectrogramme est simple.

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

C'est simple et cela prend en moyenne environ 2 ms sur une VM GCP g2. Eh bien, il y a deux problèmes principaux :

1. Habituellement, lorsque vous travaillez avec des modèles DL, vous devez exécuter le modèle sur un GPU. Cela signifie qu'une partie de votre chaîne s'exécute sur le CPU et que vous copiez ensuite les résultats sur le GPU. Pour l'inférence par lots, c'est généralement très bien puisque vous devez collecter autant de données que possible sur le GPU/transfert. Cependant, dans notre cas, nous travaillons souvent avec une image à la fois pour réduire le temps d'attente et de traitement.
2. Notre budget temps total est d'environ 33 ms/image. Cela inclut la latence de transfert du serveur API vers le serveur d'inférence ML, la copie CPU vers GPU, le prétraitement et le post-traitement pour les modèles incluant le melspectrogram. Chaque milliseconde compte lorsque vous travaillez avec un budget aussi serré. Ces deux millisecondes ont en fait contribué à avoir un flux vidéo rendu en direct fonctionnel pour Simli (enfin, il y avait de nombreuses optimisations valant chacune une milliseconde ou deux).

Rechercher des solutions en ligne

En essayant de voir comment d'autres personnes l'ont fait (heureusement, ce n'est pas un problème unique pour nous), j'ai trouvé cet article qui expliquait à la fois le fonctionnement des melspectrogrammes et fournissait une implémentation de référence qui, pour une raison quelconque, ne prenait que 1 ms (50 % amélioration). C'est un bon début mais il reste le premier problème, tout n'était pas sur le GPU. Nous utilisons PyTorch et nous comptons sur torch.compile avec le mode=reduce-overhead pour des améliorations de vitesse maximales. Cependant, un tel transfert de données peut nuire aux performances car le compilateur PyTorch ne sera pas non plus en mesure d'optimiser la fonction. La solution est un peu fastidieuse mais relativement simple, réécrivez-la à la torche. L'équipe PyTorch s'est assurée qu'une grande partie de sa syntaxe et de ses fonctionnalités sont aussi proches que possible de NumPy (avec quelques cas extrêmes qui sont généralement bien documentés, à l'exception d'un qui m'a fait perdre quelques jours mais c'est une histoire pour un autre blog) .

La réécriture de PyTorch

Il y a donc quelques étapes que nous devons suivre pour réussir à tout réécrire dans Pytorch. Les Melspectrogrammes peuvent être divisés en trois étapes :

• Calcul Transformée de Fourier à court terme
• Génération des banques de fréquences de l'échelle mel
• Génération du spectrogramme.

Il y a de bonnes et de mauvaises nouvelles. La bonne nouvelle est que toutes les fonctionnalités requises sont facilement disponibles dans pytorch ou torchaudio. La mauvaise nouvelle est que le comportement par défaut est très différent de celui de librosa, il y a donc beaucoup de configuration et d'essais et d'erreurs pour bien faire les choses. J’ai vécu ça et je partage l’info parce que je ne peux même pas souhaiter cet enfer à mon pire ennemi. Une chose que nous devons comprendre est que ce code repose fortement sur la mise en cache de certains de nos résultats pour être utilisés plus tard. Cela se fait dans une fonction d'initialisation qui prégénère tous les tableaux statiques (les banques de fréquences mel par exemple dépendent de la fréquence d'échantillonnage et du nombre de mel dont vous avez besoin). Voici notre version optimisée de la fonction melspectrogram utilisant 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)
)

Après la compilation initiale, nous avons mesuré que cette fonction prenait 350 microsecondes en utilisant un GPU Nvidia L4 (avec mise en cache de hann_window et melscale_fbanks). L'appel ajusté ressemblera à ceci :

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

Ceci fait partie d'une série d'articles sur la façon dont nous avons optimisé nos modèles pré-entraînés déployés, en optimisant les étapes de prétraitement et de post-traitement. Vous pouvez consulter https://www.simli.com/demo pour voir les modèles déployés et les avatars à latence la plus faible que nous proposons