training/train_utils_old.py

# coding: utf-8
# train_utils.py

import os
import torch
import logging
import random
import datetime
import numpy as np
from tqdm import tqdm
import csv

from torch.utils.data import DataLoader, ConcatDataset
from utils.losses import WeightedCrossEntropyLoss
from utils.measures import uar, war, mf1, wf1
from models.models import BiFormer, BiGraphFormer, BiGatedGraphFormer
from data_loading.dataset_multimodal import DatasetMultiModal
from data_loading.feature_extractor import AudioEmbeddingExtractor, TextEmbeddingExtractor
from sklearn.utils.class_weight import compute_class_weight

def custom_collate_fn(batch):
    """Собирает список образцов в единый батч, отбрасывая None (невалидные)."""
    batch = [x for x in batch if x is not None]
    if not batch:
        return None

    audios = [b["audio"] for b in batch]
    audio_tensor = torch.stack(audios)

    labels = [b["label"] for b in batch]
    label_tensor = torch.stack(labels)

    texts = [b["text"] for b in batch]

    return {
        "audio": audio_tensor,
        "label": label_tensor,
        "text": texts
    }

def get_class_weights_from_loader(train_loader, num_classes):
    """
    Вычисляет веса классов из train_loader, устойчиво к отсутствующим классам.
    Если какой-либо класс отсутствует в выборке, ему будет присвоен вес 0.0.

    :param train_loader: DataLoader с one-hot метками
    :param num_classes: Общее количество классов
    :return: np.ndarray весов длины num_classes
    """
    all_labels = []
    for batch in train_loader:
        if batch is None:
            continue
        all_labels.extend(batch["label"].argmax(dim=1).tolist())

    if not all_labels:
        raise ValueError("Нет ни одной метки в train_loader для вычисления весов классов.")

    present_classes = np.unique(all_labels)

    if len(present_classes) < num_classes:
        missing = set(range(num_classes)) - set(present_classes)
        logging.info(f"[!] Отсутствуют метки для классов: {sorted(missing)}")

    # Вычисляем веса только по тем классам, что есть
    weights_partial = compute_class_weight(
        class_weight="balanced",
        classes=present_classes,
        y=all_labels
    )

    # Собираем полный вектор весов
    full_weights = np.zeros(num_classes, dtype=np.float32)
    for cls, w in zip(present_classes, weights_partial):
        full_weights[cls] = w

    return full_weights

def make_dataset_and_loader(config, split: str, only_dataset: str = None):
    """
    Универсальная функция: объединяет датасеты, или возвращает один при only_dataset.
    """
    datasets = []

    if not hasattr(config, "datasets") or not config.datasets:
        raise ValueError("⛔ В конфиге не указана секция [datasets].")

    for dataset_name, dataset_cfg in config.datasets.items():
        if only_dataset and dataset_name != only_dataset:
            continue

        csv_path = dataset_cfg["csv_path"].format(base_dir=dataset_cfg["base_dir"], split=split)
        wav_dir  = dataset_cfg["wav_dir"].format(base_dir=dataset_cfg["base_dir"], split=split)

        logging.info(f"[{dataset_name.upper()}] Split={split}: CSV={csv_path}, WAV_DIR={wav_dir}")

        dataset = DatasetMultiModal(
            csv_path = csv_path,
            wav_dir  = wav_dir,
            emotion_columns = config.emotion_columns,
            split          = split,
            sample_rate    = config.sample_rate,
            wav_length     = config.wav_length,
            whisper_model  = config.whisper_model,
            text_column    = config.text_column,
            use_whisper_for_nontrain_if_no_text = config.use_whisper_for_nontrain_if_no_text,
            whisper_device = config.whisper_device,
            subset_size    = config.subset_size,
            merge_probability = config.merge_probability
        )

        datasets.append(dataset)

    if not datasets:
        raise ValueError(f"⚠️ Для split='{split}' не найдено ни одного подходящего датасета.")

    # Объединяем только если их несколько
    full_dataset = datasets[0] if len(datasets) == 1 else ConcatDataset(datasets)

    loader = DataLoader(
        full_dataset,
        batch_size=config.batch_size,
        shuffle=(split == "train"),
        num_workers=config.num_workers,
        collate_fn=custom_collate_fn
    )

    return full_dataset, loader

