Skip to content

training.md

Latest commit

 

History

History
304 lines (229 loc) · 9.38 KB

training.md

File metadata and controls

304 lines (229 loc) · 9.38 KB
Raw

MERC — Training Pipeline

← Back to README · Architecture · Data

Table of Contents

Overview

flowchart TD
    subgraph Init["Initialisation  (train_one_seed)"]
        SEED["Set torch + numpy seed"]
        DATA["build_dataloaders()\nnormalise on train stats"]
        WEIGHTS["compute_class_weights()\nMELD-only train labels"]
        MODEL["Instantiate MERCModel"]
        OPT["AdamW optimiser"]
        SCHED["LambdaLR: warmup-cosine"]
        SEED --> DATA --> WEIGHTS --> MODEL --> OPT --> SCHED
    end

    subgraph Loop["Per-epoch loop"]
        TRAIN["_train_epoch()\nforward · loss · backward · clip · step"]
        EVAL["_evaluate()\nno_grad · val metrics"]
        LOG["_print_epoch()\nloss · wF1 · macF1 · acc · pred dist"]
        CKPT{"val_wF1\nimproved?"}
        SAVE["torch.save(best_state)"]
        PAT["patience_counter++"]
        STOP{"patience ≥ 10?"}

        TRAIN --> EVAL --> LOG --> CKPT
        CKPT -->|yes| SAVE
        CKPT -->|no| PAT --> STOP
        STOP -->|yes| BREAK["early stop"]
        STOP -->|no| TRAIN
    end

    subgraph Test["Test evaluation"]
        LOAD["load best checkpoint"]
        TESTEVAL["_evaluate() on test_loader"]
        REPORT["per-class F1\nconfusion matrix"]
        LOAD --> TESTEVAL --> REPORT
    end

    Init --> Loop --> Test
Loading

Data Normalisation

File: merc/data/loaders.py — build_dataloaders()
Class: merc/data/preprocess.py — ModalityNormalizer

Per-modality zero-mean, unit-variance normalisation is applied to all three feature arrays. Statistics are computed from the training split only and applied to train, dev, and test:

text_mean  = text_feats[train_mask].mean(axis=0)   # (512,)
text_std   = text_feats[train_mask].std(axis=0) + 1e-8
text_feats = (text_feats - text_mean) / text_std

This ensures no leakage from dev/test into the normalisation statistics. The norm_stats dict (means and stds) is returned from build_dataloaders() so it can be saved alongside a checkpoint for deployment.

Class Weight Computation

File: merc/data/preprocess.py — compute_class_weights()
Called in: merc/train.py — train_one_seed()

Inverse-frequency weights compensate for class imbalance in the training set:

weight_c = total_samples / (num_classes × count_c)

Weights are clipped to [0.5, 5.0] to prevent extreme values.

Why MELD-only labels are used

Class weights are computed from MELD training examples only (not the combined MELD + IEMOCAP pool):

flowchart LR
    META["metadata.csv"]
    FILTER["filter: split=='train'\nAND source=='MELD'"]
    LABELS["1 901 MELD train labels"]
    CW["compute_class_weights()"]
    META --> FILTER --> LABELS --> CW
Loading

Rationale: IEMOCAP contributes no surprise, fear, or disgust examples (these map only from MELD). When the combined 6 653-sample pool is used as the denominator, disgust (58 total, 1 from IEMOCAP) receives a raw weight of 16.4× — clipped to 10 — which causes the model to catastrophically over-predict disgust on validation (where it represents only 2% of utterances). Using the 1 901-sample MELD pool brings the weights into a sensible range:

Class Combined weight (old) MELD-only weight (new)
neutral 0.51 0.50
joy 0.69 0.79
sadness 1.20 2.66
anger 0.44 1.41
surprise 3.13 1.14
fear 10.0 (clipped) 4.24
disgust 10.0 (clipped) 4.76

Loss Function

nn.CrossEntropyLoss(
    weight=class_weights,   # (7,) float tensor on device
    label_smoothing=0.1,
)

Label smoothing (ε = 0.1) spreads 10% of the target probability mass uniformly across all classes. This regularises the model and prevents overconfident predictions on the dominant class.

The effective target for the correct class becomes:

y_smooth = (1 - ε) · 1 + ε / num_classes  =  0.9 + 0.1/7  ≈  0.914

Optimiser and Scheduler

Optimiser: AdamW with weight decay applied to all parameters except biases and LayerNorm parameters.

lr = 2e-4 · weight_decay = 1e-4

Scheduler: Linear warmup followed by cosine decay. scheduler.step() is called once per epoch inside _train_epoch().

xychart-beta
    title "Learning Rate Schedule (warmup=10, max=100 epochs)"
    x-axis "Epoch" [0, 10, 20, 30, 40, 50, 60, 70, 80, 90, 100]
    y-axis "LR multiplier" 0 --> 1.05
    line [0.1, 1.0, 0.98, 0.91, 0.79, 0.65, 0.50, 0.35, 0.21, 0.09, 0.0]
