feat: Enable prediction-map generation for final linear classifier test runs

configs/data/dataset.yaml

@@ -14,7 +14,7 @@ dataset:
     test_split_filename: "split_mapping/test_split.csv"
     tiling_run_id: "fdf7550a2004474f8c7a05dc0cf1fd86"
     filter_tiles_run_id: "4e8f5d3c82124ea5a8f871a42d3ed9ba"
-    stratified_kfold_run_id: "c7eafdffa32743aa9eb6dd2bf3a185b5"
+    stratified_kfold_run_id: "850c81506684450b9af92296acfd045a"
     stratified_group_kfold_run_id: "382b41d2fa894514908e8067949c4326"
     embedding_run_id: "c325e3a5033b4077b6febb0e3e6b0bd6"
     tissue_masks_run_id: "52bc0924f8624b259819c480c7cf213f"

configs/experiment/ml/linear_classifier_stratified_group_kfold.yaml → configs/experiment/ml/linear_classifier_adamw_stratified_group_kfold.yaml

@@ -1,8 +1,13 @@
 # @package _global_
 
 defaults:
-  - /ml/linear_classifier
+  - /ml/task: kfold_linear_classifier
   - _self_
 
 kfold_strategy: stratified_group
 kfold_run_id: ${dataset.mlflow_artifacts.stratified_group_kfold_run_id}
+
+model:
+  optimizer: adamw
+  learning_rate: 1.0e-4
+  weight_decay: 0.0

configs/experiment/ml/linear_classifier_stratified_kfold.yaml → configs/experiment/ml/linear_classifier_adamw_stratified_kfold.yaml

@@ -1,8 +1,13 @@
 # @package _global_
 
 defaults:
-  - /ml/linear_classifier
+  - /ml/task: kfold_linear_classifier
   - _self_
 
 kfold_strategy: stratified
 kfold_run_id: ${dataset.mlflow_artifacts.stratified_kfold_run_id}
+
+model:
+  optimizer: adamw
+  learning_rate: 1.0e-4
+  weight_decay: 0.0

configs/experiment/ml/linear_classifier_final_adamw.yaml

@@ -0,0 +1,23 @@
+# @package _global_
+
+defaults:
+  - /ml/task: final_linear_classifier
+  - override /ml/trainer: early_stopping
+  - _self_
+
+# AdamW final: trained to convergence with the same early-stopping rule as the
+# k-fold sweep (monitor train/loss_epoch, patience 1, min_delta 1e-4), not a
+# fixed 6-epoch budget. weight_decay=1e-3 = best AdamW sweep point (flat curve).
+model:
+  optimizer: adamw
+  learning_rate: 1.0e-4
+  weight_decay: 1.0e-3
+
+trainer:
+  callbacks:
+    model_checkpoint:
+      save_last: true
+
+metadata:
+  run_name: Final Linear Classifier AdamW ${dataset.name}
+  description: "Final AdamW linear probe over frozen Virchow2 embeddings, trained on all training folds with early stopping on train/loss_epoch."

configs/experiment/ml/linear_classifier_final_lbfgs.yaml

@@ -0,0 +1,37 @@
+# @package _global_
+
+defaults:
+  - /ml/task: final_linear_classifier
+  - _self_
+
+# LBFGS final: exact solve of the convex objective on the full training batch.
+# weight_decay=1e-2 = best LBFGS sweep point. Full-batch guard requires
+# train_shuffle=false, train_drop_last=false, train_batch_size >= len(train);
+# num_workers=0 avoids the single-batch IPC deadlock.
+trainer:
+  max_epochs: 10
+
+data:
+  train_batch_size: 1000000000
+  eval_batch_size: 1024
+  train_shuffle: false
+  train_drop_last: false
+  num_workers: 0
+
+model:
+  optimizer: lbfgs
+  learning_rate: 1.0
+  weight_decay: 1.0e-2
+  lbfgs:
+    max_iter: 100
+    max_eval: null
+    tolerance_grad: 1.0e-7
+    tolerance_change: 1.0e-9
+    history_size: 100
+    line_search_fn: strong_wolfe
+    accumulate_batches: 1
+    accumulate_on_cpu: false
+
+metadata:
+  run_name: Final Linear Classifier LBFGS ${dataset.name}
+  description: "Final LBFGS linear probe over frozen Virchow2 embeddings, exact full-batch solve of the convex objective on all training folds."

