search_ofa_subnets.py

import torch
import torch.nn as nn
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 # For saving config dicts
import pandas as pd # For saving results to CSV

# Your OFA building blocks and OFAMaxViT model
from ofa_maxvit_building_blocks import OFAMaxViT, make_divisible # Ensure make_divisible is accessible
# Your data loading functions (get_ofa_val_transforms, AlbumentationsImageFolder etc.)
# For simplicity, I'll include placeholders; replace with your actual data loading.
# You need: get_ofa_val_transforms, AlbumentationsImageFolder (if used), get_ham10000_dataloaders (or just val part)
from train_ofa_maxvit import get_ofa_val_transforms, AlbumentationsImageFolder # Assuming they are in train_ofa_maxvit_ddp.py
                                                                                # Or move them to a dataset_utils.py

# --- CONFIGURATIONS ---
# These should match the supernet's initialization for architecture choices
STEM_OUT_CHANNELS = 64
GLOBAL_K_CHOICES = [3] 
GLOBAL_E_CHOICES = [4, 6] 
GLOBAL_MLP_RATIO_CHOICES = [2.0, 4.0]

# Helper to get channel choices (ensure it's available)
# def get_channel_choices(...): ... (defined as in training script)
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]
    return sorted(list(set(choices)))


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
EVAL_IMAGE_SIZE = 224
EVAL_BATCH_SIZE = 256 # Can be larger for inference
NUM_WORKERS_EVAL = 4

# Path to your trained supernet checkpoint
# SUPERNET_CHECKPOINT_PATH = "./ofa_maxvit_supernet_training/ofa_maxvit_supernet_best_val_loss.pth" # Example
SUPERNET_CHECKPOINT_PATH = "/home/dgx-s-user2/controlleddiffusion/EIS/ofa_maxvit_supernet_training/ofa_maxvit_supernet_best_val_loss.pth" # Or final one
# <<< *** MODIFY THIS TO YOUR ACTUAL BEST/FINAL SUPERNET CHECKPOINT ***

BASE_HAM10000_PATH = "/home/dgx-s-user2/controlleddiffusion/EIS/HAM10000_extracted/HAM10000_local" # <<< *** MODIFY THIS ***
SEARCH_OUTPUT_DIR = "./ofa_maxvit_subnet_search_results"
os.makedirs(SEARCH_OUTPUT_DIR, exist_ok=True)

NUM_SUBNETS_TO_SAMPLE_AND_EVAL = 1000 # Number of random subnets to evaluate (start with 100-1000)

def get_val_loader_for_search(base_dataset_path, image_size, batch_size, num_workers):
    val_dir_try1 = os.path.join(base_dataset_path, "validation")
    val_dir_try2 = os.path.join(base_dataset_path, "val")
    val_dir = val_dir_try1 if os.path.isdir(val_dir_try1) else val_dir_try2
    if not os.path.isdir(val_dir):
        raise FileNotFoundError(f"Validation directory not found at {val_dir_try1} or {val_dir_try2}")

    val_transforms = get_ofa_val_transforms(image_size) # Assuming get_ofa_val_transforms is defined
    val_dataset = AlbumentationsImageFolder(root=val_dir, transform=val_transforms) # Assuming AlbumentationsImageFolder
    
    val_loader = DataLoader(
        val_dataset, batch_size=batch_size, shuffle=False,
        num_workers=num_workers, pin_memory=True
    )
    print(f"Validation dataset: {len(val_dataset)} samples, {len(val_loader)} batches.")
    return val_loader, val_dataset.classes

def evaluate_subnet(supernet_model, active_config, val_loader, criterion, device, config_name="subnet"):
    supernet_model.set_active_subnet(active_config)
    supernet_model.eval()

    all_preds, all_labels = [], []
    total_loss = 0.0
    num_samples = 0

    with torch.no_grad():
        for images, labels in tqdm(val_loader, desc=f"Evaluating {config_name}", leave=False):
            images, labels = images.to(device), labels.to(device)
            logits = supernet_model(images)
            loss = criterion(logits, labels)
            
            total_loss += loss.item() * images.size(0)
            num_samples += images.size(0)
            
            all_preds.extend(torch.argmax(logits, dim=1).cpu().numpy())
            all_labels.extend(labels.cpu().numpy())

    avg_loss = total_loss / num_samples if num_samples > 0 else float('inf')
    macro_f1 = 0.0
    if all_labels and len(np.unique(all_labels)) > 1: # Check for diversity for F1 score
        from sklearn.metrics import f1_score as sk_f1_score # Ensure import
        macro_f1 = sk_f1_score(all_labels, all_preds, average='macro', zero_division=0)
    
    return avg_loss, macro_f1

