Align brainscore pipeline with official Brain-Score (closes #9)

benchmark.py

@@ -3,64 +3,94 @@
 import numpy as np
 import torch
 
-from utils.dataset import TVSD_Dataset
+from utils.dataset import TVSD_Dataset, TVSD_TestDataset
 from utils.load_model import load_model
-from utils.brainscore import compute_brain_score
+from utils.brainscore import score_train_test
+
+
+def _load_activations(path, device):
+    """Load saved activations and flatten to (n_stimuli, features)."""
+    a = torch.load(path, map_location=device)
+    return a.reshape(a.shape[0], -1).detach().cpu().float().numpy()
 
 
 def main(args):
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     print(f"Using device: {device}")
 
     model, model_name, _ = load_model(args.model_config)
-    # layers = [name for name, module in model.named_modules() if 'relu' not in name]
 
-    tvsd_dataset = TVSD_Dataset(root_dir=args.root_dir, monkey=args.monkey, region=args.region)
+    # Train responses (single-trial train_MUA) fit the mapping; the held-out test
+    # responses (repetition-averaged test_MUA) are what we score against. Neuron
+    # selection uses the single-trial reliability (`reliab`); the noise ceiling is
+    # the split-half + Spearman-Brown internal consistency of the 30-rep average.
+    train_ds = TVSD_Dataset(root_dir=args.root_dir, monkey=args.monkey, region=args.region)
+    test_ds = TVSD_TestDataset(
+        root_dir=args.root_dir,
+        monkey=args.monkey,
+        region=args.region,
+        recompute_reliability=True,
+        spearman_brown=True,
+        n_boot=args.ceiling_n_boot,
+        random_state=args.random_state,
+    )
+
+    single_trial_reliab = train_ds.reliability.numpy()  # (C,)
+    Y_train_all = np.asarray(train_ds.responses)  # (n_train, C)
+    Y_test_all = np.asarray(test_ds.responses).mean(axis=0)  # (n_test, C) = test_MUA
+    ceiling_all = test_ds.reliability.numpy()  # (C,) SB ceiling
+
+    mask = single_trial_reliab > args.reliability_threshold
+    Y_train_all = Y_train_all[:, mask]
+    Y_test_all = Y_test_all[:, mask]
+    ceiling = ceiling_all[mask]
+    print(
+        f"{int(mask.sum())} neuroids retained with single-trial reliability "
+        f"> {args.reliability_threshold} (median ceiling {np.nanmedian(ceiling):.3f})"
+    )
+
+    train_dir = f"{args.output_dir}/activations/TVSD_train/{model_name}"
+    test_dir = f"{args.output_dir}/activations/TVSD_test/{model_name}"
 
     layer_scores = {}
-    activation_dir = f"{args.output_dir}/activations/TVSD/{model_name}"
-    layers = os.listdir(activation_dir)
-    for layer in layers:
+    for layer in sorted(os.listdir(train_dir)):
         print(f"===== EVALUATING LAYER: {layer} =========")
-        layer_dir = f"{activation_dir}/{layer}"
-        activation_path = f"{layer_dir}/activations.pt"
-        if not os.path.exists(activation_path):
-            print(f"Activations for layer {layer} not found at {activation_path}. Skipping.")
-            continue
-        activations = torch.load(f"{layer_dir}/activations.pt", map_location=device)
-        if activations is not None:
-            print(f"Layer: {layer}, Activations shape: {activations.shape}")
-        else:
-            print(f"No activations found for layer: {layer}")
+        train_path = f"{train_dir}/{layer}/activations.pt"
+        test_path = f"{test_dir}/{layer}/activations.pt"
+        if not (os.path.exists(train_path) and os.path.exists(test_path)):
+            print(f"Missing train/test activations for layer {layer}. Skipping.")
             continue
 
-        # float32: float16 overflows StandardScaler's in-place divide to inf.
-        activations = (
-            activations.reshape(activations.shape[0], -1).detach().cpu().float().numpy()
-        )  # [B, H * W * C]
-        neural_responses, reliability = tvsd_dataset[: activations.shape[0]]
-        reliability_mask = reliability > args.reliability_threshold
-        neural_responses = neural_responses[:, reliability_mask]  # [B, C']
+        X_train = _load_activations(train_path, device)
+        X_test = _load_activations(test_path, device)
+
+        # Align rows with responses in case generation was truncated (e.g. --max_batches).
+        n_tr = min(X_train.shape[0], Y_train_all.shape[0])
+        n_te = min(X_test.shape[0], Y_test_all.shape[0])
+        X_train, Y_train = X_train[:n_tr], Y_train_all[:n_tr]
+        X_test, Y_test = X_test[:n_te], Y_test_all[:n_te]
+        print(f"train X{X_train.shape} Y{Y_train.shape} | test X{X_test.shape} Y{Y_test.shape}")
 
         if args.noise_test:
-            activations = np.random.normal(
-                size=activations.shape
-            )  # Use random noise for null results
+            X_train = np.random.normal(size=X_train.shape)
+            X_test = np.random.normal(size=X_test.shape)
         if args.permutation_test:
-            np.random.shuffle(neural_responses)
-
-        print(
-            f"{neural_responses.shape[1]} neural responses retained with reliability > {args.reliability_threshold}"
-        )
+            np.random.shuffle(Y_test)
 
-        layer_score, layer_std = compute_brain_score(
-            X=activations,
-            Y=neural_responses,
-            n_splits=args.n_splits,
+        layer_score, layer_std = score_train_test(
+            X_train,
+            Y_train,
+            X_test,
+            Y_test,
             reducer=args.reducer,
             correlation_fn=args.correlation_fn,
             pca_components=args.pca_components,
-            preprocessed=args.preprocessed,
+            skip_pca=args.skip_pca,
+            standardize=args.standardize,
+            ceiling=ceiling,
+            ceiling_normalize=args.ceiling_normalize,
+            n_boot_ci=args.n_boot_ci,
+            random_state=args.random_state,
         )
         layer_scores[layer] = {"score": layer_score, "std": layer_std}
         print(f"Score: {layer_score}, Std: {layer_std}")
@@ -108,60 +138,76 @@ def main(args):
         "--output_dir",
         type=str,
         default=f"{os.getcwd()}/outputs",
-        help="Directory to save activations.",
+        help="Directory holding generated activations and results.",
     )
     parser.add_argument(
         "--reliability_threshold",
         type=float,
         default=0.3,
-        help="Reliability threshold for neural responses.",
+        help="Keep neuroids whose single-trial reliability exceeds this (TVSD-paper selection).",
     )
     parser.add_argument(
         "--reducer",
         type=str,
         default="median",
         choices=["mean", "median"],
-        help="Reduction method for brain score.",
+        help="Reduction across neuroids.",
     )
     parser.add_argument(
         "--correlation_fn",
         type=str,
         default="pearson",
         choices=["pearson", "spearman"],
-        help="Correlation function to use for brain score computation.",
+        help="Correlation function for predictivity.",
     )
     parser.add_argument(
-        "--n_splits",
+        "--pca_components",
         type=int,
-        default=5,
-        help="Number of splits for KFold cross-validation.",
+        default=100,
+        help="PCA components to reduce features to (fit on train, applied to test).",
     )
     parser.add_argument(
-        "--pca_components",
+        "--skip_pca",
+        action="store_true",
+        help="Skip in-benchmark PCA (e.g. features were already reduced at generation).",
+    )
+    parser.add_argument(
+        "--standardize",
+        action="store_true",
+        help="Z-score features/targets (fit on train). Off by default to match Brain-Score.",
+    )
+    parser.add_argument(
+        "--ceiling_normalize",
+        action="store_true",
+        help="Normalize scores by the median noise ceiling (Brain-Score-style).",
+    )
+    parser.add_argument(
+        "--ceiling_n_boot",
         type=int,
-        default=100,
-        help="Number of PCA components to use.",
+        default=30,
+        help="Bootstrap splits for the split-half+SB noise ceiling from test reps.",
     )
     parser.add_argument(
-        "--skip_interval",
+        "--n_boot_ci",
         type=int,
-        default=1,
-        help="Skip every n-th image in the dataset.",
+        default=100,
+        help="Bootstrap resamples over test stimuli for the score's std.",
     )
     parser.add_argument(
-        "--preprocessed",
-        action="store_true",
-        help="Whether the data is preprocessed (scaled and PCA applied).",
+        "--random_state",
+        type=int,
+        default=42,
+        help="Random seed for the ceiling and bootstrap.",
     )
     parser.add_argument(
         "--noise_test",
         action="store_true",
-        help="Run with pure noise to test. Useful for debugging.",
+        help="Replace activations with pure noise (null control).",
     )
     parser.add_argument(
         "--permutation_test",
         action="store_true",
-        help="Randomly permute neural responses. Useful for debugging.",
+        help="Randomly permute test responses (null control).",
     )
 
     args = parser.parse_args()

generate_activations.py

@@ -4,7 +4,7 @@
 import torch
 from torch.utils.data import DataLoader
 
-from utils.dataset import TVSD_Dataset
+from utils.dataset import TVSD_Dataset, TVSD_TestDataset
 from utils.hooks import Activations
 from utils.load_model import load_model, resolve_transform
 
@@ -19,36 +19,49 @@ def main(args):
     layers = [name for name, module in model.named_modules() if "relu" not in name]
     hook_layers = [layers[i] for i in range(0, len(layers), hook_interval)]
 
+    # Activations are saved per split under TVSD_train / TVSD_test so the benchmark
+    # can fit on train and score on the held-out test images.
+    dataset_name = f"TVSD_{args.split}"
     activations = Activations(
         output_dir=args.output_dir,
         model_name=model_name,
-        dataset_name="TVSD",
+        dataset_name=dataset_name,
         pca_components=args.pca_components,
     )
     activations.register(model, hook_layers)
 
-    tvsd_dataset = TVSD_Dataset(root_dir=args.root_dir, monkey=args.monkey, region=args.region)
+    if args.split == "train":
+        tvsd_dataset = TVSD_Dataset(root_dir=args.root_dir, monkey=args.monkey, region=args.region)
+    else:
+        tvsd_dataset = TVSD_TestDataset(
+            root_dir=args.root_dir, monkey=args.monkey, region=args.region
+        )
     things_dataset = tvsd_dataset.get_things(things_path=args.things_path, transform=transform)
     things_loader = DataLoader(things_dataset, batch_size=args.batch_size, shuffle=False)
-    shuffled_things_loader = DataLoader(things_dataset, batch_size=args.batch_size, shuffle=True)
 
     if args.pca_components is not None:
-        print("Training IPCA models...")
-        activations.set_training_mode(True)
-
-        for i, batch in tqdm(enumerate(shuffled_things_loader), desc="Training IPCA"):
-            if args.max_pca_train_batches and i >= int(args.max_pca_train_batches):
-                break
-            activations.set_batch(i)
-            with torch.no_grad():
-                _ = model(batch.to(device))
-            activations.finalize_batch_training()
-
-        torch.cuda.empty_cache()
-        activations.save_ipca_models()
-        print("IPCA training completed.")
-
-    print("Generating activations...")
+        if args.split == "train":
+            # Fit the IPCA basis on the train images only.
+            print("Training IPCA models on train split...")
+            activations.set_training_mode(True)
+            shuffled_loader = DataLoader(things_dataset, batch_size=args.batch_size, shuffle=True)
+            for i, batch in tqdm(enumerate(shuffled_loader), desc="Training IPCA"):
+                if args.max_pca_train_batches and i >= int(args.max_pca_train_batches):
+                    break
+                activations.set_batch(i)
+                with torch.no_grad():
+                    _ = model(batch.to(device))
+                activations.finalize_batch_training()
+            torch.cuda.empty_cache()
+            activations.save_ipca_models()
+            activations.set_training_mode(False)
+            print("IPCA training completed.")
+        else:
+            # Apply the train-fit IPCA to the test split (shared basis, no leakage).
+            train_models = f"{args.output_dir}/activations/TVSD_train/{model_name}"
+            activations.load_ipca_models(train_models)
+
+    print(f"Generating {args.split} activations...")
     activations.set_training_mode(False)
 
     for i, batch in tqdm(enumerate(things_loader), desc="Generating activations"):
@@ -92,6 +105,14 @@ def main(args):
         choices=["V1", "V4", "IT"],
         help="Which brain region to use.",
     )
