File size: 7,257 Bytes
1d4bac8 90ed5f2 1d4bac8 90ed5f2 1d4bac8 90ed5f2 1d4bac8 90ed5f2 1d4bac8 90ed5f2 1d4bac8 90ed5f2 1d4bac8 90ed5f2 1d4bac8 90ed5f2 1d4bac8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 | # S2F experiments: data audit, evaluation, baselines and ablations
This folder holds the code for every number reported in the paper beyond the training of the released checkpoints:
- the dataset audit and split manifests;
- the evaluation of the released models;
- the baselines and ablation study.
Run every command from the repository root (`Shape2Force/`); a GPU is used automatically when available.
```
experiments/
layout.py dataset layout, per-set metadata (Young's modulus, TFM settings), manifest helpers
prepare_dataset.py audit legacy data vs dataset/, copy missing images, write splits/*.csv
manifest_data.py data loader driven by a split manifest (same preprocessing as data/cell_dataset.py)
ablation_models.py plain U-Net and S2F generator variants
losses.py per-image WFM-RME and weighted MSE + soft Dice losses (both used for the released model)
configs.py shared training protocol + one entry per variant
train.py train one variant
evaluate.py evaluate released checkpoints / baselines / trained runs on all manifests
run_all.py run the steps above in order
splits/ split manifests (tracked in git)
runs/ training outputs (git-ignored)
results/ audit and evaluation tables
```
## Data
`dataset/` (git-ignored, Zenodo record 10.5281/zenodo.18457879) contains:
- `singlecell/train` (236), `singlecell/test` (46);
- `spheroid/train` (129), `spheroid/test` (28).
The released single-cell model was trained on a larger split (`all_tune`: 280 training / 81 validation pairs, 1, 6 and 10 kPa). The paper also reports independent test sets. `prepare_dataset.py` reconstructs all of this from the original data tree and adds the missing images to `dataset/`:
```
dataset/singlecell_extra/<set>/<sample>/ all_tune images that are not in singlecell/ (1 and 10 kPa)
dataset/singlecell_external/<set>/<sample>/ independent test sets (DMSO, Y-27632, TGF-beta, fibronectin, ...)
```
Each manifest in `splits/` lists one sample per row:
- the sample folder relative to `dataset/`;
- condition, Young's modulus and TFM settings: colour limit of the heatmap (`ythreshold_Pa`, 500 Pa for single cells and 200 Pa for spheroids) and pixel-size setting;
- the substrate setting the original code fed to the network (`substrate_legacy`) and the nominal one (`substrate_nominal`);
- the MD5 of the bright-field image.
## Steps
```bash
# 1. audit only: what is in dataset/ vs the original tree, all_tune composition, leakage check
python -m experiments.run_all --legacy-root <path>/cellforce/dataset --steps audit
# 2. complete dataset/ and write splits/*.csv (then commit experiments/splits/)
python -m experiments.run_all --legacy-root <path>/cellforce/dataset --steps prepare
# 3. released checkpoints, substrate sweep, input-free baseline, image-only stage-1 model
python -m experiments.run_all --steps eval_released --stage1 <path>/outputs/checkpoints/last/last_checkpoint_final.pth
# 4. baselines / ablations (S2F about 7 h, U-Net about 2.5 h each on an RTX A6000; resumable), then their evaluation
python -m experiments.run_all --steps train,eval_runs --variants s2f_full,unet,unet_gan,s2f_no_gan,s2f_no_substrate,s2f_no_force
```
A single run: `python -m experiments.train --variant unet --seed 0`.
### GPUs and run time
- `--steps train` runs one job per GPU in parallel (default `--gpus 0,1`). Each job sees only its own GPU (`CUDA_VISIBLE_DEVICES`), and the next variant starts as soon as a GPU is free.
- A model is never split across GPUs: with batch size 2, that would harm batch-norm/GAN training.
- `train.py` refuses to run without CUDA unless `--allow-cpu` is given.
- Augmentation runs in `--num-workers` DataLoader processes (default 2 per job; use 0 if RAM is short).
- Finished runs (and finished phases) are skipped and interrupted runs resume from the phase's `last_checkpoint.pth`, so the command can simply be re-run.
- Run time on an RTX A6000: S2F about 3 min per epoch, i.e. about 7 h per variant (100 + 40 epochs); U-Net about 2.5 h. On two GPUs the 6 core variants take about 15 h (`s2f_no_force` reuses phase 1 of `s2f_full`).
- Progress is logged to `experiments/runs/<variant>_seed<seed>/train.log` (validation pixel r and WFM-RME every epoch) and `phase*/history.json`.
For a quick check that everything runs, add `--epochs 1 --max-batches 2` (1 epoch per phase; add `--allow-cpu` without a GPU).
## Training protocol (all variants)
Every variant is trained from scratch on `alltune_train` (280 pairs), batch size 2, 1024 × 1024 px, the augmentations of `data/augmentations.py`, Adam (β = 0.5/0.999) with cosine decay after 10 % linear warm-up, in two phases that mirror the staged training of the released model:
| phase | epochs | reconstruction loss | other terms | lr | starts from |
|---|---|---|---|---|---|
| 1 localisation | 100 | 100 × weighted MSE + soft Dice (pos_weight 3, threshold 0.2, Dice 0.5) | 1 × adversarial | 1e-4 | random weights |
| 2 magnitude | 40 | 10 × per-image WFM-RME | 1 × adversarial + 100 × MSE of mean intensity | 1e-5 | best phase-1 generator |
Each phase uses a new discriminator and optimiser and keeps the checkpoint with the lowest validation loss of that phase on `alltune_val` (81 pairs); the final model is the best phase-2 checkpoint. WFM-RME alone does not learn to localise the forces from random weights (a first protocol that used it from the start gave pixel r ≈ 0), which is also why the released model was first trained with the weighted MSE + Dice loss. Variants change one component each (`configs.py`):
| variant | change |
|---|---|
| `s2f_full` | none (from-scratch counterpart of the released model) |
| `unet` | plain U-Net (Ronneberger et al. 2015), no attention modules, no discriminator |
| `unet_gan` | plain U-Net + PatchGAN (pix2pix-style) |
| `mask_only` | binary cell mask instead of bright field (needs `cell_mask.tif` in each sample folder) |
| `s2f_no_gan` | no discriminator (λ_GAN = 0) |
| `s2f_no_force` | no stress-consistency term (λ_force = 0 in phase 2; phase 1 is identical to `s2f_full` and reused) |
| `s2f_gan_0.1`, `s2f_gan_5` | λ_GAN = 0.1 / 5 |
| `s2f_no_substrate` | bright field only (no substrate channels) |
| `s2f_nominal_substrate` | substrate channels from the nominal Young's modulus of each image |
| `s2f_no_multiscale` | no multi-scale input module |
| `s2f_bridge_no_cbam` | bottleneck without CBAM |
| `s2f_bridge_plain` | bottleneck without CBAM, Global Context Module and hierarchical attention |
| `s2f_no_attention_gates` | skip connections without attention gates |
The released checkpoint was obtained by staged fine-tuning from earlier checkpoints, with substrate channels added only in its last two stages (see Methods). `evaluate.py` reports it alongside the from-scratch variants.
## Metrics
`evaluate.py` computes every metric per image and reports the mean and SD across images:
- MSE, MS-SSIM and pixel-wise Pearson r;
- per-image WFM-RME;
- mean traction stress in Pa (heatmap value × colour limit);
- the fraction of saturated ground-truth pixels;
- inference time.
Its outputs are `results/eval/per_image.csv`, `summary.csv` and `model_params.csv`.
|