Loading 
def lr_lambda(epoch):
    if epoch < warmup_epochs:
        return (epoch + 1) / warmup_epochs          # linear: 0.1 → 1.0
    progress = (epoch - warmup_epochs) / (max_epochs - warmup_epochs)
    return 0.5 * (1.0 + cos(π · progress))          # cosine: 1.0 → 0.0

Gradient clipping is applied before each optimiser step:

nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

Training Loop

File: merc/train.py — _train_epoch()

for batch in train_loader:           # batch is a PyG Batch
    logits = model(batch)            # (N_total_utterances, 7)
    loss   = criterion(logits, batch.y)

    optimizer.zero_grad()
    loss.backward()
    clip_grad_norm_(model.parameters(), grad_clip)
    optimizer.step()

scheduler.step()                     # once per epoch

batch.y contains the flat concatenation of emotion labels for all utterances in all dialogues in the batch. The model returns one logit vector per utterance, so the shapes align directly.

Batch size = 16 dialogues per step. Each dialogue contains a variable number of utterances (typically 5–15), so the effective number of utterances per gradient step is ~80–240.

Validation and Early Stopping

File: merc/train.py — _evaluate()

Evaluation runs under torch.no_grad(). The model is set to .eval() mode (disables dropout).

Metrics computed per epoch on the dev split:

Metric Implementation
Weighted F1 f1_score(labels, preds, average="weighted")
Macro F1 f1_score(labels, preds, average="macro")
Accuracy accuracy_score(labels, preds)

Early stopping monitors val weighted F1 with patience = 10 epochs. The best checkpoint (by val wF1) is saved to merc/checkpoints/best_seed{seed}.pt and loaded before test evaluation.

flowchart TD
    EVAL["compute val_wF1"]
    CMP{"val_wF1 >\nbest_val_wF1?"}
    UPD["best_val_wF1 = val_wF1\nsave checkpoint\npatience_counter = 0"]
    INC["patience_counter += 1"]
    CHK{"patience_counter\n≥ patience (10)?"}
    CONT["continue training"]
    STOP["early stop\nload best checkpoint"]

    EVAL --> CMP
    CMP -->|yes| UPD --> CONT
    CMP -->|no| INC --> CHK
    CHK -->|no| CONT
    CHK -->|yes| STOP
Loading

Multi-Seed Evaluation

Training repeats for each seed in training.seeds (default: [42, 123, 456]). Each seed independently:

  • Sets torch.manual_seed(seed) and numpy.random.seed(seed)
  • Builds a fresh model
  • Saves its own checkpoint (best_seed{seed}.pt)
  • Returns test metrics

After all seeds complete, the script reports:

FINAL RESULTS (mean ± std across seeds)
  Weighted F1 : X.XXXX ± X.XXXX
  Macro F1    : X.XXXX ± X.XXXX
  Accuracy    : X.XXXX ± X.XXXX

  Per-class weighted F1 (mean ± std):
    neutral     : X.XXXX ± X.XXXX
    joy         : ...

Epoch Logging

Each epoch prints a one-line summary followed by the val prediction distribution:

Ep   1 | tr_loss=1.8234  tr_wF1=0.1312 | val_loss=2.1200  val_wF1=0.0390  val_macF1=0.0280  val_acc=0.0521
         val_pred_dist: neut=12  joy=45  sadn=8  ange=890  surp=6  fear=89  disg=57

The prediction distribution line is critical for diagnosing prediction collapse — if the model is predicting one or two classes for almost all val examples, this shows up immediately.

Full Config Reference

data:
  base_dir: "."
  audio_file:    "data/enhanced_npy/enhanced_audio.npy"
  text_file:     "data/enhanced_npy/enhanced_text.npy"
  video_file:    "data/enhanced_npy/enhanced_video.npy"
  metadata_file: "data/enhanced_npy/metadata.csv"
  modality_dim: 512

model:
  hidden_dim:           256    # internal hidden size throughout
  num_heads:            8      # attention heads (cross-modal + R-GAT)
  num_gnn_layers:       2      # R-GAT layers
  num_relations:        3      # temporal / same-speaker / cross-speaker
  dropout:              0.3
  num_classes:          7
  cross_speaker_window: 10     # max utterance distance for cross-speaker edges

training:
  batch_size:      16
  lr:              0.0002
  weight_decay:    0.0001
  warmup_epochs:   10
  max_epochs:      100
  patience:        10          # early stopping on val weighted F1
  grad_clip:       1.0
  label_smoothing: 0.1
  seeds:           [42, 123, 456]
  checkpoint_dir:  "merc/checkpoints"

eval:
  metrics:
    - weighted_f1
    - macro_f1
    - accuracy
    - per_class_f1