configs/experiment/ml/linear_classifier_lbfgs_stratified_group_kfold.yaml

@@ -8,7 +8,7 @@ trainer:
   max_epochs: 10
 
 data:
-  batch_size: 1000000000
+  train_batch_size: 1000000000
   train_shuffle: false
   train_drop_last: false
   num_workers: 0

configs/experiment/ml/linear_classifier_lbfgs_stratified_kfold.yaml

@@ -8,7 +8,7 @@ trainer:
   max_epochs: 10
 
 data:
-  batch_size: 1000000000
+  train_batch_size: 1000000000
   train_shuffle: false
   train_drop_last: false
 

configs/experiment/ml/linear_classifier_predict_tissue_tiles.yaml

@@ -0,0 +1,58 @@
+# @package _global_
+
+defaults:
+  - /ml/task: final_linear_classifier
+  - _self_
+
+# Predict over every test-split tile that intersects the tissue mask. This run
+# is unlabeled: it writes predictions/maps for review and does not compute test
+# metrics.
+mode: predict
+final_train_run_id: 0e2230c722134ce0985e09a18ccadf75
+checkpoint: mlflow-artifacts:/104/${final_train_run_id}/artifacts/checkpoints/last/checkpoint.ckpt
+checkpoint_weights_only: false
+
+tissue_embedding_run_id: 95a02c93c164415e94702ad5c83ccca2
+tissue_stats_run_id: ${dataset.mlflow_artifacts.tissue_stats_run_id}
+tissue_stats_artifact_path: tissue_stats
+tissue_column: tile_tissue_coverage
+tissue_min: 0.0
+
+test_embedding_uri: runs:/${tissue_embedding_run_id}/test/tiles
+test_metadata_uri: runs:/${tissue_stats_run_id}/${tissue_stats_artifact_path}/test_tiles.parquet
+
+data:
+  # num_workers MUST be 0 for this config. UnlabeledEmbeddingTilesDataset and
+  # EmbeddingTilesDataset load the entire split into memory in __init__; using
+  # worker processes (num_workers > 0) can deadlock after multiprocessing/fork.
+  num_workers: 0
+  predict:
+    _target_: ml.data.datasets.UnlabeledEmbeddingTilesDataset
+    embedding_uri: ${test_embedding_uri}
+    metadata_uri: ${test_metadata_uri}
+    tissue_column: ${tissue_column}
+    tissue_min: ${tissue_min}
+
+model:
+  optimizer: lbfgs
+  learning_rate: 1.0
+  weight_decay: 1.0e-2
+  lbfgs:
+    max_iter: 100
+    max_eval: null
+    tolerance_grad: 1.0e-7
+    tolerance_change: 1.0e-9
+    history_size: 100
+    line_search_fn: strong_wolfe
+    accumulate_batches: 1
+    accumulate_on_cpu: false
+
+metadata:
+  run_name: Predict Linear Classifier tissue tiles ${dataset.name}
+  description: "Predict final linear classifier over all held-out test tiles intersecting tissue masks for external doctor review."
+  hyperparams:
+    tissue_embedding_run_id: ${tissue_embedding_run_id}
+    tissue_stats_run_id: ${tissue_stats_run_id}
+    tissue_stats_artifact_path: ${tissue_stats_artifact_path}
+    tissue_column: ${tissue_column}
+    tissue_min: ${tissue_min}

configs/experiment/ml/linear_classifier_test_adamw.yaml

