File size: 5,070 Bytes
09dcb2a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 | from typing import Optional, Tuple, Union
import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import Wav2Vec2Model, Wav2Vec2PreTrainedModel
from transformers.modeling_outputs import SequenceClassifierOutput
from transformers.models.auto import AutoConfig, AutoModelForAudioClassification
from .configuration import Wav2Vec2DualHypersphereConfig
class Wav2Vec2DualHypersphereForAudioClassification(Wav2Vec2PreTrainedModel):
config_class = Wav2Vec2DualHypersphereConfig
def __init__(self, config: Wav2Vec2DualHypersphereConfig):
super().__init__(config)
self.num_labels = config.num_labels
self.wav2vec2 = Wav2Vec2Model(config)
proj_dim = getattr(config, "classifier_proj_size", 256)
self.projector = nn.Linear(config.hidden_size, proj_dim)
self.dropout = nn.Dropout(getattr(config, "final_dropout", 0.1))
self.classifier = nn.Linear(proj_dim, self.num_labels)
# Initialize weights and apply final processing
self.post_init()
def freeze_feature_extractor(self):
self.wav2vec2.feature_extractor._freeze_parameters()
def forward(
self,
input_values: Optional[torch.Tensor] = None,
attention_mask: Optional[torch.Tensor] = None,
output_attentions: Optional[bool] = None,
output_hidden_states: Optional[bool] = None,
return_dict: Optional[bool] = None,
labels: Optional[torch.Tensor] = None,
) -> Union[Tuple, SequenceClassifierOutput]:
r"""
labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
Labels for computing the sequence classification/regression loss.
"""
return_dict = return_dict if return_dict is not None else self.config.return_dict
outputs = self.wav2vec2(
input_values,
attention_mask=attention_mask,
output_attentions=output_attentions,
output_hidden_states=output_hidden_states,
return_dict=return_dict,
)
hidden_states = outputs[0]
# 1. Temporal Pooling with Attention Mask support
if attention_mask is not None:
mask = self.wav2vec2._get_feature_vector_attention_mask(
hidden_states.shape[1], attention_mask
)
mask = mask.unsqueeze(-1).expand_as(hidden_states)
pooled_768 = (hidden_states * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1e-9)
else:
pooled_768 = hidden_states.mean(dim=1)
# 2. First Hypersphere Projection (L2 Normalize 768-D)
norm_pooled_768 = F.normalize(pooled_768, p=2, dim=-1)
# 3. Intermediate Projection & Dropout
projected_256 = self.dropout(self.projector(norm_pooled_768))
# 4. Second Hypersphere Projection (L2 Normalize 256-D)
norm_projected_256 = F.normalize(projected_256, p=2, dim=-1)
# 5. Final Classification Logits
logits = self.classifier(norm_projected_256)
loss = None
if labels is not None:
if self.config.problem_type is None:
if self.num_labels == 1:
self.config.problem_type = "regression"
elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
self.config.problem_type = "single_label_classification"
else:
self.config.problem_type = "multi_label_classification"
if self.config.problem_type == "regression":
loss_fct = nn.MSELoss()
if self.num_labels == 1:
loss = loss_fct(logits.squeeze(), labels.squeeze())
else:
loss = loss_fct(logits, labels)
elif self.config.problem_type == "single_label_classification":
loss_fct = nn.CrossEntropyLoss()
loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
elif self.config.problem_type == "multi_label_classification":
loss_fct = nn.BCEWithLogitsLoss()
loss = loss_fct(logits, labels)
if not return_dict:
output = (logits,) + outputs[2:]
return ((loss,) + output) if loss is not None else output
return SequenceClassifierOutput(
loss=loss,
logits=logits,
hidden_states=outputs.hidden_states if hasattr(outputs, "hidden_states") else None,
attentions=outputs.attentions if hasattr(outputs, "attentions") else None,
)
# Register model & config with Auto classes
AutoConfig.register("wav2vec2-dual-hypersphere", Wav2Vec2DualHypersphereConfig)
AutoModelForAudioClassification.register(
Wav2Vec2DualHypersphereConfig, Wav2Vec2DualHypersphereForAudioClassification
)
Wav2Vec2DualHypersphereClassifier = Wav2Vec2DualHypersphereForAudioClassification
__all__ = [
"Wav2Vec2DualHypersphereConfig",
"Wav2Vec2DualHypersphereForAudioClassification",
"Wav2Vec2DualHypersphereClassifier",
] |