train_wav_vit.py

import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.nn.functional as F
import math 
from dataloader import *
from torch.optim import Adam, SGD, AdamW
from torch.utils.data import DataLoader
from sklearn.metrics import confusion_matrix
from torchvision import models


from speechbrain.lobes.models.huggingface_transformers.wav2vec2 import  Wav2Vec2, Wav2Vec2Pretrain
model_hub_w2v2 = "jonatasgrosman/wav2vec2-large-xlsr-53-english"



batch_size = 24

device = torch.device('cuda:0')

dropout = 0.1

def train_and_evaluate(model, optimizer, scheduler, train_dataloader, val_dataloader, test_dataloader, epochs, device):
    compute_loss = nn.CrossEntropyLoss()
    
    best_val_loss = float('inf')
    best_model = None

    for epoch in range(epochs):
        model.train()
        train_loss = 0
        for waveform, melspec, y in train_dataloader:
            waveform,  melspec, y = waveform.to(device), melspec.to(device), y.to(device)
            optimizer.zero_grad()
            out = model(waveform, melspec)
            loss = compute_loss(out, y)
            loss.backward()
            optimizer.step()
            scheduler.step()
            train_loss += loss.item()
        train_loss /= len(train_dataloader)

        model.eval()
        val_loss = 0
        with torch.no_grad():
            for waveform, melspec, y in val_dataloader:
                waveform, melspec, y = waveform.to(device),  melspec.to(device), y.to(device)
                out = model(waveform,  melspec)
                val_loss += compute_loss(out, y).item()
        val_loss /= len(val_dataloader) 
        print(f"Époque {epoch+1}/{epochs} - Perte d'entraînement : {train_loss:.4f}, Perte de validation : {val_loss:.4f}")        
        if val_loss < best_val_loss:
            best_val_loss = val_loss
            best_model = model.state_dict()
        
    # Évaluation sur l'ensemble de test
    model.load_state_dict(best_model)
    model.eval()
    correct = 0
    total = 0
    matrice = []
    with torch.no_grad():
        for waveform, melspec, y in test_dataloader:
            waveform,  melspec, y = waveform.to(device), melspec.to(device), y.to(device)
            out = model(waveform, melspec)
            _, predicted = torch.max(out.data, 1)
            total += y.size(0)
            correct += (predicted == y).sum().item()
            matrice.append(confusion_matrix(y.detach().cpu().numpy(),predicted.detach().cpu().numpy(),labels= np.asarray([0,1,2,3])))

    accuracy = correct / total
    CM = np.sum(np.stack(matrice,axis=0),axis=0)
    return accuracy,CM




