train_kd_student.py

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.utils.data import DataLoader
import random
import os
import numpy as np
from tqdm import tqdm
import json
import timm # For loading the teacher model

# Your OFA building blocks and OFAMaxViT model
try:
    from ofa_maxvit_building_blocks import OFAMaxViT, make_divisible # Or your actual import path
    from train_ofa_maxvit import ( # Assuming these are in train_ofa_maxvit.py or a utils file
        get_ham10000_dataloaders, 
        # STAGE_CONFIG_PARAMS_FINAL, # These define the supernet structure
        # STEM_OUT_CHANNELS,
        # GLOBAL_K_CHOICES, GLOBAL_E_CHOICES, GLOBAL_MLP_RATIO_CHOICES,
        # NUM_CLASSES_HAM10000,
        # TEACHER_MODEL_PATH, # This will be our KD teacher
        # BASE_HAM10000_PATH,
        # BATCH_SIZE as KD_BATCH_SIZE,
        # TRAIN_IMAGE_SIZE as KD_IMAGE_SIZE
    )
    # If the above configs are not easily importable, define them directly here:
    # Helper to get channel choices (ensure it's available or define it here)
    def get_channel_choices(base_channel, common_divisor=8, multipliers=[0.5, 0.75, 1.0]):
        choices = []
        for m in multipliers:
            raw_c = base_channel * m
            if raw_c < common_divisor and raw_c > 0: 
                c_candidate = common_divisor
            else:
                c_candidate = make_divisible(raw_c, common_divisor)
            if c_candidate > 0: choices.append(c_candidate)
        if not choices and base_channel > 0: choices.append(make_divisible(base_channel, common_divisor))
        elif not choices and base_channel == 0: return [0]
        if not choices and base_channel > 0:
            choices.append(make_divisible(base_channel, common_divisor))
        return sorted(list(set(choices)))

    STEM_OUT_CHANNELS = 64
    GLOBAL_K_CHOICES = [3] 
    GLOBAL_E_CHOICES = [4, 6] 
    GLOBAL_MLP_RATIO_CHOICES = [2.0, 4.0]
    STAGE_CONFIG_PARAMS_FINAL = [
        { 'C_out_stage_choices': get_channel_choices(96), 'depth_choices': [1, 2], 'stride_first_block': 2, 'num_heads_attn_max': 3, 'se_fixed_rd_channels_choices_mbconv': [16, 24, 32]},
        { 'C_out_stage_choices': get_channel_choices(192), 'depth_choices': [1, 2], 'stride_first_block': 2, 'num_heads_attn_max': 6, 'se_fixed_rd_channels_choices_mbconv': [24, 32, 48]},
        { 'C_out_stage_choices': get_channel_choices(384), 'depth_choices': [2, 3, 4, 5], 'stride_first_block': 2, 'num_heads_attn_max': 12, 'se_fixed_rd_channels_choices_mbconv': [32, 48, 64, 96]},
        { 'C_out_stage_choices': get_channel_choices(768), 'depth_choices': [1, 2], 'stride_first_block': 2, 'num_heads_attn_max': 24, 'se_fixed_rd_channels_choices_mbconv': [64, 96, 128]}
    ]
    NUM_CLASSES_HAM10000 = 7
    
except ImportError as e:
    print(f"Error importing OFAMaxViT or helper functions: {e}")
    sys.exit(1)

from sklearn.metrics import f1_score as calculate_macro_f1 # For validation

# --- KD CONFIGURATIONS ---
KD_OUTPUT_PARENT_DIR = "./ofa_maxvit_kd_students_results_100" # Parent dir for all KD student runs
os.makedirs(KD_OUTPUT_PARENT_DIR, exist_ok=True)

# Path to the JSON file containing the list of selected student configurations
STUDENT_CONFIGS_JSON_PATH = "/home/dgx-s-user2/controlleddiffusion/EIS/student_models_100.json" # <<< *** MODIFY THIS IF DIFFERENT ***

# Path to your trained OFA-MaxViT SUPERNET checkpoint
SUPERNET_CHECKPOINT_PATH = "/home/dgx-s-user2/controlleddiffusion/EIS/ofa_maxvit_supernet_training/ofa_maxvit_supernet_best_val_loss.pth" # <<< *** MODIFY THIS ***

# Path to your best Fold 4 MaxViT-Small model (THE TEACHER)
TEACHER_MODEL_PATH_KD = "/home/dgx-s-user2/controlleddiffusion/EIS/Server_Runs_KFold/maxvit_small_tf_224.in1k_k5_e20_20250523-212155/fold_4/maxvit_small_tf_224_in1k_fold4_img224_bs32_lr3e-05_e20_best_loss.pth" # <<< *** MODIFY THIS ***

BASE_HAM10000_PATH_KD = "/home/dgx-s-user2/controlleddiffusion/EIS/HAM10000_extracted/HAM10000_local" # <<< *** MODIFY THIS ***