+    parser.add_argument(
+        "--split",
+        type=str,
+        default="train",
+        choices=["train", "test"],
+        help="Which image split to generate activations for. 'test' reuses the "
+        "train-fit IPCA basis (generate train first).",
+    )
     parser.add_argument(
         "--things_path",
         type=str,

scripts/benchmark.sh

@@ -11,9 +11,12 @@
 regions=("V1" "V4" "IT")
 export PYTHONPATH="$PYTHONPATH:$(pwd)"
 module load cuda cudnn
+# Fit on train activations, score on the held-out test images. --skip_pca because
+# activations were already reduced by the train-fit IPCA at generation time (the
+# same basis is applied to the test images, so it is leak-free).
 python -u benchmark.py  --model_config $1 \
                         --monkey monkeyF \
                         --region ${regions[$SLURM_ARRAY_TASK_ID]} \
-                        --n_splits 4 \
-                        --preprocessed
-                        
+                        --skip_pca \
+                        --ceiling_normalize
+

scripts/generate_activations.sh

@@ -11,8 +11,18 @@
 
 export PYTHONPATH="$PYTHONPATH:$(pwd)"
 module load cuda cudnn
+
+# Train split first: fits the IPCA basis and saves reduced train activations.
 python -u generate_activations.py   --model_config $1 \
+                                    --split train \
                                     --monkey monkeyF \
                                     --batch_size 128 \
                                     --pca_components 100 \
                                     --max_pca_train_batches 4
+
+# Test split: reuses the train-fit IPCA (no refitting -> leak-free shared basis).
+python -u generate_activations.py   --model_config $1 \
+                                    --split test \
+                                    --monkey monkeyF \
+                                    --batch_size 128 \
+                                    --pca_components 100