Views
No views yet
(V, A, D) ∈ ℝ³ representing the utterance’s position in the emotional space. Unlike discrete emotion classifiers, this model performs a multi-target regression task, allowing for more detailed and nuanced emotion predictions from speech.1import torch
2import librosa
3import sys
4import json
5from transformers import Wav2Vec2FeatureExtractor, AutoConfig
6from transformers.models.wav2vec2.modeling_wav2vec2 import Wav2Vec2PreTrainedModel, Wav2Vec2Model
7from transformers.file_utils import ModelOutput
8from dataclasses import dataclass
9from typing import Optional, Tuple
10import torch.nn as nn
11
12@dataclass
13class SpeechClassifierOutput(ModelOutput):
14 loss: Optional[torch.FloatTensor] = None
15 logits: torch.FloatTensor = None
16 hidden_states: Optional[Tuple[torch.FloatTensor]] = None
17 attentions: Optional[Tuple[torch.FloatTensor]] = None
18
19
20class CCCLoss(nn.Module):
21 def __init__(self):
22 super().__init__()
23
24 def forward(self, gold, pred):
25 gold = gold.float()
26 pred = pred.float()
27
28 gold_mean = torch.mean(gold, dim=0)
29 pred_mean = torch.mean(pred, dim=0)
30
31 gold_var = torch.var(gold, dim=0, unbiased=False)
32 pred_var = torch.var(pred, dim=0, unbiased=False)
33
34 cov = torch.mean((gold - gold_mean) * (pred - pred_mean), dim=0)
35
36 numerator = 2 * cov
37 denominator = gold_var + pred_var + (gold_mean - pred_mean) ** 2
38
39 eps = 1e-8
40 ccc = numerator / (denominator + eps)
41
42 loss = 1.0 - ccc
43 return torch.mean(loss)
44
45class Wav2Vec2ClassificationHead(nn.Module):
46 def __init__(self, config):
47 super().__init__()
48 self.dense = nn.Linear(config.hidden_size, config.hidden_size)
49 self.dropout = nn.Dropout(config.final_dropout)
50 self.out_proj = nn.Linear(config.hidden_size, config.num_labels)
51
52 def forward(self, features, **kwargs):
53 x = features
54 x = self.dropout(x)
55 x = self.dense(x)
56 x = torch.tanh(x)
57 x = self.dropout(x)
58 x = self.out_proj(x)
59 return x
60
61class Wav2Vec2ForSpeechClassification(Wav2Vec2PreTrainedModel):
62 _tied_weights_keys = []
63
64 @property
65 def all_tied_weights_keys(self):
66 return {}
67
68 def __init__(self, config):
69 super().__init__(config)
70 self.num_labels = config.num_labels
71 self.pooling_mode = config.pooling_mode
72 self.config = config
73
74 self.wav2vec2 = Wav2Vec2Model(config)
75 self.classifier = Wav2Vec2ClassificationHead(config)
76
77 self.init_weights()
78 self.loss_fct = CCCLoss()
79
80 def freeze_feature_extractor(self):
81 self.wav2vec2.feature_extractor._freeze_parameters()
82
83 def merged_strategy(self, hidden_states, mode="mean"):
84 if mode == "mean":
85 outputs = torch.mean(hidden_states, dim=1)
86 elif mode == "sum":
87 outputs = torch.sum(hidden_states, dim=1)
88 elif mode == "max":
89 outputs = torch.max(hidden_states, dim=1)[0]
90 else:
91 raise Exception("Pooling mode not supported")
92 return outputs
93
94 def forward(
95 self,
96 input_values,
97 attention_mask=None,
98 output_attentions=None,
99 output_hidden_states=None,
100 return_dict=None,
101 labels=None,
102 ):
103 return_dict = return_dict if return_dict is not None else self.config.use_return_dict
104
105 outputs = self.wav2vec2(
106 input_values,
107 attention_mask=attention_mask,
108 output_attentions=output_attentions,
109 output_hidden_states=output_hidden_states,
110 return_dict=return_dict,
111 )
112
113 if isinstance(outputs, dict) or hasattr(outputs, 'last_hidden_state'):
114 hidden_states = outputs.last_hidden_state
115 else:
116 hidden_states = outputs[0]
117
118 # Pooling
119 hidden_states = self.merged_strategy(hidden_states, mode=self.pooling_mode)
120
121 # Logits
122 logits = self.classifier(hidden_states)
123
124 preds = torch.sigmoid(logits)
125
126 loss = None
127 if labels is not None:
128 loss = self.loss_fct(labels, preds)
129
130 if not return_dict:
131 output = (preds,) + outputs[2:]
132 return ((loss,) + output) if loss is not None else output
133
134 return SpeechClassifierOutput(
135 loss=loss,
136 logits=preds,
137 hidden_states=outputs.hidden_states,
138 attentions=outputs.attentions,
139 )
140
141model_path = 'model'
142device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
143print(f"Loading Prosody model from '{model_path}' su {device}...")
144
145config = AutoConfig.from_pretrained(model_path)
146setattr(config, 'pooling_mode', 'mean')
147
148feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(model_path)
149model = Wav2Vec2ForSpeechClassification.from_pretrained(model_path, config=config)
150model.to(device)
151model.eval()
152
153
154def get_prosody_info(audio_path):
155 print(f"Analysis file: {audio_path}")
156 speech, sr = librosa.load(audio_path, sr=16000)
157
158 inputs = feature_extractor(speech, sampling_rate=16000, return_tensors="pt", padding=True)
159
160 inputs = {key: val.to(device) for key, val in inputs.items()}
161
162 with torch.no_grad():
163 outputs = model(**inputs)
164 predictions = outputs.logits.cpu().numpy()[0]
165
166 print(predictions)
167
168 arousal = float(f"{predictions[0]:.4f}")
169 valence = float(f"{predictions[1]:.4f}")
170 dominance = float(f"{predictions[2]:.4f}")
171
172 return {
173 "arousal": arousal,
174 "valence": valence,
175 "dominance": dominance
176 }
177
178if __name__ == "__main__":
179 file_audio = "<your_audio_path>"
180 if len(sys.argv) > 1:
181 file_audio = sys.argv[1]
182
183 result = get_prosody_info(file_audio)
184
185 print("Result:")
186 print(json.dumps(result, indent=4))| Parameter | Value |
|---|---|
| Base model | facebook/wav2vec2-large-xlsr-53 |
| Batch size (train) | 32 |
| Batch size (eval) | 16 |
| Gradient accumulation | 4 steps (effective batch: 128) |
| Learning rate | 1e-4 |
| Epochs | 30 |
| Numerical precision | BF16 |
| Feature extractor | Frozen |
| Output | Multi-target regression (V, A, D) |
| Dimension | CCC | MSE |
|---|---|---|
| Arousal | 0.615 | 0.026 |
| Dominance | 0.661 | 0.026 |
| Valence | 0.522 | 0.026 |
| Mean | 0.599 | 0.026 |
1@misc{bernardini2024wav2vec2,
2 author = {Alessio Bernardini},
3 title = {Wav2Vec2.0 Italian Prosody PAD Regressor},
4 year = {2024},
5 howpublished = {\url{https://github.com/abernardini-unimi/Wav2Vec2.0-Italian-prosody}}
6}