Aqui na Simli, nos preocupamos mais com a latência. Afinal, é disso que somos: vídeo de baixa latência. Por outro lado, alguns dos algoritmos mais utilizados em Audio Machine Learning têm implementações muito lentas. Para ser claro, essas implementações geralmente são adequadas para a criação dos próprios modelos ou para inferência em lote. Mas para nós da Simli, alguns milissegundos podem significar a diferença entre uma bagunça trêmula ou um vídeo suave.
Felizmente para mim (e por procuração para você, leitor), este guia não requer muito conhecimento em matemática, pessoas muito mais inteligentes já descobriram como obter a resposta correta, estamos apenas tornando o cálculo mais eficiente. Se precisar de mais informações para entender o que é o MelSpectrogram, você pode ler este artigo. Existem várias maneiras de calcular o espectrograma, depende muito da sua aplicação. Portanto, estamos nos concentrando nos mels necessários para executar nossos modelos internos por uma questão de conveniência para o autor.

A solução comum: Librosa

Você provavelmente está aqui depois de encontrar um repositório que usa Librosa. É uma biblioteca muito útil, para ser honesto. Existem vários utilitários, maneiras fáceis de ler o áudio no disco e acesso rápido a muitas funcionalidades comumente necessárias, como reamostragem de áudio, downmixing de canais e outros. No nosso caso, estamos interessados ​​em uma funcionalidade específica: cálculo do melspectrograma. Na librosa, obter o melspectrograma é simples.

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

Simples e leva em média cerca de 2 ms em uma VM GCP g2. Bem, existem dois problemas principais:

1. Normalmente, ao trabalhar com modelos DL, você precisaria executar o modelo em uma GPU. Isso significa que parte da sua cadeia é executada na CPU e então você copia os resultados de volta para a GPU. Para inferência em lote, isso é bastante adequado, pois você deve coletar o máximo de dados que puder caber na GPU/transferência. No entanto, no nosso caso, muitas vezes trabalhamos com um quadro de cada vez para reduzir o tempo de espera e de processamento.
2. Nosso orçamento de tempo total é de aproximadamente 33 ms/quadro. Isso inclui latência de transferência do servidor API para o servidor de inferência de ML, cópia de CPU para GPU, pré-processamento e pós-processamento para os modelos, incluindo o melspectrograma. Cada milissegundo é importante quando você trabalha com um orçamento tão apertado. Esses dois milissegundos na verdade contribuíram para ter um fluxo de vídeo renderizado ao vivo funcional para Simli (bem, foram muitas otimizações, cada uma valendo um ou dois milissegundos).

Procurando soluções on-line

Enquanto tentava ver como outras pessoas fizeram isso (felizmente, este não é um problema exclusivo para nós), encontrei este artigo que explicava como funcionam os melspectrogramas e fornecia uma implementação de referência que, por algum motivo, levou apenas 1ms (50 % melhoria). É um bom começo, mas ainda há o primeiro problema: nem tudo estava na GPU. Estamos usando PyTorch e contamos com torch.compile com mode=reduce-overhead para melhorias máximas de velocidade. No entanto, transferências de dados como essa podem prejudicar o desempenho, pois o compilador PyTorch também não será capaz de otimizar a função. A solução é um pouco tediosa, mas relativamente fácil, reescreva-a rapidamente. A equipe PyTorch garantiu que grande parte de sua sintaxe e funcionalidade fosse o mais próximo possível do NumPy (com alguns casos extremos que geralmente são bem documentados, exceto um que me perdeu por alguns dias, mas isso é uma história para um blog diferente) .

A reescrita do PyTorch

Portanto, há algumas etapas que precisamos seguir para reescrever tudo com sucesso no Pytorch. Os melspectrogramas podem ser divididos em três etapas:

• Calculando a transformada de Fourier em curto espaço de tempo
• Gerando os bancos de frequência da escala mel
• Gerando o espectrograma.

Há boas e más notícias. A boa notícia é que todas as funcionalidades necessárias estão prontamente disponíveis em pytorch ou torchaudio. A má notícia é que o comportamento padrão é muito diferente do librosa, então há muita configuração e tentativa e erro para acertar. Já passei por isso e estou compartilhando a informação porque não posso nem desejar esse inferno para o meu pior inimigo. Uma coisa que precisamos entender é que esse código depende muito do armazenamento em cache de alguns de nossos resultados para serem usados ​​posteriormente. Isso é feito em uma função de inicialização que pré-gera todos os arrays estáticos (os bancos de frequência mel, por exemplo, dependem da taxa de amostragem e do número de mels necessários). Aqui está nossa versão otimizada da função 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)
)

Após a execução inicial da compilação, medimos esta função para levar 350 microssegundos usando uma GPU Nvidia L4 (com cache de hann_window e melscale_fbanks). A chamada ajustada ficará assim:

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 é parte de uma série de artigos sobre como otimizamos nossos modelos pré-treinados implantados, otimizando as etapas de pré e pós-processamento. Você pode verificar https://www.simli.com/demo para ver os modelos implantados e os avatares de menor latência que fornecemos