def run_eval(model, loader, audio_extractor, text_extractor, criterion, device="cuda"):
    """
    Оценка модели на loader'е. Возвращает (loss, uar, war, mf1, wf1).
    """
    model.eval()
    total_loss = 0.0
    total_preds = []
    total_targets = []
    total = 0

    with torch.no_grad():
        for batch in tqdm(loader):
            if batch is None:
                continue

            audio  = batch["audio"].to(device)
            labels = batch["label"].to(device)
            texts  = batch["text"]

            audio_emb = audio_extractor.extract(audio)
            text_emb  = text_extractor.extract(texts)

            logits = model(audio_emb, text_emb)
            target = labels.argmax(dim=1)

            loss = criterion(logits, target)
            bs = audio.shape[0]
            total_loss += loss.item() * bs
            total += bs

            preds = logits.argmax(dim=1)
            total_preds.extend(preds.cpu().numpy().tolist())
            total_targets.extend(target.cpu().numpy().tolist())

    avg_loss = total_loss / total

    uar_m = uar(total_targets, total_preds)
    war_m = war(total_targets, total_preds)
    mf1_m = mf1(total_targets, total_preds)
    wf1_m = wf1(total_targets, total_preds)

    return avg_loss, uar_m, war_m, mf1_m, wf1_m

