audeering/mobilenetv4-emotion

MobileNetV4 - Arousal / Dominance / Valence

Please note that this model is for research purpose only. A commercial license can be acquired with audEERING. The model expects a raw audio signal 16KHz as input, and outputs predictions for arousal, dominance and valence in a range of approximately 0...1. The model is created with distillation of Wav2Vec2 following Wav2Small paper.

CCC MSP Podcast v1.7
	Val	Dom	Aro	CPU Latency (10s audio)
Wav2Vec2	0.638	0.655	0.744	670 ms
MobileNet V4	0.497	0.619	0.720	7 ms

How To

import numpy as np
import torch
import librosa
from transformers import PreTrainedModel, PretrainedConfig
import timm
from torch import nn


signal = torch.from_numpy(
    librosa.load('test.wav', sr=16000)[0])[None, :]
device = 'cpu'


class Spectrogram(nn.Module):
    def __init__(self,
                 n_fft=144,   # num cols of DFT
                 n_time=144,  # num rows of DFT matrix
                 hop_length=72,
                 window='hann',
                 freeze_parameters=True):

        super().__init__()

        fft_window = librosa.filters.get_window(window, n_time, fftbins=True)

        fft_window = librosa.util.pad_center(fft_window, size=n_time)

        out_channels = n_fft // 2 + 1

        (x, y) = np.meshgrid(np.arange(n_time), np.arange(n_fft))
        omega = np.exp(-2 * np.pi * 1j / n_time)
        dft_matrix = np.power(omega, x * y)  # (n_fft, n_time)
        dft_matrix = dft_matrix * fft_window[None, :]
        dft_matrix = dft_matrix[0: out_channels, :]
        dft_matrix = dft_matrix[:, None, :]

        # DFT Non Square

        self.conv_real = nn.Conv1d(in_channels=1, out_channels=out_channels,
                                   kernel_size=n_fft, stride=hop_length,
                                   bias=False)

        self.conv_imag = nn.Conv1d(in_channels=1, out_channels=out_channels,
                                   kernel_size=n_fft, stride=hop_length,
                                   bias=False)

        self.conv_real.weight.data = torch.Tensor(
            np.real(dft_matrix)).contiguous()
        # (n_fft // 2 + 1, 1, n_fft)

        self.conv_imag.weight.data = torch.Tensor(
            np.imag(dft_matrix)).contiguous()
        # (n_fft // 2 + 1, 1, n_fft)

        if freeze_parameters:
            for param in self.parameters():
                param.requires_grad = False

    def forward(self, x):
        x = x[:, None, :]

        real = self.conv_real(x)
        imag = self.conv_imag(x)
        return real ** 2 + imag ** 2  # bs, mel, l


class LogmelFilterBank(nn.Module):
    
    def __init__(self,
                 sr=16000,
                 n_fft=144,
                 n_mels=52):

        super().__init__()

        W2 = librosa.filters.mel(sr=sr,
                                 n_fft=n_fft,
                                 n_mels=n_mels,
                                 fmin=0.0,
                                 fmax=sr//2).T

        self.register_buffer('melW', torch.Tensor(W2))
        self.register_buffer('amin', torch.Tensor([1e-10]))

    def forward(self, x):

        x = torch.matmul(x[:, None, :, :].transpose(2, 3), self.melW)

        x = torch.where(x > self.amin, x, self.amin)  # not in place

        x = 10 * torch.log10(x)
        return x


class MobileNetV4Emotion(PreTrainedModel):
    config_class = PretrainedConfig

    def __init__(self, config=PretrainedConfig()):
        super().__init__(config=config)
        self.spectrogram = Spectrogram()
        self.logmel = LogmelFilterBank()
        self.to_rgb = nn.Conv2d(1, 3, 1, padding=0, bias=True)
        self.v4 = timm.create_model(
            'mobilenetv4_conv_small.e1200_r224_in1k',
            num_classes=0,
            )
        self.lin = nn.Conv2d(960, 960, 1, padding=0, stride=1, bias=True)
        self.sof = nn.Conv2d(960, 960, 1, padding=0, stride=1, bias=True)
        self.arousal = nn.Linear(1920, 1)
        self.dominance = nn.Linear(1920, 1)
        self.valence = nn.Linear(1920, 1)

    def forward(self, x):
        x -= x.mean(1, keepdim=True)
        dev = ((x * x).mean(1, keepdim=True) + 1e-7).sqrt()
        x = self.logmel(self.spectrogram(x / dev))
        x = self.v4.forward_features(self.to_rgb(x))
        x = (self.lin(x) * self.sof(x).softmax(2)).sum(2) # pooling
        bs, channels, mels = x.size()
        x = x.reshape(bs, channels * mels)  # convert mels as channels
        return self.arousal(x), self.dominance(x), self.valence(x)


mobilenetv4_emotion = MobileNetV4Emotion.from_pretrained(
    'audeering/mobilenetv4-emotion').to(device).eval()
with torch.no_grad():
    logits = mobilenetv4_emotion(signal.to(device))

print(f'Arousal={logits[0]} | Dominance={logits[1]} | Valence={logits[2]}')