rpotter6298
257 lines
9.2 KiB
Python
257 lines
9.2 KiB
Python
import sys
|
|
from pathlib import Path
|
|
|
|
# Make sure we can import from v3 classes when running directly
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[3]))
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn.functional as F
|
|
import matplotlib.pyplot as plt
|
|
|
|
from v3.classes.hypertower_models import SingleEyeHT
|
|
from v3.classes.papila_builders import build_papila_data
|
|
from v3.classes.profiles import build_papila_profile
|
|
from v3.classes.loader_factory import make_loader, filter_eye_samples
|
|
from v3.classes.transforms import build_eval_transform
|
|
from v3.classes.utils import choose_device
|
|
|
|
|
|
def extract_emergent_features(model, loader, device, num_classes=2):
|
|
"""
|
|
Finds samples where the fused bridge is correct, but both individual
|
|
towers are wrong, and extracts the driving features from the fusion_dim.
|
|
|
|
Assumes `model` is a SingleEyeHT. Can be adapted for NTowerHT.
|
|
"""
|
|
model.eval()
|
|
|
|
emergent_samples = []
|
|
|
|
# Access the final linear layer weights in the HTClassifier
|
|
# HTClassifier head is: Sequential(ReLU(), Dropout(), Linear())
|
|
# Index 2 is the Linear layer
|
|
final_linear = model.bridge.classifier_fused.head[2]
|
|
final_weights = (
|
|
final_linear.weight.detach().cpu()
|
|
) # Shape: [num_classes, fusion_dim]
|
|
final_bias = final_linear.bias.detach().cpu()
|
|
|
|
with torch.no_grad():
|
|
for batch in loader:
|
|
x = batch.get("image_1")
|
|
m = batch.get("matrix_1")
|
|
y = batch.get("label_1")
|
|
|
|
if not (torch.is_tensor(x) and torch.is_tensor(m) and torch.is_tensor(y)):
|
|
continue
|
|
|
|
x, m, y = x.to(device), m.to(device), y.to(device)
|
|
|
|
# 1. Get Tower Embeddings
|
|
img_feats = model.img_tower(x)
|
|
md_feats = model.cd_tower(m)
|
|
|
|
# 2. Get Independent Tower Predictions
|
|
logits_i = model.bridge.aux_heads[0](img_feats)
|
|
logits_m = model.bridge.aux_heads[1](md_feats)
|
|
pi = logits_i.argmax(dim=1)
|
|
pm = logits_m.argmax(dim=1)
|
|
|
|
# 3. Get Fused Representation & Prediction
|
|
# _compute_fused applies the Hadamard product and optional SE gate
|
|
z_fused = model.bridge._compute_fused([img_feats, md_feats])
|
|
logits_fused = model.bridge.classifier_fused(z_fused)
|
|
pf = logits_fused.argmax(dim=1)
|
|
|
|
# 4. Find the "Aha!" Moments and "Corrections"
|
|
for i in range(len(y)):
|
|
yi = y[i].cpu().item()
|
|
is_correct = pf[i] == yi
|
|
img_wrong = pi[i] != yi
|
|
md_wrong = pm[i] != yi
|
|
|
|
if not is_correct:
|
|
continue
|
|
|
|
is_aha = img_wrong and md_wrong
|
|
is_correction = img_wrong or md_wrong
|
|
|
|
if is_correction:
|
|
category = (
|
|
"Aha!"
|
|
if is_aha
|
|
else ("Corrected Image" if img_wrong else "Corrected Clinical")
|
|
)
|
|
# Apply the ReLU that happens inside HTClassifier before the Linear layer
|
|
z_act = F.relu(z_fused[i]).cpu()
|
|
|
|
# Calculate how much each feature contributed to the correct class logit
|
|
feature_contributions = z_act * final_weights[yi]
|
|
|
|
emergent_samples.append(
|
|
{
|
|
"patient_id": (
|
|
batch.get("id_1", ["Unknown"])[i]
|
|
if "id_1" in batch
|
|
else "Unknown"
|
|
),
|
|
"target_class": yi,
|
|
"category": category,
|
|
"z_activated": z_act.numpy(),
|
|
"contributions": feature_contributions.numpy(),
|
|
"total_logit": logits_fused[i, yi].cpu().item(),
|
|
}
|
|
)
|
|
|
|
return emergent_samples
|
|
|
|
|
|
def plot_top_emergent_features(emergent_samples, top_k=10):
|
|
"""
|
|
Plots the top K contributing dimensions across all emergent success samples.
|
|
"""
|
|
if not emergent_samples:
|
|
print("No emergent success or correction samples found in this pass.")
|
|
return
|
|
|
|
ahas = [s for s in emergent_samples if s["category"] == "Aha!"]
|
|
corrected_img = [s for s in emergent_samples if s["category"] == "Corrected Image"]
|
|
corrected_clin = [
|
|
s for s in emergent_samples if s["category"] == "Corrected Clinical"
|
|
]
|
|
|
|
print(f"\nFound {len(emergent_samples)} total events of interest:")
|
|
print(f" - 'Aha!' Moments (Both wrong, Fused right): {len(ahas)}")
|
|
print(f" - Corrected Image (Image wrong, Fused right): {len(corrected_img)}")
|
|
print(
|
|
f" - Corrected Clinical (Clinical wrong, Fused right): {len(corrected_clin)}"
|
|
)
|
|
|
|
# Prioritize true Aha moments if they exist, otherwise use corrections
|
|
plot_samples = ahas if ahas else emergent_samples
|
|
plot_title = (
|
|
"Aha! Moments (Both Wrong, Fused Right)"
|
|
if ahas
|
|
else "Fusion Corrections (At least one wrong)"
|
|
)
|
|
|
|
# Average the feature contributions across samples
|
|
all_contribs = np.stack([s["contributions"] for s in plot_samples])
|
|
mean_contribs = all_contribs.mean(axis=0)
|
|
|
|
# Get indices of the top K features with the highest absolute contribution
|
|
top_indices = np.argsort(np.abs(mean_contribs))[-top_k:][::-1]
|
|
|
|
top_values = mean_contribs[top_indices]
|
|
labels = [f"Dim {idx}" for idx in top_indices]
|
|
|
|
plt.figure(figsize=(10, 6))
|
|
colors = ["green" if v > 0 else "red" for v in top_values]
|
|
|
|
plt.barh(np.arange(top_k), top_values[::-1], color=colors[::-1])
|
|
plt.yticks(np.arange(top_k), labels[::-1])
|
|
plt.xlabel("Mean Contribution to Correct Logit")
|
|
plt.title(f"Top {top_k} Features Driving {plot_title}")
|
|
plt.tight_layout()
|
|
plt.show()
|
|
|
|
print(f"\nAnalyzed {len(plot_samples)} samples for this plot.")
|
|
print("\nTop Feature Breakdown:")
|
|
for idx, val in zip(top_indices, top_values):
|
|
print(f"Dimension {idx:3d}: {val:+.4f} average logit push")
|
|
|
|
|
|
def main():
|
|
device = choose_device("auto")
|
|
print(f"Using device: {device}")
|
|
|
|
repo_root = Path(__file__).resolve().parents[3]
|
|
image_dir = repo_root / "Papila" / "FundusImages"
|
|
clinical_dir = repo_root / "Papila" / "ClinicalData"
|
|
|
|
print("Loading PAPILA data with Phase 5 settings...")
|
|
data = build_papila_data(
|
|
image_dir=str(image_dir),
|
|
clinical_dir=str(clinical_dir),
|
|
label_col="Diagnosis",
|
|
cat_cols=["Gender", "Phakic/Pseudophakic"],
|
|
iop_corr_method="ratio",
|
|
iop_drop_raw=True,
|
|
exclude_cols=["Axial_Length"],
|
|
)
|
|
# Filter to binary
|
|
data.df = data.df[data.df["Diagnosis"].isin([0, 1])].reset_index(drop=True)
|
|
|
|
print("Building dataloader (All samples)...")
|
|
profile_eye = build_papila_profile(
|
|
patient_col="Patient ID", label_col="Diagnosis", sample_mode="eye"
|
|
)
|
|
eye_samples = filter_eye_samples(
|
|
profile_eye.build_samples(df=data.df, clinical=data)
|
|
)
|
|
|
|
loader = make_loader(
|
|
eye_samples,
|
|
profile_eye.slot_descriptors(),
|
|
image_transform=build_eval_transform("refugelike"),
|
|
batch_size=16,
|
|
shuffle=False,
|
|
num_workers=4,
|
|
)
|
|
|
|
print("Building SingleEyeHT model...")
|
|
model = SingleEyeHT(
|
|
backbone="refugelike",
|
|
freeze_ratio=0.0,
|
|
augment=False,
|
|
clinical_data=data,
|
|
num_classes=2,
|
|
cd_hidden_dim=128,
|
|
fusion_dim=256,
|
|
bridge_mode="fused",
|
|
).to(device)
|
|
|
|
base_ckpt_dir = repo_root / "v3" / "results" / "phase5" / "logit_mlp_head_ckpt"
|
|
checkpoints = sorted(list(base_ckpt_dir.rglob("best_single.pt")))
|
|
|
|
if not checkpoints:
|
|
print(f"\n[!] No checkpoints found in {base_ckpt_dir}")
|
|
print("Please ensure you ran the jobs with the --save-checkpoints flag.")
|
|
return
|
|
|
|
all_emergent_data = []
|
|
|
|
for checkpoint_path in checkpoints:
|
|
print(f"\nProcessing {checkpoint_path.relative_to(repo_root)}...")
|
|
state_dict = torch.load(checkpoint_path, map_location=device)
|
|
|
|
# Backward compatibility for checkpoints saved before the N-tower bridge refactor
|
|
new_state_dict = {}
|
|
for k, v in state_dict.items():
|
|
k = k.replace("bridge.W_img.", "bridge.W.0.")
|
|
k = k.replace("bridge.W_md.", "bridge.W.1.")
|
|
k = k.replace("bridge.ln_img.", "bridge.ln.0.")
|
|
k = k.replace("bridge.ln_md.", "bridge.ln.1.")
|
|
k = k.replace("bridge.classifier_img.", "bridge.aux_heads.0.")
|
|
k = k.replace("bridge.classifier_cd.", "bridge.aux_heads.1.")
|
|
k = k.replace(
|
|
"bridge.classifier_fused.2.", "bridge.classifier_fused.head.2."
|
|
)
|
|
new_state_dict[k] = v
|
|
|
|
model.load_state_dict(new_state_dict)
|
|
|
|
emergent_data = extract_emergent_features(model, loader, device)
|
|
all_emergent_data.extend(emergent_data)
|
|
print(f" -> Found {len(emergent_data)} events of interest.")
|
|
|
|
print(
|
|
f"\nTotal aggregated events across {len(checkpoints)} checkpoints: {len(all_emergent_data)}"
|
|
)
|
|
plot_top_emergent_features(all_emergent_data, top_k=15)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|