acc_all_session = []
matrice_confusion_all_session = []
utterance_type = None
print(f"{utterance_type} (Utterance type)")
for i in range(5):
    for kk in range(2):
    # Define the sessions for training, validation, and test
        test_session = (i + 1,)
        train_sessions = tuple(j + 1 for j in range(5) if j != i)
        val_session = test_session

        # Déterminer les speakers pour la validation et le test
        all_speaker = ('Ses01F','Ses01M','Ses02F','Ses02M','Ses03F','Ses03M','Ses04F','Ses04M','Ses05F','Ses05M')
        speakers = ('Ses0'+str(i+1)+'M', 'Ses0'+str(i+1)+'F')
        if kk ==0: 
            val_speaker = speakers[i % 2]  # Alterner entre les speakers 'M' et 'F' pour chaque fold
            test_speaker = speakers[(i + 1) % 2]
        else:
            test_speaker = speakers[i % 2]  # Alterner entre les speakers 'M' et 'F' pour chaque fold
            val_speaker = speakers[(i + 1) % 2]
        
        root = '.'
        # Create the datasets
        train_dataset = IEMOCAP_MFCC(
            root=root,
            utterance_type=utterance_type, 
            sessions=train_sessions,
            transform2=mel_spectrogram,
            speakers=all_speaker
        )
        val_dataset = IEMOCAP_MFCC(
            root=root, 
            utterance_type=utterance_type,
            sessions=test_session,  
            transform2=mel_spectrogram,
            speakers=val_speaker
        )
        test_dataset = IEMOCAP_MFCC(
            root=root, 
            utterance_type=utterance_type,
            sessions=test_session, 
            transform2=mel_spectrogram,
            speakers=test_speaker,
            augmentation=False
        )

       
        class ClassificationHead(nn.Module):
            """Head for classification task."""

            def __init__(self, hidden_size,final_dropout,num_labels):
                super().__init__()
                self.dense1 = nn.Linear(hidden_size, 256)
                self.dense2 = nn.Linear(256, 256)
                self.dropout = nn.Dropout(final_dropout)
                self.out_proj = nn.Linear(256, num_labels)
                self.sigma = torch.tanh

            def forward(self, features, **kwargs):
                x = features
                x = self.dropout(x)
                x = self.dense1(x)
                x = self.sigma(x)
                x = self.dropout(x)
                x = self.dense2(x)
                x = self.sigma(x)
                x = self.dropout(x)
                x = self.out_proj(x)
                return x
            

            
        class CrossAttention(nn.Module):
            def __init__(self, n_heads, d_embed, d_cross, in_proj_bias=True, out_proj_bias=True):
                super().__init__()
                self.q_proj   = nn.Linear(d_embed, d_embed, bias=in_proj_bias)
                self.k_proj   = nn.Linear(d_cross, d_embed, bias=in_proj_bias)
                self.v_proj   = nn.Linear(d_cross, d_embed, bias=in_proj_bias)
                self.out_proj = nn.Linear(d_embed, d_embed, bias=out_proj_bias)
                self.n_heads = n_heads
                self.d_head = d_embed // n_heads
            
            def forward(self, x, y):
                input_shape = x.shape
                batch_size, sequence_length, d_embed = input_shape
                interim_shape = (batch_size, -1, self.n_heads, self.d_head)

                q = self.q_proj(x)
                k = self.k_proj(y)
                v = self.v_proj(y)

                q = q.view(interim_shape).transpose(1, 2)
                k = k.view(interim_shape).transpose(1, 2)
                v = v.view(interim_shape).transpose(1, 2)

                weight = q @ k.transpose(-1, -2)
                weight /= math.sqrt(self.d_head)
                weight = F.softmax(weight, dim=-1)

                output = weight @ v
                output = output.transpose(1, 2).contiguous()
                output = output.view(input_shape)
                output = self.out_proj(output)
                return output  

        class MultiAda(nn.Module):

            def __init__(self, dropout=0.1, device=device):
                super().__init__()
                self.n_base = n_base
                self.device = device
                self.lambda1 = lambda1
                self.lambda2 = lambda2

                self.model_wav = Wav2Vec2(model_hub_w2v2, save_path='.', freeze=False).to(device)
                
                self.vit = models.vit_b_16(weights='IMAGENET1K_V1')
                
                self.cross_attention1 = CrossAttention(16,1024,1000)
                
                self.classif = ClassificationHead(383976-1000,dropout,4) #L=6
                

            def forward(self, inputs, img): 
                
                wav_embedding = self.model_wav(inputs)
                waveform_embedding = torch.flatten(wav_embedding,start_dim=1)              
                
                vit_embedding = self.vit(img)
                
                cross_wav_vit = self.cross_attention1(wav_embedding, vit_embedding)
                cross_wav_vit = torch.flatten(cross_wav_vit,start_dim=1)
               
                complete_embedding = torch.cat([cross_wav_vit, waveform_embedding],1) 
                predictions = self.classif(complete_embedding)
                return predictions
            
            

        model = MultiAda(dropout=dropout,
                        device=device)
        model = model.to(device)

        optimizer = AdamW(model.parameters(), lr=1e-4)
        from transformers import TrainingArguments, get_linear_schedule_with_warmup
        


        scheduler = get_linear_schedule_with_warmup(
            optimizer=optimizer, 
            num_warmup_steps=160,
            num_training_steps=1600
        )


        train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
        val_dataloader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)
        test_dataloader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

        accuracy, CM = train_and_evaluate(model, optimizer, scheduler, train_dataloader, val_dataloader, test_dataloader, epochs=6, device=device)

        print(f"Précision pour le fold {i+1} et speaker test {test_speaker}: {accuracy}")
        acc_all_session.append(accuracy)
        matrice_confusion_all_session.append(CM)
        print(f"Matrice de confusion pour le fold {i+1} et speaker test {test_speaker}: {CM}")
    
print('accuracy totale: ',np.mean(acc_all_session))
print('matrice confusion cumulée: ', np.sum(np.stack(matrice_confusion_all_session,axis=0),axis=0))