def train_once(config, train_loader, dev_loaders, test_loaders, metrics_csv_path=None):
    """
    Логика обучения (train/dev/test).
    Возвращает лучшую метрику на dev и словарь метрик.
    """

    logging.info("== Запуск тренировки (train/dev/test) ==")

    csv_writer = None
    csv_file = None

    if metrics_csv_path:
        csv_file = open(metrics_csv_path, mode="w", newline="", encoding="utf-8")
        csv_writer = csv.writer(csv_file)
        csv_writer.writerow(["split", "epoch", "dataset", "loss", "uar", "war", "mf1", "wf1", "mean"])


    # Seed
    if config.random_seed > 0:
        random.seed(config.random_seed)
        torch.manual_seed(config.random_seed)
        logging.info(f"== Фиксируем random seed: {config.random_seed}")
    else:
        logging.info("== Random seed не фиксирован (0).")

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

    # Экстракторы
    audio_extractor = AudioEmbeddingExtractor(config)
    text_extractor  = TextEmbeddingExtractor(config)

    # Параметры
    hidden_dim            = config.hidden_dim
    num_classes           = len(config.emotion_columns)
    num_transformer_heads = config.num_transformer_heads
    num_graph_heads       = config.num_graph_heads
    hidden_dim_gated      = config.hidden_dim_gated
    mode                  = config.mode
    positional_encoding   = config.positional_encoding
    dropout               = config.dropout
    out_features          = config.out_features
    lr                    = config.lr
    num_epochs            = config.num_epochs
    tr_layer_number       = config.tr_layer_number
    max_patience          = config.max_patience

    dict_models = {
        'BiFormer': BiFormer,
        'BiGraphFormer': BiGraphFormer,
        'BiGatedGraphFormer': BiGatedGraphFormer,
        # 'MultiModalTransformer_v5': MultiModalTransformer_v5,
        # 'MultiModalTransformer_v4': MultiModalTransformer_v4,
        # 'MultiModalTransformer_v3': MultiModalTransformer_v3
    }

    model_cls = dict_models[config.model_name]
    model = model_cls(
        audio_dim             = config.audio_embedding_dim,
        text_dim              = config.text_embedding_dim,
        hidden_dim            = hidden_dim,
        hidden_dim_gated      = hidden_dim_gated,
        num_transformer_heads = num_transformer_heads,
        num_graph_heads       = num_graph_heads,
        seg_len               = config.max_tokens,
        mode                  = mode,
        dropout               = dropout,
        positional_encoding   = positional_encoding,
        out_features          = out_features,
        tr_layer_number       = tr_layer_number,
        device                = device,
        num_classes           = num_classes
    ).to(device)

    # Оптимизатор и лосс
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)

    class_weights = get_class_weights_from_loader(train_loader, num_classes)
    criterion = WeightedCrossEntropyLoss(class_weights)

    logging.info("Class weights: " + ", ".join(f"{name}={weight:.4f}" for name, weight in zip(config.emotion_columns, class_weights)))

    # LR Scheduler
    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
        optimizer,
        mode="max",
        factor=0.5,
        patience=2,
        min_lr=1e-7
    )

    # Early stopping по dev
    best_dev_mean = float("-inf")
    best_dev_metrics = {}
    patience_counter = 0

    for epoch in range(num_epochs):
        logging.info(f"\n=== Эпоха {epoch} ===")
        model.train()

        total_loss = 0.0
        total_samples = 0
        total_preds = []
        total_targets = []

        for batch in tqdm(train_loader):
            if batch is None:
                continue

            audio  = batch["audio"].to(device)
            labels = batch["label"].to(device)
            texts  = batch["text"]

            audio_emb = audio_extractor.extract(audio)
            text_emb  = text_extractor.extract(texts)

            logits = model(audio_emb, text_emb)
            target = labels.argmax(dim=1)
            loss   = criterion(logits, target)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            bs = audio.shape[0]
            total_loss += loss.item() * bs

            preds = logits.argmax(dim=1)
            total_preds.extend(preds.cpu().numpy().tolist())
            total_targets.extend(target.cpu().numpy().tolist())
            total_samples += bs

        train_loss = total_loss / total_samples
        uar_m = uar(total_targets, total_preds)
        war_m = war(total_targets, total_preds)
        mf1_m = mf1(total_targets, total_preds)
        wf1_m = wf1(total_targets, total_preds)
        mean_train = np.mean([uar_m, war_m, mf1_m, wf1_m])

        logging.info(
            f"[TRAIN] Loss={train_loss:.4f}, UAR={uar_m:.4f}, WAR={war_m:.4f}, "
            f"MF1={mf1_m:.4f}, WF1={wf1_m:.4f}, MEAN={mean_train:.4f}"
        )

        # --- DEV ---
        dev_means = []
        dev_metrics_by_dataset = []

        for name, loader in dev_loaders:
            d_loss, d_uar, d_war, d_mf1, d_wf1 = run_eval(
                model, loader, audio_extractor, text_extractor, criterion, device
            )
            d_mean = np.mean([d_uar, d_war, d_mf1, d_wf1])
            dev_means.append(d_mean)

            if csv_writer:
                csv_writer.writerow(["dev", epoch, name, d_loss, d_uar, d_war, d_mf1, d_wf1, d_mean])

            logging.info(
                f"[DEV:{name}] Loss={d_loss:.4f}, UAR={d_uar:.4f}, WAR={d_war:.4f}, "
                f"MF1={d_mf1:.4f}, WF1={d_wf1:.4f}, MEAN={d_mean:.4f}"
            )

            dev_metrics_by_dataset.append({
                "name": name,
                "loss": d_loss,
                "uar": d_uar,
                "war": d_war,
                "mf1": d_mf1,
                "wf1": d_wf1,
                "mean": d_mean,
            })

        mean_dev = np.mean(dev_means)
        scheduler.step(mean_dev)

        if mean_dev > best_dev_mean:
            best_dev_mean = mean_dev
            patience_counter = 0
            best_dev_metrics = {
                "mean": mean_dev
            }
            best_dev_metrics["by_dataset"] = dev_metrics_by_dataset
        else:
            patience_counter += 1
            if patience_counter >= max_patience:
                logging.info(f"Early stopping: {max_patience} эпох без улучшения.")
                break

        # --- TEST ---
        for name, loader in test_loaders:
            t_loss, t_uar, t_war, t_mf1, t_wf1 = run_eval(
                model, loader, audio_extractor, text_extractor, criterion, device
            )
            t_mean = np.mean([t_uar, t_war, t_mf1, t_wf1])
            logging.info(
                f"[TEST:{name}] Loss={t_loss:.4f}, UAR={t_uar:.4f}, WAR={t_war:.4f}, "
                f"MF1={t_mf1:.4f}, WF1={t_wf1:.4f}, MEAN={t_mean:.4f}"
            )

            if csv_writer:
                csv_writer.writerow(["test", epoch, name, t_loss, t_uar, t_war, t_mf1, t_wf1, t_mean])

    if csv_file:
        csv_file.close()

    logging.info("Тренировка завершена. Все split'ы обработаны!")
    return best_dev_mean, best_dev_metrics