#!/usr/bin/env python3 """Whisper-Medium Encoder + Linear Emotion Head 학습. Whisper encoder를 freeze하고 linear classifier head만 학습하여 benchmark_ser_models.py의 WhisperMediumAdapter에 사용할 체크포인트를 생성한다. Usage: # AI Hub 테스트 서브셋 외의 데이터로 학습 (test leakage 방지) python scripts/train_whisper_emotion_head.py \\ --train-dir data/evaluation/korean/train_audio \\ --val-dir data/evaluation/korean/val_audio \\ --output data/models/whisper_emotion_head.pt # prepare_aihub_test_subset.py의 출력으로 quick test (test leakage 주의) python scripts/train_whisper_emotion_head.py \\ --train-dir data/evaluation/korean/test_audio \\ --epochs 3 --output data/models/whisper_emotion_head.pt """ from __future__ import annotations import argparse import glob import logging import os import time from pathlib import Path import numpy as np import torch import torch.nn as nn from torch.utils.data import DataLoader, Dataset logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") logger = logging.getLogger(__name__) EVAL_LABELS = ["neutral", "joy", "sadness", "anger", "surprise", "fear"] LABEL_TO_IDX = {label: i for i, label in enumerate(EVAL_LABELS)} class EmotionAudioDataset(Dataset): """Load WAV files organized by emotion class directory.""" def __init__(self, root_dir: str, processor, max_samples_per_class: int | None = None): self.samples = [] self.processor = processor for label in EVAL_LABELS: class_dir = Path(root_dir) / label if not class_dir.exists(): logger.warning("Class directory not found: %s", class_dir) continue wavs = sorted(glob.glob(str(class_dir / "*.wav"))) if max_samples_per_class and len(wavs) > max_samples_per_class: wavs = wavs[:max_samples_per_class] for wav_path in wavs: self.samples.append({ "path": wav_path, "label": LABEL_TO_IDX[label], }) logger.info("Dataset: %d samples from %s", len(self.samples), root_dir) def __len__(self): return len(self.samples) def __getitem__(self, idx): sample = self.samples[idx] import librosa audio, sr = librosa.load(sample["path"], sr=16000) inputs = self.processor(audio, sampling_rate=16000, return_tensors="pt") features = inputs.input_features.squeeze(0) # (n_mels, T) return features, sample["label"] def collate_fn(batch): features, labels = zip(*batch) # Pad features to same length max_len = max(f.shape[1] for f in features) padded = [] for f in features: if f.shape[1] < max_len: pad = torch.zeros(f.shape[0], max_len - f.shape[1]) f = torch.cat([f, pad], dim=1) padded.append(f) return torch.stack(padded), torch.tensor(labels, dtype=torch.long) def train(args): from transformers import WhisperModel, WhisperFeatureExtractor device = torch.device(args.device) # Load Whisper encoder (frozen) logger.info("Loading Whisper-Medium encoder...") processor = WhisperFeatureExtractor.from_pretrained("openai/whisper-medium") whisper = WhisperModel.from_pretrained("openai/whisper-medium").to(device) whisper.eval() for param in whisper.parameters(): param.requires_grad = False hidden_dim = whisper.config.d_model # 1024 head = nn.Linear(hidden_dim, len(EVAL_LABELS)).to(device) # Dataset train_dataset = EmotionAudioDataset(args.train_dir, processor, args.max_samples_per_class) if len(train_dataset) == 0: logger.error("No training samples found") return train_loader = DataLoader( train_dataset, batch_size=args.batch_size, shuffle=True, collate_fn=collate_fn, num_workers=0, ) val_loader = None if args.val_dir and Path(args.val_dir).exists(): val_dataset = EmotionAudioDataset(args.val_dir, processor) if len(val_dataset) > 0: val_loader = DataLoader( val_dataset, batch_size=args.batch_size, shuffle=False, collate_fn=collate_fn, num_workers=0, ) # Optimizer optimizer = torch.optim.Adam(head.parameters(), lr=args.lr) criterion = nn.CrossEntropyLoss() # Training loop best_val_acc = 0.0 for epoch in range(args.epochs): head.train() total_loss = 0 correct = 0 total = 0 for batch_idx, (features, labels) in enumerate(train_loader): features = features.to(device) labels = labels.to(device) with torch.no_grad(): encoder_out = whisper.encoder(features) hidden = encoder_out.last_hidden_state # (B, T, D) pooled = hidden.mean(dim=1) # (B, D) logits = head(pooled) loss = criterion(logits, labels) optimizer.zero_grad() loss.backward() optimizer.step() total_loss += loss.item() * labels.size(0) preds = logits.argmax(dim=1) correct += (preds == labels).sum().item() total += labels.size(0) train_acc = correct / max(total, 1) avg_loss = total_loss / max(total, 1) logger.info("Epoch %d/%d: loss=%.4f, train_acc=%.3f", epoch + 1, args.epochs, avg_loss, train_acc) # Validation if val_loader: head.eval() val_correct = 0 val_total = 0 with torch.no_grad(): for features, labels in val_loader: features = features.to(device) labels = labels.to(device) encoder_out = whisper.encoder(features) pooled = encoder_out.last_hidden_state.mean(dim=1) logits = head(pooled) preds = logits.argmax(dim=1) val_correct += (preds == labels).sum().item() val_total += labels.size(0) val_acc = val_correct / max(val_total, 1) logger.info(" val_acc=%.3f", val_acc) if val_acc > best_val_acc: best_val_acc = val_acc save_checkpoint(head, args.output, epoch, val_acc) else: # No val set — save latest save_checkpoint(head, args.output, epoch, train_acc) logger.info("Training complete. Best checkpoint: %s", args.output) def save_checkpoint(head: nn.Linear, path: str, epoch: int, accuracy: float): Path(path).parent.mkdir(parents=True, exist_ok=True) torch.save(head.state_dict(), path) logger.info("Saved checkpoint: %s (epoch=%d, acc=%.3f)", path, epoch + 1, accuracy) def main(): parser = argparse.ArgumentParser(description="Train Whisper emotion classifier head") parser.add_argument("--train-dir", required=True, help="학습 오디오 디렉토리 ({emotion}/*.wav 구조)") parser.add_argument("--val-dir", default=None, help="검증 오디오 디렉토리 (없으면 train accuracy로 판단)") parser.add_argument("--output", default="data/models/whisper_emotion_head.pt", help="출력 체크포인트 경로") parser.add_argument("--epochs", type=int, default=10) parser.add_argument("--batch-size", type=int, default=8) parser.add_argument("--lr", type=float, default=1e-3) parser.add_argument("--device", default="cpu", choices=["cpu", "cuda"]) parser.add_argument("--max-samples-per-class", type=int, default=None, help="클래스당 최대 학습 샘플 수") args = parser.parse_args() train(args) if __name__ == "__main__": main()