def main_search():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")

    # 1. Validation DataLoader
    if not os.path.exists(BASE_HAM10000_PATH):
        print(f"ERROR: Dataset path not found: {BASE_HAM10000_PATH}")
        return
    val_loader, class_names = get_val_loader_for_search(
        BASE_HAM10000_PATH, EVAL_IMAGE_SIZE, EVAL_BATCH_SIZE, NUM_WORKERS_EVAL
    )
    num_actual_classes = len(class_names)
    if NUM_CLASSES_HAM10000 != num_actual_classes:
        print(f"Warning: Configured NUM_CLASSES ({NUM_CLASSES_HAM10000}) != detected ({num_actual_classes})")
        # Use detected for model loading if there's a mismatch, but ideally they should align.


    # 2. Load Trained OFA-MaxViT Supernet
    print(f"Loading OFA-MaxViT Supernet from: {SUPERNET_CHECKPOINT_PATH}")
    if not os.path.exists(SUPERNET_CHECKPOINT_PATH):
        print(f"ERROR: Supernet checkpoint not found: {SUPERNET_CHECKPOINT_PATH}")
        return
        
    supernet = OFAMaxViT( # Initialize with the same arch params as during training
        stem_out_channels=STEM_OUT_CHANNELS,
        stage_configs=STAGE_CONFIG_PARAMS_FINAL,
        num_classes=num_actual_classes, # Use detected num_classes
        global_k_mbconv_choices=GLOBAL_K_CHOICES,
        global_e_mbconv_choices=GLOBAL_E_CHOICES,
        global_mlp_ratio_attn_choices=GLOBAL_MLP_RATIO_CHOICES,
        # Add other global choices if your OFAMaxViT init expects them
    ).to(device)

    checkpoint = torch.load(SUPERNET_CHECKPOINT_PATH, map_location=device)
    if 'supernet_state_dict' in checkpoint:
        supernet.load_state_dict(checkpoint['supernet_state_dict'])
    elif 'model_state_dict' in checkpoint: # Some save formats
        supernet.load_state_dict(checkpoint['model_state_dict'])
    else:
        supernet.load_state_dict(checkpoint) # Assume it's just the state_dict
    print("Supernet loaded successfully.")

    criterion_eval = nn.CrossEntropyLoss().to(device)
    
    # --- Random Search Loop ---
    results = []
    print(f"Starting random search for {NUM_SUBNETS_TO_SAMPLE_AND_EVAL} subnets...")

    # Evaluate teacher (Fold 4 model) on this val set for reference, if desired
    # (Requires loading teacher model and separate evaluation)

    # Evaluate max and min subnets from the supernet first
    boundary_configs_to_eval = {
        "supernet_max": supernet.get_max_subnet_config(),
        "supernet_min": supernet.get_min_subnet_config(),
    }
    for name, b_config in boundary_configs_to_eval.items():
        print(f"Evaluating boundary subnet: {name}")
        # Simple param count for this config (very rough estimate, sum of active choices)
        # This is NOT accurate parameters. Accurate requires summing params of active layers.
        # params_estimate = sum(d for s_conf in b_config['stage_configs'] for d in [s_conf['active_depth'], s_conf['C_out_active_stage']] )
        
        loss, f1 = evaluate_subnet(supernet, b_config, val_loader, criterion_eval, device, config_name=name)
        results.append({
            'config_name': name,
            'macro_f1': f1,
            'val_loss': loss,
            'config_details': json.dumps(b_config) # Save full config
        })
        print(f"  {name}: Val Loss={loss:.4f}, Macro F1={f1:.4f}")


    for i in range(NUM_SUBNETS_TO_SAMPLE_AND_EVAL):
        active_config = supernet.sample_active_subnet_config()
        config_name = f"random_subnet_{i+1}"
        
        print(f"Evaluating {config_name} ({i+1}/{NUM_SUBNETS_TO_SAMPLE_AND_EVAL})...")
        loss, f1 = evaluate_subnet(supernet, active_config, val_loader, criterion_eval, device, config_name=config_name)
        
        results.append({
            'config_name': config_name,
            'macro_f1': f1,
            'val_loss': loss,
            'config_details': json.dumps(active_config) # Save full config
        })
        print(f"  {config_name}: Val Loss={loss:.4f}, Macro F1={f1:.4f}")

        # Save results periodically
        if (i + 1) % 20 == 0 or (i + 1) == NUM_SUBNETS_TO_SAMPLE_AND_EVAL:
            df_results = pd.DataFrame(results)
            df_results = df_results.sort_values(by='macro_f1', ascending=False)
            results_path = os.path.join(SEARCH_OUTPUT_DIR, f"subnet_search_results_upto_{i+1}.csv")
            df_results.to_csv(results_path, index=False)
            print(f"Search results saved to {results_path}")

    print("\n--- Subnet Search Finished ---")
    df_results = pd.DataFrame(results)
    df_results = df_results.sort_values(by='macro_f1', ascending=False)
    final_results_path = os.path.join(SEARCH_OUTPUT_DIR, "subnet_search_results_final.csv")
    df_results.to_csv(final_results_path, index=False)
    print(f"All search results saved to {final_results_path}")
    print("\nTop performing sampled subnets (by Macro F1):")
    print(df_results.head(10))

if __name__ == '__main__':
    # --- IMPORTANT: Set SUPERNET_CHECKPOINT_PATH and BASE_HAM10000_PATH above ---
    if SUPERNET_CHECKPOINT_PATH.startswith("./ofa_maxvit_supernet_training/ofa_maxvit_supernet"):
        if not os.path.exists(SUPERNET_CHECKPOINT_PATH): # Check if default placeholder is missing
             print(f"ERROR: Default SUPERNET_CHECKPOINT_PATH ('{SUPERNET_CHECKPOINT_PATH}') not found. Please update the path.")
             sys.exit(1)
    elif not os.path.exists(SUPERNET_CHECKPOINT_PATH):
        print(f"ERROR: Specified SUPERNET_CHECKPOINT_PATH ('{SUPERNET_CHECKPOINT_PATH}') not found. Please update the path.")
        sys.exit(1)

    if not os.path.exists(BASE_HAM10000_PATH):
        print(f"ERROR: BASE_HAM10000_PATH ('{BASE_HAM10000_PATH}') not found. Please update the path.")
        sys.exit(1)
        
    main_search()