Training Pipeline¶

The complete training process from data preparation to model saving.

Overview¶

flowchart TB
    subgraph "Phase 1: Data Preparation"
        A[User Recordings] --> B[Validate Audio]
        B --> C[TTS Synthesis]
        C --> D[Data Augmentation]
        D --> E[Hard Negative Generation]
        E --> F[Load Real Negatives]
        F --> G[Create Dataset]
    end

    subgraph "Phase 2: Training"
        G --> H[Initialize Model]
        H --> I[Training Loop]
        I --> J[Validation]
        J --> K{Early Stop?}
        K -->|No| I
        K -->|Yes| L[Best Model]
    end

    subgraph "Phase 3: Export"
        L --> M[Threshold Calibration]
        M --> N[Save PyTorch]
        N --> O[Export ONNX]
        O --> P[Save Metadata]
    end

Phase 1: Data Preparation¶

Audio Validation¶

User recordings are validated:

Positive Sample Generation¶

flowchart LR
    A[User Recordings] --> B[TTS Synthesis]
    B --> C[Speed Augmentation]
    C --> D[Pitch Augmentation]
    D --> E[Volume Augmentation]
    E --> F[Noise Augmentation]
    F --> G[5000+ Positives]

Negative Sample Generation¶

flowchart LR
    A[Wake Word] --> B[Phonetic Analysis]
    B --> C[Generate Similar Words]
    C --> D[TTS Synthesis]
    D --> E[20000+ Hard Negatives]

    F[UNAC Cache] --> G[Load Chunks]
    G --> H[10000+ Real Negatives]

Dataset Creation¶

Samples are split with unseen voice validation:

Phase 2: Training¶

Model Initialization¶

model = ASTWakeWordModel(
    freeze_base=True,  # Freeze AST
    classifier_hidden_dims=[256, 128],
    classifier_dropout=0.5,
    use_attention=True,
    use_se_block=True,
    use_tcn=True,
)

Training Configuration¶

Focal Loss¶

Focuses on hard examples:

Mixup Augmentation¶

Virtual training samples are created by mixing embeddings:

Where \(\lambda \sim Beta(\alpha, \alpha)\) and \(\alpha = 0.5\).

Training Loop¶

for epoch in range(max_epochs):
    # Training
    model.train()
    for batch in train_loader:
        embeddings = model.get_embeddings(batch.audio)

        # Apply mixup
        embeddings, labels = mixup(embeddings, batch.labels)

        # Forward pass
        logits = model.classifier(embeddings)
        loss = focal_loss(logits, labels)

        # Backward
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        scheduler.step()

    # Validation
    val_metrics = validate(model, val_loader)

    # Early stopping check
    if val_metrics.loss < best_loss:
        best_loss = val_metrics.loss
        patience_counter = 0
        save_checkpoint(model)
    else:
        patience_counter += 1
        if patience_counter >= patience:
            break

Early Stopping¶

Training stops when validation loss doesn't improve:

Phase 3: Export¶

Threshold Calibration¶

Find optimal detection threshold:

for threshold in np.arange(0.3, 0.9, 0.05):
    predictions = (probs >= threshold)
    far = false_positives / total_negatives
    frr = false_negatives / total_positives
    f1 = compute_f1(predictions, labels)

    if f1 > best_f1:
        best_threshold = threshold

Model Saving¶

PyTorch (.pt):

torch.save({
    'classifier_state_dict': model.classifier.state_dict(),
    'config': model.classifier.config,
}, f"{wake_word}.pt")

Metadata (.json):

{
  "wake_word": "jarvis",
  "threshold": 0.65,
  "accuracy": 0.971,
  "f1_score": 0.943,
  "training_config": { ... }
}

ONNX (.onnx):

torch.onnx.export(
    model,
    dummy_input,
    f"{wake_word}.onnx",
    input_names=["input_values"],
    output_names=["logits"],
    dynamic_axes={"input_values": {0: "batch"}},
)

Metrics¶

Training Metrics¶

Final Evaluation¶

Timeline¶

Training time varies significantly based on your hardware:

Reference¶

Source files:

• src/wakebuilder/models/trainer.py - ASTTrainer class
• src/wakebuilder/models/classifier.py - Model architecture
• src/wakebuilder/backend/routes/training.py - API integration