# KD Training Hyperparameters
KD_IMAGE_SIZE = 224
KD_BATCH_SIZE = 64  # Can be same or different from supernet training
NUM_EPOCHS_KD_STUDENT = 50 # Number of epochs to fine-tune each student (adjust)
LEARNING_RATE_KD_STUDENT = 3e-5 # Fine-tuning LR, usually smaller than supernet's initial LR
WEIGHT_DECAY_KD = 1e-5
KD_ALPHA = 0.3      # Weight for CE loss (student vs true labels)
                    # (1 - KD_ALPHA) will be weight for KD loss (student vs teacher)
KD_TEMPERATURE = 3.0 # Temperature for softening logits

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

def set_seed(seed=42):
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

def load_teacher_for_kd(path, num_classes, device):
    print(f"Loading KD teacher model from: {path}")
    teacher = timm.create_model('maxvit_small_tf_224.in1k', pretrained=False, num_classes=num_classes)
    teacher.load_state_dict(torch.load(path, map_location='cpu')) # Load to CPU first
    teacher.eval()
    return teacher.to(device)


def train_one_student_with_kd(student_id, student_config_dict, supernet_checkpoint_path,
                              teacher_model_instance, train_loader_kd, val_loader_kd,
                              current_student_output_dir):
    print(f"\n--- Starting KD Training for Student: {student_id} ---")
    print(f"Output directory: {current_student_output_dir}")
    os.makedirs(current_student_output_dir, exist_ok=True)

    # 1. Initialize Student (as a configured OFA-MaxViT supernet instance)
    # Each student training run gets its own instance of the supernet model
    student_model = OFAMaxViT(
        stem_out_channels=STEM_OUT_CHANNELS,
        stage_configs=STAGE_CONFIG_PARAMS_FINAL,
        num_classes=NUM_CLASSES_HAM10000,
        global_k_mbconv_choices=GLOBAL_K_CHOICES,
        global_e_mbconv_choices=GLOBAL_E_CHOICES,
        global_mlp_ratio_attn_choices=GLOBAL_MLP_RATIO_CHOICES,
    ).to(DEVICE)

    # Load weights from the TRAINED SUPERNET
    print(f"Loading base weights from supernet checkpoint: {supernet_checkpoint_path}")
    supernet_ckpt = torch.load(supernet_checkpoint_path, map_location=DEVICE)
    supernet_state_dict = supernet_ckpt.get('supernet_state_dict', supernet_ckpt.get('model_state_dict', supernet_ckpt))
    if any(key.startswith('module.') for key in supernet_state_dict.keys()):
        supernet_state_dict = {k.replace('module.', ''): v for k, v in supernet_state_dict.items()}
    student_model.load_state_dict(supernet_state_dict)
    
    # Set the specific architecture for this student
    student_model.set_active_subnet(student_config_dict)
    print(f"Student model configured to: {student_id}")
    # student_model.eval() # For complexity check if needed
    # gmacs, mparams = get_subnet_complexity(student_model, student_config_dict, ...) # Can re-verify here
    # print(f"  Approx Complexity: {mparams:.2f}M Params (Supernet state), {gmacs:.2f} GMACs")

    # 2. Optimizer and Scheduler for the student
    # We fine-tune all parameters of the supernet, but gradients only flow through active path.
    optimizer = optim.AdamW(student_model.parameters(), lr=LEARNING_RATE_KD_STUDENT, weight_decay=WEIGHT_DECAY_KD)
    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=NUM_EPOCHS_KD_STUDENT) # Simple cosine for fine-tuning

    criterion_ce = nn.CrossEntropyLoss() # No label smoothing for student fine-tuning usually, or small
    criterion_kd = nn.KLDivLoss(reduction='batchmean')

    best_val_macro_f1 = 0.0
    best_model_path = ""

    # --- KD Training Loop ---
    for epoch in range(NUM_EPOCHS_KD_STUDENT):
        student_model.train() # Student is in training mode
        teacher_model_instance.eval() # Teacher always in eval

        epoch_loss_kd = 0.0
        progress_bar_kd = tqdm(train_loader_kd, desc=f"KD Epoch {epoch+1}/{NUM_EPOCHS_KD_STUDENT} for {student_id}")

        for images, true_labels in progress_bar_kd:
            images, true_labels = images.to(DEVICE), true_labels.to(DEVICE)
            optimizer.zero_grad()

            # Student predictions (ensure it's using the active student config)
            # student_model.set_active_subnet(student_config_dict) # Should already be set
            student_logits = student_model(images)

            # Teacher predictions (soft labels)
            with torch.no_grad():
                teacher_logits = teacher_model_instance(images)
            
            # Calculate losses
            loss_ce_val = criterion_ce(student_logits, true_labels)
            loss_kd_val = criterion_kd(
                F.log_softmax(student_logits / KD_TEMPERATURE, dim=1),
                F.softmax(teacher_logits / KD_TEMPERATURE, dim=1)
            ) * (KD_TEMPERATURE * KD_TEMPERATURE)

            total_loss = KD_ALPHA * loss_ce_val + (1 - KD_ALPHA) * loss_kd_val
            
            total_loss.backward()
            optimizer.step()

            epoch_loss_kd += total_loss.item()
            progress_bar_kd.set_postfix(loss=total_loss.item(), lr=optimizer.param_groups[0]['lr'])
        
        avg_epoch_loss_kd = epoch_loss_kd / len(train_loader_kd)
        scheduler.step()
        print(f"KD Epoch {epoch+1} - Student {student_id} - Avg Loss: {avg_epoch_loss_kd:.4f}, LR: {optimizer.param_groups[0]['lr']:.6f}")

        # --- Validation for this student ---
        student_model.eval()
        # student_model.set_active_subnet(student_config_dict) # Ensure correct config for val
        val_preds, val_labels_list = [], []
        with torch.no_grad():
            for val_images, val_true_labels in tqdm(val_loader_kd, desc=f"Validating Student {student_id}"):
                val_images, val_true_labels = val_images.to(DEVICE), val_true_labels.to(DEVICE)
                val_logits = student_model(val_images)
                val_preds.extend(torch.argmax(val_logits, dim=1).cpu().numpy())
                val_labels_list.extend(val_true_labels.cpu().numpy())
        
        current_val_macro_f1 = 0.0
        if val_labels_list:
            current_val_macro_f1 = calculate_macro_f1(val_labels_list, val_preds, average='macro', zero_division=0)
            print(f"  Student {student_id} - Validation Macro F1: {current_val_macro_f1:.4f}")

            if current_val_macro_f1 > best_val_macro_f1:
                best_val_macro_f1 = current_val_macro_f1
                if best_model_path: # remove previous best
                    try: os.remove(best_model_path) 
                    except OSError: pass
                best_model_path = os.path.join(current_student_output_dir, f"{student_id}_best_f1_{best_val_macro_f1:.4f}_epoch{epoch+1}.pth")
                # Save the *supernet's state_dict* because the student IS the configured supernet
                torch.save(student_model.state_dict(), best_model_path)
                print(f"  Saved new best model for {student_id} to {best_model_path}")
    
    print(f"Finished KD for Student {student_id}. Best Val Macro F1: {best_val_macro_f1:.4f}")
    return best_val_macro_f1


