Added Transfer Learning example

README.md

@@ -51,52 +51,161 @@ Further details are available in the corresponding [**paper**](https://huggingfa
 ### Usage
 
 ```python 
- import torch
- import torch.nn as nn
- from transformers import AutoModelForAudioClassification, Wav2Vec2FeatureExtractor
+import torch
+import torch.nn as nn
+from transformers import AutoModelForAudioClassification, Wav2Vec2FeatureExtractor
 
 
- # CONFIG and MODEL SETUP
- model_name = 'amiriparian/ExHuBERT'
- feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained("facebook/hubert-base-ls960")
- model = AutoModelForAudioClassification.from_pretrained(model_name, trust_remote_code=True,
- revision="b158d45ed8578432468f3ab8d46cbe5974380812")
+# CONFIG and MODEL SETUP
+model_name = 'amiriparian/ExHuBERT'
+feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained("facebook/hubert-base-ls960")
+model = AutoModelForAudioClassification.from_pretrained(model_name, trust_remote_code=True,
+ revision="b158d45ed8578432468f3ab8d46cbe5974380812")
 
- # Freezing half of the encoder for further transfer learning
- model.freeze_og_encoder()
+# Freezing half of the encoder for further transfer learning
+model.freeze_og_encoder()
 
- sampling_rate = 16000
- device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
- model = model.to(device)
+sampling_rate = 16000
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = model.to(device)
 
 
 
- # Example application from a local audiofile
- import numpy as np
- import librosa
- import torch.nn.functional as F
- # Sample taken from the Toronto emotional speech set (TESS) https://tspace.library.utoronto.ca/handle/1807/24487
- waveform, sr_wav = librosa.load("YAF_date_angry.wav")
- # Max Padding to 3 Seconds at 16k sampling rate for the best results
- waveform = feature_extractor(waveform, sampling_rate=sampling_rate,padding = 'max_length',max_length = 48000)
- waveform = waveform['input_values'][0]
- waveform = waveform.reshape(1, -1)
- waveform = torch.from_numpy(waveform).to(device)
- with torch.no_grad():
- output = model(waveform)
- output = F.softmax(output.logits, dim = 1)
- output = output.detach().cpu().numpy().round(2)
- print(output)
+# Example application from a local audiofile
+import numpy as np
+import librosa
+import torch.nn.functional as F
+# Sample taken from the Toronto emotional speech set (TESS) https://tspace.library.utoronto.ca/handle/1807/24487
+waveform, sr_wav = librosa.load("YAF_date_angry.wav")
+# Max Padding to 3 Seconds at 16k sampling rate for the best results
+waveform = feature_extractor(waveform, sampling_rate=sampling_rate,padding = 'max_length',max_length = 48000)
+waveform = waveform['input_values'][0]
+waveform = waveform.reshape(1, -1)
+waveform = torch.from_numpy(waveform).to(device)
+with torch.no_grad():
+ output = model(waveform)
+ output = F.softmax(output.logits, dim = 1)
+ output = output.detach().cpu().numpy().round(2)
+ print(output)
 
- # [[0. 0. 0. 1. 0. 0.]]
- # Low | High Arousal
- # Neg. Neut. Pos. | Neg. Neut. Pos Valence
- # Disgust, Neutral, Kind| Anger, Surprise, Joy Example emotions
+ # [[0. 0. 0. 1. 0. 0.]]
+ # Low | High Arousal
+ # Neg. Neut. Pos. | Neg. Neut. Pos Valence
+ # Disgust, Neutral, Kind| Anger, Surprise, Joy Example emotions
 
 
 
 ```
 
+### Example of How to Train the Model for Transfer Learning
+The datasets used for showcasing are EmoDB and IEMOCAP from the HuggingFace Hub. As noted above, the model has seen both datasets before.
+
+```python
+
+import pandas as pd
+import torch
+import torch.nn as nn
+from torch.utils.data import Dataset, DataLoader
+import librosa
+import io
+from transformers import AutoModelForAudioClassification, Wav2Vec2FeatureExtractor
+
+# CONFIG and MODEL SETUP
+model_name = 'amiriparian/ExHuBERT'
+feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained("facebook/hubert-base-ls960")
+model = AutoModelForAudioClassification.from_pretrained(model_name, trust_remote_code=True,
+ revision="b158d45ed8578432468f3ab8d46cbe5974380812")
+
+# Replacing Classifier layer
+model.classifier = nn.Linear(in_features=256, out_features=7)
+# Freezing the original encoder layers and feature encoder (as in the paper) for further transfer learning
+model.freeze_og_encoder()
+model.freeze_feature_encoder()
+model.train()
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = model.to(device)
+
+# Define a custom dataset class
+class EmotionDataset(Dataset):
+ def __init__(self, dataframe, feature_extractor, max_length):
+ self.dataframe = dataframe
+ self.feature_extractor = feature_extractor
+ self.max_length = max_length
+
+ def __len__(self):
+ return len(self.dataframe)
+
+ def __getitem__(self, idx):
+ row = self.dataframe.iloc[idx]
+ # emotion = torch.tensor(row['label'], dtype=torch.int64) # For the IEMOCAP example
+ emotion = torch.tensor(row['emotion'], dtype=torch.int64) # EmoDB specific
+
+ # Decode audio bytes from the Huggingface dataset with librosa
+ audio_bytes = row['audio']['bytes']
+ audio_buffer = io.BytesIO(audio_bytes)
+ audio_data, samplerate = librosa.load(audio_buffer, sr=16000)
+
+ # Use the feature extractor to preprocess the audio. Padding/Truncating to 3 seconds gives better results
+ audio_features = self.feature_extractor(audio_data, sampling_rate=16000, return_tensors="pt", padding="max_length",
+ truncation=True, max_length=self.max_length)
+
+ audio = audio_features['input_values'].squeeze(0)
+ return audio, emotion
+
+# Load your DataFrame. Samples are shown for EmoDB and IEMOCAP from the Huggingface Hub
+df = pd.read_parquet("hf://datasets/renumics/emodb/data/train-00000-of-00001-cf0d4b1ae18136ff.parquet")
+# splits = {'session1': 'data/session1-00000-of-00001-04e11ca668d90573.parquet', 'session2': 'data/session2-00000-of-00001-f6132100b374cb18.parquet', 'session3': 'data/session3-00000-of-00001-6e102fcb5c1126b4.parquet', 'session4': 'data/session4-00000-of-00001-e39531a7c694b50d.parquet', 'session5': 'data/session5-00000-of-00001-03769060403172ce.parquet'}
+# df = pd.read_parquet("hf://datasets/Zahra99/IEMOCAP_Audio/" + splits["session1"])
+
+# Dataset and DataLoader
+dataset = EmotionDataset(df, feature_extractor, max_length=3 * 16000)
+dataloader = DataLoader(dataset, batch_size=4, shuffle=True)
+
+# Training setup
+criterion = nn.CrossEntropyLoss()
+lr = 1e-5
+non_frozen_parameters = [p for p in model.parameters() if p.requires_grad]
+optim = torch.optim.AdamW(non_frozen_parameters, lr=lr, betas=(0.9, 0.999), eps=1e-08)
+
+# Function to calculate accuracy
+def calculate_accuracy(outputs, targets):
+ _, predicted = torch.max(outputs, 1)
+ correct = (predicted == targets).sum().item()
+ return correct / targets.size(0)
+
+# Training loop
+num_epochs = 3
+for epoch in range(num_epochs):
+ model.train()
+ total_loss = 0.0
+ total_correct = 0
+ total_samples = 0
+ for batch_idx, (inputs, targets) in enumerate(dataloader):
+ inputs, targets = inputs.to(device), targets.to(device)
+
+ optim.zero_grad()
+ outputs = model(inputs).logits
+ loss = criterion(outputs, targets)
+ loss.backward()
+ optim.step()
+
+ total_loss += loss.item()
+ total_correct += (outputs.argmax(1) == targets).sum().item()
+ total_samples += targets.size(0)
+
+ epoch_loss = total_loss / len(dataloader)
+ epoch_accuracy = total_correct / total_samples
+ print(f'Epoch [{epoch + 1}/{num_epochs}], Average Loss: {epoch_loss:.4f}, Average Accuracy: {epoch_accuracy:.4f}')
+
+# Example outputs:
+# Epoch [3/3], Average Loss: 0.4572, Average Accuracy: 0.8249 for IEMOCAP
+# Epoch [3/3], Average Loss: 0.1511, Average Accuracy: 0.9850 for EmoDB
+
+
+```
+
+
 ### Citation Info