uyiosa/SER-WavLM-Multi-Attributes

The model is a recreation of 3loi/SER-Odyssey-Baseline-WavLM-Multi-Attributes for direct implementation in torch, with class definition and feed forward method. This model was recreated with the hopes of greater flexibilty of control, training/fine-tuning of model. The model was trained on the same MSP-Podcast dataset as the original, but a different smaller subset was used. The subset is evenly distributed across gender and emotion category with hopes that training would improve accuracy of valence and arousal predictions. This model is therefore a multi-attributed based model which predict arousal, dominance and valence. However, unlike the original model, I just kept the original attribute score range of 0...7 (the range the dataset follows). I will provide the evaluations later on. For now I decided to make this repo so that other people could test out my model and see what they think of the inference accuracy themselves, or retrain from scratch, modify etc. My best trained weights s of now are provided in this repo. The class definition for the model is can be found in my github.

Get class definition

git clone https://github.com/PhilipAmadasun/SER-Model-for-dimensional-attribute-prediction.git

Usage

Inference Testing

import torch
import torchaudio
from SER_Model_setup import SERModel 

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

checkpoint_path = "<model.pt file>"
checkpoint = torch.load(checkpoint_path, map_location=device)

# Create the model architecture and load weights
model = SERModel()
model.load_state_dict(checkpoint['model_state_dict'])
model.to(device)
model.eval()

audio_path = "<wav file>"
audio, sr = torchaudio.load(audio_path)

if sr != model.sample_rate:
    resampler = torchaudio.transforms.Resample(sr, model.sample_rate)
    audio = resampler(audio)
#print(audio.shape[0])

if audio.shape[0] > 1:
    audio = torch.mean(audio, dim=0, keepdim=True)

audio_len = audio.shape[-1]

# Create waveform tensor (shape: [1, audio_len])
waveform = torch.zeros(1, audio_len, dtype=torch.float32)
# print(waveform)
# print()
# print(f"waveform shape: {waveform.shape}")
# print()
waveform[0, :audio_len] = audio
# print(waveform)
# print()
# Create mask as 2D tensor: shape [1, audio_len] with ones in valid region
mask = torch.ones(1, audio_len, dtype=torch.float32)
# print(mask)
# print()
# print(f"mask shape: {mask.shape}")

# Move waveform and mask to device
waveform = waveform.to(device)
mask = mask.to(device)

# Normalize waveform using model's mean and std
mean = model.mean.to(device)
std = model.std.to(device)
waveform = (waveform - mean) / (std + 1e-6)

with torch.no_grad():
    predictions = model(waveform, mask)  # predictions shape: [1, 3]

# Extract predictions: [0,0] for arousal, [0,1] for valence, [0,2] for dominance
arousal   = predictions[0, 0].item()
valence   = predictions[0, 1].item()
dominance = predictions[0, 2].item()

print(f"Arousal: {arousal:.3f}")
print(f"Valence: {valence:.3f}")
print(f"Dominance: {dominance:.3f}")

Batch inference

import os
import glob
import torch
import torchaudio
from SER_Model_setup import SERModel  # Adjust if your model code is elsewhere

def load_model_from_checkpoint(checkpoint_path, device='cpu'):
    """
    Loads the SERModel and weights from a checkpoint, moves to device, sets eval mode.
    """
    checkpoint = torch.load(checkpoint_path, map_location=device)
    
    # Create the model architecture
    model = SERModel()
    model.load_state_dict(checkpoint['model_state_dict'])
    
    model.to(device)
    model.eval()
    return model

def batch_inference(model, file_paths, device='cpu', normalize=True):
    """
    Perform true batch inference on multiple .wav files in one forward pass.
    
    Args:
        model (SERModel): The loaded SER model in eval mode
        file_paths (list[str]): List of paths to .wav files
        device (str or torch.device): 'cpu' or 'cuda'
        normalize (bool): Whether to normalize waveforms (subtract mean, divide std)
    
    Returns:
        dict: {filename: {"arousal": float, "valence": float, "dominance": float}}
    """

    # ----------------------------------------
    # 1) Load & store all waveforms in memory
    # ----------------------------------------
    waveforms_list = []
    lengths = []
    for fp in file_paths:
        # Load audio
        audio, sr = torchaudio.load(fp)
        
        # Resample if needed
        if sr != model.sample_rate:
            resampler = torchaudio.transforms.Resample(sr, model.sample_rate)
            audio = resampler(audio)
        
        # Convert stereo -> mono if needed
        if audio.shape[0] > 1:
            audio = torch.mean(audio, dim=0, keepdim=True)
        
        # audio shape => [1, num_samples]
        lengths.append(audio.shape[-1])
        waveforms_list.append(audio)
    
    # ----------------------------------------
    # 2) Determine max length
    # ----------------------------------------
    max_len = max(lengths)
    
    # ----------------------------------------
    # 3) Pad each waveform to max length & build masks
    # ----------------------------------------
    batch_size = len(waveforms_list)
    batched_waveforms = torch.zeros(batch_size, 1, max_len, dtype=torch.float32)
    masks = torch.zeros(batch_size, max_len, dtype=torch.float32)

    for i, audio in enumerate(waveforms_list):
        cur_len = audio.shape[-1]
        batched_waveforms[i, :, :cur_len] = audio
        masks[i, :cur_len] = 1.0  # valid portion

    # ----------------------------------------
    # 4) Move batched data to device BEFORE normalization
    # ----------------------------------------
    batched_waveforms = batched_waveforms.to(device)
    masks = masks.to(device)
    
    # ----------------------------------------
    # 5) Normalize if needed (model.mean, model.std)
    # ----------------------------------------
    if normalize:
        # model.mean and model.std are buffers; ensure they're on the correct device
        mean = model.mean.to(device)
        std = model.std.to(device)
        batched_waveforms = (batched_waveforms - mean) / (std + 1e-6)
    
    # ----------------------------------------
    # 6) Single forward pass
    # ----------------------------------------
    with torch.no_grad():
        predictions = model(batched_waveforms, masks)
        # predictions shape => [batch_size, 3]
    
    # ----------------------------------------
    # 7) Build result dict
    # ----------------------------------------
    results = {}
    for i, fp in enumerate(file_paths):
        arousal   = predictions[i, 0].item()
        valence   = predictions[i, 1].item()
        dominance = predictions[i, 2].item()
        filename = os.path.basename(fp)
        results[filename] = {
            "arousal": arousal,
            "valence": valence,
            "dominance": dominance
        }
    
    return results

if __name__ == "__main__":
    # -----------------------------------------
    # Example usage
    # -----------------------------------------
    device = "cuda" if torch.cuda.is_available() else "cpu"
    
    checkpoint_path = "<weights.pt>"
    model = load_model_from_checkpoint(checkpoint_path, device=device)
    
    # Suppose you have a folder of .wav files
    wav_folder = "<directory containing .wav files>"
    wav_paths = glob.glob(os.path.join(wav_folder, "*.wav"))
    
    # Do a single pass of batch inference
    all_results = batch_inference(model, wav_paths, device=device, normalize=True)
    
    # Print results
    for fname, preds in all_results.items():
        print(f"{fname}: Arousal={preds['arousal']:.3f}, "
              f"Valence={preds['valence']:.3f}, Dominance={preds['dominance']:.3f}")