def main_kd_training():
    set_seed(42)
    print(f"Using device: {DEVICE}")

    # 1. Load Student Configurations
    if not os.path.exists(STUDENT_CONFIGS_JSON_PATH):
        print(f"ERROR: Student configurations JSON not found at {STUDENT_CONFIGS_JSON_PATH}")
        return
    with open(STUDENT_CONFIGS_JSON_PATH, 'r') as f:
        selected_student_configs_list_of_dicts = json.load(f)
    print(f"Loaded {len(selected_student_configs_list_of_dicts)} student configurations for KD.")

    # 2. DataLoaders (use the same ones for all student KD runs for consistency)
    if not os.path.exists(BASE_HAM10000_PATH_KD):
        print(f"ERROR: Dataset path for KD not found: {BASE_HAM10000_PATH_KD}")
        return
    train_loader_kd, val_loader_kd, _ = get_ham10000_dataloaders(
        base_dataset_path=BASE_HAM10000_PATH_KD,
        image_size=KD_IMAGE_SIZE,
        batch_size=KD_BATCH_SIZE,
        num_workers=NUM_WORKERS_KD
    )

    # 3. Load Teacher Model (once)
    if not os.path.exists(TEACHER_MODEL_PATH_KD):
        print(f"ERROR: Teacher model for KD not found: {TEACHER_MODEL_PATH_KD}")
        return
    teacher = load_teacher_for_kd(TEACHER_MODEL_PATH_KD, NUM_CLASSES_HAM10000, DEVICE)

    # 4. Iterate and Train Each Student
    for i, student_config in enumerate(selected_student_configs_list_of_dicts):
        student_name = f"student_{i+1}" # Or derive name from config if unique IDs exist
        
        # Create a unique output directory for this student's KD run
        current_student_run_dir = os.path.join(KD_OUTPUT_PARENT_DIR, student_name)
        
        train_one_student_with_kd(
            student_id=student_name,
            student_config_dict=student_config,
            supernet_checkpoint_path=SUPERNET_CHECKPOINT_PATH,
            teacher_model_instance=teacher,
            train_loader_kd=train_loader_kd,
            val_loader_kd=val_loader_kd,
            current_student_output_dir=current_student_run_dir
        )
    
    print("\nAll selected students have been KD trained.")


if __name__ == '__main__':
    # --- IMPORTANT: Verify all paths and hyperparameters at the top of the script ---
    if any(p.startswith("/path/to") or p.endswith("your_model.pth") for p in [STUDENT_CONFIGS_JSON_PATH, SUPERNET_CHECKPOINT_PATH, TEACHER_MODEL_PATH_KD, BASE_HAM10000_PATH_KD]):
        print("ERROR: Please update placeholder paths in the KD_CONFIGURATIONS section.")
        sys.exit(1)
    main_kd_training()