@@ -0,0 +1,32 @@
+# @package _global_
+
+defaults:
+  - /experiment/ml/linear_classifier_final_adamw
+  - _self_
+
+# Test the AdamW final checkpoint on the held-out test split. Same model
+# architecture as the final run (required for state_dict load); optimizer
+# fields are inert at test.
+#
+mode: test
+final_train_run_id: a23e478b00b04da79cfbf4d91cada8cd
+checkpoint: mlflow-artifacts:/104/${final_train_run_id}/artifacts/checkpoints/last/checkpoint.ckpt
+checkpoint_weights_only: false
+
+# num_workers MUST stay 0. EmbeddingTilesDataset loads the entire split into
+# one in-memory numpy array in __init__ and __getitem__ is pure numpy indexing
+# (no per-item IO), so workers give zero speedup; num_workers>0 forks the
+# parent (pyarrow/mlflow/fsspec thread state + large array) and deadlocks
+# before the first test batch. final_embedding_tiles defaults to 4; override here.
+data:
+  num_workers: 0
+
+trainer:
+  callbacks:
+    tiff_prediction_maps:
+      _target_: ml.callbacks.TiffPredictionMapWriter
+      slides_uri: runs:/${dataset.mlflow_artifacts.tiling_run_id}/test_split/slides.parquet
+      artifact_path: prediction_maps_tiff
+      draw_region: central_stride
+      slide_selection: all
+      max_slides: null

configs/experiment/ml/linear_classifier_test_lbfgs.yaml

@@ -0,0 +1,36 @@
+# @package _global_
+
+defaults:
+  - /experiment/ml/linear_classifier_final_lbfgs
+  - override /ml/trainer: early_stopping
+  - _self_
+
+# Test the LBFGS final checkpoint on the held-out test split. The full-batch
+# train_batch_size=1e9 is a TRAINING requirement for the convex LBFGS solve
+# only; at test there is no optimization, so use a normal batch to avoid
+# loading the whole test set as one tensor (OOM).
+#
+# num_workers MUST stay 0. EmbeddingTilesDataset loads the entire split into
+# one in-memory numpy array in __init__ and __getitem__ is pure numpy indexing
+# (no per-item IO), so workers give zero speedup; num_workers>0 forks the
+# parent (which holds pyarrow/mlflow/fsspec thread state + the large array),
+# which deadlocks before the first test batch.
+# Same model architecture as the final run (required for state_dict load).
+mode: test
+final_train_run_id: 0e2230c722134ce0985e09a18ccadf75
+checkpoint: mlflow-artifacts:/104/${final_train_run_id}/artifacts/checkpoints/last/checkpoint.ckpt
+checkpoint_weights_only: false
+
+data:
+  train_batch_size: 1024
+  num_workers: 0
+
+trainer:
+  callbacks:
+    tiff_prediction_maps:
+      _target_: ml.callbacks.TiffPredictionMapWriter
+      slides_uri: runs:/${dataset.mlflow_artifacts.tiling_run_id}/test_split/slides.parquet
+      artifact_path: prediction_maps_tiff
+      draw_region: central_stride
+      slide_selection: all
+      max_slides: null

configs/experiment/preprocessing/embeddings_virchow2_tissue_tiles_05mpp.yaml

@@ -0,0 +1,17 @@
+# @package _global_
+
+defaults:
+  - /experiment/preprocessing/embeddings_virchow2_05mpp
+  - _self_
+
+# Embeddings for every tile in the train/test tiling split that intersects the
+# tissue mask. This is the tile universe used for doctor-review prediction maps.
+splits:
+  - test
+tile_source_run_id: ${dataset.mlflow_artifacts.tissue_stats_run_id}
+tile_source_artifact_template: "tissue_stats/{split}_tiles.parquet"
+tile_filter_column: tile_tissue_coverage
+
+metadata:
+  run_name: "Embeddings: ${model} tissue tiles"
+  description: "Tile embeddings using ${model} over held-out test split tiles with tile_tissue_coverage > 0."

configs/ml.yaml

@@ -6,6 +6,7 @@ defaults:
 seed: ${random_seed:}
 mode: ???
 checkpoint: null
+checkpoint_weights_only: null
 
 trainer: {}
 

configs/ml/data/final_embedding_tiles.yaml

@@ -0,0 +1,23 @@
+# @package _global_
+
+data:
+  train_batch_size: 1024
+  num_workers: 4
+  train_shuffle: true
+  train_drop_last: false
+
+  train:
+    _target_: ml.data.datasets.EmbeddingTilesDataset
+    embedding_uri: ${train_embedding_uri}
+    metadata_uri: ${train_metadata_uri}
+    class_indices: ${class_indices}
+    thresholds: ${thresholds}
+    tissue_prop_min: ${tissue_prop_min}
+
+  test:
+    _target_: ml.data.datasets.EmbeddingTilesDataset
+    embedding_uri: ${test_embedding_uri}
+    metadata_uri: ${test_metadata_uri}
+    class_indices: ${class_indices}
+    thresholds: ${thresholds}
+    tissue_prop_min: ${tissue_prop_min}

