run_kd_pipeline.py

"""
Complete Knowledge Distillation Pipeline
This script runs the full pipeline: train teacher → train student baseline → knowledge distillation
"""

import numpy as np
import os
import tensorflow as tf
from config import *
from models.knowledge_distillation import (
    load_teacher_model_logits,
    train_student_model_with_kd,
)
from models.teacher import train_and_evaluate as train_teacher


def setup_directories():
    """Creates necessary directories defined in config."""
    os.makedirs(MODELS_DIR, exist_ok=True)
    os.makedirs(RESULTS_DIR, exist_ok=True)
    print(f"Directories checked: {MODELS_DIR}, {RESULTS_DIR}")


def main():
    # Set constant seed for testing purpose
    # Original study utilized random seed
    tf.random.set_seed(RANDOM_SEED)
    np.random.seed(RANDOM_SEED)

    # Setup the directory if one does not exist
    setup_directories()

    # Enable memory growth for GPU
    physical_devices = tf.config.list_physical_devices("GPU")
    if physical_devices:
        for device in physical_devices:
            tf.config.experimental.set_memory_growth(device, True)

    print("STEP 1: Training Teacher Model (Dual-Branch Transformer)")
    teacher_model, teacher_history = train_teacher()

    print("STEP 2: Loading Teacher Model")
    teacher_model = load_teacher_model_logits()

    print("STEP 3: Knowledge Distillation Training")

    student_configs = [
        # Model Type, Save Path for KD Model
        ('cnn-s1', STUDENT_S1_KD_NAME),
        ('cnn-s2', STUDENT_S2_KD_NAME)
    ]

    for model_type, student_model_kd_name in student_configs:
        print(f"\n\n{'=' * 50}")
        print(f"Running KD Pipeline for Student: {model_type.upper()}")
        print(f"{'=' * 50}")

        # Note: We pass model_type and the specific save path here
        distilled_student = train_student_model_with_kd(
            teacher_model=teacher_model,
            model_type=model_type,
            student_kd_name=student_model_kd_name
        )

    print("\nPipeline Complete!")
    print(f"Models saved: {STUDENT_S1_KD_NAME}, {STUDENT_S2_KD_NAME}")


if __name__ == "__main__":
    main()