configs/ml/data/embedding.yaml → configs/ml/data/kfold_embedding_tiles.yaml

@@ -1,7 +1,7 @@
 # @package _global_
 
 data:
-  batch_size: 1024
+  train_batch_size: 1024
   num_workers: 4
   train_shuffle: true
   train_drop_last: true

configs/ml/model/linear_classifier.yaml

@@ -11,15 +11,7 @@ model:
 
   class_indices: ${class_indices}
 
-  optimizer: adamw
-  learning_rate: 1.0e-4
-  weight_decay: 0.0
-  lbfgs:
-    max_iter: 100
-    max_eval: null
-    tolerance_grad: 1.0e-7
-    tolerance_change: 1.0e-9
-    history_size: 100
-    line_search_fn: strong_wolfe
-    accumulate_batches: 1
-    accumulate_on_cpu: false
+  optimizer: ???
+  learning_rate: ???
+  weight_decay: ???
+  lbfgs: null

configs/ml/task/final_linear_classifier.yaml

@@ -0,0 +1,48 @@
+# @package _global_
+
+defaults:
+  - /data: dataset
+  - /class_mapping: collapse_alterations_to_other
+  - /ml/trainer: final_with_prediction_maps
+  - /ml/data: final_embedding_tiles
+  - /ml/model: linear_classifier
+  - _self_
+
+mode: fit
+
+embedding_run_id: ${dataset.mlflow_artifacts.embedding_run_id}
+kfold_strategy: stratified
+kfold_run_id: ${dataset.mlflow_artifacts.stratified_kfold_run_id}
+filter_tiles_run_id: ${dataset.mlflow_artifacts.filter_tiles_run_id}
+
+train_embedding_uri: runs:/${embedding_run_id}/train/tiles
+test_embedding_uri: runs:/${embedding_run_id}/test/tiles
+train_metadata_uri: runs:/${kfold_run_id}/kfold_split/kfold_tiles.parquet
+test_metadata_uri: runs:/${filter_tiles_run_id}/filter_tiles/test_tiles.parquet
+
+tissue_prop_min: 0.2
+thresholds:
+  Nerve: 0.0
+  Blood: 0.0
+  Connective-Tissue: 0.4
+  Fat: 0.6
+  Epithelium: 0.2
+  Muscle: 0.5
+  Other: 0.5
+
+mlflow_artifact_path: linear_classifier_final
+
+metadata:
+  run_name: Final Linear Classifier ${dataset.name}
+  description: "Final linear probe over frozen Virchow2 embeddings trained on all training folds for ${trainer.max_epochs} epochs."
+  hyperparams:
+    embedding_run_id: ${embedding_run_id}
+    kfold_strategy: ${kfold_strategy}
+    kfold_run_id: ${kfold_run_id}
+    filter_tiles_run_id: ${filter_tiles_run_id}
+    tissue_prop_min: ${tissue_prop_min}
+    thresholds: ${thresholds}
+    learning_rate: ${model.learning_rate}
+    weight_decay: ${model.weight_decay}
+    batch_size: ${data.train_batch_size}
+    max_epochs: ${trainer.max_epochs}

configs/ml/linear_classifier.yaml → configs/ml/task/kfold_linear_classifier.yaml

@@ -3,8 +3,8 @@
 defaults:
   - /data: dataset
   - /class_mapping: collapse_alterations_to_other
-  - /ml/trainer: default
-  - /ml/data: embedding
+  - /ml/trainer: early_stopping
+  - /ml/data: kfold_embedding_tiles
   - /ml/model: linear_classifier
   - _self_
 
@@ -49,6 +49,6 @@ metadata:
     learning_rate: ${model.learning_rate}
     weight_decay: ${model.weight_decay}
     lbfgs: ${model.lbfgs}
-    batch_size: ${data.batch_size}
+    batch_size: ${data.train_batch_size}
     train_shuffle: ${data.train_shuffle}
     train_drop_last: ${data.train_drop_last}