diff --git a/.gitattributes b/.gitattributes index a6344aac8c09253b3b630fb776ae94478aa0275b..2bf9b9fbb61b5001482b4851446795b29da7397d 100644 --- a/.gitattributes +++ b/.gitattributes @@ -33,3 +33,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text *.zip filter=lfs diff=lfs merge=lfs -text *.zst filter=lfs diff=lfs merge=lfs -text *tfevents* filter=lfs diff=lfs merge=lfs -text +fomo_tune_baseline/output/task1/task1.sif filter=lfs diff=lfs merge=lfs -text +fomo_tune_baseline/output/task3/task3.sif filter=lfs diff=lfs merge=lfs -text +fomo_tune_baseline/output/task5/task5.sif filter=lfs diff=lfs merge=lfs -text diff --git a/fomo_tune_baseline/build.sh b/fomo_tune_baseline/build.sh new file mode 100644 index 0000000000000000000000000000000000000000..4a5b2abb12dbd201b05763feeccc7c3980848e98 --- /dev/null +++ b/fomo_tune_baseline/build.sh @@ -0,0 +1,31 @@ +#!/usr/bin/env bash +# Package each trained run into its challenge .sif. Run on the login node, after launch.sh: +# apptainer lives there and on n-6 only, and a build needs no GPU driver. +# +# Slow. Apptainer always re-runs %post, so each of the three re-downloads ~3G of wheels. + +set -euo pipefail + +ROOT="/data/connor/nanobrain.1" +cd $ROOT + +EXP_DIR="experiments/fomo_tune_baseline" +OUT_DIR="${EXP_DIR}/output" + +runs=(task1 task5 task3) + +for name in "${runs[@]}"; do + # build.py names the sif after `task` in the run's saved config, which is the run name here + sif="${OUT_DIR}/${name}/${name}.sif" + + if [[ -f "${sif}" ]]; then + echo "sif ${sif} exists; skipping" + continue + fi + + echo "=== ${name} ===" + uv run --no-sync python -m fomo_tune.build "${OUT_DIR}/${name}" +done + +echo "=== sifs ===" +ls -lh "${OUT_DIR}"/*/*.sif diff --git a/fomo_tune_baseline/launch.sh b/fomo_tune_baseline/launch.sh new file mode 100644 index 0000000000000000000000000000000000000000..79e67b364a5ad3b6de9a3ca0799c3606df6653c9 --- /dev/null +++ b/fomo_tune_baseline/launch.sh @@ -0,0 +1,43 @@ +#!/usr/bin/env bash +#SBATCH --job-name=fomo_tune +#SBATCH --nodes=1 +#SBATCH --ntasks-per-node=1 +#SBATCH --gpus-per-task=1 +#SBATCH --time=2:00:00 +#SBATCH --partition=main +#SBATCH --output=slurms/slurm-%j.out +#SBATCH --account=sophont + +set -euo pipefail + +ROOT="/data/connor/nanobrain.1" +cd $ROOT + +EXP_DIR="experiments/fomo_tune_baseline" +OUT_DIR="${EXP_DIR}/output" + +# name, module. Cheapest first, so a broken environment fails in 90s. +# One job, run sequentially: the whole grid is ~12 minutes and each task loads the same +# 3.9G checkpoint, so an array would only add queue time and cache contention. +runs=( + "task1 main_task1" + "task5 main_task5" + "task3 main_task3" +) + +for run in "${runs[@]}"; do + read -r name module <<<"${run}" + + if [[ -f "${OUT_DIR}/${name}/metrics.json" ]]; then + echo "result ${name} exists; skipping" + continue + fi + + echo "=== ${name} ===" + uv run --no-sync python -m "fomo_tune.${module}" train \ + output_root="${OUT_DIR}" \ + name="${name}" +done + +echo "=== results ===" +cat "${OUT_DIR}"/*/metrics.json diff --git a/fomo_tune_baseline/output/task1/build/Apptainer.def b/fomo_tune_baseline/output/task1/build/Apptainer.def new file mode 100644 index 0000000000000000000000000000000000000000..98a8bdfc1fefdbffdaed0dedbc7166a3da28b472 --- /dev/null +++ b/fomo_tune_baseline/output/task1/build/Apptainer.def @@ -0,0 +1,31 @@ +Bootstrap: docker +From: python:3.11-slim + +# NOT buildable where it sits: the %files paths below are relative to the build cwd, which is the +# staging dir `build.py` writes. Build it with `python -m fomo_tune.build `, not by +# pointing apptainer at this file. +# +# Versions are pinned to the training environment: numpy, scikit-learn and joblib because they +# unpickle `head.joblib`, torch because that is what the checkpoint was written by. + +%files + fomo_tune /app/fomo_tune + smri_mae /app/smri_mae + model /app/model + predict.py /app/predict.py + +%post + pip install --no-cache-dir \ + torch==2.8.0 \ + numpy==2.4.6 \ + nibabel==5.4.2 \ + einops==0.8.2 \ + jaxtyping==0.3.10 \ + timm==1.0.27 \ + huggingface-hub==0.36.2 \ + scikit-learn==1.8.0 \ + joblib==1.5.3 \ + omegaconf==2.3.0 + +%runscript + exec python /app/predict.py "$@" diff --git a/fomo_tune_baseline/output/task1/build/fomo_tune/README.md b/fomo_tune_baseline/output/task1/build/fomo_tune/README.md new file mode 100644 index 0000000000000000000000000000000000000000..aa7a6cdadfe314f9d1b0fc2d48ef67b85032d3c5 --- /dev/null +++ b/fomo_tune_baseline/output/task1/build/fomo_tune/README.md @@ -0,0 +1,239 @@ +# fomo_tune + +The five FOMO26 challenge tasks, one script each, tuned independently. + +This is a spinoff of `nanobrain.eval`, which scored every backbone on every task through one fixed +probe. That was the right shape for a benchmark and the wrong shape for a competition: here we care +about one backbone (sMRI MAE) and five scores, and each task wants a different method. **Nothing +here imports `nanobrain.eval`, and it should stay that way** — this package may be shared with +people who won't get the eval suite. + +## Layout + +| File | | +|---|---| +| `datasets.py` | core, **frozen**. One `load_fomo_task()` per task, streaming the challenge zips into an HF dataset. Raw niftis, no resampling — the backbone transform does that. | +| `backbone.py` | core, **frozen**. `load_backbone(ckpt_path) -> (SmriMaeBackbone, SmriMaeTransform)`. Frozen sMRI MAE encoder; the transform canonicalizes to RAS, rescales to 1mm, fits to the pretraining shape, z-scores in a mean-threshold brain mask. | +| `utils.py` | core. `set_seed`, `git_sha`, `setup_logging`. | +| `main_task.py` | shell. One task, end to end. Task 1 is the worked example; copy it. | +| `build.py` + `Apptainer.def` | shell. Package a run dir into the challenge `.sif`. Shared by every task. | + +`datasets.py` and `backbone.py` are settled and their caches are warm. Treat them as read-only: +new work goes in `main_task.py`. If one of them genuinely needs to change, that is a +conversation first, because it invalidates every score already recorded. + +## The pattern + +`main_task1.py` is in three sections, and the split is the point of the whole design. + +**`Task1Method` — the part we tune.** Features, head, hyperparameters, anything that might move +the score. Its interface is: + +```python +method.fit(rows) # rows are dataset records: subject, label, images +method.predict(images) # -> the challenge's output for one subject +method.save(model_dir) # config.yaml + head.joblib +Task1Method.load(model_dir, **overrides) +``` + +**The protocol — fixed.** Pool out-of-fold predictions over all subjects, bootstrap subjects for +the CI. No repeats, no stratification; the bootstrap is the only variance estimate. Splitting is +per-task but fixed within a task — leave-one-out where n is tiny (task 1, n=21), **20-fold** where +it isn't (tasks 3 and 5), which is close enough to LOO without paying for 494 refits. Once a task's +scheme is set, hold it or scores stop being comparable across iterations. That is also why +`cross_validate` seeds its shuffle at 0 rather than from `cfg.seed`: the folds are part of the +protocol, so tuning the run's seed must not silently redraw them. + +**Two entrypoints.** `train` runs the protocol then fits a head on all subjects and saves it; +`predict` is the challenge CLI. Both go through `Method.predict`, which is why every fold +exercises the code the submission will run. + +That last point is the load-bearing one. `predict` is not a wrapper written at packaging time — it +is the same call cross-validation already made once per held-out subject. When you add a task, keep +that property. + +```bash +uv run python -m fomo_tune.main_task1 train modalities=[dwi_b1000,flair] name=task1_dwi_flair +uv run python -m fomo_tune.main_task1 predict --model-dir output/fomo_tune/task1_dwi/model \ + --adc adc.nii.gz --dwi dwi.nii.gz --flair flair.nii.gz --output prob.txt +``` + +`train` takes omegaconf dotlist overrides against the `Config` dataclass at the top of the file. +It writes `config.yaml`, `log.txt`, `metrics.json`, and `model/` into `{output_root}/{name}/`. + +## Status + +Tasks 1, 5 and 3 are drafted and verified. Task 1 is also packaged — its container passes the +challenge validator; 5 and 3 have not been built yet. **Tasks 2 and 4 are tabled** — both are +segmentation, both need `predict` to emit a nifti on the input grid, and neither is worth opening +until the classification and regression tasks are settled. + +All three on `vitl_fomo300`, one H100, wall being the cross-validation loop: + +| run | result | wall | +|---|---|---| +| `task1_dwi`, dwi_b1000, n=21, LOO | AUROC **0.990** [0.944, 1.000] | 25s | +| `task5_t1w`, t1w, n=48, 20-fold | AUROC **0.984** [0.953, 1.000] | 73s | +| `task3_t1w`, t1w, n=494, 20-fold | r **0.962** [0.956, 0.968], MAE **3.71y** [3.45, 3.97] | 260s | + +**Task 3's row is one fold-seed stale.** It was measured before `cross_validate` froze its shuffle +at 0, so it is a 20-fold run with `random_state=4466`. Task 1 (LOO) and task 5 are current. The +re-run is cheap — 260s on a GPU — it just has not been done. Expect a shift of the same order task +5 saw when its folds moved (0.948 → 0.984, i.e. inside the CI but not negligible). + +Task 1's earlier probe sweep got 0.954 [0.861, 1.000] on the same checkpoint +(`experiments/eval_global_0728`), so it roughly reproduces — the gap is LOO vs 5×5 stratified CV, +one interpolation instead of two, and a head selected on AUROC instead of balanced accuracy. + +Two checks worth repeating per task — `.claude/scratch/verify_task1.py` and +`.claude/scratch/verify_task35.py ` do both: +- features are **bit-identical** whether the nifti comes from the HF dataset wrapper or from + `nib.load` off disk, so CV numbers transfer to the container +- the `predict` CLI agrees with the in-process method + +## What changes per task + +Counts and modalities, read from the local zips: + +| Task | n | Inputs | Output | Split | Notes | +|---|---|---|---|---|---| +| 1 infarct | 21 | adc, dwi_b1000, flair (+t2s/swi) | probability | LOO | done | +| 5 polymicrogyria | 48 | t1w | probability | 20-fold | done | +| 3 brain age | 494 | t1w | age in years | 20-fold | done — RidgeCV head, scored by **Pearson r and MAE**, each with its own bootstrap CI | +| 2 meningioma | 23 | dwi_b1000, flair (+t2s/swi) | mask, input grid | — | tabled | +| 4 trigeminal | 40 | t2w | mask, labels 1=nerve 2=vessel | — | tabled | + +Tasks 5 and 3 diverge from task 1 only where that table says. `cross_validate` over a shuffled +`KFold` replaces `leave_one_out`; both take one modality, so `features` loses the +concat-over-modalities loop and `Config` loses `modalities`; the challenge CLI flag is `--t1` for +both, and it is `--t1` for task 3 too even though the file in the zip is `t1w.nii.gz`. + +Task 3 is the first regression, so its `score` loops over the two metrics rather than returning +one, and it drops task 1's guard against bootstrap resamples with fewer than two distinct labels. +The analogous degenerate case for regression is a resample with no spread in `y`, where Pearson r +is undefined rather than merely unstable — at n=494 it does not happen. + +When tasks 2 and 4 come back: `predict` must write a nifti on the input's grid, and the method +needs localized features rather than a pooled vector — `backbone.forward` returns `patch_coords` +in world mm for exactly that. Task 4's label order (1=nerve, 2=vessel) is still a guess and needs +confirming against the challenge data before per-class numbers mean anything. + +## Gotchas + +**Raw niftis are on disk** at `data/fomo_eval/Task_/preprocessed//ses-01/`, which is the +easy way to exercise `predict` on a real file rather than one written out of the dataset: + +```bash +uv run python -m fomo_tune.main_task1 predict \ + --model-dir output/fomo_tune/task1_dwi/model \ + --adc data/fomo_eval/Task_1/preprocessed/sub-20/ses-01/adc.nii.gz \ + --dwi data/fomo_eval/Task_1/preprocessed/sub-20/ses-01/dwi_b1000.nii.gz \ + --flair data/fomo_eval/Task_1/preprocessed/sub-20/ses-01/flair.nii.gz \ + --output /tmp/prob.txt +``` + +Task 5 breaks the naming: `Task_5/preprocessed/sub_01/ses_01/t1.nii.gz` — underscores throughout, +and `t1` not `t1w`. `datasets.py` already handles it; anything you write by hand won't. + +```bash +uv run python -m fomo_tune.main_task5 predict --model-dir output/fomo_tune/task5_t1w/model \ + --t1 data/fomo_eval/Task_5/preprocessed/sub_01/ses_01/t1.nii.gz --output /tmp/prob.txt +``` + +**Volumes are wildly anisotropic.** Task 1's DWI is 0.46×0.46×**5.6**mm, so the transform +upsamples z by 5.6× to reach 1mm iso. Nothing is wrong, but don't read the 1mm grid as real +resolution. + +**The backbone never saw skull or neck.** Pretraining used a SynthSeg brain mask; the transform +substitutes a mean-intensity threshold, which keeps both. Known fidelity gap — see +`.claude/memory/smri-mae-preprocessing-gap.md`. + +**Probabilities are not calibrated.** `LogisticRegressionCV` on ~20 samples × 1024 features shrinks +hard; task 1's out-of-fold probabilities all land in 0.48–0.52 with near-perfect ranking. Fine for +AUROC, which is what the challenge scores, but don't read them as probabilities. Task 5's do span +0–1, which is n=48 rather than n=21 and not evidence of calibration. + +**n is tiny, so the CI is the result.** Task 1's is ~0.06 wide at the top of the range. Most tuning +deltas you chase will be inside it. `.claude/NOTES.md` thread 1 has the longer argument. + +**GPUs need an allocation** — the login node has no driver. See the `gpu-session` skill. + +## Submission + +`build.py` packages a run dir into the `.sif` the challenge wants. One command, taking the run dir +the shipped head was saved into: + +```bash +uv run python -m fomo_tune.build output/fomo_tune/task1_dwi +``` + +It stages `/app`, then builds from there: + +``` +/app/predict.py # shim: calls fomo_tune.main_task predict +/app/model/config.yaml # from the run dir +/app/model/head.joblib # from the run dir +/app/model/backbone.pth # stripped checkpoint, --ckpt-path points here +``` + +**Both `build.py` and `Apptainer.def` are shared across tasks**, which they can be because nothing +in staging or in the dependency list is task-specific. The one thing that does vary is the module +the shim imports, and that comes from `task` in the run's saved config — so a run dir knows which +task it belongs to, and `build.py` never needs telling. + +`predict.py` is **generated at build time** rather than checked in. It is eight lines whose whole +meaning is the container layout staged around it, so there is nowhere outside a container to run +it. This does not weaken the point above about `predict` not being written at packaging time: the +logic still lives in `main_task.py`, exercised once per fold, and the shim only picks the +subcommand and two paths. + +**`Apptainer.def` is not buildable where it sits.** Its `%files` paths are relative to the build +cwd, which is the staging dir. Pointing `apptainer build` at it in the repo fails confusingly; go +through `build.py`. + +The run dir deliberately does *not* carry backbone weights — that checkpoint is 3.9G and would be +copied on every run. `--ckpt-path` overrides what `config.yaml` recorded, so the saved config stays +a faithful record of what trained rather than being rewritten at build time. + +**The staged checkpoint is stripped to `model` and `args`**, which is 3.9G → 1.3G because the rest +is optimizer state inference never reads. `load_backbone` needs no change for this, and `predict` +gives a bit-identical probability either way (0.524739 on `sub-20`, checked on GPU). + +**The base image is `python:3.11-slim`, not a CUDA image.** The PyPI torch wheel *is* the cu128 +build and vendors the whole CUDA userspace as `nvidia-*` packages, so all the container needs from +the host is the driver, which `--nv`/`--nvccli` binds in. That keeps the SIF at 5.3G (4.0G of +image, 1.3G of weights) against roughly double for `pytorch/pytorch` and far more for NGC. +Versions are pinned to the training environment +mostly so `head.joblib` unpickles against the numpy/sklearn that wrote it. + +Apptainer caches the bootstrap layers but **always re-runs `%post`**, so every build re-downloads +~3G of wheels. If that gets annoying, bake a deps-only base SIF and `Bootstrap: localimage` off it. + +### Validating + +`third_party/container-validator` is the challenge's own validator, test niftis included: + +```bash +python third_party/container-validator/container_validator/validate.py \ + --task task1 --sif output/fomo_tune/task1_dwi/task1.sif +``` + +It runs `python /app/predict.py --flair /input/… --adc … --dwi … --swi … --output /output/.txt` +inside an `apptainer instance` with `/input`, `/output` and `/tmp` bound — which is exactly the +shim's contract, so nothing in `predict.py` is guessing at the interface. + +One thing it does that is easy to miss: it takes GPU via `--nvccli` rather than `--nv`, and one of +its tests runs `nvidia-smi -L` **inside** the container. `python:3.11-slim` ships no `nvidia-smi`, +so that test passes only because `--nvccli` injects the host one — a CUDA base image would hide +that dependency rather than remove it. + +**The `task1_dwi` container passes all 20 validator tests**, and `predict` inside it returns +0.524739 on `sub-20`, identical to the same call outside the container. So the packaging is +verified end to end, not just built. + +**Run it on a compute node with apptainer, which as of 2026-08-11 means `n-6`** — `salloc +--nodelist=n-6`. The other nodes fail the validator's preflight. The login node has apptainer but +no driver, and +`predict` there dies inside `can_use_cudnn_attention` — the jagged-SDPA path reaches into CUDA even +when the tensors are on CPU, so a driver-less host fails at the forward pass rather than falling +back. That is the CPU gap worth remembering; it is not a container problem. diff --git a/fomo_tune_baseline/output/task1/build/fomo_tune/backbone.py b/fomo_tune_baseline/output/task1/build/fomo_tune/backbone.py new file mode 100644 index 0000000000000000000000000000000000000000..4a1275a32e0f9da0963043225e8b3aefcbe8edc6 --- /dev/null +++ b/fomo_tune_baseline/output/task1/build/fomo_tune/backbone.py @@ -0,0 +1,153 @@ +import inspect + +import nibabel as nib +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange +from torch import Tensor + +import smri_mae.model_mae as models_mae + + +class SmriMaeBackbone(nn.Module): + grid_coords: Tensor + + def __init__(self, encoder: models_mae.MaskedEncoder): + super().__init__() + self.encoder = encoder + self.img_size = self.encoder.patchify.img_size + + grid_size = self.encoder.patchify.grid_size + patch_size = np.array(self.encoder.patchify.patch_size) + grid_coords = rearrange(np.indices(grid_size), "c x y z -> (x y z) c") + grid_coords = grid_coords * patch_size + (patch_size - 1) / 2 + grid_coords = torch.as_tensor(grid_coords, dtype=torch.float32) + self.register_buffer("grid_coords", grid_coords) + + def forward(self, batch: dict[str, Tensor]) -> dict[str, Tensor]: + images = batch["image"] + mask = batch["mask"] + affine = batch["affine"] + + B, C, X, Y, Z = images.shape + assert (X, Y, Z) == self.img_size, f"expected {self.img_size}, got {(X, Y, Z)}" + + _, _, patch_embeds, _, patch_ids, token_mask = self.encoder(images, mask=mask) + + # [B, L, 3] world xyz coords of embeddings + patch_coords = self.grid_coords[patch_ids, :] + rot = affine[:, :3, :3] + trans = affine[:, :3, 3] + patch_coords = patch_coords @ rot.transpose(1, 2) + trans[:, None, :] + + return { + "patch_embeds": patch_embeds, + "patch_ids": patch_ids, + "token_mask": token_mask, + "patch_coords": patch_coords, + } + + +class SmriMaeTransform: + def __init__( + self, + img_size: tuple[int, int, int] = (208, 240, 208), + spacing: tuple[float, float, float] = (1.0, 1.0, 1.0), + ): + self.img_size = img_size + self.spacing = spacing + + def __call__(self, img: nib.Nifti1Image) -> dict[str, Tensor]: + # repack image to handle incomplete hf Nifti interface + img = nib.Nifti1Image(img.dataobj, img.affine, img.header) + img = nib.as_closest_canonical(img) + + data = torch.from_numpy(np.ascontiguousarray(img.get_fdata(dtype=np.float32))) + affine = np.asarray(img.affine) + + spacing = img.header.get_zooms() + if max(abs(s - s_) for s, s_ in zip(spacing, self.spacing)) > 0.05: + data, affine = rescale(data, affine, spacing, self.spacing) + + data, affine = fit_to_shape(data, affine, target_shape=self.img_size) + + # mean threshold, not the SynthSeg mask used in pretraining, so skull and neck stay in + mask = data > data.mean() + brain = data[mask] + mean = brain.mean() + # population std (correction=0) to match the pretraining normalization + std = brain.std(correction=0).clamp_min(1e-6) + data = torch.where(mask, (data - mean) / std, 0.0) + + return { + "image": data.unsqueeze(0), + "mask": mask.unsqueeze(0), + "affine": torch.as_tensor(affine, dtype=torch.float32), + } + + +def rescale( + x: torch.Tensor, + affine: np.ndarray, + spacing: tuple[float, ...], + target_spacing: tuple[float, ...] = (1.0, 1.0, 1.0), +) -> tuple[torch.Tensor, np.ndarray]: + scales = tuple([current / target for current, target in zip(spacing, target_spacing)]) + resampled = F.interpolate(x[None, None], scale_factor=scales, mode="trilinear").squeeze(0, 1) + + # align_corners=False reads output voxel j from input voxel (j + 0.5) / scale - 0.5 + scale = np.asarray(scales, dtype=float) + step = np.diag([*(1 / scale), 1.0]) + step[:3, 3] = 0.5 / scale - 0.5 + return resampled, affine @ step + + +def fit_to_shape( + x: torch.Tensor, affine: np.ndarray, target_shape: tuple[int, ...] +) -> tuple[torch.Tensor, np.ndarray]: + """Centre the volume in `target_shape`, padding the short axes and cropping the long ones.""" + pads = [target - size for size, target in zip(x.shape, target_shape)] + padding = [side for pad in reversed(pads) for side in (pad // 2, pad - pad // 2)] + + # a crop is a negative pad, so output voxel k came from input voxel k - pad // 2 either way + step = np.eye(4) + step[:3, 3] = [-(pad // 2) for pad in pads] + return F.pad(x, padding), affine @ step + + +def resolve_ckpt(ckpt_path: str) -> str: + """A local path for a checkpoint, downloading it if it is an hf://// URI.""" + from huggingface_hub import hf_hub_download + + if ckpt_path.startswith("hf://"): + org, repo, *rest = ckpt_path.removeprefix("hf://").split("/") + return hf_hub_download(f"{org}/{repo}", "/".join(rest)) + + return ckpt_path + + +def load_backbone(ckpt_path: str) -> tuple[SmriMaeBackbone, SmriMaeTransform]: + path = resolve_ckpt(ckpt_path) + ckpt = torch.load(path, map_location="cpu", weights_only=True, mmap=True) + args = ckpt["args"] + + model_fn = models_mae.__dict__[args["model"]] + model: models_mae.MaskedAutoencoderViT = model_fn( + img_size=args["img_size"], + in_chans=args.get("in_chans", 1), + patch_size=args["patch_size"], + # older checkpoints carry training flags the current model_mae no longer takes + **filter_kwargs(models_mae.MaskedAutoencoderViT, args.get("model_kwargs") or {}), + ) + model.load_state_dict(ckpt["model"]) + backbone = SmriMaeBackbone(model.encoder) + transform = SmriMaeTransform(img_size=args["img_size"]) + return backbone, transform + + +def filter_kwargs(func, kwargs): + signature = inspect.signature(func) + kwargs = {k: v for k, v in kwargs.items() if k in signature.parameters} + return kwargs diff --git a/fomo_tune_baseline/output/task1/build/fomo_tune/datasets.py b/fomo_tune_baseline/output/task1/build/fomo_tune/datasets.py new file mode 100644 index 0000000000000000000000000000000000000000..64c2ac2da798086e88cbb335601ae850c2dad249 --- /dev/null +++ b/fomo_tune_baseline/output/task1/build/fomo_tune/datasets.py @@ -0,0 +1,205 @@ +import os +import shutil +import tempfile +import zipfile +from collections.abc import Generator +from contextlib import contextmanager +from pathlib import Path + +import fsspec +from datasets import Dataset, Features, Nifti, Value + +FOMO_EVAL_BASE_URL = os.getenv( + "FOMO_EVAL_BASE_URL", + "https://sid.erda.dk/share_redirect/fmeuvo1EdF", +) +FOMO_EVAL_TASK5_URL = os.getenv( + "FOMO_EVAL_TASK5_URL", + "https://huggingface.co/datasets/medarc/smri-fm/resolve/main/fomo_eval/Task_5.zip", +) + + +@contextmanager +def open_zip(url: str) -> Generator[zipfile.ZipFile, None, None]: + """Open a task zip, copying a remote url to a temp file first.""" + with tempfile.TemporaryDirectory() as tmp: + local = Path(url) + if not local.exists(): + local = Path(tmp) / "task.zip" + with fsspec.open(url) as src, local.open("wb") as dst: + shutil.copyfileobj(src, dst) + with zipfile.ZipFile(local) as zf: + yield zf + + +def subject_ids(zf: zipfile.ZipFile) -> list[str]: + return sorted({name.split("/")[2] for name in zf.namelist() if name.endswith(".nii.gz")}) + + +# ---- Task 1: acute infarct (classification; positives also carry a lesion mask) -------- + + +def load_fomo_task1() -> Dataset: + # No 4th modality: it is swi on 16 subjects and t2s on the other 5. + suffixes = ("adc", "dwi_b1000", "flair") + features = Features( + { + "subject": Value("string"), + "label": Value("int32"), + **{suffix: Nifti() for suffix in suffixes}, + } + ) + dataset = Dataset.from_generator( + _fomo_task1_generator, + features=features, + gen_kwargs={"suffixes": suffixes}, + writer_batch_size=16, + ) + return dataset + + +def _fomo_task1_generator(suffixes: tuple[str, ...]): + url = f"{FOMO_EVAL_BASE_URL}/Task_1.zip" + with open_zip(url) as zf: + for sub in subject_ids(zf): + label = int(zf.read(f"Task_1/labels/{sub}/ses-01/label.txt").strip()) + sample = {"subject": sub, "label": label} + for suffix in suffixes: + name = f"Task_1/preprocessed/{sub}/ses-01/{suffix}.nii.gz" + sample[suffix] = {"path": None, "bytes": zf.read(name)} + yield sample + + +# ---- Task 2: meningioma segmentation --------------------------------------------------- + + +def load_fomo_task2() -> Dataset: + # No 4th modality: it is t2s on 15 subjects and swi on the other 8. + suffixes = ("dwi_b1000", "flair") + features = Features( + { + "subject": Value("string"), + **{suffix: Nifti() for suffix in suffixes}, + "seg": Nifti(), + } + ) + dataset = Dataset.from_generator( + _fomo_task2_generator, + features=features, + gen_kwargs={"suffixes": suffixes}, + writer_batch_size=16, + ) + return dataset + + +def _fomo_task2_generator(suffixes: tuple[str, ...]): + url = f"{FOMO_EVAL_BASE_URL}/Task_2.zip" + with open_zip(url) as zf: + for sub in subject_ids(zf): + sample = {"subject": sub} + for suffix in suffixes: + name = f"Task_2/preprocessed/{sub}/ses-01/{suffix}.nii.gz" + sample[suffix] = {"path": None, "bytes": zf.read(name)} + # Seg is on the image grid (shapes match) but its affine differs by up to 0.03mm. + name = f"Task_2/labels/{sub}/ses-01/seg.nii.gz" + sample["seg"] = {"path": None, "bytes": zf.read(name)} + yield sample + + +# ---- Task 3: brain age regression ------------------------------------------------------ + + +def load_fomo_task3() -> Dataset: + features = Features( + { + "subject": Value("string"), + "age": Value("int32"), + "t1w": Nifti(), + } + ) + dataset = Dataset.from_generator( + _fomo_task3_generator, + features=features, + writer_batch_size=16, + ) + return dataset + + +def _fomo_task3_generator(): + url = f"{FOMO_EVAL_BASE_URL}/Task_3.zip" + with open_zip(url) as zf: + for sub in subject_ids(zf): + age = int(zf.read(f"Task_3/labels/{sub}/ses-01/labels.txt").strip()) + image_gz = zf.read(f"Task_3/preprocessed/{sub}/ses-01/t1w.nii.gz") + sample = { + "subject": sub, + "age": age, + "t1w": {"path": None, "bytes": image_gz}, + } + yield sample + + +# ---- Task 4: trigeminal nerve/vessel segmentation -------------------------------------- + + +def load_fomo_task4() -> Dataset: + # Volumes are uncropped 0.5mm near-iso, ~360x512x512; crop before feeding a model. + features = Features( + { + "subject": Value("string"), + "t2w": Nifti(), + "seg": Nifti(), + } + ) + dataset = Dataset.from_generator( + _fomo_task4_generator, + features=features, + writer_batch_size=16, + ) + return dataset + + +def _fomo_task4_generator(): + url = f"{FOMO_EVAL_BASE_URL}/Task_4.zip" + with open_zip(url) as zf: + for sub in subject_ids(zf): + image_gz = zf.read(f"Task_4/preprocessed/{sub}/ses-01/t2w.nii.gz") + seg_gz = zf.read(f"Task_4/labels/{sub}/ses-01/seg.nii.gz") + sample = { + "subject": sub, + "t2w": {"path": None, "bytes": image_gz}, + "seg": {"path": None, "bytes": seg_gz}, + } + yield sample + + +# ---- Task 5: polymicrogyria classification --------------------------------------------- + + +def load_fomo_task5() -> Dataset: + features = Features( + { + "subject": Value("string"), + "label": Value("int32"), + "t1w": Nifti(), + } + ) + dataset = Dataset.from_generator( + _fomo_task5_generator, + features=features, + writer_batch_size=16, + ) + return dataset + + +def _fomo_task5_generator(): + with open_zip(FOMO_EVAL_TASK5_URL) as zf: + for sub in subject_ids(zf): + label = int(zf.read(f"Task_5/labels/{sub}/ses_01/labels.txt").strip()) + image_gz = zf.read(f"Task_5/preprocessed/{sub}/ses_01/t1.nii.gz") + sample = { + "subject": sub, + "label": label, + "t1w": {"path": None, "bytes": image_gz}, + } + yield sample diff --git a/fomo_tune_baseline/output/task1/build/fomo_tune/main_task1.py b/fomo_tune_baseline/output/task1/build/fomo_tune/main_task1.py new file mode 100644 index 0000000000000000000000000000000000000000..bfbc47d49ceeed6fc43dc52a0a007d933ff86d4b --- /dev/null +++ b/fomo_tune_baseline/output/task1/build/fomo_tune/main_task1.py @@ -0,0 +1,253 @@ +"""FOMO task 1: acute infarct classification, scored by AUROC as the challenge scores it. + +`Task1Method` is the part we tune -- features, head, hyperparameters. The protocol below it is +fixed so scores stay comparable across iterations: leave one subject out, pool the out-of-fold +predictions, bootstrap subjects for the CI. + +`train` runs that protocol then fits and saves a head; `predict` is the challenge contract, +modality paths in and one probability out. Both go through `Task1Method.predict`, so every fold +exercises the path the submission will run. +""" + +import argparse +import json +import logging +import time +from dataclasses import dataclass, field +from pathlib import Path + +import joblib +import nibabel as nib +import numpy as np +import torch +from omegaconf import OmegaConf +from sklearn.linear_model import LogisticRegressionCV +from sklearn.metrics import roc_auc_score +from sklearn.pipeline import make_pipeline +from sklearn.preprocessing import StandardScaler + +from fomo_tune.backbone import load_backbone +from fomo_tune.utils import git_sha, set_seed, setup_logging + +logger = logging.getLogger("fomo_tune") + +Images = dict[str, nib.Nifti1Image] + + +@dataclass +class Config: + task: str = "task1" + ckpt_path: str = ( + "/data/mihir-stuff/smri-pretrained/pretrain_full_90_10_h100/checkpoint-last.pth" + ) + modalities: list[str] = field(default_factory=lambda: ["dwi_b1000"]) + output_root: str = "output/fomo_tune" + name: str = "task1" + device: str = "cuda" + seed: int = 4466 + + +# ---- method: the part we tune ----------------------------------------------------------- + + +class Task1Method: + """Frozen sMRI MAE, mean-pooled tokens per modality concatenated, logistic head.""" + + def __init__(self, cfg: Config): + self.cfg = cfg + self.backbone, self.transform = load_backbone(cfg.ckpt_path) + self.device = torch.device(cfg.device) + self.backbone.to(self.device).eval().requires_grad_(False) + self.modalities = list(cfg.modalities) + self.cache: dict[str, np.ndarray] = {} + self.head = None + + @torch.inference_mode() + def features(self, images: Images) -> np.ndarray: + """(D,) per subject. A pure function of the images, so training and inference agree.""" + pooled = [] + for modality in self.modalities: + sample = self.transform(images[modality]) + batch = {key: value[None].to(self.device) for key, value in sample.items()} + + with torch.autocast("cuda", torch.bfloat16, enabled=self.device.type == "cuda"): + out = self.backbone(batch) + + patch_embeds = out["patch_embeds"] + token_mask = out["token_mask"].bool().unsqueeze(-1) + embed = (patch_embeds * token_mask).sum(dim=1) / token_mask.sum(dim=1) + pooled.append(embed[0].float().cpu()) + + return torch.cat(pooled).numpy() + + def cached_features(self, row: dict) -> np.ndarray: + if row["subject"] not in self.cache: + self.cache[row["subject"]] = self.features(row) + return self.cache[row["subject"]] + + def fit(self, rows: list[dict]) -> None: + X = np.stack([self.cached_features(row) for row in rows]) + y = np.array([row["label"] for row in rows]) + + clf = LogisticRegressionCV( + Cs=10, + class_weight="balanced", + scoring="roc_auc", + max_iter=1000, + l1_ratios=(0,), + use_legacy_attributes=False, + ) + self.head = make_pipeline(StandardScaler(), clf) + self.head.fit(X, y) + self.positive = list(self.head.classes_).index(1) + + def predict(self, images: Images) -> float: + """Positive-class probability. Indexes `classes_` rather than assuming column 1, which + would silently score the wrong class if the label order differed.""" + X = self.features(images)[None] + probs = self.head.predict_proba(X)[0] + return float(probs[self.positive]) + + def save(self, model_dir: Path) -> None: + """Everything `load` needs but the backbone weights, which stay wherever `ckpt_path` + points -- a few hundred KB, so a run saves one without copying a 3.7G checkpoint.""" + model_dir.mkdir(parents=True, exist_ok=True) + OmegaConf.save(self.cfg, model_dir / "config.yaml") + joblib.dump({"head": self.head, "positive": self.positive}, model_dir / "head.joblib") + + @classmethod + def load(cls, model_dir: Path, **overrides) -> "Task1Method": + """Rebuild a fitted method from `save`. Overrides are Config fields, for what differs + between here and the container -- the backbone path, the device.""" + cfg = OmegaConf.merge( + OmegaConf.structured(Config), OmegaConf.load(model_dir / "config.yaml"), overrides + ) + method = cls(cfg) + state = joblib.load(model_dir / "head.joblib") + method.head, method.positive = state["head"], state["positive"] + return method + + +# ---- protocol: the part we hold fixed --------------------------------------------------- + +# Every image the task ships. The method picks which of them it wants, as at inference, where +# the challenge hands over all four modalities whether or not a model uses them. +IMAGE_COLS = ("adc", "dwi_b1000", "flair") + + +def leave_one_out(rows: list[dict], method: Task1Method) -> tuple[np.ndarray, np.ndarray]: + """Out-of-fold score for every subject, each predicted by a head fit on the other n-1.""" + y = np.array([row["label"] for row in rows]) + oof = np.zeros(len(rows), dtype=float) + start = time.perf_counter() + for held_out, row in enumerate(rows): + method.fit([r for r in rows if r["subject"] != row["subject"]]) + oof[held_out] = method.predict({key: row[key] for key in IMAGE_COLS}) + logger.info( + f"fold {held_out + 1}/{len(rows)} {row['subject']} " + f"y={y[held_out]} p={oof[held_out]:.3f} ({time.perf_counter() - start:.0f}s)" + ) + return y, oof + + +def score( + y: np.ndarray, oof: np.ndarray, seed: int = 0, n_boot: int = 2000, alpha: float = 0.05 +) -> dict: + """AUROC, the challenge metric, plus a percentile CI resampling subjects with replacement.""" + rng = np.random.default_rng(seed) + samples = [] + for _ in range(n_boot): + rows = rng.integers(0, len(y), size=len(y)) + if len(np.unique(y[rows])) < 2: + continue + samples.append(roc_auc_score(y[rows], oof[rows])) + + low, high = np.percentile(samples, [100 * alpha / 2, 100 * (1 - alpha / 2)]) + return { + "auroc": float(roc_auc_score(y, oof)), + "auroc_ci_low": float(low), + "auroc_ci_high": float(high), + } + + +# ---- entrypoints ------------------------------------------------------------------------ + + +def train(args: argparse.Namespace) -> None: + # imported here, not at the top, so the container needs no dataset stack to run `predict` + from fomo_tune.datasets import load_fomo_task1 + + cfg = OmegaConf.merge(OmegaConf.structured(Config), OmegaConf.from_dotlist(args.overrides)) + run_dir = Path(cfg.output_root) / cfg.name + run_dir.mkdir(parents=True, exist_ok=True) + + setup_logging(run_dir) + set_seed(cfg.seed) + logger.info(f"run {cfg.name} (git {git_sha()})") + logger.info(f"config:\n{OmegaConf.to_yaml(cfg).rstrip()}") + OmegaConf.save(cfg, run_dir / "config.yaml") + + # decoded once: leave-one-out revisits every subject n times, and the niftis are small + rows = list(load_fomo_task1()) + logger.info(f"dataset: {len(rows)} subjects, {sum(r['label'] for r in rows)} positive") + + method = Task1Method(cfg) + start = time.perf_counter() + y, oof = leave_one_out(rows, method) + run_time = time.perf_counter() - start + summary = score(y, oof) + + # the shipped head sees all n subjects, so it is not any of the models scored above + method.fit(rows) + method.save(run_dir / "model") + + record = {"name": cfg.name, **summary, "run_time": round(run_time, 1)} + (run_dir / "metrics.json").write_text(json.dumps(record) + "\n") + scores = " ".join(f"{k}={v:.4f}" for k, v in summary.items()) + logger.info(f"result: {scores} ({run_time:.0f}s)") + + +def predict(args: argparse.Namespace) -> None: + """The challenge contract: modality paths in, one probability written to `--output`. + + `/app/predict.py` in the container is a shim over this, so what scores the submission is the + code leave-one-out already ran, not something generated at build time. + """ + overrides = {"device": args.device} + if args.ckpt_path: + overrides["ckpt_path"] = args.ckpt_path + method = Task1Method.load(args.model_dir, **overrides) + + # every image the challenge hands over, as in `leave_one_out`; the method takes what it uses + paths = {"adc": args.adc, "dwi_b1000": args.dwi, "flair": args.flair} + probability = method.predict({key: nib.load(path) for key, path in paths.items()}) + + args.output.write_text(f"{probability:.6f}\n") + + +def main() -> None: + parser = argparse.ArgumentParser() + modes = parser.add_subparsers(required=True) + + train_parser = modes.add_parser("train", help="leave-one-out over the task, then fit and save") + train_parser.add_argument("overrides", nargs="*", help="config overrides, e.g. device=cpu") + train_parser.set_defaults(run=train) + + predict_parser = modes.add_parser("predict", help="one subject, one probability") + for flag in ("--flair", "--adc", "--dwi"): + predict_parser.add_argument(flag, type=Path, required=True) + # accepted and ignored: the 4th modality is swi on some subjects and t2s on others + for flag in ("--t2s", "--swi"): + predict_parser.add_argument(flag, type=Path) + predict_parser.add_argument("--output", type=Path, required=True) + predict_parser.add_argument("--model-dir", type=Path, default=Path("/app/model")) + predict_parser.add_argument("--ckpt-path", help="overrides the trained config's backbone path") + predict_parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") + predict_parser.set_defaults(run=predict) + + args = parser.parse_args() + args.run(args) + + +if __name__ == "__main__": + main() diff --git a/fomo_tune_baseline/output/task1/build/fomo_tune/main_task3.py b/fomo_tune_baseline/output/task1/build/fomo_tune/main_task3.py new file mode 100644 index 0000000000000000000000000000000000000000..02db5413c10a0d44eb248e7f2739a436c45806c2 --- /dev/null +++ b/fomo_tune_baseline/output/task1/build/fomo_tune/main_task3.py @@ -0,0 +1,241 @@ +"""FOMO task 3: brain age regression, scored by Pearson r and MAE as the challenge scores it. + +`Task3Method` is the part we tune -- features, head, hyperparameters. The protocol below it is +fixed so scores stay comparable across iterations: 20-fold over the 494 subjects, pool the +out-of-fold predictions, bootstrap subjects for the CI. + +`train` runs that protocol then fits and saves a head; `predict` is the challenge contract, one t1 +path in and one age out. Both go through `Task3Method.predict`, so every fold exercises the path +the submission will run. +""" + +import argparse +import json +import logging +import time +from dataclasses import dataclass +from pathlib import Path + +import joblib +import nibabel as nib +import numpy as np +import torch +from omegaconf import OmegaConf +from sklearn.linear_model import RidgeCV +from sklearn.model_selection import KFold +from sklearn.pipeline import make_pipeline +from sklearn.preprocessing import StandardScaler + +from fomo_tune.backbone import load_backbone +from fomo_tune.utils import git_sha, set_seed, setup_logging + +logger = logging.getLogger("fomo_tune") + +Images = dict[str, nib.Nifti1Image] + + +@dataclass +class Config: + task: str = "task3" + ckpt_path: str = ( + "/data/mihir-stuff/smri-pretrained/pretrain_full_90_10_h100/checkpoint-last.pth" + ) + output_root: str = "output/fomo_tune" + name: str = "task3" + device: str = "cuda" + seed: int = 4466 + + +# ---- method: the part we tune ----------------------------------------------------------- + + +class Task3Method: + """Frozen sMRI MAE, mean-pooled tokens over the t1w, ridge head.""" + + def __init__(self, cfg: Config): + self.cfg = cfg + self.backbone, self.transform = load_backbone(cfg.ckpt_path) + self.device = torch.device(cfg.device) + self.backbone.to(self.device).eval().requires_grad_(False) + self.cache: dict[str, np.ndarray] = {} + self.head = None + + @torch.inference_mode() + def features(self, images: Images) -> np.ndarray: + """(D,) per subject. A pure function of the images, so training and inference agree.""" + sample = self.transform(images["t1w"]) + batch = {key: value[None].to(self.device) for key, value in sample.items()} + + with torch.autocast("cuda", torch.bfloat16, enabled=self.device.type == "cuda"): + out = self.backbone(batch) + + patch_embeds = out["patch_embeds"] + token_mask = out["token_mask"].bool().unsqueeze(-1) + embed = (patch_embeds * token_mask).sum(dim=1) / token_mask.sum(dim=1) + return embed[0].float().cpu().numpy() + + def cached_features(self, row: dict) -> np.ndarray: + if row["subject"] not in self.cache: + self.cache[row["subject"]] = self.features(row) + return self.cache[row["subject"]] + + def fit(self, rows: list[dict]) -> None: + X = np.stack([self.cached_features(row) for row in rows]) + y = np.array([row["age"] for row in rows], dtype=float) + + # RidgeCV picks alpha by its own efficient leave-one-out, so the fold's own split is + # never touched by model selection + self.head = make_pipeline(StandardScaler(), RidgeCV(alphas=np.logspace(-3, 6, 19))) + self.head.fit(X, y) + + def predict(self, images: Images) -> float: + """Age in years.""" + X = self.features(images)[None] + return float(self.head.predict(X)[0]) + + def save(self, model_dir: Path) -> None: + """Everything `load` needs but the backbone weights, which stay wherever `ckpt_path` + points -- a few hundred KB, so a run saves one without copying a 3.7G checkpoint.""" + model_dir.mkdir(parents=True, exist_ok=True) + OmegaConf.save(self.cfg, model_dir / "config.yaml") + joblib.dump(self.head, model_dir / "head.joblib") + + @classmethod + def load(cls, model_dir: Path, **overrides) -> "Task3Method": + """Rebuild a fitted method from `save`. Overrides are Config fields, for what differs + between here and the container -- the backbone path, the device.""" + cfg = OmegaConf.merge( + OmegaConf.structured(Config), OmegaConf.load(model_dir / "config.yaml"), overrides + ) + method = cls(cfg) + method.head = joblib.load(model_dir / "head.joblib") + return method + + +# ---- protocol: the part we hold fixed --------------------------------------------------- + +# Every image the task ships. The method picks which of them it wants, as at inference, where the +# challenge hands over the modalities whether or not a model uses them. +IMAGE_COLS = ("t1w",) + + +def cross_validate( + rows: list[dict], method: Task3Method, seed: int = 0, n_folds: int = 20 +) -> tuple[np.ndarray, np.ndarray]: + """Out-of-fold age for every subject, each predicted by a head fit on the other folds.""" + y = np.array([row["age"] for row in rows], dtype=float) + oof = np.zeros(len(rows), dtype=float) + folds = KFold(n_splits=n_folds, shuffle=True, random_state=seed) + start = time.perf_counter() + for fold, (train, test) in enumerate(folds.split(rows)): + method.fit([rows[i] for i in train]) + for i in test: + oof[i] = method.predict({key: rows[i][key] for key in IMAGE_COLS}) + logger.info( + f"fold {fold + 1}/{n_folds} n={len(test)} mae={np.abs(y[test] - oof[test]).mean():.2f} " + f"({time.perf_counter() - start:.0f}s)" + ) + return y, oof + + +def metrics(y: np.ndarray, oof: np.ndarray) -> dict: + return { + "pearson_r": float(np.corrcoef(y, oof)[0, 1]), + "mae": float(np.abs(y - oof).mean()), + } + + +def score( + y: np.ndarray, oof: np.ndarray, seed: int = 0, n_boot: int = 2000, alpha: float = 0.05 +) -> dict: + """Both challenge metrics, each with a percentile CI resampling subjects with replacement.""" + rng = np.random.default_rng(seed) + resamples = rng.integers(0, len(y), size=(n_boot, len(y))) + + summary = {} + for name, point in metrics(y, oof).items(): + samples = [metrics(y[rows], oof[rows])[name] for rows in resamples] + low, high = np.percentile(samples, [100 * alpha / 2, 100 * (1 - alpha / 2)]) + summary[name] = point + summary[f"{name}_ci_low"] = float(low) + summary[f"{name}_ci_high"] = float(high) + return summary + + +# ---- entrypoints ------------------------------------------------------------------------ + + +def train(args: argparse.Namespace) -> None: + # imported here, not at the top, so the container needs no dataset stack to run `predict` + from fomo_tune.datasets import load_fomo_task3 + + cfg = OmegaConf.merge(OmegaConf.structured(Config), OmegaConf.from_dotlist(args.overrides)) + run_dir = Path(cfg.output_root) / cfg.name + run_dir.mkdir(parents=True, exist_ok=True) + + setup_logging(run_dir) + set_seed(cfg.seed) + logger.info(f"run {cfg.name} (git {git_sha()})") + logger.info(f"config:\n{OmegaConf.to_yaml(cfg).rstrip()}") + OmegaConf.save(cfg, run_dir / "config.yaml") + + rows = list(load_fomo_task3()) + ages = np.array([row["age"] for row in rows]) + logger.info( + f"dataset: {len(rows)} subjects, age {ages.min()}-{ages.max()} mean {ages.mean():.1f}" + ) + + method = Task3Method(cfg) + start = time.perf_counter() + y, oof = cross_validate(rows, method) + run_time = time.perf_counter() - start + summary = score(y, oof) + + # the shipped head sees all n subjects, so it is not any of the models scored above + method.fit(rows) + method.save(run_dir / "model") + + record = {"name": cfg.name, **summary, "run_time": round(run_time, 1)} + (run_dir / "metrics.json").write_text(json.dumps(record) + "\n") + scores = " ".join(f"{k}={v:.4f}" for k, v in summary.items()) + logger.info(f"result: {scores} ({run_time:.0f}s)") + + +def predict(args: argparse.Namespace) -> None: + """The challenge contract: a t1 path in, one age written to `--output`. + + `/app/predict.py` in the container is a shim over this, so what scores the submission is the + code cross-validation already ran, not something generated at build time. + """ + overrides = {"device": args.device} + if args.ckpt_path: + overrides["ckpt_path"] = args.ckpt_path + method = Task3Method.load(args.model_dir, **overrides) + + age = method.predict({"t1w": nib.load(args.t1)}) + + args.output.write_text(f"{age:.6f}\n") + + +def main() -> None: + parser = argparse.ArgumentParser() + modes = parser.add_subparsers(required=True) + + train_parser = modes.add_parser("train", help="cross-validate over the task, then fit and save") + train_parser.add_argument("overrides", nargs="*", help="config overrides, e.g. device=cpu") + train_parser.set_defaults(run=train) + + predict_parser = modes.add_parser("predict", help="one subject, one age in years") + predict_parser.add_argument("--t1", type=Path, required=True) + predict_parser.add_argument("--output", type=Path, required=True) + predict_parser.add_argument("--model-dir", type=Path, default=Path("/app/model")) + predict_parser.add_argument("--ckpt-path", help="overrides the trained config's backbone path") + predict_parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") + predict_parser.set_defaults(run=predict) + + args = parser.parse_args() + args.run(args) + + +if __name__ == "__main__": + main() diff --git a/fomo_tune_baseline/output/task1/build/fomo_tune/main_task5.py b/fomo_tune_baseline/output/task1/build/fomo_tune/main_task5.py new file mode 100644 index 0000000000000000000000000000000000000000..d8cccbecf130497c483b3a4b7665e835ca428873 --- /dev/null +++ b/fomo_tune_baseline/output/task1/build/fomo_tune/main_task5.py @@ -0,0 +1,245 @@ +"""FOMO task 5: polymicrogyria classification, scored by AUROC as the challenge scores it. + +`Task5Method` is the part we tune -- features, head, hyperparameters. The protocol below it is +fixed so scores stay comparable across iterations: 20-fold over the 48 subjects, pool the +out-of-fold predictions, bootstrap subjects for the CI. + +`train` runs that protocol then fits and saves a head; `predict` is the challenge contract, one t1 +path in and one probability out. Both go through `Task5Method.predict`, so every fold exercises +the path the submission will run. +""" + +import argparse +import json +import logging +import time +from dataclasses import dataclass +from pathlib import Path + +import joblib +import nibabel as nib +import numpy as np +import torch +from omegaconf import OmegaConf +from sklearn.linear_model import LogisticRegressionCV +from sklearn.metrics import roc_auc_score +from sklearn.model_selection import KFold +from sklearn.pipeline import make_pipeline +from sklearn.preprocessing import StandardScaler + +from fomo_tune.backbone import load_backbone +from fomo_tune.utils import git_sha, set_seed, setup_logging + +logger = logging.getLogger("fomo_tune") + +Images = dict[str, nib.Nifti1Image] + + +@dataclass +class Config: + task: str = "task5" + ckpt_path: str = ( + "/data/mihir-stuff/smri-pretrained/pretrain_full_90_10_h100/checkpoint-last.pth" + ) + output_root: str = "output/fomo_tune" + name: str = "task5" + device: str = "cuda" + seed: int = 4466 + + +# ---- method: the part we tune ----------------------------------------------------------- + + +class Task5Method: + """Frozen sMRI MAE, mean-pooled tokens over the t1w, logistic head.""" + + def __init__(self, cfg: Config): + self.cfg = cfg + self.backbone, self.transform = load_backbone(cfg.ckpt_path) + self.device = torch.device(cfg.device) + self.backbone.to(self.device).eval().requires_grad_(False) + self.cache: dict[str, np.ndarray] = {} + self.head = None + + @torch.inference_mode() + def features(self, images: Images) -> np.ndarray: + """(D,) per subject. A pure function of the images, so training and inference agree.""" + sample = self.transform(images["t1w"]) + batch = {key: value[None].to(self.device) for key, value in sample.items()} + + with torch.autocast("cuda", torch.bfloat16, enabled=self.device.type == "cuda"): + out = self.backbone(batch) + + patch_embeds = out["patch_embeds"] + token_mask = out["token_mask"].bool().unsqueeze(-1) + embed = (patch_embeds * token_mask).sum(dim=1) / token_mask.sum(dim=1) + return embed[0].float().cpu().numpy() + + def cached_features(self, row: dict) -> np.ndarray: + if row["subject"] not in self.cache: + self.cache[row["subject"]] = self.features(row) + return self.cache[row["subject"]] + + def fit(self, rows: list[dict]) -> None: + X = np.stack([self.cached_features(row) for row in rows]) + y = np.array([row["label"] for row in rows]) + + clf = LogisticRegressionCV( + Cs=10, + class_weight="balanced", + scoring="roc_auc", + max_iter=1000, + l1_ratios=(0,), + use_legacy_attributes=False, + ) + self.head = make_pipeline(StandardScaler(), clf) + self.head.fit(X, y) + self.positive = list(self.head.classes_).index(1) + + def predict(self, images: Images) -> float: + """Positive-class probability. Indexes `classes_` rather than assuming column 1, which + would silently score the wrong class if the label order differed.""" + X = self.features(images)[None] + probs = self.head.predict_proba(X)[0] + return float(probs[self.positive]) + + def save(self, model_dir: Path) -> None: + """Everything `load` needs but the backbone weights, which stay wherever `ckpt_path` + points -- a few hundred KB, so a run saves one without copying a 3.7G checkpoint.""" + model_dir.mkdir(parents=True, exist_ok=True) + OmegaConf.save(self.cfg, model_dir / "config.yaml") + joblib.dump({"head": self.head, "positive": self.positive}, model_dir / "head.joblib") + + @classmethod + def load(cls, model_dir: Path, **overrides) -> "Task5Method": + """Rebuild a fitted method from `save`. Overrides are Config fields, for what differs + between here and the container -- the backbone path, the device.""" + cfg = OmegaConf.merge( + OmegaConf.structured(Config), OmegaConf.load(model_dir / "config.yaml"), overrides + ) + method = cls(cfg) + state = joblib.load(model_dir / "head.joblib") + method.head, method.positive = state["head"], state["positive"] + return method + + +# ---- protocol: the part we hold fixed --------------------------------------------------- + +# Every image the task ships. The method picks which of them it wants, as at inference, where the +# challenge hands over the modalities whether or not a model uses them. +IMAGE_COLS = ("t1w",) + + +def cross_validate( + rows: list[dict], method: Task5Method, seed: int = 0, n_folds: int = 20 +) -> tuple[np.ndarray, np.ndarray]: + """Out-of-fold score for every subject, each predicted by a head fit on the other folds.""" + y = np.array([row["label"] for row in rows]) + oof = np.zeros(len(rows), dtype=float) + folds = KFold(n_splits=n_folds, shuffle=True, random_state=seed) + start = time.perf_counter() + for fold, (train, test) in enumerate(folds.split(rows)): + method.fit([rows[i] for i in train]) + for i in test: + oof[i] = method.predict({key: rows[i][key] for key in IMAGE_COLS}) + logger.info( + f"fold {fold + 1}/{n_folds} n={len(test)} y={y[test]} " + f"p={np.round(oof[test], 3)} ({time.perf_counter() - start:.0f}s)" + ) + return y, oof + + +def score( + y: np.ndarray, oof: np.ndarray, seed: int = 0, n_boot: int = 2000, alpha: float = 0.05 +) -> dict: + """AUROC, the challenge metric, plus a percentile CI resampling subjects with replacement.""" + rng = np.random.default_rng(seed) + samples = [] + for _ in range(n_boot): + rows = rng.integers(0, len(y), size=len(y)) + if len(np.unique(y[rows])) < 2: + continue + samples.append(roc_auc_score(y[rows], oof[rows])) + + low, high = np.percentile(samples, [100 * alpha / 2, 100 * (1 - alpha / 2)]) + return { + "auroc": float(roc_auc_score(y, oof)), + "auroc_ci_low": float(low), + "auroc_ci_high": float(high), + } + + +# ---- entrypoints ------------------------------------------------------------------------ + + +def train(args: argparse.Namespace) -> None: + # imported here, not at the top, so the container needs no dataset stack to run `predict` + from fomo_tune.datasets import load_fomo_task5 + + cfg = OmegaConf.merge(OmegaConf.structured(Config), OmegaConf.from_dotlist(args.overrides)) + run_dir = Path(cfg.output_root) / cfg.name + run_dir.mkdir(parents=True, exist_ok=True) + + setup_logging(run_dir) + set_seed(cfg.seed) + logger.info(f"run {cfg.name} (git {git_sha()})") + logger.info(f"config:\n{OmegaConf.to_yaml(cfg).rstrip()}") + OmegaConf.save(cfg, run_dir / "config.yaml") + + rows = list(load_fomo_task5()) + logger.info(f"dataset: {len(rows)} subjects, {sum(r['label'] for r in rows)} positive") + + method = Task5Method(cfg) + start = time.perf_counter() + y, oof = cross_validate(rows, method) + run_time = time.perf_counter() - start + summary = score(y, oof) + + # the shipped head sees all n subjects, so it is not any of the models scored above + method.fit(rows) + method.save(run_dir / "model") + + record = {"name": cfg.name, **summary, "run_time": round(run_time, 1)} + (run_dir / "metrics.json").write_text(json.dumps(record) + "\n") + scores = " ".join(f"{k}={v:.4f}" for k, v in summary.items()) + logger.info(f"result: {scores} ({run_time:.0f}s)") + + +def predict(args: argparse.Namespace) -> None: + """The challenge contract: a t1 path in, one probability written to `--output`. + + `/app/predict.py` in the container is a shim over this, so what scores the submission is the + code cross-validation already ran, not something generated at build time. + """ + overrides = {"device": args.device} + if args.ckpt_path: + overrides["ckpt_path"] = args.ckpt_path + method = Task5Method.load(args.model_dir, **overrides) + + probability = method.predict({"t1w": nib.load(args.t1)}) + + args.output.write_text(f"{probability:.6f}\n") + + +def main() -> None: + parser = argparse.ArgumentParser() + modes = parser.add_subparsers(required=True) + + train_parser = modes.add_parser("train", help="cross-validate over the task, then fit and save") + train_parser.add_argument("overrides", nargs="*", help="config overrides, e.g. device=cpu") + train_parser.set_defaults(run=train) + + predict_parser = modes.add_parser("predict", help="one subject, one probability") + predict_parser.add_argument("--t1", type=Path, required=True) + predict_parser.add_argument("--output", type=Path, required=True) + predict_parser.add_argument("--model-dir", type=Path, default=Path("/app/model")) + predict_parser.add_argument("--ckpt-path", help="overrides the trained config's backbone path") + predict_parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") + predict_parser.set_defaults(run=predict) + + args = parser.parse_args() + args.run(args) + + +if __name__ == "__main__": + main() diff --git a/fomo_tune_baseline/output/task1/build/fomo_tune/utils.py b/fomo_tune_baseline/output/task1/build/fomo_tune/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..57a5c6d0396d8811e763254a987b331cf41ee1f1 --- /dev/null +++ b/fomo_tune_baseline/output/task1/build/fomo_tune/utils.py @@ -0,0 +1,33 @@ +import logging +import random +import subprocess +import sys +from pathlib import Path + +import numpy as np +import torch + +logger = logging.getLogger("fomo_tune") + + +def set_seed(seed: int) -> None: + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + + +def git_sha() -> str: + kwargs = dict(cwd=Path(__file__).parent, capture_output=True, text=True, check=True) + sha = subprocess.run(["git", "rev-parse", "--short", "HEAD"], **kwargs).stdout.strip() + dirty = subprocess.run(["git", "status", "--porcelain", "-uno"], **kwargs).stdout.strip() + return f"{sha}-dirty" if dirty else sha + + +def setup_logging(run_dir: Path) -> None: + handlers = [logging.StreamHandler(sys.stdout), logging.FileHandler(run_dir / "log.txt")] + logger.setLevel(logging.INFO) + logger.handlers.clear() + for handler in handlers: + handler.setFormatter(logging.Formatter("%(asctime)s %(message)s", datefmt="%H:%M:%S")) + logger.addHandler(handler) + logger.propagate = False diff --git a/fomo_tune_baseline/output/task1/build/model/backbone.pth b/fomo_tune_baseline/output/task1/build/model/backbone.pth new file mode 100644 index 0000000000000000000000000000000000000000..793b42f5cea0b055534c29cd18c21f3200873219 --- /dev/null +++ b/fomo_tune_baseline/output/task1/build/model/backbone.pth @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:aacf2582e9464e7fd4bd3d5ed41a5b9aa770e5b04a24b2a82c1c4ba5fb5f6e7a +size 1389681124 diff --git a/fomo_tune_baseline/output/task1/build/model/config.yaml b/fomo_tune_baseline/output/task1/build/model/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..5dca875152bb8d3adf4bf758bd4681bf19e1d83e --- /dev/null +++ b/fomo_tune_baseline/output/task1/build/model/config.yaml @@ -0,0 +1,8 @@ +task: task1 +ckpt_path: /data/mihir-stuff/smri-pretrained/pretrain_full_90_10_h100/checkpoint-last.pth +modalities: +- dwi_b1000 +output_root: experiments/fomo_tune_baseline/output +name: task1 +device: cuda +seed: 4466 diff --git a/fomo_tune_baseline/output/task1/build/model/head.joblib b/fomo_tune_baseline/output/task1/build/model/head.joblib new file mode 100644 index 0000000000000000000000000000000000000000..90ed37cfc3dfd4a13e1f754990808e15c9b1f7e0 --- /dev/null +++ b/fomo_tune_baseline/output/task1/build/model/head.joblib @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f4a4826f5ea9f3f74279b798caf6a5970a3c63b02b70c4d0828ee5325d29e0dd +size 445215 diff --git a/fomo_tune_baseline/output/task1/build/predict.py b/fomo_tune_baseline/output/task1/build/predict.py new file mode 100644 index 0000000000000000000000000000000000000000..6ecbf03e1ba0f3fba34984bd334f8d44befdc132 --- /dev/null +++ b/fomo_tune_baseline/output/task1/build/predict.py @@ -0,0 +1,16 @@ + +import sys + +from fomo_tune.main_task1 import main + +sys.argv = [ + sys.argv[0], + "predict", + *sys.argv[1:], + "--model-dir", + "/app/model", + "--ckpt-path", + "/app/model/backbone.pth", +] + +main() diff --git a/fomo_tune_baseline/output/task1/build/smri_mae/config/default_pretrain.yaml b/fomo_tune_baseline/output/task1/build/smri_mae/config/default_pretrain.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c16290b1e2c24e4d206f7e32a7b9c2ee5266067e --- /dev/null +++ b/fomo_tune_baseline/output/task1/build/smri_mae/config/default_pretrain.yaml @@ -0,0 +1,98 @@ +# Name of the run. Used for output directory suffix and wandb. +name: pretrain + +# Description of the run. Goes in wandb notes. +notes: null + +# Root output directory. +# The run writes to checkpoints/ when name is set. +output_dir: checkpoints + +# Standard 3D structural MRI volume size. +img_size: [208, 240, 208] +patch_size: 8 + +# Masking. +mask_ratio: 0.80 +pred_mask_ratio: null +pad_to_multiple: 32 + +# Model. +model: mae_vit_large +model_kwargs: + # target normalization: null/none, global, slice, or patch. + target_norm: none + + no_decode_pos: false + mask_drop_scale: false + + class_token: true + reg_tokens: 0 + no_embed_class: false + + decoder_depth: 4 + drop_path_rate: 0.0 + +# Datasets. +datasets: + fomo_train: + url: datasets/FOMO_with_dwi/shard.{000000..001620}.tar + samples_per_epoch: 243200 + shuffle: true + buffer_size: 8000 + drop_last: true + + fomo_val: + url: datasets/FOMO_with_dwi/shard.{001621..001800}.tar + samples_per_epoch: 26880 + shuffle: false + buffer_size: 1000 + drop_last: true + +train_dataset: fomo_train +eval_datasets: + - fomo_val + +# Data loader. +num_workers: 4 +prefetch_factor: 2 +presend_cuda: true + +# Optimization. +epochs: 100 +batch_size: 64 +accum_iter: 1 + +base_lr: 0.001 +min_lr: 1e-6 +warmup_epochs: 10 +weight_decay: 0.05 +betas: [0.9, 0.95] +clip_grad: 1.0 + +amp: true +amp_dtype: bfloat16 + +# Checkpointing. +ckpt: null +resume: false +auto_resume: true +start_epoch: 0 +checkpoint_period: 10 +max_checkpoints: 5 + +# Evaluation. +eval_period: 10 + +# Sync checkpoints to an R2 bucket using the AWS CLI. Set to a URL to enable. +r2_sync: null + +device: cuda +distributed: false +seed: 7338 +eval_seed: 7338 +debug: false + +wandb: false +wandb_entity: null +wandb_project: smri-fm diff --git a/fomo_tune_baseline/output/task1/build/smri_mae/main_pretrain.py b/fomo_tune_baseline/output/task1/build/smri_mae/main_pretrain.py new file mode 100644 index 0000000000000000000000000000000000000000..5470551d5b56bf4653fad4c2a23888c417c0c6e8 --- /dev/null +++ b/fomo_tune_baseline/output/task1/build/smri_mae/main_pretrain.py @@ -0,0 +1,486 @@ +# Copyright (c) Sophont, Inc +# This source code is licensed under the Apache License, Version 2.0 +# +# References: +# deit: https://github.com/facebookresearch/deit/blob/main/main.py +# capi: https://github.com/facebookresearch/capi/blob/main/train_capi.py + +import argparse +import datetime +import json +import math +import random +import subprocess +import time +from contextlib import nullcontext +from functools import partial +from itertools import islice +from pathlib import Path +from typing import Iterable, Sequence + +import torch +import torch.nn as nn +import wandb +import webdataset as wds +from omegaconf import DictConfig, OmegaConf +from PIL import Image + +from matplotlib import pyplot as plt +from torch import Tensor + +import data.mri_data as mri_data +import smri_mae.model_mae as models_mae +import smri_mae.utils as ut +import smri_mae.visualization as vis + +DEFAULT_CONFIG = Path(__file__).parent / "config/default_pretrain.yaml" + +MODELS_DICT = models_mae.__dict__ + + +def main(args: DictConfig): + # setup + ut.init_distributed_mode(args) + global_rank = ut.get_rank() + is_master = global_rank == 0 + world_size = ut.get_world_size() + device = torch.device(args.device) + ut.configure_flash_sdpa() + ut.random_seed(args.seed, rank=global_rank) + + if args.name and not args.output_dir.endswith(args.name): + args.output_dir = f"{args.output_dir}/{args.name}" + output_dir = Path(args.output_dir) + + if is_master: + output_dir.mkdir(parents=True, exist_ok=True) + out_cfg_path = output_dir / "config.yaml" + if out_cfg_path.exists(): + prev_cfg = OmegaConf.load(out_cfg_path) + assert args == prev_cfg, "current config doesn't match previous config" + else: + OmegaConf.save(args, out_cfg_path) + + if args.wandb: + wandb.init( + entity=args.wandb_entity, + project=args.wandb_project, + name=args.name, + notes=args.notes, + config=OmegaConf.to_container(args), + ) + + ut.setup_for_distributed(log_path=output_dir / "log.txt") + + print("pretraining 3D ViTMAE") + print(f"start: {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}") + print(f"cwd: {Path.cwd()}") + print(ut.get_sha()) + print("config:", OmegaConf.to_yaml(args), sep="\n") + + # data loaders + train_loader, eval_loaders = create_data_loaders(args) + + # model + model = MODELS_DICT[args.model]( + img_size=args.img_size, + in_chans=args.get("in_chans", 1), + patch_size=args.patch_size, + **(args.get("model_kwargs") or {}), + ) + model.to(device) + print("model:", model, sep="\n") + num_params = sum(p.numel() for p in model.parameters() if p.requires_grad) + print(f"num params: {num_params / 1e6:.1f}M") + + model_without_ddp = model + if args.distributed: + model = torch.nn.parallel.DistributedDataParallel( + model, + device_ids=[args.gpu], + gradient_as_bucket_view=True, + ) + model_without_ddp = model.module + + # optimizer + total_batch_size = args.batch_size * args.accum_iter * world_size + print( + f"total batch size: {total_batch_size} = " + f"{args.batch_size} bs per gpu x {args.accum_iter} accum x {world_size} gpus" + ) + + if not args.get("lr"): + args.lr = args.base_lr * total_batch_size / 256 + print(f"lr: {args.lr:.2e} = {args.base_lr:.2e} x {total_batch_size} / 256") + else: + print(f"lr: {args.lr:.2e}") + + param_groups = ut.get_param_groups(model) + ut.update_lr(param_groups, args.lr) + ut.update_wd(param_groups, args.weight_decay) + # cast or else it corrupts the checkpoint + betas = tuple(args.betas) if args.betas is not None else None + optimizer = torch.optim.AdamW(param_groups, betas=betas, fused=True) + + epoch_num_batches = len(train_loader) + steps_per_epoch = math.ceil(epoch_num_batches / args.accum_iter) + total_steps = args.epochs * steps_per_epoch + warmup_steps = args.warmup_epochs * steps_per_epoch + lr_schedule = ut.WarmupThenCosine( + base_value=args.lr, + final_value=args.min_lr, + total_iters=total_steps, + warmup_iters=warmup_steps, + ) + print(f"full schedule: epochs = {args.epochs} (steps = {total_steps})") + print(f"warmup: epochs = {args.warmup_epochs} (steps = {warmup_steps})") + + # loss scaling not needed for bfloat16 (according to timm) + if args.amp and args.amp_dtype != "bfloat16": + loss_scaler = torch.GradScaler(device.type) + else: + loss_scaler = None + + # load checkpoint/resume training + ut.load_model(args, model_without_ddp, optimizer, loss_scaler) + + print(f"start training for {args.epochs} epochs") + start_time = time.monotonic() + for epoch in range(args.start_epoch, args.epochs): + train_stats = train_one_epoch( + args, + model, + train_loader, + optimizer, + loss_scaler, + lr_schedule, + epoch, + device, + ) + eval_stats = {} + eval_plots = {} + eval_period = args.get("eval_period", 1) + if eval_period and (epoch % eval_period == 0 or epoch == args.epochs - 1): + for name, loader in eval_loaders.items(): + stats, plots = evaluate( + args, + model, + loader, + epoch, + device, + eval_name=name, + ) + eval_stats.update(stats) + eval_plots.update(plots) + + merged_stats = {"epoch": epoch, **train_stats, **eval_stats} + if is_master: + with (output_dir / "log.json").open("a") as f: + print(json.dumps(merged_stats), file=f) + + for plot_name, img in eval_plots.items(): + plot_name = plot_name.replace("/", "__") + img.save(output_dir / f"{plot_name}__{epoch:05d}.png") + + ut.save_model(args, epoch, model_without_ddp, optimizer, loss_scaler) + sync_checkpoints_to_r2(args, output_dir) + + if args.distributed: + torch.distributed.destroy_process_group() + + total_time = time.monotonic() - start_time + print(f"done! training time: {datetime.timedelta(seconds=int(total_time))}") + + +def create_data_loaders(args: DictConfig): + data_loaders = {} + dataset_names = [args.train_dataset] + args.eval_datasets + + for dataset_name in dataset_names: + dataset_config = args.datasets[dataset_name].copy() + drop_last = dataset_config.pop("drop_last") + is_train = dataset_name == args.train_dataset + + print(f"loading dataset: {dataset_name}\n\n{OmegaConf.to_yaml(dataset_config)}") + shuffle = dataset_config["shuffle"] + samples_per_epoch = dataset_config.pop("samples_per_epoch") + dataset = mri_data.make_sparse_wds_dataset( + dataset_config["url"], + shuffle=shuffle, + buffer_size=dataset_config["buffer_size"], + ) + num_workers = int(args.num_workers) + loader_kwargs = { + "batch_size": args.batch_size, + "collate_fn": partial(mri_data.collate, include_meta=not is_train), + "shuffle": False, + "num_workers": num_workers, + "persistent_workers": num_workers > 0, + "pin_memory": True, + "drop_last": drop_last, + "prefetch_factor": args.prefetch_factor, + } + loader = wds.WebLoader(dataset, **loader_kwargs) + num_batches = samples_per_epoch // (ut.get_world_size() * args.batch_size) + loader = loader.with_epoch(num_batches) + loader = loader.with_length(num_batches, silent=True) + + data_loaders[dataset_name] = loader + + train_loader = data_loaders.pop(args.train_dataset) + return train_loader, data_loaders + + +def sync_checkpoints_to_r2(args: DictConfig, output_dir: Path) -> None: + r2_sync_url = args.get("r2_sync") + if not r2_sync_url or not ut.is_main_process(): + return + + cmd = ["aws", "s3", "sync", str(output_dir), str(r2_sync_url), "--profile", "r2"] + print(f"syncing checkpoints to R2: {output_dir} -> {r2_sync_url}") + subprocess.run(cmd, check=True) + + +def train_one_epoch( + args: DictConfig, + model: nn.Module, + data_loader: Iterable, + optimizer: torch.optim.Optimizer, + loss_scaler: torch.GradScaler | None, + lr_schedule: Sequence[float], + epoch: int, + device: torch.device, +): + model.train() + + metric_logger = ut.MetricLogger(delimiter=" ") + metric_logger.add_meter("lr", ut.SmoothedValue(window_size=1, fmt="{value:.6f}")) + metric_logger.add_meter("grad", ut.SmoothedValue()) + header = f"Train: [{epoch}]" + log_wandb = args.wandb and ut.is_main_process() + + epoch_num_batches = len(data_loader) + steps_per_epoch = math.ceil(epoch_num_batches / args.accum_iter) + + print_freq = args.get("print_freq", 100) if not args.debug else 1 + num_batches = epoch_num_batches if not args.debug else 10 + amp_dtype = getattr(torch, args.amp_dtype) + use_cuda = device.type == "cuda" + if use_cuda and args.presend_cuda: + data_loader = ut.pre_send_to_cuda_wrapper( + data_loader, device, dtype_map={torch.float16: amp_dtype} + ) + + optimizer.zero_grad() + + for batch_idx, batch in enumerate( + metric_logger.log_every(data_loader, print_freq, header, total_steps=num_batches) + ): + if use_cuda and not args.presend_cuda: + batch = ut.send_data(batch, device, dtype_map={torch.float16: amp_dtype}) + + batch_step = batch_idx + 1 + log_step = batch_step % print_freq == 0 or batch_step == num_batches + update_in_epoch = batch_idx // args.accum_iter + group_size = min(args.accum_iter, num_batches - update_in_epoch * args.accum_iter) + need_update = batch_step % args.accum_iter == 0 or batch_step == num_batches + global_step = epoch * steps_per_epoch + update_in_epoch + lr = lr_schedule[global_step] + if need_update: + ut.update_lr(optimizer.param_groups, lr) + + images, img_mask = mri_data.densify_sparse_image_batch( + batch["image_values"], + batch["img_mask"], + (int(args.get("in_chans", 1)), *args.img_size), + dtype=amp_dtype, + ) + + sync_context = model.no_sync() if args.distributed and not need_update else nullcontext() + with sync_context: + with torch.autocast(device_type=device.type, dtype=amp_dtype, enabled=args.amp): + loss = model( + images, + img_mask=img_mask, + mask_ratio=args.mask_ratio, + pred_mask_ratio=args.pred_mask_ratio, + pad_to_multiple=args.pad_to_multiple, + with_state=False, + ) + + loss_for_log = loss.detach() + torch._assert_async(torch.isfinite(loss_for_log), "non-finite loss") + + grad_norm = ut.backward_step( + loss / group_size, + optimizer, + scaler=loss_scaler, + need_update=need_update, + max_norm=args.clip_grad, + ) + + if need_update and log_step: + loss_value = loss_for_log.item() + grad_norm_value = grad_norm.item() + metric_logger.update(loss=loss_value, lr=lr, grad=grad_norm_value) + if log_wandb: + wandb.log( + { + "train/loss": loss_value, + "train/lr": lr, + "train/grad": grad_norm_value, + }, + step=int(1000 * (epoch + batch_step / epoch_num_batches)), + ) + + # gather the stats from all processes + metric_logger.synchronize_between_processes() + print("Averaged stats:", metric_logger) + return {f"train/{k}": meter.global_avg for k, meter in metric_logger.meters.items()} + + +@torch.inference_mode() +def evaluate( + args: DictConfig, + model: nn.Module, + data_loader: Iterable, + epoch: int, + device: torch.device, + eval_name: str, +): + model.eval() + + metric_logger = ut.MetricLogger(delimiter=" ") + header = f"Eval ({eval_name}): [{epoch}]" + is_master = ut.is_main_process() + log_wandb = args.wandb and is_master + + epoch_num_batches = len(data_loader) + if epoch_num_batches <= 0: + raise ValueError(f"eval loader {eval_name!r} has zero batches") + + print_freq = args.get("print_freq", 100) if not args.debug else 1 + num_batches = epoch_num_batches if not args.debug else 10 + num_batches = min(num_batches, epoch_num_batches) + eval_seed = int(args.get("eval_seed", args.seed)) + ut.get_rank() + example_step = random.Random(eval_seed).randint(1, num_batches) + amp_dtype = getattr(torch, args.amp_dtype) + use_cuda = device.type == "cuda" + rng_state = ut.capture_rng_state() + torch.set_rng_state(torch.Generator().manual_seed(eval_seed).get_state()) + if use_cuda: + torch.cuda.manual_seed(eval_seed) + if use_cuda and args.presend_cuda: + data_loader = ut.pre_send_to_cuda_wrapper( + data_loader, device, dtype_map={torch.float16: amp_dtype} + ) + + eval_batches = islice(data_loader, num_batches) + for batch_idx, batch in enumerate( + metric_logger.log_every(eval_batches, print_freq, header, total_steps=num_batches) + ): + if use_cuda and not args.presend_cuda: + batch = ut.send_data(batch, device, dtype_map={torch.float16: amp_dtype}) + + batch_step = batch_idx + 1 + + images, img_mask = mri_data.densify_sparse_image_batch( + batch["image_values"], + batch["img_mask"], + (int(args.get("in_chans", 1)), *args.img_size), + dtype=amp_dtype, + ) + + with torch.autocast(device_type=device.type, dtype=amp_dtype, enabled=args.amp): + loss, state = model( + images, + img_mask=img_mask, + mask_ratio=args.mask_ratio, + pred_mask_ratio=args.pred_mask_ratio, + pad_to_multiple=args.pad_to_multiple, + ) + + loss_value = loss.detach().float().item() + finite = torch.tensor(int(math.isfinite(loss_value)), dtype=torch.int32, device=device) + if args.distributed: + torch.distributed.all_reduce(finite, op=torch.distributed.ReduceOp.MIN) + if not finite.item(): + raise RuntimeError("non-finite validation loss detected") + metric_logger.meters["loss"].update(loss_value, n=int(batch["img_mask"].shape[0])) + + if is_master and batch_step == example_step: + example_batch = {"image": images[:1], "img_mask": img_mask[:1]} + if "meta" in batch: + example_batch["meta"] = batch["meta"][:1] + example_state = { + "pred_images": state["pred_images"][:1], + "pred_mask": state["pred_mask"][:1], + } + example_data = { + "batch": ut.send_data(example_batch, "cpu"), + "state": ut.send_data(example_state, "cpu"), + } + + # gather the stats from all processes + metric_logger.synchronize_between_processes() + print(f"Averaged stats ({eval_name}):", metric_logger) + stats = {f"eval/{eval_name}/{k}": meter.global_avg for k, meter in metric_logger.meters.items()} + + plots = {} + if is_master: + print(f"Making plots ({eval_name}): example={example_step}") + plots = make_plots(args, **example_data) + plots = {f"eval/{eval_name}/{k}": img for k, img in plots.items()} + + if log_wandb: + wandb.log(stats, step=1000 * (epoch + 1)) + wandb.log( + {k: wandb.Image(img, caption=f"example={example_step}") for k, img in plots.items()}, + step=1000 * (epoch + 1), + ) + ut.restore_rng_state(rng_state) + return stats, plots + + +def make_plots( + args: DictConfig, + batch: dict[str, Tensor], + state: dict[str, Tensor], +) -> dict[str, Image.Image]: + fig_kwargs = args.get("fig_kwargs", {}) + + images = batch["image"] + img_mask = batch.get("img_mask") + if img_mask is not None: + img_mask = img_mask.expand_as(images) + + raw_mean, raw_std = vis.raw_stats_from_batch(batch) + + plots = {} + mask_pred_fig = vis.plot_mask_pred( + target=images, + pred=state["pred_images"], + pred_mask=state["pred_mask"], + img_mask=img_mask, + patch_size=args.patch_size, + raw_mean=raw_mean, + raw_std=raw_std, + **ut.filter_kwargs(vis.plot_mask_pred, fig_kwargs), + ) + plots["mask_pred"] = vis.fig2pil(mask_pred_fig) + plt.close(mask_pred_fig) + + return plots + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--cfg-path", type=str, default=None) + parser.add_argument("--overrides", type=str, default=None, nargs="+") + args = parser.parse_args() + cfg = OmegaConf.load(DEFAULT_CONFIG) + if args.cfg_path: + cfg = OmegaConf.unsafe_merge(cfg, OmegaConf.load(args.cfg_path)) + if args.overrides: + cfg = OmegaConf.unsafe_merge(cfg, OmegaConf.from_dotlist(args.overrides)) + main(cfg) diff --git a/fomo_tune_baseline/output/task1/build/smri_mae/masking.py b/fomo_tune_baseline/output/task1/build/smri_mae/masking.py new file mode 100644 index 0000000000000000000000000000000000000000..28f9d53cbbd8d32a3923f9b1a6655c92cd526bf7 --- /dev/null +++ b/fomo_tune_baseline/output/task1/build/smri_mae/masking.py @@ -0,0 +1,80 @@ +import torch +from jaxtyping import Float, Int +from torch import Tensor + + +def pad_patch_mask( + patch_mask: Float[Tensor, "B N"], + mask_ratio: float, + shuffle: bool = False, + generator: torch.Generator | None = None, + pad_to_multiple: int | None = None, +) -> tuple[Float[Tensor, "B N"], Int[Tensor, "B L"], Tensor]: + """ + Select each row's own mask-ratio count, then pad ids to the batch max length. + + Returns: + - selected patch mask [B, N] + - padded selected patch ids [B, Lpad] + - token mask [B, Lpad], true for real ids and false for padding + """ + if not 0.0 <= mask_ratio <= 1.0: + raise ValueError(f"mask_ratio must be in [0, 1], got {mask_ratio}") + + B, N = patch_mask.shape + device = patch_mask.device + patch_mask = patch_mask.to(dtype=torch.bool) + + valid_counts = patch_mask.sum(dim=1) + num_keep = torch.floor(valid_counts.to(torch.float64) * (1.0 - mask_ratio)).to(torch.long) + if not shuffle: + selected = patch_mask & (patch_mask.cumsum(dim=1) <= num_keep.unsqueeze(1)) + padded_ids, token_mask = patch_ids_from_mask( + selected, + pad_to_multiple=pad_to_multiple, + ) + return selected, padded_ids, token_mask + + # One masked sort directly produces random valid IDs. The previous + # shuffle/select/inverse-shuffle path required two full argsorts plus a + # dynamic nonzero/scatter solely to recover the same selected set. + noise = torch.rand(B, N, generator=generator, device=device) + noise.masked_fill_(~patch_mask, torch.inf) + shuffled_ids = torch.argsort(noise, dim=1) + + max_count = int(num_keep.max().item()) + if pad_to_multiple is not None: + if pad_to_multiple <= 0: + raise ValueError(f"pad_to_multiple must be positive, got {pad_to_multiple}") + max_count = (max_count + pad_to_multiple - 1) // pad_to_multiple * pad_to_multiple + padded_ids = shuffled_ids[:, :max_count] + token_mask = torch.arange(max_count, device=device).unsqueeze(0) < num_keep.unsqueeze(1) + selected = torch.zeros_like(patch_mask).scatter_(1, padded_ids, token_mask) + return selected, padded_ids, token_mask + + +def patch_ids_from_mask( + patch_mask: Tensor, + pad_to_multiple: int | None = None, +) -> tuple[Int[Tensor, "B L"], Tensor]: + """Return optionally rounded patch IDs and their token-validity mask.""" + if pad_to_multiple is not None and pad_to_multiple <= 0: + raise ValueError(f"pad_to_multiple must be positive, got {pad_to_multiple}") + + patch_mask = patch_mask.to(dtype=torch.bool) + B, N = patch_mask.shape + device = patch_mask.device + counts = patch_mask.sum(dim=1) + max_count = int(counts.max().item()) + if pad_to_multiple is not None: + max_count = (max_count + pad_to_multiple - 1) // pad_to_multiple * pad_to_multiple + + patch_ids = torch.zeros((B, max_count), dtype=torch.long, device=device) + token_mask = torch.arange(max_count, device=device).unsqueeze(0) < counts.unsqueeze(1) + if max_count == 0: + return patch_ids, token_mask + + batch_ids, selected_ids = patch_mask.nonzero(as_tuple=True) + slot_ids = patch_mask.cumsum(dim=1)[batch_ids, selected_ids].to(torch.long) - 1 + patch_ids[batch_ids, slot_ids] = selected_ids + return patch_ids, token_mask diff --git a/fomo_tune_baseline/output/task1/build/smri_mae/model_mae.py b/fomo_tune_baseline/output/task1/build/smri_mae/model_mae.py new file mode 100644 index 0000000000000000000000000000000000000000..f9a4d884784e53f55ae557c46dd4afa2eb70e7d3 --- /dev/null +++ b/fomo_tune_baseline/output/task1/build/smri_mae/model_mae.py @@ -0,0 +1,916 @@ +# Copyright (c) Sophont, Inc +# This source code is licensed under the Apache License, Version 2.0 +# +# References: +# capi: https://github.com/facebookresearch/capi/blob/main/model.py +# timm: https://github.com/huggingface/pytorch-image-models/blob/v1.0.20/timm/models/vision_transformer.py + +""" +From-scratch re-implementation of the original MAE model. + +MaskedEncoder: standard ViT with masking +MaskedDecoder: standard self-attention MAE decoder +MaskedAutoEncoderViT: full MAE model for 3D structural MRI volumes +""" + +from collections.abc import Sequence +from typing import Literal + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch import Tensor +from torch.utils.checkpoint import checkpoint +from huggingface_hub import PyTorchModelHubMixin +from jaxtyping import Float, Int + +from .modules import ( + AbsolutePosEmbed, + Block, + LayerNorm, + Normalize, + JaggedBatch, + Patchify3D, + SeparablePosEmbed, + SinCosPosEmbed3D, + unpack_tokens, +) +from .masking import pad_patch_mask + + +class MaskedEncoder(nn.Module): + """ + Masked transformer encoder. + """ + + def __init__( + self, + patchify: nn.Module, + patch_embed: nn.Module, + pos_embed: nn.Module, + depth: int = 12, + embed_dim: int = 768, + num_heads: int = 12, + qkv_bias: bool = True, + proj_bias: bool = True, + mlp_ratio: int | float = 4, + class_token: bool = True, + reg_tokens: int = 0, + no_embed_class: bool = False, + final_norm: bool = True, + drop_path_rate: float = 0.0, + mask_drop_scale: bool = False, + ): + super().__init__() + self.num_prefix_tokens = int(class_token) + reg_tokens + self.num_reg_tokens = reg_tokens + self.has_class_token = class_token + self.no_embed_class = no_embed_class + + # scale inputs by 1 / observed rate (like dropout) + self.mask_drop_scale = mask_drop_scale + + # inject tokenization modules, so that the encoder doesn't specifically need to + # know how the data are tokenized, while still implementing a complete + # self-contained model. + self.patchify = patchify + self.patch_embed = patch_embed + self.pos_embed = pos_embed + + R = reg_tokens + self.cls_token = nn.Parameter(torch.empty(1, 1, embed_dim)) if class_token else None + self.reg_token = nn.Parameter(torch.empty(1, R, embed_dim)) if reg_tokens else None + + if not no_embed_class: + self.cls_token_pos = nn.Parameter(torch.empty(1, 1, embed_dim)) if class_token else None + self.reg_token_pos = nn.Parameter(torch.empty(1, R, embed_dim)) if reg_tokens else None + else: + self.cls_token_pos = self.reg_token_pos = None + + # stochastic depth decay rule + dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)] + + self.blocks = nn.ModuleList( + [ + Block( + dim=embed_dim, + num_heads=num_heads, + qkv_bias=qkv_bias, + proj_bias=proj_bias, + mlp_ratio=mlp_ratio, + drop_path=dpr[ii], + ) + for ii in range(depth) + ] + ) + + self.norm = LayerNorm(embed_dim) if final_norm else nn.Identity() + + self.reset_parameters() + + def extra_repr(self): + return ( + f"class_token={self.has_class_token}, reg_tokens={self.num_reg_tokens}, " + f"no_embed_class={self.no_embed_class}, mask_drop_scale={self.mask_drop_scale}" + ) + + def reset_parameters(self) -> None: + for p in [self.cls_token, self.cls_token_pos, self.reg_token, self.reg_token_pos]: + if p is not None: + nn.init.trunc_normal_(p, std=0.02) + + def cat_tokens(self, x: Tensor) -> Tensor: + # prepend cls and reg tokens with optional learned position embedding + # the cls and reg pos embedding is ofc redundant, but included in many other + # implementations. + B, _, _ = x.shape + + to_cat = [] + if self.has_class_token: + cls_token = self.cls_token + if not self.no_embed_class: + cls_token = cls_token + self.cls_token_pos + to_cat.append(cls_token.expand(B, -1, -1)) + + if self.num_reg_tokens: + reg_token = self.reg_token + if not self.no_embed_class: + reg_token = reg_token + self.reg_token_pos + to_cat.append(reg_token.expand(B, -1, -1)) + + if to_cat: + x = torch.cat(to_cat + [x], dim=1) + return x + + def cat_token_mask(self, token_mask: Tensor, batch_size: int) -> Tensor: + if self.num_prefix_tokens: + prefix_mask = torch.ones( + (batch_size, self.num_prefix_tokens), + dtype=torch.bool, + device=token_mask.device, + ) + token_mask = torch.cat([prefix_mask, token_mask], dim=1) + return token_mask + + def chunk_tokens(self, x: Tensor) -> tuple[Tensor | None, Tensor | None, Tensor]: + cls_offset = int(self.has_class_token) + cls = x[:, :cls_offset] if self.has_class_token else None + if self.num_reg_tokens: + reg = x[:, cls_offset : self.num_prefix_tokens, :] + else: + reg = None + patch = x[:, self.num_prefix_tokens :, :] + return cls, reg, patch + + def forward( + self, + x: Tensor, + mask: Tensor | None = None, + mask_ratio: float | None = None, + pad_to_multiple: int | None = None, + ) -> tuple[ + Float[Tensor, "B 1 D"] | None, + Float[Tensor, "B R D"] | None, + Float[Tensor, "B L D"], + Tensor | None, + Int[Tensor, "B L"] | None, + Tensor | None, + ]: + """ + x: input data shape [B, C, D, H, W] + mask: visible mask, 1 = visible, 0 = invisible. broadcastable shape + mask_ratio: mask ratio for uniform random masking + + returns: + - cls_embeds: [B, 1, D] + - reg_embeds: [B, R, D] + - patch_embeds: [B, L, D], where L is the number of visible patches + - mask: observed mask, 1 = observed, 0 = unobserved. same shape as input + - mask_ids: indices of visible patches [B L] + - token_mask: valid token mask for padded per-sample masking [B L] + """ + # apply mask to the input + if mask is not None: + mask = mask.to(device=x.device, dtype=torch.bool).expand_as(x) + x = x.masked_fill(~mask, 0) + + # patchify input + x = self.patchify(x) + B, N, P = x.shape + + # patchify mask and apply dropout style scaling + if mask is not None: + mask_patches = self.patchify(mask) + patch_num_obs = mask_patches.sum(dim=-1) + patch_mask = patch_num_obs > 0 + if self.mask_drop_scale: + patch_num_obs = patch_num_obs.to(x.dtype) + x = x * (P / patch_num_obs.unsqueeze(-1).clamp(min=1.0)) + elif mask_ratio is not None: + patch_mask = torch.ones((B, N), dtype=torch.bool, device=x.device) + mask_patches = patch_mask.unsqueeze(-1).expand(-1, -1, P) + else: + patch_mask = mask_patches = None + + # patch and position embed + x = self.patch_embed(x) + x = self.pos_embed(x) + + if mask is not None or mask_ratio is not None: + mask_ratio = 0.0 if mask_ratio is None else mask_ratio + patch_mask, mask_ids, token_mask = pad_patch_mask( + patch_mask, + mask_ratio=mask_ratio, + shuffle=mask_ratio > 0, + pad_to_multiple=pad_to_multiple, + ) + + mask_patches = mask_patches & patch_mask.unsqueeze(-1) + mask = self.patchify.unpatchify(mask_patches) + x = x.gather(1, mask_ids.unsqueeze(-1).expand(-1, -1, x.shape[-1])) + else: + mask_ids = None + token_mask = None + + cls_embeds, reg_embeds, patch_embeds = self.forward_patch_embeds( + x, + token_mask=token_mask, + ) + return cls_embeds, reg_embeds, patch_embeds, mask, mask_ids, token_mask + + def forward_patch_embeds( + self, + x: Float[Tensor, "B L D"], + token_mask: Tensor | None = None, + ) -> tuple[ + Float[Tensor, "B 1 D"] | None, + Float[Tensor, "B R D"] | None, + Float[Tensor, "B L D"], + ]: + B = x.shape[0] + if token_mask is None: + token_mask = torch.ones(x.shape[:2], dtype=torch.bool, device=x.device) + x = self.cat_tokens(x) + token_mask = self.cat_token_mask(token_mask, B) + jagged_batch = JaggedBatch.from_mask(token_mask) + x = x[token_mask] + for block in self.blocks: + x = block(x, jagged_batch=jagged_batch) + x = self.norm(x) + x = unpack_tokens(x, token_mask) + + cls_embeds, reg_embeds, patch_embeds = self.chunk_tokens(x) + return cls_embeds, reg_embeds, patch_embeds + + def forward_visible_ids( + self, + x: Tensor, + visible_ids: Int[Tensor, "B L"], + img_mask: Tensor | None = None, + ) -> tuple[ + Float[Tensor, "B 1 D"] | None, + Float[Tensor, "B R D"] | None, + Float[Tensor, "B L D"], + ]: + if img_mask is not None: + img_mask = img_mask.to(device=x.device, dtype=torch.bool).expand_as(x) + x = x.masked_fill(~img_mask, 0) + + x = self.patchify(x) + if self.mask_drop_scale and img_mask is not None: + mask_patches = self.patchify(img_mask) + patch_num_obs = mask_patches.sum(dim=-1).to(x.dtype) + x = x * (self.patchify.patch_dim / patch_num_obs.unsqueeze(-1).clamp(min=1.0)) + x = self.patch_embed(x) + x = self.pos_embed(x) + visible_ids = visible_ids.to(device=x.device) + x = x.gather(1, visible_ids.unsqueeze(-1).expand(-1, -1, x.shape[-1])) + return self.forward_patch_embeds(x) + + def forward_embedding( + self, + x: Tensor, + mask: Tensor | None = None, + mask_ratio: float | None = None, + ): + cls_embeds, reg_embeds, patch_embeds, *_ = self.forward( + x, + mask=mask, + mask_ratio=mask_ratio, + ) + return cls_embeds, reg_embeds, patch_embeds + + +class MaskedDecoder(nn.Module): + """Self-attention MAE decoder supporting sparse subset decoding via pred_ids.""" + + def __init__( + self, + pos_embed: nn.Module, + head: nn.Module | None = None, + input_dim: int | None = None, + depth: int = 12, + embed_dim: int = 768, + num_heads: int = 12, + qkv_bias: bool = True, + proj_bias: bool = True, + mlp_ratio: int | float = 4, + class_token: bool = True, + no_embed_class: bool = False, + final_norm: bool = True, + ): + super().__init__() + input_dim = embed_dim if input_dim is None else input_dim + self.has_class_token = class_token + self.no_embed_class = no_embed_class + + self.cls_token = nn.Parameter(torch.empty(1, 1, embed_dim)) if class_token else None + self.cls_token_pos = ( + nn.Parameter(torch.empty(1, 1, embed_dim)) + if class_token and not no_embed_class + else None + ) + self.mask_token = nn.Parameter(torch.empty(1, 1, embed_dim)) + + # decoder position embedding, encodes query position information into masks + self.pos_embed = pos_embed + + self.proj = nn.Identity() if input_dim == embed_dim else nn.Linear(input_dim, embed_dim) + + self.blocks = nn.ModuleList( + [ + Block( + dim=embed_dim, + num_heads=num_heads, + qkv_bias=qkv_bias, + proj_bias=proj_bias, + mlp_ratio=mlp_ratio, + ) + for _ in range(depth) + ] + ) + + self.norm = LayerNorm(embed_dim) if final_norm else nn.Identity() + + # optional injected prediction head + self.head = nn.Identity() if head is None else head + + self.reset_parameters() + + def extra_repr(self): + return f"class_token={self.has_class_token}, no_embed_class={self.no_embed_class}" + + def reset_parameters(self) -> None: + # official mae initializes decoder cls token to zeros + # although perhaps this was an oversight + if self.cls_token is not None: + nn.init.zeros_(self.cls_token) + if self.cls_token_pos is not None: + nn.init.trunc_normal_(self.cls_token_pos, std=0.02) + nn.init.trunc_normal_(self.mask_token, std=0.02) + + def cat_tokens(self, x: Tensor) -> Tensor: + if not self.has_class_token: + return x + cls_token = self.cls_token + if not self.no_embed_class: + cls_token = cls_token + self.cls_token_pos + return torch.cat([cls_token.expand(x.shape[0], -1, -1), x], dim=1) + + def cat_token_mask(self, token_mask: Tensor, batch_size: int) -> Tensor: + if self.has_class_token: + cls_mask = torch.ones( + (batch_size, 1), + dtype=torch.bool, + device=token_mask.device, + ) + token_mask = torch.cat([cls_mask, token_mask], dim=1) + return token_mask + + def chunk_tokens(self, x: Tensor) -> tuple[Tensor | None, Tensor]: + cls_offset = int(self.has_class_token) + cls = x[:, :cls_offset] if self.has_class_token else None + patch = x[:, cls_offset:, :] + return cls, patch + + def forward( + self, + embeds: Float[Tensor, "B L D"], + embed_ids: Int[Tensor, "B L"] | None = None, + pred_ids: Int[Tensor, "B Q"] | None = None, + embed_token_mask: Tensor | None = None, + pred_token_mask: Tensor | None = None, + packed_output: bool = False, + ) -> Float[Tensor, "B Q P"] | Float[Tensor, "T P"]: + """ + embeds: input patch embeddings. + embed_ids: optional patch indices for input embeddings. If not provided, no + position will be added to the embeddings. + pred_ids: patch indices of query mask positions. If None, decode *all* patches. + + returns: + - pred [B, Q, P] where Q is the number of prediction patches and P is the output + dimension + """ + B, L, _ = embeds.shape + + Q = self.pos_embed.num_patches if pred_ids is None else pred_ids.shape[1] + mask = self.mask_token.expand(B, Q, -1) + mask = self.pos_embed(mask, pos_ids=pred_ids) + + embeds = self.proj(embeds) + + if embed_ids is not None: + embeds = self.pos_embed(embeds, pos_ids=embed_ids) + if embed_token_mask is None: + embed_token_mask = torch.ones((B, L), dtype=torch.bool, device=embeds.device) + if pred_token_mask is None: + pred_token_mask = torch.ones((B, Q), dtype=torch.bool, device=embeds.device) + x = torch.cat([embeds, mask], dim=1) + token_mask = torch.cat([embed_token_mask, pred_token_mask], dim=1) + + x = self.cat_tokens(x) + token_mask = self.cat_token_mask(token_mask, B) + jagged_batch = JaggedBatch.from_mask(token_mask) + x = x[token_mask] + # Keep headroom for rare maximum-length PSP batches. + checkpoint_start = max(0, len(self.blocks) - 2) + for block_index, block in enumerate(self.blocks): + if self.training and torch.is_grad_enabled() and block_index >= checkpoint_start: + x = checkpoint(block, x, jagged_batch, use_reentrant=False) + else: + x = block(x, jagged_batch=jagged_batch) + + x = self.norm(x) + if packed_output: + pred_offset = int(self.has_class_token) + L + prediction_mask = F.pad(pred_token_mask, (pred_offset, 0)) + return self.head(x[prediction_mask[token_mask]]) + + x = unpack_tokens(x, token_mask) + _, x = self.chunk_tokens(x) + + pred = x[:, L:] + pred = pred.masked_fill(~pred_token_mask.unsqueeze(-1), 0) + pred = self.head(pred) + return pred + + +class MaskedAutoencoderViT(nn.Module, PyTorchModelHubMixin): + def __init__( + self, + img_size: int | tuple[int, int, int] = (208, 240, 208), + patch_size: int | tuple[int, int, int] = (16, 16, 16), + in_chans: int = 1, + depth: int = 12, + embed_dim: int = 768, + num_heads: int = 12, + decoder_depth: int = 4, + decoder_embed_dim: int | None = 512, + decoder_num_heads: int | None = 16, # default from mae, head dim = 32 + qkv_bias: bool = True, + proj_bias: bool = True, + mlp_ratio: int | float = 4, + class_token: bool = True, + reg_tokens: int = 0, + no_embed_class: bool = False, + drop_path_rate: float = 0.0, + mask_drop_scale: bool = False, + no_decode_pos: bool = False, + pos_embed: Literal["abs", "sep", "sincos"] = "sincos", + target_norm: Literal["none", "global", "slice", "patch"] | None = None, + ): + super().__init__() + img_size = _to_3d_tuple(img_size, "img_size") + patch_size = _to_3d_tuple(patch_size, "patch_size") + + self.no_decode_pos = no_decode_pos # don't pos encode embeddings in decoder + + # patchify reshapes input into sequence of flattened patches, shape [B, N, P] + ndim = 3 + patchify = Patchify3D(img_size, patch_size, in_chans=in_chans) + + # linear patch embedding P -> D + patch_embed = nn.Linear(patchify.patch_dim, embed_dim) + + # position embedding + # separable position embedding decouples the first spatial axis from the + # others. Fixed sin/cos embeddings are the default for sMRI volumes. + if pos_embed == "sincos": + pos_embed_layer = SinCosPosEmbed3D + else: + pos_embed_layer = {"abs": AbsolutePosEmbed, "sep": SeparablePosEmbed}[pos_embed] + pos_embed = pos_embed_layer(embed_dim, patchify.grid_size) + + # encoder. for inference, this model can be extracted and used like a regular vit + self.encoder = MaskedEncoder( + patchify=patchify, + patch_embed=patch_embed, + pos_embed=pos_embed, + depth=depth, + embed_dim=embed_dim, + num_heads=num_heads, + qkv_bias=qkv_bias, + proj_bias=proj_bias, + mlp_ratio=mlp_ratio, + class_token=class_token, + reg_tokens=reg_tokens, + no_embed_class=no_embed_class, + drop_path_rate=drop_path_rate, + mask_drop_scale=mask_drop_scale, + ) + + self.pred_patchify = patchify + + # fall back to encoder architecture width + decoder_embed_dim = decoder_embed_dim or embed_dim + decoder_num_heads = decoder_num_heads or num_heads + + decoder_pos_embed = pos_embed_layer(decoder_embed_dim, self.pred_patchify.grid_size) + # we might want to try tying the weights of the prediction head to the patch + # embedding at some point. + decoder_head = nn.Linear(decoder_embed_dim, self.pred_patchify.patch_dim) + + self.decoder = MaskedDecoder( + pos_embed=decoder_pos_embed, + head=decoder_head, + input_dim=embed_dim, + depth=decoder_depth, + embed_dim=decoder_embed_dim, + num_heads=decoder_num_heads, + qkv_bias=qkv_bias, + proj_bias=proj_bias, + mlp_ratio=mlp_ratio, + class_token=class_token, + no_embed_class=no_embed_class, + ) + + # mae style target normalization + # dim is relative to an unflattened embedding tensor of shape [B, *grid_size, D] + if target_norm not in {"none", None}: + norm_dim = { + "global": tuple(range(1, ndim + 2)), # full sequence + "slice": tuple(range(2, ndim + 2)), # each depth slice along first dim + "patch": -1, # normalize each patch independently (mae pix norm loss) + }[target_norm] + self.target_norm = Normalize(self.pred_patchify.grid_size, dim=norm_dim) + else: + self.target_norm = None + + self.init_weights() + + def extra_repr(self): + return f"no_decode_pos={self.no_decode_pos}" + + def init_weights(self): + self.apply(_init_weights) + + def prepare_targets(self, images: Tensor, img_mask: Tensor | None): + """ + images: [B, C, D, H, W] + img_mask: mask of valid data. only used for computing correct normalization + stats. same shape as images. + """ + targets_patches = self.pred_patchify(images) # [B, N, P] + + # target normalization + if self.target_norm is not None: + # full image data mask used for normalization stats only + if img_mask is not None: + img_mask_patches = self.pred_patchify(img_mask) + else: + img_mask_patches = None + targets_patches, *targets_stats = self.target_norm(targets_patches, img_mask_patches) + else: + targets_stats = None + + return targets_patches, targets_stats + + def prepare_masks( + self, + img_mask: Tensor, + visible_mask: Tensor | None, + pred_mask: Tensor | None, + device: torch.device, + ): + img_mask = img_mask.to(device=device, dtype=torch.bool) + + if visible_mask is None: + visible_mask = img_mask + else: + visible_mask = img_mask & visible_mask.to(device=device, dtype=torch.bool) + + if pred_mask is None: + pred_mask = img_mask + else: + pred_mask = img_mask & pred_mask.to(device=device, dtype=torch.bool) + + return img_mask, visible_mask, pred_mask + + def prepare_pred_mask( + self, + visible_mask: Tensor, + pred_mask: Tensor | None = None, + pred_mask_ratio: float | None = None, + pad_to_multiple: int | None = None, + ): + """ + prepare prediction mask by removing visible content + visible_mask: [B, C, D, H, W], 1 = visible, 0 = invisible + pred_mask: same shape, 1 = predict, 0 = don't predict + """ + if pred_mask is None: + pred_mask = torch.ones_like(visible_mask) + + pred_mask = pred_mask & ~visible_mask + + pred_mask_patches = self.pred_patchify(pred_mask) + pred_patch_mask = pred_mask_patches.any(dim=-1) + # Optionally subsample the prediction candidates. + mask_ratio = 0.0 if pred_mask_ratio is None else pred_mask_ratio + pred_patch_mask, pred_ids, pred_token_mask = pad_patch_mask( + pred_patch_mask, + mask_ratio=mask_ratio, + # With per-sample padding every candidate is retained when the ratio + # is zero, so randomizing their order is pure overhead. + shuffle=mask_ratio > 0, + pad_to_multiple=pad_to_multiple, + ) + pred_mask_patches = pred_mask_patches & pred_patch_mask.unsqueeze(-1) + return pred_mask_patches, pred_ids, pred_token_mask + + def forward_decoder( + self, + patch_embeds: Float[Tensor, "B L D"], + visible_ids: Int[Tensor, "B L"], + pred_ids: Int[Tensor, "B Q"] | None, + visible_token_mask: Tensor | None = None, + pred_token_mask: Tensor | None = None, + packed_output: bool = False, + ) -> Float[Tensor, "B Q P"] | Float[Tensor, "T P"]: + return self.decoder( + patch_embeds, + embed_ids=None if self.no_decode_pos else visible_ids, + pred_ids=pred_ids, + embed_token_mask=visible_token_mask, + pred_token_mask=pred_token_mask, + packed_output=packed_output, + ) + + def forward_loss( + self, + preds: Float[Tensor, "T P"], + targets_patches: Float[Tensor, "B N P"], + pred_mask_patches: Float[Tensor, "B N P"], + pred_ids: Int[Tensor, "B Q"], + pred_token_mask: Tensor, + ) -> Tensor: + """Average valid-voxel MSE within each scan, then average across scans.""" + batch_ids, slot_ids = pred_token_mask.nonzero(as_tuple=True) + patch_ids = pred_ids[batch_ids, slot_ids] + targets = targets_patches[batch_ids, patch_ids] + voxel_mask = pred_mask_patches[batch_ids, patch_ids] + + patch_errors = ((preds - targets) ** 2 * voxel_mask).sum(dim=1) + patch_voxels = voxel_mask.sum(dim=1).to(dtype=patch_errors.dtype) + batch_size = targets_patches.shape[0] + scan_errors = patch_errors.new_zeros(batch_size).scatter_add_(0, batch_ids, patch_errors) + scan_voxels = patch_voxels.new_zeros(batch_size).scatter_add_(0, batch_ids, patch_voxels) + return (scan_errors / scan_voxels).mean() + + @torch.no_grad() + def forward_pred_images( + self, + preds: Float[Tensor, "B Q P"], + pred_ids: Int[Tensor, "B Q"], + pred_token_mask: Tensor | None = None, + img_mask: Tensor | None = None, + targets_stats: tuple[Tensor, Tensor] | None = None, + ) -> Tensor: + B, _, P = preds.shape + N = self.pred_patchify.num_patches + if pred_token_mask is not None: + preds = preds.masked_fill(~pred_token_mask.unsqueeze(-1), 0) + + preds = torch.zeros((B, N, P), dtype=preds.dtype, device=preds.device).scatter_add_( + 1, pred_ids.unsqueeze(-1).expand(-1, -1, P), preds + ) + + if targets_stats is not None: + targets_mean, targets_std = targets_stats + preds = preds * targets_std + targets_mean + + pred_images = self.pred_patchify.unpatchify(preds) + if img_mask is not None: + pred_images = pred_images.masked_fill(~img_mask, 0) + return pred_images + + def forward( + self, + images: Tensor, + img_mask: Tensor, + mask_ratio: float, + pred_mask_ratio: float | None = None, + pad_to_multiple: int | None = None, + with_state: bool = True, + ) -> Tensor | tuple[Tensor, dict]: + img_mask, visible_mask, pred_mask = self.prepare_masks( + img_mask, + None, + None, + device=images.device, + ) + targets_patches, targets_stats = self.prepare_targets(images, img_mask) + + ( + cls_embeds, + reg_embeds, + patch_embeds, + visible_mask, + visible_ids, + visible_token_mask, + ) = self.encoder( + images, + mask=visible_mask, + mask_ratio=mask_ratio, + pad_to_multiple=pad_to_multiple, + ) + + pred_mask_patches, pred_ids, pred_token_mask = self.prepare_pred_mask( + visible_mask, + pred_mask=pred_mask, + pred_mask_ratio=pred_mask_ratio, + pad_to_multiple=pad_to_multiple, + ) + + preds = self.forward_decoder( + patch_embeds, + visible_ids, + pred_ids, + visible_token_mask=visible_token_mask, + pred_token_mask=pred_token_mask, + packed_output=not with_state, + ) + + loss_preds = preds if not with_state else preds[pred_token_mask] + loss = self.forward_loss( + loss_preds, + targets_patches, + pred_mask_patches, + pred_ids, + pred_token_mask, + ) + + if not with_state: + return loss + + pred_mask = self.pred_patchify.unpatchify(pred_mask_patches) + pred_images = self.forward_pred_images( + preds, + pred_ids, + pred_token_mask=pred_token_mask, + img_mask=img_mask, + targets_stats=targets_stats, + ) + + state = { + "targets_patches": targets_patches, + "targets_stats": targets_stats, + "patch_embeds": patch_embeds, + "cls_embeds": cls_embeds, + "reg_embeds": reg_embeds, + "visible_mask": visible_mask, + "visible_ids": visible_ids, + "visible_token_mask": visible_token_mask, + "pred_mask": pred_mask, + "pred_ids": pred_ids, + "pred_token_mask": pred_token_mask, + "preds": preds, + "pred_images": pred_images, + } + return loss, state + + def forward_embedding( + self, + x: Tensor, + mask: Tensor | None = None, + mask_ratio: float | None = None, + ): + return self.encoder.forward_embedding(x, mask=mask, mask_ratio=mask_ratio) + + +class MaskedViT(MaskedEncoder, PyTorchModelHubMixin): + def __init__( + self, + img_size: int | tuple[int, int, int] = (208, 240, 208), + in_chans: int = 1, + patch_size: int | tuple[int, int, int] = (16, 16, 16), + depth: int = 12, + embed_dim: int = 768, + num_heads: int = 12, + qkv_bias: bool = True, + proj_bias: bool = True, + mlp_ratio: int | float = 4, + class_token: bool = True, + reg_tokens: int = 0, + no_embed_class: bool = False, + final_norm: bool = True, + drop_path_rate: float = 0.0, + mask_drop_scale: bool = False, + pos_embed: Literal["abs", "sep", "sincos"] = "sincos", + ): + img_size = _to_3d_tuple(img_size, "img_size") + patch_size = _to_3d_tuple(patch_size, "patch_size") + + patchify = Patchify3D(img_size, patch_size, in_chans=in_chans) + patch_embed = nn.Linear(patchify.patch_dim, embed_dim) + if pos_embed == "sincos": + pos_embed_layer = SinCosPosEmbed3D + else: + pos_embed_layer = {"abs": AbsolutePosEmbed, "sep": SeparablePosEmbed}[pos_embed] + pos_embed = pos_embed_layer(embed_dim, patchify.grid_size) + + super().__init__( + patchify=patchify, + patch_embed=patch_embed, + pos_embed=pos_embed, + depth=depth, + embed_dim=embed_dim, + num_heads=num_heads, + qkv_bias=qkv_bias, + proj_bias=proj_bias, + mlp_ratio=mlp_ratio, + class_token=class_token, + reg_tokens=reg_tokens, + no_embed_class=no_embed_class, + final_norm=final_norm, + drop_path_rate=drop_path_rate, + mask_drop_scale=mask_drop_scale, + ) + + self.init_weights() + + def init_weights(self): + self.apply(_init_weights) + + +def _to_3d_tuple(value: int | Sequence[int], name: str) -> tuple[int, int, int]: + if isinstance(value, int): + return (value, value, value) + if len(value) != 3: + raise ValueError(f"{name} must have exactly 3 spatial dimensions, got {tuple(value)}") + return tuple(int(item) for item in value) + + +# JAX ViT xavier uniform init +# https://github.com/facebookresearch/capi/blob/main/model.py +def _init_weights(m: nn.Module) -> None: + if isinstance(m, nn.Linear): + nn.init.xavier_uniform_(m.weight) + if m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.LayerNorm) and m.elementwise_affine: + nn.init.constant_(m.weight, 1.0) + if m.bias is not None: + nn.init.constant_(m.bias, 0) + + +def _create_vit(**kwargs): + model = MaskedViT(**kwargs) + return model + + +def _create_mae_vit(**kwargs): + model = MaskedAutoencoderViT(**kwargs) + return model + + +def mae_vit_small(**kwargs): + model_args = dict(embed_dim=384, depth=12, num_heads=6) + return _create_mae_vit(**model_args, **kwargs) + + +def mae_vit_base(**kwargs): + model_args = dict(embed_dim=768, depth=12, num_heads=12) + return _create_mae_vit(**model_args, **kwargs) + + +def mae_vit_large(**kwargs): + model_args = dict(embed_dim=1024, depth=24, num_heads=16) + return _create_mae_vit(**model_args, **kwargs) + + +def mae_vit_huge(**kwargs): + model_args = dict(embed_dim=1280, depth=32, num_heads=16) + return _create_mae_vit(**model_args, **kwargs) + + +# "patch embed" baseline model, depth 0 ViT (hah) +def patch_embed_small(**kwargs): + model_args = dict(embed_dim=384, depth=0) + return _create_vit(**model_args, **kwargs) + + +def patch_embed_base(**kwargs): + model_args = dict(embed_dim=768, depth=0) + return _create_vit(**model_args, **kwargs) diff --git a/fomo_tune_baseline/output/task1/build/smri_mae/modules.py b/fomo_tune_baseline/output/task1/build/smri_mae/modules.py new file mode 100644 index 0000000000000000000000000000000000000000..828aecb1bd57182f1690feb4e943395db1bdd515 --- /dev/null +++ b/fomo_tune_baseline/output/task1/build/smri_mae/modules.py @@ -0,0 +1,453 @@ +# This source code is licensed under the Apache License, Version 2.0 +# +# References: +# capi: https://github.com/facebookresearch/capi/blob/main/model.py +# timm: https://github.com/huggingface/pytorch-image-models/blob/v1.0.20/timm/models/vision_transformer.py +# vjepa2: https://github.com/facebookresearch/vjepa2/blob/main/src/models/utils/pos_embs.py + +import math +from functools import partial +from typing import NamedTuple, Type + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch import Tensor +from einops import rearrange +from jaxtyping import Float, Int +from timm.layers import DropPath, to_3tuple + +Layer = Type[nn.Module] + + +class JaggedBatch(NamedTuple): + """Sequence boundaries and cached launch metadata for jagged attention.""" + + offsets: Tensor + max_seqlen: int + + @classmethod + def from_mask(cls, mask: Tensor) -> "JaggedBatch": + mask = mask.to(dtype=torch.bool) + counts = mask.sum(dim=1) + return cls( + offsets=F.pad(counts.cumsum(dim=0), (1, 0)), + max_seqlen=mask.shape[1], + ) + + def as_nested(self, tokens: Tensor) -> Tensor: + # Cached conservative bounds avoid min/max reductions and GPU-to-CPU + # synchronization when Flash SDPA inspects the jagged sequence lengths. + return torch.nested.nested_tensor_from_jagged( + tokens, + self.offsets, + min_seqlen=1, + max_seqlen=self.max_seqlen, + ).transpose(1, 2) + + +def unpack_tokens(tokens: Tensor, token_mask: Tensor) -> Tensor: + """Restore packed values to a padded batch, filling invalid slots with zero.""" + output = tokens.new_zeros((*token_mask.shape, *tokens.shape[1:])) + return output.index_put((token_mask,), tokens) + + +def jagged_scaled_dot_product_attention( + query: Tensor, + key: Tensor, + value: Tensor, + jagged_batch: JaggedBatch, +) -> Tensor: + """Run SDPA on a packed batch of variable-length sequences.""" + output_jagged = F.scaled_dot_product_attention( + jagged_batch.as_nested(query), + jagged_batch.as_nested(key), + jagged_batch.as_nested(value), + ) + return output_jagged.transpose(1, 2).values() + + +class Attention(nn.Module): + def __init__( + self, + dim: int, + num_heads: int, + qkv_bias: bool = False, + proj_bias: bool = False, + ) -> None: + super().__init__() + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.qkv = nn.Linear(dim, 3 * dim, bias=qkv_bias) + self.proj = nn.Linear(dim, dim, bias=proj_bias) + + def extra_repr(self): + return f"num_heads={self.num_heads}" + + def forward( + self, + x: Float[Tensor, "L D"], + jagged_batch: JaggedBatch, + ) -> Float[Tensor, "L D"]: + L, D = x.shape + h, dh = self.num_heads, self.head_dim + + qkv = self.qkv(x).reshape(L, 3, h, dh) + q, k, v = qkv.unbind(1) + + x = jagged_scaled_dot_product_attention( + q, + k, + v, + jagged_batch=jagged_batch, + ) + x = x.reshape(L, D) + x = self.proj(x) + return x + + +class Mlp(nn.Module): + def __init__( + self, + dim: int, + mlp_ratio: int | float = 4, + bias: bool = False, + ) -> None: + super().__init__() + hidden_features = int(dim * mlp_ratio) + self.fc1 = nn.Linear(dim, hidden_features, bias=bias) + self.act = nn.GELU() + self.fc2 = nn.Linear(hidden_features, dim, bias=bias) + + def forward(self, x: Float[Tensor, "... D"]) -> Float[Tensor, "... D"]: + x = self.fc1(x) + x = self.act(x) + x = self.fc2(x) + return x + + +# timm default eps=1e-6 +LayerNorm = partial(nn.LayerNorm, eps=1e-6) + + +class Block(nn.Module): + def __init__( + self, + dim: int, + num_heads: int, + qkv_bias: bool = False, + proj_bias: bool = False, + mlp_ratio: int | float = 4, + drop_path: float = 0.0, + norm_layer: Layer = LayerNorm, + ) -> None: + super().__init__() + self.norm1 = norm_layer(dim) + self.attn = Attention( + dim=dim, + num_heads=num_heads, + qkv_bias=qkv_bias, + proj_bias=proj_bias, + ) + self.drop_path1 = DropPath(drop_path) if drop_path > 0 else nn.Identity() + + self.norm2 = norm_layer(dim) + self.mlp = Mlp( + dim=dim, + mlp_ratio=mlp_ratio, + bias=proj_bias, + ) + self.drop_path2 = DropPath(drop_path) if drop_path > 0 else nn.Identity() + + def forward( + self, + x: Float[Tensor, "L D"], + jagged_batch: JaggedBatch, + ) -> Float[Tensor, "L D"]: + x = x + self.drop_path1( + self.attn( + self.norm1(x), + jagged_batch=jagged_batch, + ) + ) + x = x + self.drop_path2(self.mlp(self.norm2(x))) + return x + + +# Patching and position embedding modules + + +class Patchify3D(nn.Module): + def __init__( + self, + img_size: int | tuple[int, int, int], + patch_size: int | tuple[int, int, int], + in_chans: int = 3, + ) -> None: + super().__init__() + self.img_size = to_3tuple(img_size) + self.patch_size = to_3tuple(patch_size) + self.in_chans = in_chans + + T, H, W = self.img_size + p_t, p_h, p_w = self.patch_size + if T % p_t or H % p_h or W % p_w: + raise ValueError( + f"img_size {self.img_size} must be divisible by patch_size {self.patch_size}" + ) + self.grid_size = (T // p_t, H // p_h, W // p_w) + self.num_patches = math.prod(self.grid_size) + self.patch_dim = in_chans * math.prod(self.patch_size) + + def forward(self, x: Float[Tensor, "B C T H W"]) -> Float[Tensor, "B N P"]: + x = patchify3d(x, self.patch_size) + return x + + def unpatchify(self, x: Float[Tensor, "B N P"]) -> Float[Tensor, "B C T H W"]: + x = unpatchify3d(x, patch_size=self.patch_size, img_size=self.img_size) + return x + + def extra_repr(self): + return f"{self.img_size}, {self.patch_size}, in_chans={self.in_chans}" + + +def patchify3d(x: Tensor, patch_size: tuple[int, int, int]) -> Tensor: + p_t, p_h, p_w = to_3tuple(patch_size) + B, C, T, H, W = x.shape + x = rearrange(x, "b c (t u) (h p) (w q) -> b (t h w) (c u p q)", u=p_t, p=p_h, q=p_w) + return x + + +def unpatchify3d( + x: Tensor, + patch_size: tuple[int, int, int], + img_size: tuple[int, int, int], +) -> Tensor: + B, N, P = x.shape + p_t, p_h, p_w = to_3tuple(patch_size) + T, H, W = to_3tuple(img_size) + x = rearrange( + x, + "b (t h w) (c u p q) -> b c (t u) (h p) (w q)", + t=T // p_t, + h=H // p_h, + w=W // p_w, + u=p_t, + p=p_h, + q=p_w, + ) + return x + + +class AbsolutePosEmbed(nn.Module): + def __init__(self, embed_dim: int, grid_size: tuple[int, ...]) -> None: + super().__init__() + self.embed_dim = embed_dim + self.grid_size = grid_size + self.num_patches = math.prod(grid_size) + + self.weight = nn.Parameter(torch.empty(self.num_patches, embed_dim)) + self.reset_parameters() + + def reset_parameters(self): + nn.init.trunc_normal_(self.weight, std=0.02) + + def forward( + self, + x: Float[Tensor, "B L D"], + pos_ids: Int[Tensor, "B L"] | None = None, + ) -> Float[Tensor, "B L D"]: + x = apply_pos_embed(x, self.weight, pos_ids=pos_ids) + return x + + def extra_repr(self): + return f"{self.embed_dim}, {self.grid_size}" + + +class SeparablePosEmbed(nn.Module): + def __init__(self, embed_dim: int, grid_size: tuple[int, ...]) -> None: + super().__init__() + self.embed_dim = embed_dim + self.grid_size = grid_size + self.num_patches = math.prod(grid_size) + + N_t, *grid_size_spatial = grid_size + N_s = math.prod(grid_size_spatial) + self.weight_spatial = nn.Parameter(torch.empty(1, N_s, embed_dim)) + self.weight_temporal = nn.Parameter(torch.empty(N_t, 1, embed_dim)) + self.reset_parameters() + + def reset_parameters(self): + nn.init.trunc_normal_(self.weight_spatial, std=0.02) + nn.init.trunc_normal_(self.weight_temporal, std=0.02) + + def forward( + self, + x: Float[Tensor, "B L D"], + pos_ids: Int[Tensor, "B L"] | None = None, + ) -> Float[Tensor, "B L D"]: + B, N, D = x.shape + weight = (self.weight_temporal + self.weight_spatial).flatten(0, 1) # [N, D] + x = apply_pos_embed(x, weight, pos_ids=pos_ids) + return x + + def extra_repr(self): + return f"{self.embed_dim}, {self.grid_size}" + + +class SinCosPosEmbed3D(nn.Module): + def __init__(self, embed_dim: int, grid_size: tuple[int, int, int]) -> None: + super().__init__() + self.embed_dim = embed_dim + self.grid_size = grid_size + self.num_patches = math.prod(grid_size) + + N_t, N_h, N_w = grid_size + weight = get_3d_sincos_pos_embed( + embed_dim=embed_dim, + grid_size=(N_h, N_w), + grid_depth=N_t, + uniform_power=True, + ) + self.weight = nn.Parameter(torch.from_numpy(weight).float(), requires_grad=False) + + def forward( + self, + x: Float[Tensor, "B L D"], + pos_ids: Int[Tensor, "B L"] | None = None, + ) -> Float[Tensor, "B L D"]: + x = apply_pos_embed(x, self.weight, pos_ids=pos_ids) + return x + + def extra_repr(self): + return f"{self.embed_dim}, {self.grid_size}" + + +# sincos pos embed utils from vjepa2, but fixed the confusing meshgrid indexing + + +def get_3d_sincos_pos_embed(embed_dim, grid_size, grid_depth, cls_token=False, uniform_power=False): + """ + grid_size: tuple of int of the grid height and width + grid_depth: int of the grid depth + returns: + pos_embed: [grid_depth*grid_height*grid_width, embed_dim] (w/o cls_token) + or [1+grid_depth*grid_height*grid_width, embed_dim] (w/ cls_token) + """ + grid_d = np.arange(grid_depth, dtype=float) + grid_h = np.arange(grid_size[0], dtype=float) + grid_w = np.arange(grid_size[1], dtype=float) + grid_d, grid_h, grid_w = np.meshgrid(grid_d, grid_h, grid_w, indexing="ij") + + if not uniform_power: + h_embed_dim = embed_dim // 4 + w_embed_dim = embed_dim // 4 + d_embed_dim = embed_dim // 2 + else: + h_embed_dim = w_embed_dim = d_embed_dim = int(np.ceil(embed_dim / 6) * 2) + + emb_h = get_1d_sincos_pos_embed_from_grid(h_embed_dim, grid_h) # (T*H*W, D1) + emb_w = get_1d_sincos_pos_embed_from_grid(w_embed_dim, grid_w) # (T*H*W, D2) + emb_d = get_1d_sincos_pos_embed_from_grid(d_embed_dim, grid_d) # (T*H*W, D3) + pos_embed = np.concatenate([emb_d, emb_h, emb_w], axis=1) + pos_embed = pos_embed[:, :embed_dim] + if cls_token: + pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0) + return pos_embed + + +def get_1d_sincos_pos_embed_from_grid(embed_dim, pos): + """ + embed_dim: output dimension for each position + pos: a list of positions to be encoded: size (M,) + returns: (M, D) + """ + assert embed_dim % 2 == 0 + omega = np.arange(embed_dim // 2, dtype=float) + omega /= embed_dim / 2.0 + omega = 1.0 / 10000**omega # (D/2,) + + pos = pos.reshape(-1) # (M,) + out = np.einsum("m,d->md", pos, omega) # (M, D/2), outer product + + emb_sin = np.sin(out) # (M, D/2) + emb_cos = np.cos(out) # (M, D/2) + + emb = np.concatenate([emb_sin, emb_cos], axis=1) # (M, D) + return emb + + +def apply_pos_embed( + x: Float[Tensor, "B L D"], + weight: Float[Tensor, "N D"], + pos_ids: Int[Tensor, "B L"] | None = None, +) -> Float[Tensor, "B L D"]: + B, L, D = x.shape + weight = weight.expand(B, -1, -1) + if pos_ids is not None: + weight = weight.gather(1, pos_ids.unsqueeze(-1).expand(-1, -1, D)) + x = x + weight + return x + + +# (masked) normalization used for MAE target normalization + + +class Normalize(nn.Module): + def __init__( + self, + grid_size: tuple[int, ...], + dim: int | tuple[int, ...] | None = -1, + eps: float = 1e-6, + ) -> None: + super().__init__() + self.grid_size = grid_size + self.dim = dim + self.eps = eps + + def forward(self, x: Tensor, mask: Tensor | None = None) -> tuple[Tensor, Tensor, Tensor]: + """ + Normalize input sequence along dim(s) after reshaping to grid. + Returns tuple of (x, mean, std). + """ + B, N, D = x.shape + x = x.reshape((B, *self.grid_size, D)) + if mask is not None: + mask = mask.reshape((B, *self.grid_size, D)) + x, mean, std = masked_normalize(x, mask, dim=self.dim, eps=self.eps) + else: + x, mean, std = normalize(x, dim=self.dim, eps=self.eps) + mean = mean.expand_as(x).reshape(B, N, D) + std = std.expand_as(x).reshape(B, N, D) + x = x.reshape(B, N, D) + return x, mean, std + + def extra_repr(self): + return f"{self.grid_size}, dim={self.dim}" + + +def masked_normalize( + x: Tensor, + mask: Tensor, + dim: int | tuple[int, ...] | None = -1, + eps: float = 1e-6, +) -> tuple[Tensor, Tensor, Tensor]: + num_obs = mask.sum(dim=dim, keepdim=True).clamp(min=1) + mean = (mask * x).sum(dim=dim, keepdim=True) / num_obs + var = (mask * (x - mean) ** 2).sum(dim=dim, keepdim=True) / num_obs + std = (var + eps) ** 0.5 + x = mask * (x - mean) / std + return x, mean, std + + +def normalize( + x: Tensor, + dim: int | tuple[int, ...] | None = -1, + eps: float = 1e-6, +) -> tuple[Tensor, Tensor, Tensor]: + mean = x.mean(dim=dim, keepdim=True) + var = torch.var(x, dim=dim, keepdim=True, unbiased=False) + std = (var + eps) ** 0.5 + x = (x - mean) / std + return x, mean, std diff --git a/fomo_tune_baseline/output/task1/build/smri_mae/utils.py b/fomo_tune_baseline/output/task1/build/smri_mae/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..0f6274f813429552c9ddfccfe4033ae677b52342 --- /dev/null +++ b/fomo_tune_baseline/output/task1/build/smri_mae/utils.py @@ -0,0 +1,581 @@ +# Copyright (c) Sophont, Inc +# This source code is licensed under the Apache License, Version 2.0 +# +# References: +# deit: https://github.com/facebookresearch/deit/blob/main/utils.py +# beit3: https://github.com/microsoft/unilm/blob/master/beit3/utils.py +# capi: https://github.com/facebookresearch/capi/blob/main/utils.py +# dinov2: https://github.com/facebookresearch/dinov2/blob/main/dinov2/utils/param_groups.py +# timm: https://github.com/huggingface/pytorch-image-models/blob/main/timm/utils/cuda.py +# dino: https://github.com/facebookresearch/dino/blob/main/utils.py + +import datetime +import inspect +import math +import os +import random +import subprocess +import time +from collections import defaultdict, deque +from omegaconf import OmegaConf +from pathlib import Path + +import numpy as np +import torch +import torch.distributed as dist +import torch.nn as nn +from torch import Tensor +from torch.amp import GradScaler +from torch.optim import Optimizer + + +# these very useful utils copied from deit with only minor changes +# thanks to the original authors, wherever you are + + +def configure_flash_sdpa() -> None: + """Use Flash Attention exclusively for CUDA SDPA.""" + torch.backends.cuda.enable_flash_sdp(True) + torch.backends.cuda.enable_mem_efficient_sdp(False) + torch.backends.cuda.enable_math_sdp(False) + torch.backends.cuda.enable_cudnn_sdp(False) + print("SDPA backend: flash") + + +class SmoothedValue: + """Track a series of values and provide access to smoothed values over a + window or the global series average. + """ + + def __init__(self, window_size=20, fmt=None): + if fmt is None: + fmt = "{median:.4f} ({global_avg:.4f})" + self.deque = deque(maxlen=window_size) + self.total = 0.0 + self.count = 0 + self.fmt = fmt + + def update(self, value, n=1): + value = float(value) + if math.isfinite(value): + self.deque.append(value) + self.count += n + self.total += value * n + + def synchronize_between_processes(self): + """ + Warning: does not synchronize the deque! + """ + if not is_dist_avail_and_initialized(): + return + t = torch.tensor([self.count, self.total], dtype=torch.float64, device="cuda") + dist.barrier() + dist.all_reduce(t) + t = t.tolist() + self.count = int(t[0]) + self.total = t[1] + + @property + def median(self): + if not self.count: + return float("nan") + d = torch.tensor(list(self.deque)) + return d.median().item() + + @property + def avg(self): + if not self.count: + return float("nan") + d = torch.tensor(list(self.deque), dtype=torch.float32) + return d.mean().item() + + @property + def global_avg(self): + if not self.count: + return float("nan") + return self.total / self.count + + @property + def max(self): + if not self.count: + return float("nan") + return max(self.deque) + + @property + def value(self): + if not self.count: + return float("nan") + return self.deque[-1] + + def __str__(self): + return self.fmt.format( + median=self.median, + avg=self.avg, + global_avg=self.global_avg, + max=self.max, + value=self.value, + ) + + +class MetricLogger: + def __init__(self, delimiter="\t"): + self.meters = defaultdict(SmoothedValue) + self.delimiter = delimiter + + def update(self, **kwargs): + for k, v in kwargs.items(): + if v is None: + continue + if isinstance(v, (torch.Tensor, np.generic)): + v = v.item() + assert isinstance(v, (float, int)) + self.meters[k].update(v) + + def __getattr__(self, attr): + if attr in self.meters: + return self.meters[attr] + if attr in self.__dict__: + return self.__dict__[attr] + raise AttributeError("'{}' object has no attribute '{}'".format(type(self).__name__, attr)) + + def __str__(self): + loss_str = [] + for name, meter in self.meters.items(): + loss_str.append("{}: {}".format(name, str(meter))) + return self.delimiter.join(loss_str) + + def synchronize_between_processes(self): + for meter in self.meters.values(): + meter.synchronize_between_processes() + + def add_meter(self, name, meter): + self.meters[name] = meter + + def log_every(self, iterable, print_freq, header=None, total_steps=None): + i = 0 + total_steps = total_steps or len(iterable) + if not header: + header = "" + start_time = time.time() + end = time.time() + iter_time = SmoothedValue(fmt="{avg:.4f}") + data_time = SmoothedValue(fmt="{avg:.4f}") + space_fmt = ":" + str(len(str(total_steps))) + "d" + log_msg = [ + header, + "[{0" + space_fmt + "}/{1}]", + "eta: {eta}", + "{meters}", + "time: {time}", + "data: {data}", + ] + if torch.cuda.is_available(): + log_msg.append("max mem: {memory:.0f}") + log_msg = self.delimiter.join(log_msg) + MB = 1024.0 * 1024.0 + for obj in iterable: + if i >= total_steps: + break + data_time.update(time.time() - end) + yield obj + iter_time.update(time.time() - end) + if i % print_freq == 0 or i == total_steps - 1: + eta_seconds = iter_time.global_avg * (total_steps - i) + eta_string = str(datetime.timedelta(seconds=int(eta_seconds))) + if torch.cuda.is_available(): + print( + log_msg.format( + i, + total_steps, + eta=eta_string, + meters=str(self), + time=str(iter_time), + data=str(data_time), + memory=torch.cuda.max_memory_allocated() / MB, + ) + ) + + else: + print( + log_msg.format( + i, + total_steps, + eta=eta_string, + meters=str(self), + time=str(iter_time), + data=str(data_time), + ) + ) + i += 1 + end = time.time() + total_time = time.time() - start_time + total_time_str = str(datetime.timedelta(seconds=int(total_time))) + print( + "{} Total time: {} ({:.4f} s / it)".format( + header, total_time_str, total_time / total_steps + ) + ) + + +def setup_for_distributed(log_path=None): + """ + This function disables printing when not in master process + """ + import builtins as __builtin__ + + builtin_print = __builtin__.print + + is_master = is_main_process() + + def print(*args, **kwargs): + force = kwargs.pop("force", False) + if is_master or force: + builtin_print(*args, **kwargs) + # tee to log file + if log_path and "file" not in kwargs: + with open(log_path, "a") as f: + builtin_print(*args, file=f, **kwargs) + + __builtin__.print = print + + +def is_dist_avail_and_initialized(): + if not dist.is_available(): + return False + if not dist.is_initialized(): + return False + return True + + +def get_world_size(): + if not is_dist_avail_and_initialized(): + return 1 + return dist.get_world_size() + + +def get_rank(): + if not is_dist_avail_and_initialized(): + return 0 + return dist.get_rank() + + +def is_main_process(): + return get_rank() == 0 + + +def save_on_master(obj, path, *args, **kwargs): + if is_main_process(): + path = Path(path) + tmp_path = path.with_name(f".{path.name}.tmp-{os.getpid()}") + try: + torch.save(obj, tmp_path, *args, **kwargs) + os.replace(tmp_path, path) + except Exception: + tmp_path.unlink(missing_ok=True) + raise + + +def init_distributed_mode(args): + # removed slurm block, can add if we use slurm + if "RANK" in os.environ and "WORLD_SIZE" in os.environ: + args.rank = int(os.environ["RANK"]) + args.world_size = int(os.environ["WORLD_SIZE"]) + args.gpu = int(os.environ["LOCAL_RANK"]) + else: + args.distributed = False + return + + args.distributed = True + + torch.cuda.set_device(args.gpu) + args.dist_backend = "nccl" + print(f"| distributed init (rank {args.rank})") + torch.distributed.init_process_group( + backend=args.dist_backend, + world_size=args.world_size, + rank=args.rank, + device_id=args.gpu, + ) + torch.distributed.barrier() + + +# checkpoint saving utils adapted from beit3 + + +def capture_rng_state() -> dict: + state = {"torch": torch.get_rng_state()} + if torch.cuda.is_available(): + state["cuda"] = torch.cuda.get_rng_state() + return state + + +def restore_rng_state(state: dict) -> None: + torch.set_rng_state(state["torch"]) + if "cuda" in state and torch.cuda.is_available(): + torch.cuda.set_rng_state(state["cuda"]) + + +def _all_rank_rng_states() -> list[dict]: + local_state = capture_rng_state() + if not is_dist_avail_and_initialized(): + return [local_state] + states = [None] * get_world_size() + dist.all_gather_object(states, local_state) + return states + + +def save_model(args, epoch, model_without_ddp, optimizer, loss_scaler): + output_dir = Path(args.output_dir) + checkpoint_path = output_dir / f"checkpoint-{epoch:05d}.pth" + last_checkpoint_path = output_dir / "checkpoint-last.pth" + if epoch % args.checkpoint_period != 0 and epoch != args.epochs - 1: + return + + to_save = { + "model": model_without_ddp.state_dict(), + "optimizer": optimizer.state_dict(), + "epoch": epoch, + "scaler": None if loss_scaler is None else loss_scaler.state_dict(), + "args": OmegaConf.to_container(args), + "rng_states": _all_rank_rng_states(), + } + + print(f"saving checkpoint {last_checkpoint_path}") + save_on_master(to_save, last_checkpoint_path) + print(f"saving checkpoint {checkpoint_path}") + save_on_master(to_save, checkpoint_path) + + if args.max_checkpoints and is_main_process(): + all_checkpoints = sorted(output_dir.glob("checkpoint-[0-9]*.pth")) + del_count = max(0, len(all_checkpoints) - args.max_checkpoints) + for checkpoint_path in all_checkpoints[:del_count]: + print(f"removing checkpoint {checkpoint_path}") + checkpoint_path.unlink() + + +def load_model(args, model_without_ddp, optimizer, loss_scaler): + auto_resume = getattr(args, "auto_resume", True) + output_dir = Path(args.output_dir) + + last_checkpoint_path = output_dir / "checkpoint-last.pth" + if auto_resume and last_checkpoint_path.exists(): + args.ckpt = str(last_checkpoint_path) + args.resume = True + + if args.ckpt: + ckpt = torch.load(args.ckpt, map_location="cpu", weights_only=True) + model_without_ddp.load_state_dict(ckpt["model"]) + print(f"loaded model state from checkpoint {args.ckpt}") + + if args.resume: + optimizer.load_state_dict(ckpt["optimizer"]) + if loss_scaler is not None: + loss_scaler.load_state_dict(ckpt["scaler"]) + args.start_epoch = ckpt["epoch"] + 1 + rng_states = ckpt.get("rng_states") + if rng_states is not None: + if len(rng_states) != get_world_size(): + raise ValueError( + "checkpoint RNG state world size does not match current world size: " + f"{len(rng_states)} != {get_world_size()}" + ) + restore_rng_state(rng_states[get_rank()]) + print(f"restored RNG state for rank {get_rank()}") + print(f"loaded optimizer state, resuming training from {args.start_epoch}") + + +# optimization utils + + +# from capi +class WarmupThenCosine: + def __init__( + self, + base_value: float, + final_value: float, + total_iters: int, + warmup_iters: int = 0, + start_warmup_value: float = 0.0, + freeze_iters: int = 0, + truncate_cos: float = 1.0, + ): + super().__init__() + self.final_value = final_value + self.total_iters = total_iters + + freeze_schedule = np.zeros(freeze_iters) + + warmup_schedule = np.linspace(start_warmup_value, base_value, warmup_iters) + + iters = np.arange(total_iters - warmup_iters - freeze_iters) + schedule = final_value + 0.5 * (base_value - final_value) * ( + 1 + np.cos(np.pi * truncate_cos * iters / len(iters)) + ) + self.schedule = np.concatenate((freeze_schedule, warmup_schedule, schedule)) + assert len(self.schedule) == self.total_iters + + def __getitem__(self, it: int) -> float: + if it >= self.total_iters: + return self.final_value + # cast to float or else it can corrupt the checkpoint + return float(self.schedule[it]) + + +# adapted from timm backward logic +# https://github.com/huggingface/pytorch-image-models/blob/main/timm/utils/cuda.py +def backward_step( + loss: Tensor, + optimizer: Optimizer, + scaler: GradScaler = None, + need_update: bool = True, + max_norm: float | None = None, +) -> Tensor | None: + if scaler is not None: + scaler.scale(loss).backward() + else: + loss.backward() + + if need_update: + if scaler is not None: + scaler.unscale_(optimizer) + + total_norm = clip_grad(optimizer, max_norm) + + if scaler is not None: + scaler.step(optimizer) + scaler.update() + else: + optimizer.step() + optimizer.zero_grad() + else: + total_norm = None + return total_norm + + +def clip_grad(optimizer: Optimizer, max_norm: float | None = None) -> Tensor: + params = [p for group in optimizer.param_groups for p in group["params"]] + if max_norm: + total_norm = nn.utils.clip_grad_norm_(params, max_norm, error_if_nonfinite=False) + else: + grads = [p.grad for p in params if p.grad is not None] + total_norm = nn.utils.get_total_norm(grads, error_if_nonfinite=False) + torch._assert_async(torch.isfinite(total_norm), "non-finite gradient norm") + return total_norm + + +# from dinov2 with some minor changes +def get_param_groups(model, patch_embed_lr_mult=1.0): + # no lr decay, we could add this later if needed + all_params = [] + + for name, param in model.named_parameters(): + if not param.requires_grad: + continue + d = {"param": param, "lr_multiplier": 1.0, "wd_multiplier": 1.0, "name": name} + + if name.endswith(".bias") or "norm" in name or "gamma" in name: + d["wd_multiplier"] = 0.0 + + if "patch_embed" in name: + d["lr_multiplier"] = d["lr_multiplier"] * patch_embed_lr_mult + + all_params.append(d) + + param_groups = _fuse_param_groups(all_params) + return param_groups + + +def _fuse_param_groups(all_param_groups): + fused_param_groups = defaultdict(lambda: {"params": []}) + for d in all_param_groups: + keys = sorted(set(d.keys()) - {"param", "name"}) + identifier = "_".join(f"{k}{d[k]}" for k in keys) + for k in keys: + fused_param_groups[identifier][k] = d[k] + fused_param_groups[identifier]["params"].append(d["param"]) + + param_groups = list(fused_param_groups.values()) + return param_groups + + +def update_lr(param_groups, lr: float): + for group in param_groups: + group["lr"] = lr * group["lr_multiplier"] + + +def update_wd(param_groups, weight_decay: float | None = None): + for group in param_groups: + group["weight_decay"] = weight_decay * group["wd_multiplier"] + + +# moving data to cuda utils copied from capi +# added device argument + + +def send_data(x, device=None, dtype_map=None): + if device is None: + device = torch.device("cuda") + else: + device = torch.device(device) + + if isinstance(x, torch.Tensor): + dtype = dtype_map.get(x.dtype) if dtype_map else None + return x.to(device=device, dtype=dtype, non_blocking=True) + if isinstance(x, dict): + return {k: send_data(v, device=device, dtype_map=dtype_map) for k, v in x.items()} + if isinstance(x, list): + return [send_data(v, device=device, dtype_map=dtype_map) for v in x] + return x + + +def pre_send_to_cuda_wrapper(generator, device=None, dtype_map=None): + """From apex""" + data = None + stream = torch.cuda.Stream(device) + for next_data in generator: + with torch.cuda.stream(stream): + next_data = send_data(next_data, device=device, dtype_map=dtype_map) + if data is not None: + yield data + torch.cuda.current_stream(device).wait_stream(stream) + data = next_data + if data is not None: + yield data + + +# other misc utils + + +# from dino +def get_sha(): + cwd = os.path.dirname(os.path.abspath(__file__)) + + def _run(command): + return subprocess.check_output(command, cwd=cwd).decode("ascii").strip() + + sha = "N/A" + diff = "clean" + branch = "N/A" + try: + sha = _run(["git", "rev-parse", "HEAD"]) + diff = _run(["git", "diff-index", "HEAD"]) + diff = "has uncommitted changes" if diff else "clean" + branch = _run(["git", "rev-parse", "--abbrev-ref", "HEAD"]) + except Exception: + pass + message = f"sha: {sha}, status: {diff}, branch: {branch}" + return message + + +# from timm +def random_seed(seed=42, rank=0): + torch.manual_seed(seed + rank) + np.random.seed(seed + rank) + random.seed(seed + rank) + + +# mine :) +def filter_kwargs(func, kwargs): + sigature = inspect.signature(func) + kwargs = {k: v for k, v in kwargs.items() if k in sigature.parameters} + return kwargs diff --git a/fomo_tune_baseline/output/task1/build/smri_mae/visualization.py b/fomo_tune_baseline/output/task1/build/smri_mae/visualization.py new file mode 100644 index 0000000000000000000000000000000000000000..f721fe06e93dc1a2dfb4a679069029bdf55afc6b --- /dev/null +++ b/fomo_tune_baseline/output/task1/build/smri_mae/visualization.py @@ -0,0 +1,400 @@ +from collections.abc import Mapping +from io import BytesIO + +import torch + +from matplotlib import patches +from matplotlib import pyplot as plt +from PIL import Image +from torch import Tensor + +VIEW_NAMES = { + "sagittal": "Sagittal", + "saggital": "Sagittal", + "axial": "Axial", + "coronal": "Coronal", +} + + +def fig2pil(fig) -> Image.Image: + buffer = BytesIO() + fig.savefig(buffer, format="png", dpi=fig.dpi, facecolor=fig.get_facecolor()) + buffer.seek(0) + image = Image.open(buffer).convert("RGB") + buffer.close() + return image + + +def raw_stats_from_batch(batch: dict) -> tuple[Tensor | None, Tensor | None]: + metas = batch.get("meta") + if not metas: + return None, None + + means = [] + stds = [] + for meta in metas: + try: + mean = meta["raw_mean"] + std = meta["raw_std"] + except (KeyError, TypeError): + return None, None + if mean in ("", None) or std in ("", None): + return None, None + means.append(float(mean)) + stds.append(float(std)) + return torch.tensor(means), torch.tensor(stds) + + +def plot_mask_pred( + target: Tensor, + pred: Tensor, + pred_mask: Tensor | None = None, + img_mask: Tensor | None = None, + sample_idx: int = 0, + channel_idx: int = 0, + slice_idx: int | Mapping[str, int] | None = None, + patch_size: int | tuple[int, int, int] = 16, + views: tuple[str, ...] = ("sagittal", "axial", "coronal"), + cmap: str = "gray", + figsize: tuple[float, float] | None = None, + mask_style: str = "blank", + raw_mean: float | Tensor | None = None, + raw_std: float | Tensor | None = None, +): + target_vol = _select_volume(target, sample_idx=sample_idx, channel_idx=channel_idx) + pred_vol = _select_volume(pred, sample_idx=sample_idx, channel_idx=channel_idx) + if raw_mean is not None and raw_std is not None: + raw_mean = _select_scalar(raw_mean, sample_idx=sample_idx) + raw_std = _select_scalar(raw_std, sample_idx=sample_idx) + target_vol = target_vol * raw_std + raw_mean + pred_vol = pred_vol * raw_std + raw_mean + pred_mask_vol = ( + torch.zeros_like(target_vol) + if pred_mask is None + else _select_volume(pred_mask, sample_idx=sample_idx, channel_idx=channel_idx) > 0 + ) + img_mask_vol = None + if img_mask is not None: + img_mask_vol = _select_volume(img_mask, sample_idx=sample_idx, channel_idx=channel_idx) > 0 + + composite_vol = _prediction_composite(target_vol, pred_vol, pred_mask_vol) + vmin, vmax = _intensity_limits(target_vol, img_mask_vol) + + patch_size = _as_3tuple(patch_size) + view_items = [] + for view in views: + view_key = view.lower() + if view_key not in VIEW_NAMES: + raise ValueError(f"unknown MRI view {view!r}; expected one of {tuple(VIEW_NAMES)}") + target_slice = _extract_view_slice(target_vol, view_key, slice_idx) + composite_slice = _extract_view_slice(composite_vol, view_key, slice_idx) + mask_slice = _extract_view_slice(pred_mask_vol.float(), view_key, slice_idx) > 0 + img_mask_slice = None + if img_mask_vol is not None: + img_mask_slice = _extract_view_slice(img_mask_vol.float(), view_key, slice_idx) > 0 + view_items.append( + { + "key": view_key, + "title": VIEW_NAMES[view_key], + "target": _masked_input_display(target_slice, mask_slice, img_mask_slice, vmin), + "composite": _apply_display_mask(composite_slice, img_mask_slice, vmin), + "actual": _apply_display_mask(target_slice, img_mask_slice, vmin), + "mask": mask_slice, + "img_mask": img_mask_slice, + "patch_rc": _view_patch_size(view_key, patch_size), + } + ) + + _crop_view_items(view_items) + + fig, axes, layout = _make_figure_canvas(view_items, figsize=figsize) + for item, x in zip(view_items, layout["col_centers"]): + fig.text( + x, + layout["title_y"], + item["title"], + ha="center", + va="center", + color="#f8fafc", + fontsize=7, + ) + for label, y in zip(("Masked", "Pred", "Actual"), layout["row_centers"]): + fig.text(layout["label_x"], y, label, ha="right", va="center", color="#cbd5e1", fontsize=6) + + for item, top_ax, middle_ax, bottom_ax in zip(view_items, axes[0], axes[1], axes[2]): + top_ax.imshow( + item["target"], + cmap=cmap, + vmin=vmin, + vmax=vmax, + interpolation="nearest", + origin="upper", + ) + if mask_style == "boxes": + _draw_patch_boxes(top_ax, item["mask"], item["patch_rc"]) + elif mask_style != "blank": + raise ValueError("mask_style must be 'blank' or 'boxes'") + _style_axis(top_ax) + + middle_ax.imshow( + item["composite"], + cmap=cmap, + vmin=vmin, + vmax=vmax, + interpolation="nearest", + origin="upper", + ) + _style_axis(middle_ax) + + bottom_ax.imshow( + item["actual"], + cmap=cmap, + vmin=vmin, + vmax=vmax, + interpolation="nearest", + origin="upper", + ) + _style_axis(bottom_ax) + return fig + + +def _select_volume( + x: Tensor, + sample_idx: int = 0, + channel_idx: int = 0, +) -> Tensor: + x = x.detach().float().cpu() + if x.ndim == 5: + return x[sample_idx, channel_idx] + if x.ndim == 4: + return x[sample_idx] + if x.ndim == 3: + return x + raise ValueError(f"expected a 3D volume tensor, got shape {tuple(x.shape)}") + + +def _select_scalar(value: float | Tensor, sample_idx: int = 0) -> float: + if isinstance(value, Tensor): + value = value.detach().float().cpu() + if value.ndim > 0: + value = value.reshape(-1)[sample_idx] + return float(value) + return float(value) + + +def _prediction_composite(target: Tensor, pred: Tensor, pred_mask: Tensor) -> Tensor: + pred_mask = pred_mask.to(dtype=target.dtype) + return target * (1 - pred_mask) + pred * pred_mask + + +def _extract_view_slice( + volume: Tensor, + view: str, + slice_idx: int | Mapping[str, int] | None = None, +) -> Tensor: + if isinstance(slice_idx, Mapping): + slice_idx = slice_idx.get(view) + + if view in ("sagittal", "saggital"): + idx = _resolve_slice_idx(volume.shape[0], slice_idx) + return volume[idx, :, :].transpose(0, 1).flip(0) + if view == "axial": + idx = _resolve_slice_idx(volume.shape[2], slice_idx) + return volume[:, :, idx].transpose(0, 1).flip(0) + if view == "coronal": + idx = _resolve_slice_idx(volume.shape[1], slice_idx) + return volume[:, idx, :].transpose(0, 1).flip(0) + raise ValueError(f"unknown MRI view {view!r}") + + +def _resolve_slice_idx(size: int, slice_idx: int | None) -> int: + idx = size // 2 if slice_idx is None else int(slice_idx) + if idx < 0: + idx += size + if idx < 0 or idx >= size: + raise IndexError(f"slice index {idx} is out of bounds for axis with size {size}") + return idx + + +def _intensity_limits(volume: Tensor, mask: Tensor | None = None) -> tuple[float, float]: + values = volume[mask] if mask is not None and mask.any() else volume.flatten() + values = values[torch.isfinite(values)] + if values.numel() == 0: + return 0.0, 1.0 + if values.numel() < 32: + vmin = values.min() + vmax = values.max() + else: + vmin, vmax = torch.quantile(values, torch.tensor([0.005, 0.995])) + if torch.isclose(vmin, vmax): + delta = max(abs(float(vmin)) * 0.05, 1.0) + return float(vmin) - delta, float(vmax) + delta + return float(vmin), float(vmax) + + +def _apply_display_mask(image: Tensor, mask: Tensor | None, fill_value: float) -> Tensor: + if mask is None: + return image + return torch.where(mask, image, torch.full_like(image, fill_value)) + + +def _masked_input_display( + image: Tensor, + pred_mask: Tensor, + img_mask: Tensor | None, + fill_value: float, +) -> Tensor: + display = torch.where(pred_mask, torch.full_like(image, fill_value), image) + return _apply_display_mask(display, img_mask, fill_value) + + +def _crop_view_items(view_items: list[dict]) -> None: + for item in view_items: + mask = item["img_mask"] + if mask is None: + mask = item["actual"] != item["actual"].min() + row_slice, col_slice = _content_crop(mask, item["patch_rc"]) + for key in ("target", "composite", "actual", "mask"): + item[key] = item[key][row_slice, col_slice] + if item["img_mask"] is not None: + item["img_mask"] = item["img_mask"][row_slice, col_slice] + + +def _content_crop(mask: Tensor, patch_size: tuple[int, int]) -> tuple[slice, slice]: + mask = mask.detach().cpu() > 0 + if not mask.any(): + return slice(None), slice(None) + + rows, cols = mask.nonzero(as_tuple=True) + patch_h, patch_w = patch_size + height, width = mask.shape + row0 = max((int(rows.min()) // patch_h - 1) * patch_h, 0) + row1 = min((int(rows.max()) // patch_h + 2) * patch_h, height) + col0 = max((int(cols.min()) // patch_w - 1) * patch_w, 0) + col1 = min((int(cols.max()) // patch_w + 2) * patch_w, width) + return slice(row0, row1), slice(col0, col1) + + +def _as_3tuple(value: int | tuple[int, int, int]) -> tuple[int, int, int]: + if isinstance(value, int): + return (value, value, value) + if len(value) != 3: + raise ValueError(f"expected a 3-tuple patch size, got {value!r}") + return tuple(int(v) for v in value) + + +def _view_patch_size(view: str, patch_size: tuple[int, int, int]) -> tuple[int, int]: + p_x, p_y, p_z = patch_size + if view in ("sagittal", "saggital"): + return p_z, p_y + if view == "axial": + return p_y, p_x + if view == "coronal": + return p_z, p_x + raise ValueError(f"unknown MRI view {view!r}") + + +def _make_figure_canvas( + view_items: list[dict[str, Tensor | str | tuple[int, int]]], + figsize: tuple[float, float] | None = None, +): + dpi = 160 + left = 58 + right = 6 + top = 18 + bottom = 8 + row_gap = 14 + col_gap = 8 + widths = [int(item["target"].shape[1]) for item in view_items] + heights = [int(item["target"].shape[0]) for item in view_items] + row_h = max(heights) + num_rows = 3 + fig_w = left + right + sum(widths) + col_gap * (len(widths) - 1) + fig_h = top + bottom + row_h * num_rows + row_gap * (num_rows - 1) + + scale = 1.35 + if figsize is not None: + requested_w = figsize[0] * dpi + requested_h = figsize[1] * dpi + scale = max(requested_w / fig_w, requested_h / fig_h) + figsize = (fig_w * scale / dpi, fig_h * scale / dpi) + fig = plt.figure(figsize=figsize, dpi=dpi, facecolor="#0b0f14") + + axes = [[] for _ in range(num_rows)] + col_centers = [] + x = left + for width, height in zip(widths, heights): + col_centers.append((x + width / 2) / fig_w) + ys = [ + bottom + (num_rows - row - 1) * (row_h + row_gap) + (row_h - height) / 2 + for row in range(num_rows) + ] + for row, y in enumerate(ys): + axes[row].append( + fig.add_axes( + [ + x / fig_w, + y / fig_h, + width / fig_w, + height / fig_h, + ], + facecolor="black", + ) + ) + x += width + col_gap + row_centers = [ + (bottom + (num_rows - row - 1) * (row_h + row_gap) + row_h / 2) / fig_h + for row in range(num_rows) + ] + layout = { + "col_centers": col_centers, + "row_centers": row_centers, + "label_x": (left - 8) / fig_w, + "title_y": (fig_h - top / 2) / fig_h, + } + + return fig, axes, layout + + +def _style_axis(ax) -> None: + ax.set_xticks([]) + ax.set_yticks([]) + for spine in ax.spines.values(): + spine.set_visible(False) + + +def _draw_patch_boxes( + ax, + mask: Tensor, + patch_size: tuple[int, int], + color: str = "#facc15", +) -> None: + for col, row, width, height in _patch_rectangles(mask, patch_size): + ax.add_patch( + patches.Rectangle( + (col - 0.5, row - 0.5), + width, + height, + fill=False, + edgecolor=color, + linewidth=0.75, + alpha=0.95, + ) + ) + + +def _patch_rectangles( + mask: Tensor, + patch_size: tuple[int, int], +) -> list[tuple[int, int, int, int]]: + mask = mask.detach().cpu() > 0 + patch_h, patch_w = patch_size + height, width = mask.shape + rectangles = [] + for row in range(0, height, patch_h): + box_h = min(patch_h, height - row) + for col in range(0, width, patch_w): + box_w = min(patch_w, width - col) + if mask[row : row + box_h, col : col + box_w].any(): + rectangles.append((col, row, box_w, box_h)) + return rectangles diff --git a/fomo_tune_baseline/output/task1/config.yaml b/fomo_tune_baseline/output/task1/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..5dca875152bb8d3adf4bf758bd4681bf19e1d83e --- /dev/null +++ b/fomo_tune_baseline/output/task1/config.yaml @@ -0,0 +1,8 @@ +task: task1 +ckpt_path: /data/mihir-stuff/smri-pretrained/pretrain_full_90_10_h100/checkpoint-last.pth +modalities: +- dwi_b1000 +output_root: experiments/fomo_tune_baseline/output +name: task1 +device: cuda +seed: 4466 diff --git a/fomo_tune_baseline/output/task1/log.txt b/fomo_tune_baseline/output/task1/log.txt new file mode 100644 index 0000000000000000000000000000000000000000..d8a6290875009d93d74e3fb74fa0d5157375504c --- /dev/null +++ b/fomo_tune_baseline/output/task1/log.txt @@ -0,0 +1,33 @@ +16:31:38 run task1 (git 1df2e5d) +16:31:38 config: +task: task1 +ckpt_path: /data/mihir-stuff/smri-pretrained/pretrain_full_90_10_h100/checkpoint-last.pth +modalities: +- dwi_b1000 +output_root: experiments/fomo_tune_baseline/output +name: task1 +device: cuda +seed: 4466 +16:31:40 dataset: 21 subjects, 13 positive +16:31:48 fold 1/21 sub-01 y=1 p=0.512 (4s) +16:31:49 fold 2/21 sub-02 y=1 p=0.514 (4s) +16:31:49 fold 3/21 sub-03 y=0 p=0.486 (5s) +16:31:49 fold 4/21 sub-04 y=1 p=0.499 (5s) +16:31:50 fold 5/21 sub-05 y=1 p=0.507 (5s) +16:31:50 fold 6/21 sub-06 y=1 p=0.493 (6s) +16:31:50 fold 7/21 sub-07 y=0 p=0.489 (6s) +16:31:51 fold 8/21 sub-08 y=0 p=0.490 (6s) +16:31:51 fold 9/21 sub-09 y=1 p=0.523 (7s) +16:31:52 fold 10/21 sub-10 y=1 p=0.525 (7s) +16:31:52 fold 11/21 sub-11 y=0 p=0.479 (7s) +16:31:52 fold 12/21 sub-12 y=0 p=0.474 (8s) +16:31:53 fold 13/21 sub-13 y=1 p=0.516 (8s) +16:31:53 fold 14/21 sub-14 y=1 p=0.536 (9s) +16:31:53 fold 15/21 sub-15 y=0 p=0.486 (9s) +16:31:54 fold 16/21 sub-16 y=0 p=0.482 (9s) +16:31:54 fold 17/21 sub-17 y=1 p=0.520 (10s) +16:31:54 fold 18/21 sub-18 y=1 p=0.518 (10s) +16:31:55 fold 19/21 sub-19 y=1 p=0.520 (10s) +16:31:55 fold 20/21 sub-20 y=1 p=0.519 (11s) +16:31:56 fold 21/21 sub-21 y=0 p=0.498 (11s) +16:31:57 result: auroc=0.9904 auroc_ci_low=0.9444 auroc_ci_high=1.0000 (11s) diff --git a/fomo_tune_baseline/output/task1/metrics.json b/fomo_tune_baseline/output/task1/metrics.json new file mode 100644 index 0000000000000000000000000000000000000000..bc2f4616637ebf45df2eba86b2bf79b47e50ed1f --- /dev/null +++ b/fomo_tune_baseline/output/task1/metrics.json @@ -0,0 +1 @@ +{"name": "task1", "auroc": 0.9903846153846154, "auroc_ci_low": 0.9444444444444445, "auroc_ci_high": 1.0, "run_time": 11.2} diff --git a/fomo_tune_baseline/output/task1/model/config.yaml b/fomo_tune_baseline/output/task1/model/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..5dca875152bb8d3adf4bf758bd4681bf19e1d83e --- /dev/null +++ b/fomo_tune_baseline/output/task1/model/config.yaml @@ -0,0 +1,8 @@ +task: task1 +ckpt_path: /data/mihir-stuff/smri-pretrained/pretrain_full_90_10_h100/checkpoint-last.pth +modalities: +- dwi_b1000 +output_root: experiments/fomo_tune_baseline/output +name: task1 +device: cuda +seed: 4466 diff --git a/fomo_tune_baseline/output/task1/model/head.joblib b/fomo_tune_baseline/output/task1/model/head.joblib new file mode 100644 index 0000000000000000000000000000000000000000..90ed37cfc3dfd4a13e1f754990808e15c9b1f7e0 --- /dev/null +++ b/fomo_tune_baseline/output/task1/model/head.joblib @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f4a4826f5ea9f3f74279b798caf6a5970a3c63b02b70c4d0828ee5325d29e0dd +size 445215 diff --git a/fomo_tune_baseline/output/task1/task1.sif b/fomo_tune_baseline/output/task1/task1.sif new file mode 100644 index 0000000000000000000000000000000000000000..a9bb31f019054cf3985b44a1de07730923f6b466 --- /dev/null +++ b/fomo_tune_baseline/output/task1/task1.sif @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:0d9c2a0d53bf894c9788fe0fd45a6bb0e5a5f7c4c7554fd4b1c488397d35fd93 +size 5329231872 diff --git a/fomo_tune_baseline/output/task3/build/Apptainer.def b/fomo_tune_baseline/output/task3/build/Apptainer.def new file mode 100644 index 0000000000000000000000000000000000000000..98a8bdfc1fefdbffdaed0dedbc7166a3da28b472 --- /dev/null +++ b/fomo_tune_baseline/output/task3/build/Apptainer.def @@ -0,0 +1,31 @@ +Bootstrap: docker +From: python:3.11-slim + +# NOT buildable where it sits: the %files paths below are relative to the build cwd, which is the +# staging dir `build.py` writes. Build it with `python -m fomo_tune.build `, not by +# pointing apptainer at this file. +# +# Versions are pinned to the training environment: numpy, scikit-learn and joblib because they +# unpickle `head.joblib`, torch because that is what the checkpoint was written by. + +%files + fomo_tune /app/fomo_tune + smri_mae /app/smri_mae + model /app/model + predict.py /app/predict.py + +%post + pip install --no-cache-dir \ + torch==2.8.0 \ + numpy==2.4.6 \ + nibabel==5.4.2 \ + einops==0.8.2 \ + jaxtyping==0.3.10 \ + timm==1.0.27 \ + huggingface-hub==0.36.2 \ + scikit-learn==1.8.0 \ + joblib==1.5.3 \ + omegaconf==2.3.0 + +%runscript + exec python /app/predict.py "$@" diff --git a/fomo_tune_baseline/output/task3/build/fomo_tune/README.md b/fomo_tune_baseline/output/task3/build/fomo_tune/README.md new file mode 100644 index 0000000000000000000000000000000000000000..aa7a6cdadfe314f9d1b0fc2d48ef67b85032d3c5 --- /dev/null +++ b/fomo_tune_baseline/output/task3/build/fomo_tune/README.md @@ -0,0 +1,239 @@ +# fomo_tune + +The five FOMO26 challenge tasks, one script each, tuned independently. + +This is a spinoff of `nanobrain.eval`, which scored every backbone on every task through one fixed +probe. That was the right shape for a benchmark and the wrong shape for a competition: here we care +about one backbone (sMRI MAE) and five scores, and each task wants a different method. **Nothing +here imports `nanobrain.eval`, and it should stay that way** — this package may be shared with +people who won't get the eval suite. + +## Layout + +| File | | +|---|---| +| `datasets.py` | core, **frozen**. One `load_fomo_task()` per task, streaming the challenge zips into an HF dataset. Raw niftis, no resampling — the backbone transform does that. | +| `backbone.py` | core, **frozen**. `load_backbone(ckpt_path) -> (SmriMaeBackbone, SmriMaeTransform)`. Frozen sMRI MAE encoder; the transform canonicalizes to RAS, rescales to 1mm, fits to the pretraining shape, z-scores in a mean-threshold brain mask. | +| `utils.py` | core. `set_seed`, `git_sha`, `setup_logging`. | +| `main_task.py` | shell. One task, end to end. Task 1 is the worked example; copy it. | +| `build.py` + `Apptainer.def` | shell. Package a run dir into the challenge `.sif`. Shared by every task. | + +`datasets.py` and `backbone.py` are settled and their caches are warm. Treat them as read-only: +new work goes in `main_task.py`. If one of them genuinely needs to change, that is a +conversation first, because it invalidates every score already recorded. + +## The pattern + +`main_task1.py` is in three sections, and the split is the point of the whole design. + +**`Task1Method` — the part we tune.** Features, head, hyperparameters, anything that might move +the score. Its interface is: + +```python +method.fit(rows) # rows are dataset records: subject, label, images +method.predict(images) # -> the challenge's output for one subject +method.save(model_dir) # config.yaml + head.joblib +Task1Method.load(model_dir, **overrides) +``` + +**The protocol — fixed.** Pool out-of-fold predictions over all subjects, bootstrap subjects for +the CI. No repeats, no stratification; the bootstrap is the only variance estimate. Splitting is +per-task but fixed within a task — leave-one-out where n is tiny (task 1, n=21), **20-fold** where +it isn't (tasks 3 and 5), which is close enough to LOO without paying for 494 refits. Once a task's +scheme is set, hold it or scores stop being comparable across iterations. That is also why +`cross_validate` seeds its shuffle at 0 rather than from `cfg.seed`: the folds are part of the +protocol, so tuning the run's seed must not silently redraw them. + +**Two entrypoints.** `train` runs the protocol then fits a head on all subjects and saves it; +`predict` is the challenge CLI. Both go through `Method.predict`, which is why every fold +exercises the code the submission will run. + +That last point is the load-bearing one. `predict` is not a wrapper written at packaging time — it +is the same call cross-validation already made once per held-out subject. When you add a task, keep +that property. + +```bash +uv run python -m fomo_tune.main_task1 train modalities=[dwi_b1000,flair] name=task1_dwi_flair +uv run python -m fomo_tune.main_task1 predict --model-dir output/fomo_tune/task1_dwi/model \ + --adc adc.nii.gz --dwi dwi.nii.gz --flair flair.nii.gz --output prob.txt +``` + +`train` takes omegaconf dotlist overrides against the `Config` dataclass at the top of the file. +It writes `config.yaml`, `log.txt`, `metrics.json`, and `model/` into `{output_root}/{name}/`. + +## Status + +Tasks 1, 5 and 3 are drafted and verified. Task 1 is also packaged — its container passes the +challenge validator; 5 and 3 have not been built yet. **Tasks 2 and 4 are tabled** — both are +segmentation, both need `predict` to emit a nifti on the input grid, and neither is worth opening +until the classification and regression tasks are settled. + +All three on `vitl_fomo300`, one H100, wall being the cross-validation loop: + +| run | result | wall | +|---|---|---| +| `task1_dwi`, dwi_b1000, n=21, LOO | AUROC **0.990** [0.944, 1.000] | 25s | +| `task5_t1w`, t1w, n=48, 20-fold | AUROC **0.984** [0.953, 1.000] | 73s | +| `task3_t1w`, t1w, n=494, 20-fold | r **0.962** [0.956, 0.968], MAE **3.71y** [3.45, 3.97] | 260s | + +**Task 3's row is one fold-seed stale.** It was measured before `cross_validate` froze its shuffle +at 0, so it is a 20-fold run with `random_state=4466`. Task 1 (LOO) and task 5 are current. The +re-run is cheap — 260s on a GPU — it just has not been done. Expect a shift of the same order task +5 saw when its folds moved (0.948 → 0.984, i.e. inside the CI but not negligible). + +Task 1's earlier probe sweep got 0.954 [0.861, 1.000] on the same checkpoint +(`experiments/eval_global_0728`), so it roughly reproduces — the gap is LOO vs 5×5 stratified CV, +one interpolation instead of two, and a head selected on AUROC instead of balanced accuracy. + +Two checks worth repeating per task — `.claude/scratch/verify_task1.py` and +`.claude/scratch/verify_task35.py ` do both: +- features are **bit-identical** whether the nifti comes from the HF dataset wrapper or from + `nib.load` off disk, so CV numbers transfer to the container +- the `predict` CLI agrees with the in-process method + +## What changes per task + +Counts and modalities, read from the local zips: + +| Task | n | Inputs | Output | Split | Notes | +|---|---|---|---|---|---| +| 1 infarct | 21 | adc, dwi_b1000, flair (+t2s/swi) | probability | LOO | done | +| 5 polymicrogyria | 48 | t1w | probability | 20-fold | done | +| 3 brain age | 494 | t1w | age in years | 20-fold | done — RidgeCV head, scored by **Pearson r and MAE**, each with its own bootstrap CI | +| 2 meningioma | 23 | dwi_b1000, flair (+t2s/swi) | mask, input grid | — | tabled | +| 4 trigeminal | 40 | t2w | mask, labels 1=nerve 2=vessel | — | tabled | + +Tasks 5 and 3 diverge from task 1 only where that table says. `cross_validate` over a shuffled +`KFold` replaces `leave_one_out`; both take one modality, so `features` loses the +concat-over-modalities loop and `Config` loses `modalities`; the challenge CLI flag is `--t1` for +both, and it is `--t1` for task 3 too even though the file in the zip is `t1w.nii.gz`. + +Task 3 is the first regression, so its `score` loops over the two metrics rather than returning +one, and it drops task 1's guard against bootstrap resamples with fewer than two distinct labels. +The analogous degenerate case for regression is a resample with no spread in `y`, where Pearson r +is undefined rather than merely unstable — at n=494 it does not happen. + +When tasks 2 and 4 come back: `predict` must write a nifti on the input's grid, and the method +needs localized features rather than a pooled vector — `backbone.forward` returns `patch_coords` +in world mm for exactly that. Task 4's label order (1=nerve, 2=vessel) is still a guess and needs +confirming against the challenge data before per-class numbers mean anything. + +## Gotchas + +**Raw niftis are on disk** at `data/fomo_eval/Task_/preprocessed//ses-01/`, which is the +easy way to exercise `predict` on a real file rather than one written out of the dataset: + +```bash +uv run python -m fomo_tune.main_task1 predict \ + --model-dir output/fomo_tune/task1_dwi/model \ + --adc data/fomo_eval/Task_1/preprocessed/sub-20/ses-01/adc.nii.gz \ + --dwi data/fomo_eval/Task_1/preprocessed/sub-20/ses-01/dwi_b1000.nii.gz \ + --flair data/fomo_eval/Task_1/preprocessed/sub-20/ses-01/flair.nii.gz \ + --output /tmp/prob.txt +``` + +Task 5 breaks the naming: `Task_5/preprocessed/sub_01/ses_01/t1.nii.gz` — underscores throughout, +and `t1` not `t1w`. `datasets.py` already handles it; anything you write by hand won't. + +```bash +uv run python -m fomo_tune.main_task5 predict --model-dir output/fomo_tune/task5_t1w/model \ + --t1 data/fomo_eval/Task_5/preprocessed/sub_01/ses_01/t1.nii.gz --output /tmp/prob.txt +``` + +**Volumes are wildly anisotropic.** Task 1's DWI is 0.46×0.46×**5.6**mm, so the transform +upsamples z by 5.6× to reach 1mm iso. Nothing is wrong, but don't read the 1mm grid as real +resolution. + +**The backbone never saw skull or neck.** Pretraining used a SynthSeg brain mask; the transform +substitutes a mean-intensity threshold, which keeps both. Known fidelity gap — see +`.claude/memory/smri-mae-preprocessing-gap.md`. + +**Probabilities are not calibrated.** `LogisticRegressionCV` on ~20 samples × 1024 features shrinks +hard; task 1's out-of-fold probabilities all land in 0.48–0.52 with near-perfect ranking. Fine for +AUROC, which is what the challenge scores, but don't read them as probabilities. Task 5's do span +0–1, which is n=48 rather than n=21 and not evidence of calibration. + +**n is tiny, so the CI is the result.** Task 1's is ~0.06 wide at the top of the range. Most tuning +deltas you chase will be inside it. `.claude/NOTES.md` thread 1 has the longer argument. + +**GPUs need an allocation** — the login node has no driver. See the `gpu-session` skill. + +## Submission + +`build.py` packages a run dir into the `.sif` the challenge wants. One command, taking the run dir +the shipped head was saved into: + +```bash +uv run python -m fomo_tune.build output/fomo_tune/task1_dwi +``` + +It stages `/app`, then builds from there: + +``` +/app/predict.py # shim: calls fomo_tune.main_task predict +/app/model/config.yaml # from the run dir +/app/model/head.joblib # from the run dir +/app/model/backbone.pth # stripped checkpoint, --ckpt-path points here +``` + +**Both `build.py` and `Apptainer.def` are shared across tasks**, which they can be because nothing +in staging or in the dependency list is task-specific. The one thing that does vary is the module +the shim imports, and that comes from `task` in the run's saved config — so a run dir knows which +task it belongs to, and `build.py` never needs telling. + +`predict.py` is **generated at build time** rather than checked in. It is eight lines whose whole +meaning is the container layout staged around it, so there is nowhere outside a container to run +it. This does not weaken the point above about `predict` not being written at packaging time: the +logic still lives in `main_task.py`, exercised once per fold, and the shim only picks the +subcommand and two paths. + +**`Apptainer.def` is not buildable where it sits.** Its `%files` paths are relative to the build +cwd, which is the staging dir. Pointing `apptainer build` at it in the repo fails confusingly; go +through `build.py`. + +The run dir deliberately does *not* carry backbone weights — that checkpoint is 3.9G and would be +copied on every run. `--ckpt-path` overrides what `config.yaml` recorded, so the saved config stays +a faithful record of what trained rather than being rewritten at build time. + +**The staged checkpoint is stripped to `model` and `args`**, which is 3.9G → 1.3G because the rest +is optimizer state inference never reads. `load_backbone` needs no change for this, and `predict` +gives a bit-identical probability either way (0.524739 on `sub-20`, checked on GPU). + +**The base image is `python:3.11-slim`, not a CUDA image.** The PyPI torch wheel *is* the cu128 +build and vendors the whole CUDA userspace as `nvidia-*` packages, so all the container needs from +the host is the driver, which `--nv`/`--nvccli` binds in. That keeps the SIF at 5.3G (4.0G of +image, 1.3G of weights) against roughly double for `pytorch/pytorch` and far more for NGC. +Versions are pinned to the training environment +mostly so `head.joblib` unpickles against the numpy/sklearn that wrote it. + +Apptainer caches the bootstrap layers but **always re-runs `%post`**, so every build re-downloads +~3G of wheels. If that gets annoying, bake a deps-only base SIF and `Bootstrap: localimage` off it. + +### Validating + +`third_party/container-validator` is the challenge's own validator, test niftis included: + +```bash +python third_party/container-validator/container_validator/validate.py \ + --task task1 --sif output/fomo_tune/task1_dwi/task1.sif +``` + +It runs `python /app/predict.py --flair /input/… --adc … --dwi … --swi … --output /output/.txt` +inside an `apptainer instance` with `/input`, `/output` and `/tmp` bound — which is exactly the +shim's contract, so nothing in `predict.py` is guessing at the interface. + +One thing it does that is easy to miss: it takes GPU via `--nvccli` rather than `--nv`, and one of +its tests runs `nvidia-smi -L` **inside** the container. `python:3.11-slim` ships no `nvidia-smi`, +so that test passes only because `--nvccli` injects the host one — a CUDA base image would hide +that dependency rather than remove it. + +**The `task1_dwi` container passes all 20 validator tests**, and `predict` inside it returns +0.524739 on `sub-20`, identical to the same call outside the container. So the packaging is +verified end to end, not just built. + +**Run it on a compute node with apptainer, which as of 2026-08-11 means `n-6`** — `salloc +--nodelist=n-6`. The other nodes fail the validator's preflight. The login node has apptainer but +no driver, and +`predict` there dies inside `can_use_cudnn_attention` — the jagged-SDPA path reaches into CUDA even +when the tensors are on CPU, so a driver-less host fails at the forward pass rather than falling +back. That is the CPU gap worth remembering; it is not a container problem. diff --git a/fomo_tune_baseline/output/task3/build/fomo_tune/backbone.py b/fomo_tune_baseline/output/task3/build/fomo_tune/backbone.py new file mode 100644 index 0000000000000000000000000000000000000000..4a1275a32e0f9da0963043225e8b3aefcbe8edc6 --- /dev/null +++ b/fomo_tune_baseline/output/task3/build/fomo_tune/backbone.py @@ -0,0 +1,153 @@ +import inspect + +import nibabel as nib +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange +from torch import Tensor + +import smri_mae.model_mae as models_mae + + +class SmriMaeBackbone(nn.Module): + grid_coords: Tensor + + def __init__(self, encoder: models_mae.MaskedEncoder): + super().__init__() + self.encoder = encoder + self.img_size = self.encoder.patchify.img_size + + grid_size = self.encoder.patchify.grid_size + patch_size = np.array(self.encoder.patchify.patch_size) + grid_coords = rearrange(np.indices(grid_size), "c x y z -> (x y z) c") + grid_coords = grid_coords * patch_size + (patch_size - 1) / 2 + grid_coords = torch.as_tensor(grid_coords, dtype=torch.float32) + self.register_buffer("grid_coords", grid_coords) + + def forward(self, batch: dict[str, Tensor]) -> dict[str, Tensor]: + images = batch["image"] + mask = batch["mask"] + affine = batch["affine"] + + B, C, X, Y, Z = images.shape + assert (X, Y, Z) == self.img_size, f"expected {self.img_size}, got {(X, Y, Z)}" + + _, _, patch_embeds, _, patch_ids, token_mask = self.encoder(images, mask=mask) + + # [B, L, 3] world xyz coords of embeddings + patch_coords = self.grid_coords[patch_ids, :] + rot = affine[:, :3, :3] + trans = affine[:, :3, 3] + patch_coords = patch_coords @ rot.transpose(1, 2) + trans[:, None, :] + + return { + "patch_embeds": patch_embeds, + "patch_ids": patch_ids, + "token_mask": token_mask, + "patch_coords": patch_coords, + } + + +class SmriMaeTransform: + def __init__( + self, + img_size: tuple[int, int, int] = (208, 240, 208), + spacing: tuple[float, float, float] = (1.0, 1.0, 1.0), + ): + self.img_size = img_size + self.spacing = spacing + + def __call__(self, img: nib.Nifti1Image) -> dict[str, Tensor]: + # repack image to handle incomplete hf Nifti interface + img = nib.Nifti1Image(img.dataobj, img.affine, img.header) + img = nib.as_closest_canonical(img) + + data = torch.from_numpy(np.ascontiguousarray(img.get_fdata(dtype=np.float32))) + affine = np.asarray(img.affine) + + spacing = img.header.get_zooms() + if max(abs(s - s_) for s, s_ in zip(spacing, self.spacing)) > 0.05: + data, affine = rescale(data, affine, spacing, self.spacing) + + data, affine = fit_to_shape(data, affine, target_shape=self.img_size) + + # mean threshold, not the SynthSeg mask used in pretraining, so skull and neck stay in + mask = data > data.mean() + brain = data[mask] + mean = brain.mean() + # population std (correction=0) to match the pretraining normalization + std = brain.std(correction=0).clamp_min(1e-6) + data = torch.where(mask, (data - mean) / std, 0.0) + + return { + "image": data.unsqueeze(0), + "mask": mask.unsqueeze(0), + "affine": torch.as_tensor(affine, dtype=torch.float32), + } + + +def rescale( + x: torch.Tensor, + affine: np.ndarray, + spacing: tuple[float, ...], + target_spacing: tuple[float, ...] = (1.0, 1.0, 1.0), +) -> tuple[torch.Tensor, np.ndarray]: + scales = tuple([current / target for current, target in zip(spacing, target_spacing)]) + resampled = F.interpolate(x[None, None], scale_factor=scales, mode="trilinear").squeeze(0, 1) + + # align_corners=False reads output voxel j from input voxel (j + 0.5) / scale - 0.5 + scale = np.asarray(scales, dtype=float) + step = np.diag([*(1 / scale), 1.0]) + step[:3, 3] = 0.5 / scale - 0.5 + return resampled, affine @ step + + +def fit_to_shape( + x: torch.Tensor, affine: np.ndarray, target_shape: tuple[int, ...] +) -> tuple[torch.Tensor, np.ndarray]: + """Centre the volume in `target_shape`, padding the short axes and cropping the long ones.""" + pads = [target - size for size, target in zip(x.shape, target_shape)] + padding = [side for pad in reversed(pads) for side in (pad // 2, pad - pad // 2)] + + # a crop is a negative pad, so output voxel k came from input voxel k - pad // 2 either way + step = np.eye(4) + step[:3, 3] = [-(pad // 2) for pad in pads] + return F.pad(x, padding), affine @ step + + +def resolve_ckpt(ckpt_path: str) -> str: + """A local path for a checkpoint, downloading it if it is an hf://// URI.""" + from huggingface_hub import hf_hub_download + + if ckpt_path.startswith("hf://"): + org, repo, *rest = ckpt_path.removeprefix("hf://").split("/") + return hf_hub_download(f"{org}/{repo}", "/".join(rest)) + + return ckpt_path + + +def load_backbone(ckpt_path: str) -> tuple[SmriMaeBackbone, SmriMaeTransform]: + path = resolve_ckpt(ckpt_path) + ckpt = torch.load(path, map_location="cpu", weights_only=True, mmap=True) + args = ckpt["args"] + + model_fn = models_mae.__dict__[args["model"]] + model: models_mae.MaskedAutoencoderViT = model_fn( + img_size=args["img_size"], + in_chans=args.get("in_chans", 1), + patch_size=args["patch_size"], + # older checkpoints carry training flags the current model_mae no longer takes + **filter_kwargs(models_mae.MaskedAutoencoderViT, args.get("model_kwargs") or {}), + ) + model.load_state_dict(ckpt["model"]) + backbone = SmriMaeBackbone(model.encoder) + transform = SmriMaeTransform(img_size=args["img_size"]) + return backbone, transform + + +def filter_kwargs(func, kwargs): + signature = inspect.signature(func) + kwargs = {k: v for k, v in kwargs.items() if k in signature.parameters} + return kwargs diff --git a/fomo_tune_baseline/output/task3/build/fomo_tune/datasets.py b/fomo_tune_baseline/output/task3/build/fomo_tune/datasets.py new file mode 100644 index 0000000000000000000000000000000000000000..64c2ac2da798086e88cbb335601ae850c2dad249 --- /dev/null +++ b/fomo_tune_baseline/output/task3/build/fomo_tune/datasets.py @@ -0,0 +1,205 @@ +import os +import shutil +import tempfile +import zipfile +from collections.abc import Generator +from contextlib import contextmanager +from pathlib import Path + +import fsspec +from datasets import Dataset, Features, Nifti, Value + +FOMO_EVAL_BASE_URL = os.getenv( + "FOMO_EVAL_BASE_URL", + "https://sid.erda.dk/share_redirect/fmeuvo1EdF", +) +FOMO_EVAL_TASK5_URL = os.getenv( + "FOMO_EVAL_TASK5_URL", + "https://huggingface.co/datasets/medarc/smri-fm/resolve/main/fomo_eval/Task_5.zip", +) + + +@contextmanager +def open_zip(url: str) -> Generator[zipfile.ZipFile, None, None]: + """Open a task zip, copying a remote url to a temp file first.""" + with tempfile.TemporaryDirectory() as tmp: + local = Path(url) + if not local.exists(): + local = Path(tmp) / "task.zip" + with fsspec.open(url) as src, local.open("wb") as dst: + shutil.copyfileobj(src, dst) + with zipfile.ZipFile(local) as zf: + yield zf + + +def subject_ids(zf: zipfile.ZipFile) -> list[str]: + return sorted({name.split("/")[2] for name in zf.namelist() if name.endswith(".nii.gz")}) + + +# ---- Task 1: acute infarct (classification; positives also carry a lesion mask) -------- + + +def load_fomo_task1() -> Dataset: + # No 4th modality: it is swi on 16 subjects and t2s on the other 5. + suffixes = ("adc", "dwi_b1000", "flair") + features = Features( + { + "subject": Value("string"), + "label": Value("int32"), + **{suffix: Nifti() for suffix in suffixes}, + } + ) + dataset = Dataset.from_generator( + _fomo_task1_generator, + features=features, + gen_kwargs={"suffixes": suffixes}, + writer_batch_size=16, + ) + return dataset + + +def _fomo_task1_generator(suffixes: tuple[str, ...]): + url = f"{FOMO_EVAL_BASE_URL}/Task_1.zip" + with open_zip(url) as zf: + for sub in subject_ids(zf): + label = int(zf.read(f"Task_1/labels/{sub}/ses-01/label.txt").strip()) + sample = {"subject": sub, "label": label} + for suffix in suffixes: + name = f"Task_1/preprocessed/{sub}/ses-01/{suffix}.nii.gz" + sample[suffix] = {"path": None, "bytes": zf.read(name)} + yield sample + + +# ---- Task 2: meningioma segmentation --------------------------------------------------- + + +def load_fomo_task2() -> Dataset: + # No 4th modality: it is t2s on 15 subjects and swi on the other 8. + suffixes = ("dwi_b1000", "flair") + features = Features( + { + "subject": Value("string"), + **{suffix: Nifti() for suffix in suffixes}, + "seg": Nifti(), + } + ) + dataset = Dataset.from_generator( + _fomo_task2_generator, + features=features, + gen_kwargs={"suffixes": suffixes}, + writer_batch_size=16, + ) + return dataset + + +def _fomo_task2_generator(suffixes: tuple[str, ...]): + url = f"{FOMO_EVAL_BASE_URL}/Task_2.zip" + with open_zip(url) as zf: + for sub in subject_ids(zf): + sample = {"subject": sub} + for suffix in suffixes: + name = f"Task_2/preprocessed/{sub}/ses-01/{suffix}.nii.gz" + sample[suffix] = {"path": None, "bytes": zf.read(name)} + # Seg is on the image grid (shapes match) but its affine differs by up to 0.03mm. + name = f"Task_2/labels/{sub}/ses-01/seg.nii.gz" + sample["seg"] = {"path": None, "bytes": zf.read(name)} + yield sample + + +# ---- Task 3: brain age regression ------------------------------------------------------ + + +def load_fomo_task3() -> Dataset: + features = Features( + { + "subject": Value("string"), + "age": Value("int32"), + "t1w": Nifti(), + } + ) + dataset = Dataset.from_generator( + _fomo_task3_generator, + features=features, + writer_batch_size=16, + ) + return dataset + + +def _fomo_task3_generator(): + url = f"{FOMO_EVAL_BASE_URL}/Task_3.zip" + with open_zip(url) as zf: + for sub in subject_ids(zf): + age = int(zf.read(f"Task_3/labels/{sub}/ses-01/labels.txt").strip()) + image_gz = zf.read(f"Task_3/preprocessed/{sub}/ses-01/t1w.nii.gz") + sample = { + "subject": sub, + "age": age, + "t1w": {"path": None, "bytes": image_gz}, + } + yield sample + + +# ---- Task 4: trigeminal nerve/vessel segmentation -------------------------------------- + + +def load_fomo_task4() -> Dataset: + # Volumes are uncropped 0.5mm near-iso, ~360x512x512; crop before feeding a model. + features = Features( + { + "subject": Value("string"), + "t2w": Nifti(), + "seg": Nifti(), + } + ) + dataset = Dataset.from_generator( + _fomo_task4_generator, + features=features, + writer_batch_size=16, + ) + return dataset + + +def _fomo_task4_generator(): + url = f"{FOMO_EVAL_BASE_URL}/Task_4.zip" + with open_zip(url) as zf: + for sub in subject_ids(zf): + image_gz = zf.read(f"Task_4/preprocessed/{sub}/ses-01/t2w.nii.gz") + seg_gz = zf.read(f"Task_4/labels/{sub}/ses-01/seg.nii.gz") + sample = { + "subject": sub, + "t2w": {"path": None, "bytes": image_gz}, + "seg": {"path": None, "bytes": seg_gz}, + } + yield sample + + +# ---- Task 5: polymicrogyria classification --------------------------------------------- + + +def load_fomo_task5() -> Dataset: + features = Features( + { + "subject": Value("string"), + "label": Value("int32"), + "t1w": Nifti(), + } + ) + dataset = Dataset.from_generator( + _fomo_task5_generator, + features=features, + writer_batch_size=16, + ) + return dataset + + +def _fomo_task5_generator(): + with open_zip(FOMO_EVAL_TASK5_URL) as zf: + for sub in subject_ids(zf): + label = int(zf.read(f"Task_5/labels/{sub}/ses_01/labels.txt").strip()) + image_gz = zf.read(f"Task_5/preprocessed/{sub}/ses_01/t1.nii.gz") + sample = { + "subject": sub, + "label": label, + "t1w": {"path": None, "bytes": image_gz}, + } + yield sample diff --git a/fomo_tune_baseline/output/task3/build/fomo_tune/main_task1.py b/fomo_tune_baseline/output/task3/build/fomo_tune/main_task1.py new file mode 100644 index 0000000000000000000000000000000000000000..bfbc47d49ceeed6fc43dc52a0a007d933ff86d4b --- /dev/null +++ b/fomo_tune_baseline/output/task3/build/fomo_tune/main_task1.py @@ -0,0 +1,253 @@ +"""FOMO task 1: acute infarct classification, scored by AUROC as the challenge scores it. + +`Task1Method` is the part we tune -- features, head, hyperparameters. The protocol below it is +fixed so scores stay comparable across iterations: leave one subject out, pool the out-of-fold +predictions, bootstrap subjects for the CI. + +`train` runs that protocol then fits and saves a head; `predict` is the challenge contract, +modality paths in and one probability out. Both go through `Task1Method.predict`, so every fold +exercises the path the submission will run. +""" + +import argparse +import json +import logging +import time +from dataclasses import dataclass, field +from pathlib import Path + +import joblib +import nibabel as nib +import numpy as np +import torch +from omegaconf import OmegaConf +from sklearn.linear_model import LogisticRegressionCV +from sklearn.metrics import roc_auc_score +from sklearn.pipeline import make_pipeline +from sklearn.preprocessing import StandardScaler + +from fomo_tune.backbone import load_backbone +from fomo_tune.utils import git_sha, set_seed, setup_logging + +logger = logging.getLogger("fomo_tune") + +Images = dict[str, nib.Nifti1Image] + + +@dataclass +class Config: + task: str = "task1" + ckpt_path: str = ( + "/data/mihir-stuff/smri-pretrained/pretrain_full_90_10_h100/checkpoint-last.pth" + ) + modalities: list[str] = field(default_factory=lambda: ["dwi_b1000"]) + output_root: str = "output/fomo_tune" + name: str = "task1" + device: str = "cuda" + seed: int = 4466 + + +# ---- method: the part we tune ----------------------------------------------------------- + + +class Task1Method: + """Frozen sMRI MAE, mean-pooled tokens per modality concatenated, logistic head.""" + + def __init__(self, cfg: Config): + self.cfg = cfg + self.backbone, self.transform = load_backbone(cfg.ckpt_path) + self.device = torch.device(cfg.device) + self.backbone.to(self.device).eval().requires_grad_(False) + self.modalities = list(cfg.modalities) + self.cache: dict[str, np.ndarray] = {} + self.head = None + + @torch.inference_mode() + def features(self, images: Images) -> np.ndarray: + """(D,) per subject. A pure function of the images, so training and inference agree.""" + pooled = [] + for modality in self.modalities: + sample = self.transform(images[modality]) + batch = {key: value[None].to(self.device) for key, value in sample.items()} + + with torch.autocast("cuda", torch.bfloat16, enabled=self.device.type == "cuda"): + out = self.backbone(batch) + + patch_embeds = out["patch_embeds"] + token_mask = out["token_mask"].bool().unsqueeze(-1) + embed = (patch_embeds * token_mask).sum(dim=1) / token_mask.sum(dim=1) + pooled.append(embed[0].float().cpu()) + + return torch.cat(pooled).numpy() + + def cached_features(self, row: dict) -> np.ndarray: + if row["subject"] not in self.cache: + self.cache[row["subject"]] = self.features(row) + return self.cache[row["subject"]] + + def fit(self, rows: list[dict]) -> None: + X = np.stack([self.cached_features(row) for row in rows]) + y = np.array([row["label"] for row in rows]) + + clf = LogisticRegressionCV( + Cs=10, + class_weight="balanced", + scoring="roc_auc", + max_iter=1000, + l1_ratios=(0,), + use_legacy_attributes=False, + ) + self.head = make_pipeline(StandardScaler(), clf) + self.head.fit(X, y) + self.positive = list(self.head.classes_).index(1) + + def predict(self, images: Images) -> float: + """Positive-class probability. Indexes `classes_` rather than assuming column 1, which + would silently score the wrong class if the label order differed.""" + X = self.features(images)[None] + probs = self.head.predict_proba(X)[0] + return float(probs[self.positive]) + + def save(self, model_dir: Path) -> None: + """Everything `load` needs but the backbone weights, which stay wherever `ckpt_path` + points -- a few hundred KB, so a run saves one without copying a 3.7G checkpoint.""" + model_dir.mkdir(parents=True, exist_ok=True) + OmegaConf.save(self.cfg, model_dir / "config.yaml") + joblib.dump({"head": self.head, "positive": self.positive}, model_dir / "head.joblib") + + @classmethod + def load(cls, model_dir: Path, **overrides) -> "Task1Method": + """Rebuild a fitted method from `save`. Overrides are Config fields, for what differs + between here and the container -- the backbone path, the device.""" + cfg = OmegaConf.merge( + OmegaConf.structured(Config), OmegaConf.load(model_dir / "config.yaml"), overrides + ) + method = cls(cfg) + state = joblib.load(model_dir / "head.joblib") + method.head, method.positive = state["head"], state["positive"] + return method + + +# ---- protocol: the part we hold fixed --------------------------------------------------- + +# Every image the task ships. The method picks which of them it wants, as at inference, where +# the challenge hands over all four modalities whether or not a model uses them. +IMAGE_COLS = ("adc", "dwi_b1000", "flair") + + +def leave_one_out(rows: list[dict], method: Task1Method) -> tuple[np.ndarray, np.ndarray]: + """Out-of-fold score for every subject, each predicted by a head fit on the other n-1.""" + y = np.array([row["label"] for row in rows]) + oof = np.zeros(len(rows), dtype=float) + start = time.perf_counter() + for held_out, row in enumerate(rows): + method.fit([r for r in rows if r["subject"] != row["subject"]]) + oof[held_out] = method.predict({key: row[key] for key in IMAGE_COLS}) + logger.info( + f"fold {held_out + 1}/{len(rows)} {row['subject']} " + f"y={y[held_out]} p={oof[held_out]:.3f} ({time.perf_counter() - start:.0f}s)" + ) + return y, oof + + +def score( + y: np.ndarray, oof: np.ndarray, seed: int = 0, n_boot: int = 2000, alpha: float = 0.05 +) -> dict: + """AUROC, the challenge metric, plus a percentile CI resampling subjects with replacement.""" + rng = np.random.default_rng(seed) + samples = [] + for _ in range(n_boot): + rows = rng.integers(0, len(y), size=len(y)) + if len(np.unique(y[rows])) < 2: + continue + samples.append(roc_auc_score(y[rows], oof[rows])) + + low, high = np.percentile(samples, [100 * alpha / 2, 100 * (1 - alpha / 2)]) + return { + "auroc": float(roc_auc_score(y, oof)), + "auroc_ci_low": float(low), + "auroc_ci_high": float(high), + } + + +# ---- entrypoints ------------------------------------------------------------------------ + + +def train(args: argparse.Namespace) -> None: + # imported here, not at the top, so the container needs no dataset stack to run `predict` + from fomo_tune.datasets import load_fomo_task1 + + cfg = OmegaConf.merge(OmegaConf.structured(Config), OmegaConf.from_dotlist(args.overrides)) + run_dir = Path(cfg.output_root) / cfg.name + run_dir.mkdir(parents=True, exist_ok=True) + + setup_logging(run_dir) + set_seed(cfg.seed) + logger.info(f"run {cfg.name} (git {git_sha()})") + logger.info(f"config:\n{OmegaConf.to_yaml(cfg).rstrip()}") + OmegaConf.save(cfg, run_dir / "config.yaml") + + # decoded once: leave-one-out revisits every subject n times, and the niftis are small + rows = list(load_fomo_task1()) + logger.info(f"dataset: {len(rows)} subjects, {sum(r['label'] for r in rows)} positive") + + method = Task1Method(cfg) + start = time.perf_counter() + y, oof = leave_one_out(rows, method) + run_time = time.perf_counter() - start + summary = score(y, oof) + + # the shipped head sees all n subjects, so it is not any of the models scored above + method.fit(rows) + method.save(run_dir / "model") + + record = {"name": cfg.name, **summary, "run_time": round(run_time, 1)} + (run_dir / "metrics.json").write_text(json.dumps(record) + "\n") + scores = " ".join(f"{k}={v:.4f}" for k, v in summary.items()) + logger.info(f"result: {scores} ({run_time:.0f}s)") + + +def predict(args: argparse.Namespace) -> None: + """The challenge contract: modality paths in, one probability written to `--output`. + + `/app/predict.py` in the container is a shim over this, so what scores the submission is the + code leave-one-out already ran, not something generated at build time. + """ + overrides = {"device": args.device} + if args.ckpt_path: + overrides["ckpt_path"] = args.ckpt_path + method = Task1Method.load(args.model_dir, **overrides) + + # every image the challenge hands over, as in `leave_one_out`; the method takes what it uses + paths = {"adc": args.adc, "dwi_b1000": args.dwi, "flair": args.flair} + probability = method.predict({key: nib.load(path) for key, path in paths.items()}) + + args.output.write_text(f"{probability:.6f}\n") + + +def main() -> None: + parser = argparse.ArgumentParser() + modes = parser.add_subparsers(required=True) + + train_parser = modes.add_parser("train", help="leave-one-out over the task, then fit and save") + train_parser.add_argument("overrides", nargs="*", help="config overrides, e.g. device=cpu") + train_parser.set_defaults(run=train) + + predict_parser = modes.add_parser("predict", help="one subject, one probability") + for flag in ("--flair", "--adc", "--dwi"): + predict_parser.add_argument(flag, type=Path, required=True) + # accepted and ignored: the 4th modality is swi on some subjects and t2s on others + for flag in ("--t2s", "--swi"): + predict_parser.add_argument(flag, type=Path) + predict_parser.add_argument("--output", type=Path, required=True) + predict_parser.add_argument("--model-dir", type=Path, default=Path("/app/model")) + predict_parser.add_argument("--ckpt-path", help="overrides the trained config's backbone path") + predict_parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") + predict_parser.set_defaults(run=predict) + + args = parser.parse_args() + args.run(args) + + +if __name__ == "__main__": + main() diff --git a/fomo_tune_baseline/output/task3/build/fomo_tune/main_task3.py b/fomo_tune_baseline/output/task3/build/fomo_tune/main_task3.py new file mode 100644 index 0000000000000000000000000000000000000000..02db5413c10a0d44eb248e7f2739a436c45806c2 --- /dev/null +++ b/fomo_tune_baseline/output/task3/build/fomo_tune/main_task3.py @@ -0,0 +1,241 @@ +"""FOMO task 3: brain age regression, scored by Pearson r and MAE as the challenge scores it. + +`Task3Method` is the part we tune -- features, head, hyperparameters. The protocol below it is +fixed so scores stay comparable across iterations: 20-fold over the 494 subjects, pool the +out-of-fold predictions, bootstrap subjects for the CI. + +`train` runs that protocol then fits and saves a head; `predict` is the challenge contract, one t1 +path in and one age out. Both go through `Task3Method.predict`, so every fold exercises the path +the submission will run. +""" + +import argparse +import json +import logging +import time +from dataclasses import dataclass +from pathlib import Path + +import joblib +import nibabel as nib +import numpy as np +import torch +from omegaconf import OmegaConf +from sklearn.linear_model import RidgeCV +from sklearn.model_selection import KFold +from sklearn.pipeline import make_pipeline +from sklearn.preprocessing import StandardScaler + +from fomo_tune.backbone import load_backbone +from fomo_tune.utils import git_sha, set_seed, setup_logging + +logger = logging.getLogger("fomo_tune") + +Images = dict[str, nib.Nifti1Image] + + +@dataclass +class Config: + task: str = "task3" + ckpt_path: str = ( + "/data/mihir-stuff/smri-pretrained/pretrain_full_90_10_h100/checkpoint-last.pth" + ) + output_root: str = "output/fomo_tune" + name: str = "task3" + device: str = "cuda" + seed: int = 4466 + + +# ---- method: the part we tune ----------------------------------------------------------- + + +class Task3Method: + """Frozen sMRI MAE, mean-pooled tokens over the t1w, ridge head.""" + + def __init__(self, cfg: Config): + self.cfg = cfg + self.backbone, self.transform = load_backbone(cfg.ckpt_path) + self.device = torch.device(cfg.device) + self.backbone.to(self.device).eval().requires_grad_(False) + self.cache: dict[str, np.ndarray] = {} + self.head = None + + @torch.inference_mode() + def features(self, images: Images) -> np.ndarray: + """(D,) per subject. A pure function of the images, so training and inference agree.""" + sample = self.transform(images["t1w"]) + batch = {key: value[None].to(self.device) for key, value in sample.items()} + + with torch.autocast("cuda", torch.bfloat16, enabled=self.device.type == "cuda"): + out = self.backbone(batch) + + patch_embeds = out["patch_embeds"] + token_mask = out["token_mask"].bool().unsqueeze(-1) + embed = (patch_embeds * token_mask).sum(dim=1) / token_mask.sum(dim=1) + return embed[0].float().cpu().numpy() + + def cached_features(self, row: dict) -> np.ndarray: + if row["subject"] not in self.cache: + self.cache[row["subject"]] = self.features(row) + return self.cache[row["subject"]] + + def fit(self, rows: list[dict]) -> None: + X = np.stack([self.cached_features(row) for row in rows]) + y = np.array([row["age"] for row in rows], dtype=float) + + # RidgeCV picks alpha by its own efficient leave-one-out, so the fold's own split is + # never touched by model selection + self.head = make_pipeline(StandardScaler(), RidgeCV(alphas=np.logspace(-3, 6, 19))) + self.head.fit(X, y) + + def predict(self, images: Images) -> float: + """Age in years.""" + X = self.features(images)[None] + return float(self.head.predict(X)[0]) + + def save(self, model_dir: Path) -> None: + """Everything `load` needs but the backbone weights, which stay wherever `ckpt_path` + points -- a few hundred KB, so a run saves one without copying a 3.7G checkpoint.""" + model_dir.mkdir(parents=True, exist_ok=True) + OmegaConf.save(self.cfg, model_dir / "config.yaml") + joblib.dump(self.head, model_dir / "head.joblib") + + @classmethod + def load(cls, model_dir: Path, **overrides) -> "Task3Method": + """Rebuild a fitted method from `save`. Overrides are Config fields, for what differs + between here and the container -- the backbone path, the device.""" + cfg = OmegaConf.merge( + OmegaConf.structured(Config), OmegaConf.load(model_dir / "config.yaml"), overrides + ) + method = cls(cfg) + method.head = joblib.load(model_dir / "head.joblib") + return method + + +# ---- protocol: the part we hold fixed --------------------------------------------------- + +# Every image the task ships. The method picks which of them it wants, as at inference, where the +# challenge hands over the modalities whether or not a model uses them. +IMAGE_COLS = ("t1w",) + + +def cross_validate( + rows: list[dict], method: Task3Method, seed: int = 0, n_folds: int = 20 +) -> tuple[np.ndarray, np.ndarray]: + """Out-of-fold age for every subject, each predicted by a head fit on the other folds.""" + y = np.array([row["age"] for row in rows], dtype=float) + oof = np.zeros(len(rows), dtype=float) + folds = KFold(n_splits=n_folds, shuffle=True, random_state=seed) + start = time.perf_counter() + for fold, (train, test) in enumerate(folds.split(rows)): + method.fit([rows[i] for i in train]) + for i in test: + oof[i] = method.predict({key: rows[i][key] for key in IMAGE_COLS}) + logger.info( + f"fold {fold + 1}/{n_folds} n={len(test)} mae={np.abs(y[test] - oof[test]).mean():.2f} " + f"({time.perf_counter() - start:.0f}s)" + ) + return y, oof + + +def metrics(y: np.ndarray, oof: np.ndarray) -> dict: + return { + "pearson_r": float(np.corrcoef(y, oof)[0, 1]), + "mae": float(np.abs(y - oof).mean()), + } + + +def score( + y: np.ndarray, oof: np.ndarray, seed: int = 0, n_boot: int = 2000, alpha: float = 0.05 +) -> dict: + """Both challenge metrics, each with a percentile CI resampling subjects with replacement.""" + rng = np.random.default_rng(seed) + resamples = rng.integers(0, len(y), size=(n_boot, len(y))) + + summary = {} + for name, point in metrics(y, oof).items(): + samples = [metrics(y[rows], oof[rows])[name] for rows in resamples] + low, high = np.percentile(samples, [100 * alpha / 2, 100 * (1 - alpha / 2)]) + summary[name] = point + summary[f"{name}_ci_low"] = float(low) + summary[f"{name}_ci_high"] = float(high) + return summary + + +# ---- entrypoints ------------------------------------------------------------------------ + + +def train(args: argparse.Namespace) -> None: + # imported here, not at the top, so the container needs no dataset stack to run `predict` + from fomo_tune.datasets import load_fomo_task3 + + cfg = OmegaConf.merge(OmegaConf.structured(Config), OmegaConf.from_dotlist(args.overrides)) + run_dir = Path(cfg.output_root) / cfg.name + run_dir.mkdir(parents=True, exist_ok=True) + + setup_logging(run_dir) + set_seed(cfg.seed) + logger.info(f"run {cfg.name} (git {git_sha()})") + logger.info(f"config:\n{OmegaConf.to_yaml(cfg).rstrip()}") + OmegaConf.save(cfg, run_dir / "config.yaml") + + rows = list(load_fomo_task3()) + ages = np.array([row["age"] for row in rows]) + logger.info( + f"dataset: {len(rows)} subjects, age {ages.min()}-{ages.max()} mean {ages.mean():.1f}" + ) + + method = Task3Method(cfg) + start = time.perf_counter() + y, oof = cross_validate(rows, method) + run_time = time.perf_counter() - start + summary = score(y, oof) + + # the shipped head sees all n subjects, so it is not any of the models scored above + method.fit(rows) + method.save(run_dir / "model") + + record = {"name": cfg.name, **summary, "run_time": round(run_time, 1)} + (run_dir / "metrics.json").write_text(json.dumps(record) + "\n") + scores = " ".join(f"{k}={v:.4f}" for k, v in summary.items()) + logger.info(f"result: {scores} ({run_time:.0f}s)") + + +def predict(args: argparse.Namespace) -> None: + """The challenge contract: a t1 path in, one age written to `--output`. + + `/app/predict.py` in the container is a shim over this, so what scores the submission is the + code cross-validation already ran, not something generated at build time. + """ + overrides = {"device": args.device} + if args.ckpt_path: + overrides["ckpt_path"] = args.ckpt_path + method = Task3Method.load(args.model_dir, **overrides) + + age = method.predict({"t1w": nib.load(args.t1)}) + + args.output.write_text(f"{age:.6f}\n") + + +def main() -> None: + parser = argparse.ArgumentParser() + modes = parser.add_subparsers(required=True) + + train_parser = modes.add_parser("train", help="cross-validate over the task, then fit and save") + train_parser.add_argument("overrides", nargs="*", help="config overrides, e.g. device=cpu") + train_parser.set_defaults(run=train) + + predict_parser = modes.add_parser("predict", help="one subject, one age in years") + predict_parser.add_argument("--t1", type=Path, required=True) + predict_parser.add_argument("--output", type=Path, required=True) + predict_parser.add_argument("--model-dir", type=Path, default=Path("/app/model")) + predict_parser.add_argument("--ckpt-path", help="overrides the trained config's backbone path") + predict_parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") + predict_parser.set_defaults(run=predict) + + args = parser.parse_args() + args.run(args) + + +if __name__ == "__main__": + main() diff --git a/fomo_tune_baseline/output/task3/build/fomo_tune/main_task5.py b/fomo_tune_baseline/output/task3/build/fomo_tune/main_task5.py new file mode 100644 index 0000000000000000000000000000000000000000..d8cccbecf130497c483b3a4b7665e835ca428873 --- /dev/null +++ b/fomo_tune_baseline/output/task3/build/fomo_tune/main_task5.py @@ -0,0 +1,245 @@ +"""FOMO task 5: polymicrogyria classification, scored by AUROC as the challenge scores it. + +`Task5Method` is the part we tune -- features, head, hyperparameters. The protocol below it is +fixed so scores stay comparable across iterations: 20-fold over the 48 subjects, pool the +out-of-fold predictions, bootstrap subjects for the CI. + +`train` runs that protocol then fits and saves a head; `predict` is the challenge contract, one t1 +path in and one probability out. Both go through `Task5Method.predict`, so every fold exercises +the path the submission will run. +""" + +import argparse +import json +import logging +import time +from dataclasses import dataclass +from pathlib import Path + +import joblib +import nibabel as nib +import numpy as np +import torch +from omegaconf import OmegaConf +from sklearn.linear_model import LogisticRegressionCV +from sklearn.metrics import roc_auc_score +from sklearn.model_selection import KFold +from sklearn.pipeline import make_pipeline +from sklearn.preprocessing import StandardScaler + +from fomo_tune.backbone import load_backbone +from fomo_tune.utils import git_sha, set_seed, setup_logging + +logger = logging.getLogger("fomo_tune") + +Images = dict[str, nib.Nifti1Image] + + +@dataclass +class Config: + task: str = "task5" + ckpt_path: str = ( + "/data/mihir-stuff/smri-pretrained/pretrain_full_90_10_h100/checkpoint-last.pth" + ) + output_root: str = "output/fomo_tune" + name: str = "task5" + device: str = "cuda" + seed: int = 4466 + + +# ---- method: the part we tune ----------------------------------------------------------- + + +class Task5Method: + """Frozen sMRI MAE, mean-pooled tokens over the t1w, logistic head.""" + + def __init__(self, cfg: Config): + self.cfg = cfg + self.backbone, self.transform = load_backbone(cfg.ckpt_path) + self.device = torch.device(cfg.device) + self.backbone.to(self.device).eval().requires_grad_(False) + self.cache: dict[str, np.ndarray] = {} + self.head = None + + @torch.inference_mode() + def features(self, images: Images) -> np.ndarray: + """(D,) per subject. A pure function of the images, so training and inference agree.""" + sample = self.transform(images["t1w"]) + batch = {key: value[None].to(self.device) for key, value in sample.items()} + + with torch.autocast("cuda", torch.bfloat16, enabled=self.device.type == "cuda"): + out = self.backbone(batch) + + patch_embeds = out["patch_embeds"] + token_mask = out["token_mask"].bool().unsqueeze(-1) + embed = (patch_embeds * token_mask).sum(dim=1) / token_mask.sum(dim=1) + return embed[0].float().cpu().numpy() + + def cached_features(self, row: dict) -> np.ndarray: + if row["subject"] not in self.cache: + self.cache[row["subject"]] = self.features(row) + return self.cache[row["subject"]] + + def fit(self, rows: list[dict]) -> None: + X = np.stack([self.cached_features(row) for row in rows]) + y = np.array([row["label"] for row in rows]) + + clf = LogisticRegressionCV( + Cs=10, + class_weight="balanced", + scoring="roc_auc", + max_iter=1000, + l1_ratios=(0,), + use_legacy_attributes=False, + ) + self.head = make_pipeline(StandardScaler(), clf) + self.head.fit(X, y) + self.positive = list(self.head.classes_).index(1) + + def predict(self, images: Images) -> float: + """Positive-class probability. Indexes `classes_` rather than assuming column 1, which + would silently score the wrong class if the label order differed.""" + X = self.features(images)[None] + probs = self.head.predict_proba(X)[0] + return float(probs[self.positive]) + + def save(self, model_dir: Path) -> None: + """Everything `load` needs but the backbone weights, which stay wherever `ckpt_path` + points -- a few hundred KB, so a run saves one without copying a 3.7G checkpoint.""" + model_dir.mkdir(parents=True, exist_ok=True) + OmegaConf.save(self.cfg, model_dir / "config.yaml") + joblib.dump({"head": self.head, "positive": self.positive}, model_dir / "head.joblib") + + @classmethod + def load(cls, model_dir: Path, **overrides) -> "Task5Method": + """Rebuild a fitted method from `save`. Overrides are Config fields, for what differs + between here and the container -- the backbone path, the device.""" + cfg = OmegaConf.merge( + OmegaConf.structured(Config), OmegaConf.load(model_dir / "config.yaml"), overrides + ) + method = cls(cfg) + state = joblib.load(model_dir / "head.joblib") + method.head, method.positive = state["head"], state["positive"] + return method + + +# ---- protocol: the part we hold fixed --------------------------------------------------- + +# Every image the task ships. The method picks which of them it wants, as at inference, where the +# challenge hands over the modalities whether or not a model uses them. +IMAGE_COLS = ("t1w",) + + +def cross_validate( + rows: list[dict], method: Task5Method, seed: int = 0, n_folds: int = 20 +) -> tuple[np.ndarray, np.ndarray]: + """Out-of-fold score for every subject, each predicted by a head fit on the other folds.""" + y = np.array([row["label"] for row in rows]) + oof = np.zeros(len(rows), dtype=float) + folds = KFold(n_splits=n_folds, shuffle=True, random_state=seed) + start = time.perf_counter() + for fold, (train, test) in enumerate(folds.split(rows)): + method.fit([rows[i] for i in train]) + for i in test: + oof[i] = method.predict({key: rows[i][key] for key in IMAGE_COLS}) + logger.info( + f"fold {fold + 1}/{n_folds} n={len(test)} y={y[test]} " + f"p={np.round(oof[test], 3)} ({time.perf_counter() - start:.0f}s)" + ) + return y, oof + + +def score( + y: np.ndarray, oof: np.ndarray, seed: int = 0, n_boot: int = 2000, alpha: float = 0.05 +) -> dict: + """AUROC, the challenge metric, plus a percentile CI resampling subjects with replacement.""" + rng = np.random.default_rng(seed) + samples = [] + for _ in range(n_boot): + rows = rng.integers(0, len(y), size=len(y)) + if len(np.unique(y[rows])) < 2: + continue + samples.append(roc_auc_score(y[rows], oof[rows])) + + low, high = np.percentile(samples, [100 * alpha / 2, 100 * (1 - alpha / 2)]) + return { + "auroc": float(roc_auc_score(y, oof)), + "auroc_ci_low": float(low), + "auroc_ci_high": float(high), + } + + +# ---- entrypoints ------------------------------------------------------------------------ + + +def train(args: argparse.Namespace) -> None: + # imported here, not at the top, so the container needs no dataset stack to run `predict` + from fomo_tune.datasets import load_fomo_task5 + + cfg = OmegaConf.merge(OmegaConf.structured(Config), OmegaConf.from_dotlist(args.overrides)) + run_dir = Path(cfg.output_root) / cfg.name + run_dir.mkdir(parents=True, exist_ok=True) + + setup_logging(run_dir) + set_seed(cfg.seed) + logger.info(f"run {cfg.name} (git {git_sha()})") + logger.info(f"config:\n{OmegaConf.to_yaml(cfg).rstrip()}") + OmegaConf.save(cfg, run_dir / "config.yaml") + + rows = list(load_fomo_task5()) + logger.info(f"dataset: {len(rows)} subjects, {sum(r['label'] for r in rows)} positive") + + method = Task5Method(cfg) + start = time.perf_counter() + y, oof = cross_validate(rows, method) + run_time = time.perf_counter() - start + summary = score(y, oof) + + # the shipped head sees all n subjects, so it is not any of the models scored above + method.fit(rows) + method.save(run_dir / "model") + + record = {"name": cfg.name, **summary, "run_time": round(run_time, 1)} + (run_dir / "metrics.json").write_text(json.dumps(record) + "\n") + scores = " ".join(f"{k}={v:.4f}" for k, v in summary.items()) + logger.info(f"result: {scores} ({run_time:.0f}s)") + + +def predict(args: argparse.Namespace) -> None: + """The challenge contract: a t1 path in, one probability written to `--output`. + + `/app/predict.py` in the container is a shim over this, so what scores the submission is the + code cross-validation already ran, not something generated at build time. + """ + overrides = {"device": args.device} + if args.ckpt_path: + overrides["ckpt_path"] = args.ckpt_path + method = Task5Method.load(args.model_dir, **overrides) + + probability = method.predict({"t1w": nib.load(args.t1)}) + + args.output.write_text(f"{probability:.6f}\n") + + +def main() -> None: + parser = argparse.ArgumentParser() + modes = parser.add_subparsers(required=True) + + train_parser = modes.add_parser("train", help="cross-validate over the task, then fit and save") + train_parser.add_argument("overrides", nargs="*", help="config overrides, e.g. device=cpu") + train_parser.set_defaults(run=train) + + predict_parser = modes.add_parser("predict", help="one subject, one probability") + predict_parser.add_argument("--t1", type=Path, required=True) + predict_parser.add_argument("--output", type=Path, required=True) + predict_parser.add_argument("--model-dir", type=Path, default=Path("/app/model")) + predict_parser.add_argument("--ckpt-path", help="overrides the trained config's backbone path") + predict_parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") + predict_parser.set_defaults(run=predict) + + args = parser.parse_args() + args.run(args) + + +if __name__ == "__main__": + main() diff --git a/fomo_tune_baseline/output/task3/build/fomo_tune/utils.py b/fomo_tune_baseline/output/task3/build/fomo_tune/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..57a5c6d0396d8811e763254a987b331cf41ee1f1 --- /dev/null +++ b/fomo_tune_baseline/output/task3/build/fomo_tune/utils.py @@ -0,0 +1,33 @@ +import logging +import random +import subprocess +import sys +from pathlib import Path + +import numpy as np +import torch + +logger = logging.getLogger("fomo_tune") + + +def set_seed(seed: int) -> None: + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + + +def git_sha() -> str: + kwargs = dict(cwd=Path(__file__).parent, capture_output=True, text=True, check=True) + sha = subprocess.run(["git", "rev-parse", "--short", "HEAD"], **kwargs).stdout.strip() + dirty = subprocess.run(["git", "status", "--porcelain", "-uno"], **kwargs).stdout.strip() + return f"{sha}-dirty" if dirty else sha + + +def setup_logging(run_dir: Path) -> None: + handlers = [logging.StreamHandler(sys.stdout), logging.FileHandler(run_dir / "log.txt")] + logger.setLevel(logging.INFO) + logger.handlers.clear() + for handler in handlers: + handler.setFormatter(logging.Formatter("%(asctime)s %(message)s", datefmt="%H:%M:%S")) + logger.addHandler(handler) + logger.propagate = False diff --git a/fomo_tune_baseline/output/task3/build/model/backbone.pth b/fomo_tune_baseline/output/task3/build/model/backbone.pth new file mode 100644 index 0000000000000000000000000000000000000000..793b42f5cea0b055534c29cd18c21f3200873219 --- /dev/null +++ b/fomo_tune_baseline/output/task3/build/model/backbone.pth @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:aacf2582e9464e7fd4bd3d5ed41a5b9aa770e5b04a24b2a82c1c4ba5fb5f6e7a +size 1389681124 diff --git a/fomo_tune_baseline/output/task3/build/model/config.yaml b/fomo_tune_baseline/output/task3/build/model/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..42af3932e7182587c4a9c2a47db74aeb12d18aa3 --- /dev/null +++ b/fomo_tune_baseline/output/task3/build/model/config.yaml @@ -0,0 +1,6 @@ +task: task3 +ckpt_path: /data/mihir-stuff/smri-pretrained/pretrain_full_90_10_h100/checkpoint-last.pth +output_root: experiments/fomo_tune_baseline/output +name: task3 +device: cuda +seed: 4466 diff --git a/fomo_tune_baseline/output/task3/build/model/head.joblib b/fomo_tune_baseline/output/task3/build/model/head.joblib new file mode 100644 index 0000000000000000000000000000000000000000..7ade23f9dc38a75d598a1a8dfeb4930ff2d4b1a9 --- /dev/null +++ b/fomo_tune_baseline/output/task3/build/model/head.joblib @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e263a95e347fbaf86b4edc44717e55891711bf37a7a1309c7f04753af46a4dbf +size 29906 diff --git a/fomo_tune_baseline/output/task3/build/predict.py b/fomo_tune_baseline/output/task3/build/predict.py new file mode 100644 index 0000000000000000000000000000000000000000..9a5e95f26f14fe78d637164dd0b5361073c03e5a --- /dev/null +++ b/fomo_tune_baseline/output/task3/build/predict.py @@ -0,0 +1,16 @@ + +import sys + +from fomo_tune.main_task3 import main + +sys.argv = [ + sys.argv[0], + "predict", + *sys.argv[1:], + "--model-dir", + "/app/model", + "--ckpt-path", + "/app/model/backbone.pth", +] + +main() diff --git a/fomo_tune_baseline/output/task3/build/smri_mae/config/default_pretrain.yaml b/fomo_tune_baseline/output/task3/build/smri_mae/config/default_pretrain.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c16290b1e2c24e4d206f7e32a7b9c2ee5266067e --- /dev/null +++ b/fomo_tune_baseline/output/task3/build/smri_mae/config/default_pretrain.yaml @@ -0,0 +1,98 @@ +# Name of the run. Used for output directory suffix and wandb. +name: pretrain + +# Description of the run. Goes in wandb notes. +notes: null + +# Root output directory. +# The run writes to checkpoints/ when name is set. +output_dir: checkpoints + +# Standard 3D structural MRI volume size. +img_size: [208, 240, 208] +patch_size: 8 + +# Masking. +mask_ratio: 0.80 +pred_mask_ratio: null +pad_to_multiple: 32 + +# Model. +model: mae_vit_large +model_kwargs: + # target normalization: null/none, global, slice, or patch. + target_norm: none + + no_decode_pos: false + mask_drop_scale: false + + class_token: true + reg_tokens: 0 + no_embed_class: false + + decoder_depth: 4 + drop_path_rate: 0.0 + +# Datasets. +datasets: + fomo_train: + url: datasets/FOMO_with_dwi/shard.{000000..001620}.tar + samples_per_epoch: 243200 + shuffle: true + buffer_size: 8000 + drop_last: true + + fomo_val: + url: datasets/FOMO_with_dwi/shard.{001621..001800}.tar + samples_per_epoch: 26880 + shuffle: false + buffer_size: 1000 + drop_last: true + +train_dataset: fomo_train +eval_datasets: + - fomo_val + +# Data loader. +num_workers: 4 +prefetch_factor: 2 +presend_cuda: true + +# Optimization. +epochs: 100 +batch_size: 64 +accum_iter: 1 + +base_lr: 0.001 +min_lr: 1e-6 +warmup_epochs: 10 +weight_decay: 0.05 +betas: [0.9, 0.95] +clip_grad: 1.0 + +amp: true +amp_dtype: bfloat16 + +# Checkpointing. +ckpt: null +resume: false +auto_resume: true +start_epoch: 0 +checkpoint_period: 10 +max_checkpoints: 5 + +# Evaluation. +eval_period: 10 + +# Sync checkpoints to an R2 bucket using the AWS CLI. Set to a URL to enable. +r2_sync: null + +device: cuda +distributed: false +seed: 7338 +eval_seed: 7338 +debug: false + +wandb: false +wandb_entity: null +wandb_project: smri-fm diff --git a/fomo_tune_baseline/output/task3/build/smri_mae/main_pretrain.py b/fomo_tune_baseline/output/task3/build/smri_mae/main_pretrain.py new file mode 100644 index 0000000000000000000000000000000000000000..5470551d5b56bf4653fad4c2a23888c417c0c6e8 --- /dev/null +++ b/fomo_tune_baseline/output/task3/build/smri_mae/main_pretrain.py @@ -0,0 +1,486 @@ +# Copyright (c) Sophont, Inc +# This source code is licensed under the Apache License, Version 2.0 +# +# References: +# deit: https://github.com/facebookresearch/deit/blob/main/main.py +# capi: https://github.com/facebookresearch/capi/blob/main/train_capi.py + +import argparse +import datetime +import json +import math +import random +import subprocess +import time +from contextlib import nullcontext +from functools import partial +from itertools import islice +from pathlib import Path +from typing import Iterable, Sequence + +import torch +import torch.nn as nn +import wandb +import webdataset as wds +from omegaconf import DictConfig, OmegaConf +from PIL import Image + +from matplotlib import pyplot as plt +from torch import Tensor + +import data.mri_data as mri_data +import smri_mae.model_mae as models_mae +import smri_mae.utils as ut +import smri_mae.visualization as vis + +DEFAULT_CONFIG = Path(__file__).parent / "config/default_pretrain.yaml" + +MODELS_DICT = models_mae.__dict__ + + +def main(args: DictConfig): + # setup + ut.init_distributed_mode(args) + global_rank = ut.get_rank() + is_master = global_rank == 0 + world_size = ut.get_world_size() + device = torch.device(args.device) + ut.configure_flash_sdpa() + ut.random_seed(args.seed, rank=global_rank) + + if args.name and not args.output_dir.endswith(args.name): + args.output_dir = f"{args.output_dir}/{args.name}" + output_dir = Path(args.output_dir) + + if is_master: + output_dir.mkdir(parents=True, exist_ok=True) + out_cfg_path = output_dir / "config.yaml" + if out_cfg_path.exists(): + prev_cfg = OmegaConf.load(out_cfg_path) + assert args == prev_cfg, "current config doesn't match previous config" + else: + OmegaConf.save(args, out_cfg_path) + + if args.wandb: + wandb.init( + entity=args.wandb_entity, + project=args.wandb_project, + name=args.name, + notes=args.notes, + config=OmegaConf.to_container(args), + ) + + ut.setup_for_distributed(log_path=output_dir / "log.txt") + + print("pretraining 3D ViTMAE") + print(f"start: {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}") + print(f"cwd: {Path.cwd()}") + print(ut.get_sha()) + print("config:", OmegaConf.to_yaml(args), sep="\n") + + # data loaders + train_loader, eval_loaders = create_data_loaders(args) + + # model + model = MODELS_DICT[args.model]( + img_size=args.img_size, + in_chans=args.get("in_chans", 1), + patch_size=args.patch_size, + **(args.get("model_kwargs") or {}), + ) + model.to(device) + print("model:", model, sep="\n") + num_params = sum(p.numel() for p in model.parameters() if p.requires_grad) + print(f"num params: {num_params / 1e6:.1f}M") + + model_without_ddp = model + if args.distributed: + model = torch.nn.parallel.DistributedDataParallel( + model, + device_ids=[args.gpu], + gradient_as_bucket_view=True, + ) + model_without_ddp = model.module + + # optimizer + total_batch_size = args.batch_size * args.accum_iter * world_size + print( + f"total batch size: {total_batch_size} = " + f"{args.batch_size} bs per gpu x {args.accum_iter} accum x {world_size} gpus" + ) + + if not args.get("lr"): + args.lr = args.base_lr * total_batch_size / 256 + print(f"lr: {args.lr:.2e} = {args.base_lr:.2e} x {total_batch_size} / 256") + else: + print(f"lr: {args.lr:.2e}") + + param_groups = ut.get_param_groups(model) + ut.update_lr(param_groups, args.lr) + ut.update_wd(param_groups, args.weight_decay) + # cast or else it corrupts the checkpoint + betas = tuple(args.betas) if args.betas is not None else None + optimizer = torch.optim.AdamW(param_groups, betas=betas, fused=True) + + epoch_num_batches = len(train_loader) + steps_per_epoch = math.ceil(epoch_num_batches / args.accum_iter) + total_steps = args.epochs * steps_per_epoch + warmup_steps = args.warmup_epochs * steps_per_epoch + lr_schedule = ut.WarmupThenCosine( + base_value=args.lr, + final_value=args.min_lr, + total_iters=total_steps, + warmup_iters=warmup_steps, + ) + print(f"full schedule: epochs = {args.epochs} (steps = {total_steps})") + print(f"warmup: epochs = {args.warmup_epochs} (steps = {warmup_steps})") + + # loss scaling not needed for bfloat16 (according to timm) + if args.amp and args.amp_dtype != "bfloat16": + loss_scaler = torch.GradScaler(device.type) + else: + loss_scaler = None + + # load checkpoint/resume training + ut.load_model(args, model_without_ddp, optimizer, loss_scaler) + + print(f"start training for {args.epochs} epochs") + start_time = time.monotonic() + for epoch in range(args.start_epoch, args.epochs): + train_stats = train_one_epoch( + args, + model, + train_loader, + optimizer, + loss_scaler, + lr_schedule, + epoch, + device, + ) + eval_stats = {} + eval_plots = {} + eval_period = args.get("eval_period", 1) + if eval_period and (epoch % eval_period == 0 or epoch == args.epochs - 1): + for name, loader in eval_loaders.items(): + stats, plots = evaluate( + args, + model, + loader, + epoch, + device, + eval_name=name, + ) + eval_stats.update(stats) + eval_plots.update(plots) + + merged_stats = {"epoch": epoch, **train_stats, **eval_stats} + if is_master: + with (output_dir / "log.json").open("a") as f: + print(json.dumps(merged_stats), file=f) + + for plot_name, img in eval_plots.items(): + plot_name = plot_name.replace("/", "__") + img.save(output_dir / f"{plot_name}__{epoch:05d}.png") + + ut.save_model(args, epoch, model_without_ddp, optimizer, loss_scaler) + sync_checkpoints_to_r2(args, output_dir) + + if args.distributed: + torch.distributed.destroy_process_group() + + total_time = time.monotonic() - start_time + print(f"done! training time: {datetime.timedelta(seconds=int(total_time))}") + + +def create_data_loaders(args: DictConfig): + data_loaders = {} + dataset_names = [args.train_dataset] + args.eval_datasets + + for dataset_name in dataset_names: + dataset_config = args.datasets[dataset_name].copy() + drop_last = dataset_config.pop("drop_last") + is_train = dataset_name == args.train_dataset + + print(f"loading dataset: {dataset_name}\n\n{OmegaConf.to_yaml(dataset_config)}") + shuffle = dataset_config["shuffle"] + samples_per_epoch = dataset_config.pop("samples_per_epoch") + dataset = mri_data.make_sparse_wds_dataset( + dataset_config["url"], + shuffle=shuffle, + buffer_size=dataset_config["buffer_size"], + ) + num_workers = int(args.num_workers) + loader_kwargs = { + "batch_size": args.batch_size, + "collate_fn": partial(mri_data.collate, include_meta=not is_train), + "shuffle": False, + "num_workers": num_workers, + "persistent_workers": num_workers > 0, + "pin_memory": True, + "drop_last": drop_last, + "prefetch_factor": args.prefetch_factor, + } + loader = wds.WebLoader(dataset, **loader_kwargs) + num_batches = samples_per_epoch // (ut.get_world_size() * args.batch_size) + loader = loader.with_epoch(num_batches) + loader = loader.with_length(num_batches, silent=True) + + data_loaders[dataset_name] = loader + + train_loader = data_loaders.pop(args.train_dataset) + return train_loader, data_loaders + + +def sync_checkpoints_to_r2(args: DictConfig, output_dir: Path) -> None: + r2_sync_url = args.get("r2_sync") + if not r2_sync_url or not ut.is_main_process(): + return + + cmd = ["aws", "s3", "sync", str(output_dir), str(r2_sync_url), "--profile", "r2"] + print(f"syncing checkpoints to R2: {output_dir} -> {r2_sync_url}") + subprocess.run(cmd, check=True) + + +def train_one_epoch( + args: DictConfig, + model: nn.Module, + data_loader: Iterable, + optimizer: torch.optim.Optimizer, + loss_scaler: torch.GradScaler | None, + lr_schedule: Sequence[float], + epoch: int, + device: torch.device, +): + model.train() + + metric_logger = ut.MetricLogger(delimiter=" ") + metric_logger.add_meter("lr", ut.SmoothedValue(window_size=1, fmt="{value:.6f}")) + metric_logger.add_meter("grad", ut.SmoothedValue()) + header = f"Train: [{epoch}]" + log_wandb = args.wandb and ut.is_main_process() + + epoch_num_batches = len(data_loader) + steps_per_epoch = math.ceil(epoch_num_batches / args.accum_iter) + + print_freq = args.get("print_freq", 100) if not args.debug else 1 + num_batches = epoch_num_batches if not args.debug else 10 + amp_dtype = getattr(torch, args.amp_dtype) + use_cuda = device.type == "cuda" + if use_cuda and args.presend_cuda: + data_loader = ut.pre_send_to_cuda_wrapper( + data_loader, device, dtype_map={torch.float16: amp_dtype} + ) + + optimizer.zero_grad() + + for batch_idx, batch in enumerate( + metric_logger.log_every(data_loader, print_freq, header, total_steps=num_batches) + ): + if use_cuda and not args.presend_cuda: + batch = ut.send_data(batch, device, dtype_map={torch.float16: amp_dtype}) + + batch_step = batch_idx + 1 + log_step = batch_step % print_freq == 0 or batch_step == num_batches + update_in_epoch = batch_idx // args.accum_iter + group_size = min(args.accum_iter, num_batches - update_in_epoch * args.accum_iter) + need_update = batch_step % args.accum_iter == 0 or batch_step == num_batches + global_step = epoch * steps_per_epoch + update_in_epoch + lr = lr_schedule[global_step] + if need_update: + ut.update_lr(optimizer.param_groups, lr) + + images, img_mask = mri_data.densify_sparse_image_batch( + batch["image_values"], + batch["img_mask"], + (int(args.get("in_chans", 1)), *args.img_size), + dtype=amp_dtype, + ) + + sync_context = model.no_sync() if args.distributed and not need_update else nullcontext() + with sync_context: + with torch.autocast(device_type=device.type, dtype=amp_dtype, enabled=args.amp): + loss = model( + images, + img_mask=img_mask, + mask_ratio=args.mask_ratio, + pred_mask_ratio=args.pred_mask_ratio, + pad_to_multiple=args.pad_to_multiple, + with_state=False, + ) + + loss_for_log = loss.detach() + torch._assert_async(torch.isfinite(loss_for_log), "non-finite loss") + + grad_norm = ut.backward_step( + loss / group_size, + optimizer, + scaler=loss_scaler, + need_update=need_update, + max_norm=args.clip_grad, + ) + + if need_update and log_step: + loss_value = loss_for_log.item() + grad_norm_value = grad_norm.item() + metric_logger.update(loss=loss_value, lr=lr, grad=grad_norm_value) + if log_wandb: + wandb.log( + { + "train/loss": loss_value, + "train/lr": lr, + "train/grad": grad_norm_value, + }, + step=int(1000 * (epoch + batch_step / epoch_num_batches)), + ) + + # gather the stats from all processes + metric_logger.synchronize_between_processes() + print("Averaged stats:", metric_logger) + return {f"train/{k}": meter.global_avg for k, meter in metric_logger.meters.items()} + + +@torch.inference_mode() +def evaluate( + args: DictConfig, + model: nn.Module, + data_loader: Iterable, + epoch: int, + device: torch.device, + eval_name: str, +): + model.eval() + + metric_logger = ut.MetricLogger(delimiter=" ") + header = f"Eval ({eval_name}): [{epoch}]" + is_master = ut.is_main_process() + log_wandb = args.wandb and is_master + + epoch_num_batches = len(data_loader) + if epoch_num_batches <= 0: + raise ValueError(f"eval loader {eval_name!r} has zero batches") + + print_freq = args.get("print_freq", 100) if not args.debug else 1 + num_batches = epoch_num_batches if not args.debug else 10 + num_batches = min(num_batches, epoch_num_batches) + eval_seed = int(args.get("eval_seed", args.seed)) + ut.get_rank() + example_step = random.Random(eval_seed).randint(1, num_batches) + amp_dtype = getattr(torch, args.amp_dtype) + use_cuda = device.type == "cuda" + rng_state = ut.capture_rng_state() + torch.set_rng_state(torch.Generator().manual_seed(eval_seed).get_state()) + if use_cuda: + torch.cuda.manual_seed(eval_seed) + if use_cuda and args.presend_cuda: + data_loader = ut.pre_send_to_cuda_wrapper( + data_loader, device, dtype_map={torch.float16: amp_dtype} + ) + + eval_batches = islice(data_loader, num_batches) + for batch_idx, batch in enumerate( + metric_logger.log_every(eval_batches, print_freq, header, total_steps=num_batches) + ): + if use_cuda and not args.presend_cuda: + batch = ut.send_data(batch, device, dtype_map={torch.float16: amp_dtype}) + + batch_step = batch_idx + 1 + + images, img_mask = mri_data.densify_sparse_image_batch( + batch["image_values"], + batch["img_mask"], + (int(args.get("in_chans", 1)), *args.img_size), + dtype=amp_dtype, + ) + + with torch.autocast(device_type=device.type, dtype=amp_dtype, enabled=args.amp): + loss, state = model( + images, + img_mask=img_mask, + mask_ratio=args.mask_ratio, + pred_mask_ratio=args.pred_mask_ratio, + pad_to_multiple=args.pad_to_multiple, + ) + + loss_value = loss.detach().float().item() + finite = torch.tensor(int(math.isfinite(loss_value)), dtype=torch.int32, device=device) + if args.distributed: + torch.distributed.all_reduce(finite, op=torch.distributed.ReduceOp.MIN) + if not finite.item(): + raise RuntimeError("non-finite validation loss detected") + metric_logger.meters["loss"].update(loss_value, n=int(batch["img_mask"].shape[0])) + + if is_master and batch_step == example_step: + example_batch = {"image": images[:1], "img_mask": img_mask[:1]} + if "meta" in batch: + example_batch["meta"] = batch["meta"][:1] + example_state = { + "pred_images": state["pred_images"][:1], + "pred_mask": state["pred_mask"][:1], + } + example_data = { + "batch": ut.send_data(example_batch, "cpu"), + "state": ut.send_data(example_state, "cpu"), + } + + # gather the stats from all processes + metric_logger.synchronize_between_processes() + print(f"Averaged stats ({eval_name}):", metric_logger) + stats = {f"eval/{eval_name}/{k}": meter.global_avg for k, meter in metric_logger.meters.items()} + + plots = {} + if is_master: + print(f"Making plots ({eval_name}): example={example_step}") + plots = make_plots(args, **example_data) + plots = {f"eval/{eval_name}/{k}": img for k, img in plots.items()} + + if log_wandb: + wandb.log(stats, step=1000 * (epoch + 1)) + wandb.log( + {k: wandb.Image(img, caption=f"example={example_step}") for k, img in plots.items()}, + step=1000 * (epoch + 1), + ) + ut.restore_rng_state(rng_state) + return stats, plots + + +def make_plots( + args: DictConfig, + batch: dict[str, Tensor], + state: dict[str, Tensor], +) -> dict[str, Image.Image]: + fig_kwargs = args.get("fig_kwargs", {}) + + images = batch["image"] + img_mask = batch.get("img_mask") + if img_mask is not None: + img_mask = img_mask.expand_as(images) + + raw_mean, raw_std = vis.raw_stats_from_batch(batch) + + plots = {} + mask_pred_fig = vis.plot_mask_pred( + target=images, + pred=state["pred_images"], + pred_mask=state["pred_mask"], + img_mask=img_mask, + patch_size=args.patch_size, + raw_mean=raw_mean, + raw_std=raw_std, + **ut.filter_kwargs(vis.plot_mask_pred, fig_kwargs), + ) + plots["mask_pred"] = vis.fig2pil(mask_pred_fig) + plt.close(mask_pred_fig) + + return plots + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--cfg-path", type=str, default=None) + parser.add_argument("--overrides", type=str, default=None, nargs="+") + args = parser.parse_args() + cfg = OmegaConf.load(DEFAULT_CONFIG) + if args.cfg_path: + cfg = OmegaConf.unsafe_merge(cfg, OmegaConf.load(args.cfg_path)) + if args.overrides: + cfg = OmegaConf.unsafe_merge(cfg, OmegaConf.from_dotlist(args.overrides)) + main(cfg) diff --git a/fomo_tune_baseline/output/task3/build/smri_mae/masking.py b/fomo_tune_baseline/output/task3/build/smri_mae/masking.py new file mode 100644 index 0000000000000000000000000000000000000000..28f9d53cbbd8d32a3923f9b1a6655c92cd526bf7 --- /dev/null +++ b/fomo_tune_baseline/output/task3/build/smri_mae/masking.py @@ -0,0 +1,80 @@ +import torch +from jaxtyping import Float, Int +from torch import Tensor + + +def pad_patch_mask( + patch_mask: Float[Tensor, "B N"], + mask_ratio: float, + shuffle: bool = False, + generator: torch.Generator | None = None, + pad_to_multiple: int | None = None, +) -> tuple[Float[Tensor, "B N"], Int[Tensor, "B L"], Tensor]: + """ + Select each row's own mask-ratio count, then pad ids to the batch max length. + + Returns: + - selected patch mask [B, N] + - padded selected patch ids [B, Lpad] + - token mask [B, Lpad], true for real ids and false for padding + """ + if not 0.0 <= mask_ratio <= 1.0: + raise ValueError(f"mask_ratio must be in [0, 1], got {mask_ratio}") + + B, N = patch_mask.shape + device = patch_mask.device + patch_mask = patch_mask.to(dtype=torch.bool) + + valid_counts = patch_mask.sum(dim=1) + num_keep = torch.floor(valid_counts.to(torch.float64) * (1.0 - mask_ratio)).to(torch.long) + if not shuffle: + selected = patch_mask & (patch_mask.cumsum(dim=1) <= num_keep.unsqueeze(1)) + padded_ids, token_mask = patch_ids_from_mask( + selected, + pad_to_multiple=pad_to_multiple, + ) + return selected, padded_ids, token_mask + + # One masked sort directly produces random valid IDs. The previous + # shuffle/select/inverse-shuffle path required two full argsorts plus a + # dynamic nonzero/scatter solely to recover the same selected set. + noise = torch.rand(B, N, generator=generator, device=device) + noise.masked_fill_(~patch_mask, torch.inf) + shuffled_ids = torch.argsort(noise, dim=1) + + max_count = int(num_keep.max().item()) + if pad_to_multiple is not None: + if pad_to_multiple <= 0: + raise ValueError(f"pad_to_multiple must be positive, got {pad_to_multiple}") + max_count = (max_count + pad_to_multiple - 1) // pad_to_multiple * pad_to_multiple + padded_ids = shuffled_ids[:, :max_count] + token_mask = torch.arange(max_count, device=device).unsqueeze(0) < num_keep.unsqueeze(1) + selected = torch.zeros_like(patch_mask).scatter_(1, padded_ids, token_mask) + return selected, padded_ids, token_mask + + +def patch_ids_from_mask( + patch_mask: Tensor, + pad_to_multiple: int | None = None, +) -> tuple[Int[Tensor, "B L"], Tensor]: + """Return optionally rounded patch IDs and their token-validity mask.""" + if pad_to_multiple is not None and pad_to_multiple <= 0: + raise ValueError(f"pad_to_multiple must be positive, got {pad_to_multiple}") + + patch_mask = patch_mask.to(dtype=torch.bool) + B, N = patch_mask.shape + device = patch_mask.device + counts = patch_mask.sum(dim=1) + max_count = int(counts.max().item()) + if pad_to_multiple is not None: + max_count = (max_count + pad_to_multiple - 1) // pad_to_multiple * pad_to_multiple + + patch_ids = torch.zeros((B, max_count), dtype=torch.long, device=device) + token_mask = torch.arange(max_count, device=device).unsqueeze(0) < counts.unsqueeze(1) + if max_count == 0: + return patch_ids, token_mask + + batch_ids, selected_ids = patch_mask.nonzero(as_tuple=True) + slot_ids = patch_mask.cumsum(dim=1)[batch_ids, selected_ids].to(torch.long) - 1 + patch_ids[batch_ids, slot_ids] = selected_ids + return patch_ids, token_mask diff --git a/fomo_tune_baseline/output/task3/build/smri_mae/model_mae.py b/fomo_tune_baseline/output/task3/build/smri_mae/model_mae.py new file mode 100644 index 0000000000000000000000000000000000000000..f9a4d884784e53f55ae557c46dd4afa2eb70e7d3 --- /dev/null +++ b/fomo_tune_baseline/output/task3/build/smri_mae/model_mae.py @@ -0,0 +1,916 @@ +# Copyright (c) Sophont, Inc +# This source code is licensed under the Apache License, Version 2.0 +# +# References: +# capi: https://github.com/facebookresearch/capi/blob/main/model.py +# timm: https://github.com/huggingface/pytorch-image-models/blob/v1.0.20/timm/models/vision_transformer.py + +""" +From-scratch re-implementation of the original MAE model. + +MaskedEncoder: standard ViT with masking +MaskedDecoder: standard self-attention MAE decoder +MaskedAutoEncoderViT: full MAE model for 3D structural MRI volumes +""" + +from collections.abc import Sequence +from typing import Literal + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch import Tensor +from torch.utils.checkpoint import checkpoint +from huggingface_hub import PyTorchModelHubMixin +from jaxtyping import Float, Int + +from .modules import ( + AbsolutePosEmbed, + Block, + LayerNorm, + Normalize, + JaggedBatch, + Patchify3D, + SeparablePosEmbed, + SinCosPosEmbed3D, + unpack_tokens, +) +from .masking import pad_patch_mask + + +class MaskedEncoder(nn.Module): + """ + Masked transformer encoder. + """ + + def __init__( + self, + patchify: nn.Module, + patch_embed: nn.Module, + pos_embed: nn.Module, + depth: int = 12, + embed_dim: int = 768, + num_heads: int = 12, + qkv_bias: bool = True, + proj_bias: bool = True, + mlp_ratio: int | float = 4, + class_token: bool = True, + reg_tokens: int = 0, + no_embed_class: bool = False, + final_norm: bool = True, + drop_path_rate: float = 0.0, + mask_drop_scale: bool = False, + ): + super().__init__() + self.num_prefix_tokens = int(class_token) + reg_tokens + self.num_reg_tokens = reg_tokens + self.has_class_token = class_token + self.no_embed_class = no_embed_class + + # scale inputs by 1 / observed rate (like dropout) + self.mask_drop_scale = mask_drop_scale + + # inject tokenization modules, so that the encoder doesn't specifically need to + # know how the data are tokenized, while still implementing a complete + # self-contained model. + self.patchify = patchify + self.patch_embed = patch_embed + self.pos_embed = pos_embed + + R = reg_tokens + self.cls_token = nn.Parameter(torch.empty(1, 1, embed_dim)) if class_token else None + self.reg_token = nn.Parameter(torch.empty(1, R, embed_dim)) if reg_tokens else None + + if not no_embed_class: + self.cls_token_pos = nn.Parameter(torch.empty(1, 1, embed_dim)) if class_token else None + self.reg_token_pos = nn.Parameter(torch.empty(1, R, embed_dim)) if reg_tokens else None + else: + self.cls_token_pos = self.reg_token_pos = None + + # stochastic depth decay rule + dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)] + + self.blocks = nn.ModuleList( + [ + Block( + dim=embed_dim, + num_heads=num_heads, + qkv_bias=qkv_bias, + proj_bias=proj_bias, + mlp_ratio=mlp_ratio, + drop_path=dpr[ii], + ) + for ii in range(depth) + ] + ) + + self.norm = LayerNorm(embed_dim) if final_norm else nn.Identity() + + self.reset_parameters() + + def extra_repr(self): + return ( + f"class_token={self.has_class_token}, reg_tokens={self.num_reg_tokens}, " + f"no_embed_class={self.no_embed_class}, mask_drop_scale={self.mask_drop_scale}" + ) + + def reset_parameters(self) -> None: + for p in [self.cls_token, self.cls_token_pos, self.reg_token, self.reg_token_pos]: + if p is not None: + nn.init.trunc_normal_(p, std=0.02) + + def cat_tokens(self, x: Tensor) -> Tensor: + # prepend cls and reg tokens with optional learned position embedding + # the cls and reg pos embedding is ofc redundant, but included in many other + # implementations. + B, _, _ = x.shape + + to_cat = [] + if self.has_class_token: + cls_token = self.cls_token + if not self.no_embed_class: + cls_token = cls_token + self.cls_token_pos + to_cat.append(cls_token.expand(B, -1, -1)) + + if self.num_reg_tokens: + reg_token = self.reg_token + if not self.no_embed_class: + reg_token = reg_token + self.reg_token_pos + to_cat.append(reg_token.expand(B, -1, -1)) + + if to_cat: + x = torch.cat(to_cat + [x], dim=1) + return x + + def cat_token_mask(self, token_mask: Tensor, batch_size: int) -> Tensor: + if self.num_prefix_tokens: + prefix_mask = torch.ones( + (batch_size, self.num_prefix_tokens), + dtype=torch.bool, + device=token_mask.device, + ) + token_mask = torch.cat([prefix_mask, token_mask], dim=1) + return token_mask + + def chunk_tokens(self, x: Tensor) -> tuple[Tensor | None, Tensor | None, Tensor]: + cls_offset = int(self.has_class_token) + cls = x[:, :cls_offset] if self.has_class_token else None + if self.num_reg_tokens: + reg = x[:, cls_offset : self.num_prefix_tokens, :] + else: + reg = None + patch = x[:, self.num_prefix_tokens :, :] + return cls, reg, patch + + def forward( + self, + x: Tensor, + mask: Tensor | None = None, + mask_ratio: float | None = None, + pad_to_multiple: int | None = None, + ) -> tuple[ + Float[Tensor, "B 1 D"] | None, + Float[Tensor, "B R D"] | None, + Float[Tensor, "B L D"], + Tensor | None, + Int[Tensor, "B L"] | None, + Tensor | None, + ]: + """ + x: input data shape [B, C, D, H, W] + mask: visible mask, 1 = visible, 0 = invisible. broadcastable shape + mask_ratio: mask ratio for uniform random masking + + returns: + - cls_embeds: [B, 1, D] + - reg_embeds: [B, R, D] + - patch_embeds: [B, L, D], where L is the number of visible patches + - mask: observed mask, 1 = observed, 0 = unobserved. same shape as input + - mask_ids: indices of visible patches [B L] + - token_mask: valid token mask for padded per-sample masking [B L] + """ + # apply mask to the input + if mask is not None: + mask = mask.to(device=x.device, dtype=torch.bool).expand_as(x) + x = x.masked_fill(~mask, 0) + + # patchify input + x = self.patchify(x) + B, N, P = x.shape + + # patchify mask and apply dropout style scaling + if mask is not None: + mask_patches = self.patchify(mask) + patch_num_obs = mask_patches.sum(dim=-1) + patch_mask = patch_num_obs > 0 + if self.mask_drop_scale: + patch_num_obs = patch_num_obs.to(x.dtype) + x = x * (P / patch_num_obs.unsqueeze(-1).clamp(min=1.0)) + elif mask_ratio is not None: + patch_mask = torch.ones((B, N), dtype=torch.bool, device=x.device) + mask_patches = patch_mask.unsqueeze(-1).expand(-1, -1, P) + else: + patch_mask = mask_patches = None + + # patch and position embed + x = self.patch_embed(x) + x = self.pos_embed(x) + + if mask is not None or mask_ratio is not None: + mask_ratio = 0.0 if mask_ratio is None else mask_ratio + patch_mask, mask_ids, token_mask = pad_patch_mask( + patch_mask, + mask_ratio=mask_ratio, + shuffle=mask_ratio > 0, + pad_to_multiple=pad_to_multiple, + ) + + mask_patches = mask_patches & patch_mask.unsqueeze(-1) + mask = self.patchify.unpatchify(mask_patches) + x = x.gather(1, mask_ids.unsqueeze(-1).expand(-1, -1, x.shape[-1])) + else: + mask_ids = None + token_mask = None + + cls_embeds, reg_embeds, patch_embeds = self.forward_patch_embeds( + x, + token_mask=token_mask, + ) + return cls_embeds, reg_embeds, patch_embeds, mask, mask_ids, token_mask + + def forward_patch_embeds( + self, + x: Float[Tensor, "B L D"], + token_mask: Tensor | None = None, + ) -> tuple[ + Float[Tensor, "B 1 D"] | None, + Float[Tensor, "B R D"] | None, + Float[Tensor, "B L D"], + ]: + B = x.shape[0] + if token_mask is None: + token_mask = torch.ones(x.shape[:2], dtype=torch.bool, device=x.device) + x = self.cat_tokens(x) + token_mask = self.cat_token_mask(token_mask, B) + jagged_batch = JaggedBatch.from_mask(token_mask) + x = x[token_mask] + for block in self.blocks: + x = block(x, jagged_batch=jagged_batch) + x = self.norm(x) + x = unpack_tokens(x, token_mask) + + cls_embeds, reg_embeds, patch_embeds = self.chunk_tokens(x) + return cls_embeds, reg_embeds, patch_embeds + + def forward_visible_ids( + self, + x: Tensor, + visible_ids: Int[Tensor, "B L"], + img_mask: Tensor | None = None, + ) -> tuple[ + Float[Tensor, "B 1 D"] | None, + Float[Tensor, "B R D"] | None, + Float[Tensor, "B L D"], + ]: + if img_mask is not None: + img_mask = img_mask.to(device=x.device, dtype=torch.bool).expand_as(x) + x = x.masked_fill(~img_mask, 0) + + x = self.patchify(x) + if self.mask_drop_scale and img_mask is not None: + mask_patches = self.patchify(img_mask) + patch_num_obs = mask_patches.sum(dim=-1).to(x.dtype) + x = x * (self.patchify.patch_dim / patch_num_obs.unsqueeze(-1).clamp(min=1.0)) + x = self.patch_embed(x) + x = self.pos_embed(x) + visible_ids = visible_ids.to(device=x.device) + x = x.gather(1, visible_ids.unsqueeze(-1).expand(-1, -1, x.shape[-1])) + return self.forward_patch_embeds(x) + + def forward_embedding( + self, + x: Tensor, + mask: Tensor | None = None, + mask_ratio: float | None = None, + ): + cls_embeds, reg_embeds, patch_embeds, *_ = self.forward( + x, + mask=mask, + mask_ratio=mask_ratio, + ) + return cls_embeds, reg_embeds, patch_embeds + + +class MaskedDecoder(nn.Module): + """Self-attention MAE decoder supporting sparse subset decoding via pred_ids.""" + + def __init__( + self, + pos_embed: nn.Module, + head: nn.Module | None = None, + input_dim: int | None = None, + depth: int = 12, + embed_dim: int = 768, + num_heads: int = 12, + qkv_bias: bool = True, + proj_bias: bool = True, + mlp_ratio: int | float = 4, + class_token: bool = True, + no_embed_class: bool = False, + final_norm: bool = True, + ): + super().__init__() + input_dim = embed_dim if input_dim is None else input_dim + self.has_class_token = class_token + self.no_embed_class = no_embed_class + + self.cls_token = nn.Parameter(torch.empty(1, 1, embed_dim)) if class_token else None + self.cls_token_pos = ( + nn.Parameter(torch.empty(1, 1, embed_dim)) + if class_token and not no_embed_class + else None + ) + self.mask_token = nn.Parameter(torch.empty(1, 1, embed_dim)) + + # decoder position embedding, encodes query position information into masks + self.pos_embed = pos_embed + + self.proj = nn.Identity() if input_dim == embed_dim else nn.Linear(input_dim, embed_dim) + + self.blocks = nn.ModuleList( + [ + Block( + dim=embed_dim, + num_heads=num_heads, + qkv_bias=qkv_bias, + proj_bias=proj_bias, + mlp_ratio=mlp_ratio, + ) + for _ in range(depth) + ] + ) + + self.norm = LayerNorm(embed_dim) if final_norm else nn.Identity() + + # optional injected prediction head + self.head = nn.Identity() if head is None else head + + self.reset_parameters() + + def extra_repr(self): + return f"class_token={self.has_class_token}, no_embed_class={self.no_embed_class}" + + def reset_parameters(self) -> None: + # official mae initializes decoder cls token to zeros + # although perhaps this was an oversight + if self.cls_token is not None: + nn.init.zeros_(self.cls_token) + if self.cls_token_pos is not None: + nn.init.trunc_normal_(self.cls_token_pos, std=0.02) + nn.init.trunc_normal_(self.mask_token, std=0.02) + + def cat_tokens(self, x: Tensor) -> Tensor: + if not self.has_class_token: + return x + cls_token = self.cls_token + if not self.no_embed_class: + cls_token = cls_token + self.cls_token_pos + return torch.cat([cls_token.expand(x.shape[0], -1, -1), x], dim=1) + + def cat_token_mask(self, token_mask: Tensor, batch_size: int) -> Tensor: + if self.has_class_token: + cls_mask = torch.ones( + (batch_size, 1), + dtype=torch.bool, + device=token_mask.device, + ) + token_mask = torch.cat([cls_mask, token_mask], dim=1) + return token_mask + + def chunk_tokens(self, x: Tensor) -> tuple[Tensor | None, Tensor]: + cls_offset = int(self.has_class_token) + cls = x[:, :cls_offset] if self.has_class_token else None + patch = x[:, cls_offset:, :] + return cls, patch + + def forward( + self, + embeds: Float[Tensor, "B L D"], + embed_ids: Int[Tensor, "B L"] | None = None, + pred_ids: Int[Tensor, "B Q"] | None = None, + embed_token_mask: Tensor | None = None, + pred_token_mask: Tensor | None = None, + packed_output: bool = False, + ) -> Float[Tensor, "B Q P"] | Float[Tensor, "T P"]: + """ + embeds: input patch embeddings. + embed_ids: optional patch indices for input embeddings. If not provided, no + position will be added to the embeddings. + pred_ids: patch indices of query mask positions. If None, decode *all* patches. + + returns: + - pred [B, Q, P] where Q is the number of prediction patches and P is the output + dimension + """ + B, L, _ = embeds.shape + + Q = self.pos_embed.num_patches if pred_ids is None else pred_ids.shape[1] + mask = self.mask_token.expand(B, Q, -1) + mask = self.pos_embed(mask, pos_ids=pred_ids) + + embeds = self.proj(embeds) + + if embed_ids is not None: + embeds = self.pos_embed(embeds, pos_ids=embed_ids) + if embed_token_mask is None: + embed_token_mask = torch.ones((B, L), dtype=torch.bool, device=embeds.device) + if pred_token_mask is None: + pred_token_mask = torch.ones((B, Q), dtype=torch.bool, device=embeds.device) + x = torch.cat([embeds, mask], dim=1) + token_mask = torch.cat([embed_token_mask, pred_token_mask], dim=1) + + x = self.cat_tokens(x) + token_mask = self.cat_token_mask(token_mask, B) + jagged_batch = JaggedBatch.from_mask(token_mask) + x = x[token_mask] + # Keep headroom for rare maximum-length PSP batches. + checkpoint_start = max(0, len(self.blocks) - 2) + for block_index, block in enumerate(self.blocks): + if self.training and torch.is_grad_enabled() and block_index >= checkpoint_start: + x = checkpoint(block, x, jagged_batch, use_reentrant=False) + else: + x = block(x, jagged_batch=jagged_batch) + + x = self.norm(x) + if packed_output: + pred_offset = int(self.has_class_token) + L + prediction_mask = F.pad(pred_token_mask, (pred_offset, 0)) + return self.head(x[prediction_mask[token_mask]]) + + x = unpack_tokens(x, token_mask) + _, x = self.chunk_tokens(x) + + pred = x[:, L:] + pred = pred.masked_fill(~pred_token_mask.unsqueeze(-1), 0) + pred = self.head(pred) + return pred + + +class MaskedAutoencoderViT(nn.Module, PyTorchModelHubMixin): + def __init__( + self, + img_size: int | tuple[int, int, int] = (208, 240, 208), + patch_size: int | tuple[int, int, int] = (16, 16, 16), + in_chans: int = 1, + depth: int = 12, + embed_dim: int = 768, + num_heads: int = 12, + decoder_depth: int = 4, + decoder_embed_dim: int | None = 512, + decoder_num_heads: int | None = 16, # default from mae, head dim = 32 + qkv_bias: bool = True, + proj_bias: bool = True, + mlp_ratio: int | float = 4, + class_token: bool = True, + reg_tokens: int = 0, + no_embed_class: bool = False, + drop_path_rate: float = 0.0, + mask_drop_scale: bool = False, + no_decode_pos: bool = False, + pos_embed: Literal["abs", "sep", "sincos"] = "sincos", + target_norm: Literal["none", "global", "slice", "patch"] | None = None, + ): + super().__init__() + img_size = _to_3d_tuple(img_size, "img_size") + patch_size = _to_3d_tuple(patch_size, "patch_size") + + self.no_decode_pos = no_decode_pos # don't pos encode embeddings in decoder + + # patchify reshapes input into sequence of flattened patches, shape [B, N, P] + ndim = 3 + patchify = Patchify3D(img_size, patch_size, in_chans=in_chans) + + # linear patch embedding P -> D + patch_embed = nn.Linear(patchify.patch_dim, embed_dim) + + # position embedding + # separable position embedding decouples the first spatial axis from the + # others. Fixed sin/cos embeddings are the default for sMRI volumes. + if pos_embed == "sincos": + pos_embed_layer = SinCosPosEmbed3D + else: + pos_embed_layer = {"abs": AbsolutePosEmbed, "sep": SeparablePosEmbed}[pos_embed] + pos_embed = pos_embed_layer(embed_dim, patchify.grid_size) + + # encoder. for inference, this model can be extracted and used like a regular vit + self.encoder = MaskedEncoder( + patchify=patchify, + patch_embed=patch_embed, + pos_embed=pos_embed, + depth=depth, + embed_dim=embed_dim, + num_heads=num_heads, + qkv_bias=qkv_bias, + proj_bias=proj_bias, + mlp_ratio=mlp_ratio, + class_token=class_token, + reg_tokens=reg_tokens, + no_embed_class=no_embed_class, + drop_path_rate=drop_path_rate, + mask_drop_scale=mask_drop_scale, + ) + + self.pred_patchify = patchify + + # fall back to encoder architecture width + decoder_embed_dim = decoder_embed_dim or embed_dim + decoder_num_heads = decoder_num_heads or num_heads + + decoder_pos_embed = pos_embed_layer(decoder_embed_dim, self.pred_patchify.grid_size) + # we might want to try tying the weights of the prediction head to the patch + # embedding at some point. + decoder_head = nn.Linear(decoder_embed_dim, self.pred_patchify.patch_dim) + + self.decoder = MaskedDecoder( + pos_embed=decoder_pos_embed, + head=decoder_head, + input_dim=embed_dim, + depth=decoder_depth, + embed_dim=decoder_embed_dim, + num_heads=decoder_num_heads, + qkv_bias=qkv_bias, + proj_bias=proj_bias, + mlp_ratio=mlp_ratio, + class_token=class_token, + no_embed_class=no_embed_class, + ) + + # mae style target normalization + # dim is relative to an unflattened embedding tensor of shape [B, *grid_size, D] + if target_norm not in {"none", None}: + norm_dim = { + "global": tuple(range(1, ndim + 2)), # full sequence + "slice": tuple(range(2, ndim + 2)), # each depth slice along first dim + "patch": -1, # normalize each patch independently (mae pix norm loss) + }[target_norm] + self.target_norm = Normalize(self.pred_patchify.grid_size, dim=norm_dim) + else: + self.target_norm = None + + self.init_weights() + + def extra_repr(self): + return f"no_decode_pos={self.no_decode_pos}" + + def init_weights(self): + self.apply(_init_weights) + + def prepare_targets(self, images: Tensor, img_mask: Tensor | None): + """ + images: [B, C, D, H, W] + img_mask: mask of valid data. only used for computing correct normalization + stats. same shape as images. + """ + targets_patches = self.pred_patchify(images) # [B, N, P] + + # target normalization + if self.target_norm is not None: + # full image data mask used for normalization stats only + if img_mask is not None: + img_mask_patches = self.pred_patchify(img_mask) + else: + img_mask_patches = None + targets_patches, *targets_stats = self.target_norm(targets_patches, img_mask_patches) + else: + targets_stats = None + + return targets_patches, targets_stats + + def prepare_masks( + self, + img_mask: Tensor, + visible_mask: Tensor | None, + pred_mask: Tensor | None, + device: torch.device, + ): + img_mask = img_mask.to(device=device, dtype=torch.bool) + + if visible_mask is None: + visible_mask = img_mask + else: + visible_mask = img_mask & visible_mask.to(device=device, dtype=torch.bool) + + if pred_mask is None: + pred_mask = img_mask + else: + pred_mask = img_mask & pred_mask.to(device=device, dtype=torch.bool) + + return img_mask, visible_mask, pred_mask + + def prepare_pred_mask( + self, + visible_mask: Tensor, + pred_mask: Tensor | None = None, + pred_mask_ratio: float | None = None, + pad_to_multiple: int | None = None, + ): + """ + prepare prediction mask by removing visible content + visible_mask: [B, C, D, H, W], 1 = visible, 0 = invisible + pred_mask: same shape, 1 = predict, 0 = don't predict + """ + if pred_mask is None: + pred_mask = torch.ones_like(visible_mask) + + pred_mask = pred_mask & ~visible_mask + + pred_mask_patches = self.pred_patchify(pred_mask) + pred_patch_mask = pred_mask_patches.any(dim=-1) + # Optionally subsample the prediction candidates. + mask_ratio = 0.0 if pred_mask_ratio is None else pred_mask_ratio + pred_patch_mask, pred_ids, pred_token_mask = pad_patch_mask( + pred_patch_mask, + mask_ratio=mask_ratio, + # With per-sample padding every candidate is retained when the ratio + # is zero, so randomizing their order is pure overhead. + shuffle=mask_ratio > 0, + pad_to_multiple=pad_to_multiple, + ) + pred_mask_patches = pred_mask_patches & pred_patch_mask.unsqueeze(-1) + return pred_mask_patches, pred_ids, pred_token_mask + + def forward_decoder( + self, + patch_embeds: Float[Tensor, "B L D"], + visible_ids: Int[Tensor, "B L"], + pred_ids: Int[Tensor, "B Q"] | None, + visible_token_mask: Tensor | None = None, + pred_token_mask: Tensor | None = None, + packed_output: bool = False, + ) -> Float[Tensor, "B Q P"] | Float[Tensor, "T P"]: + return self.decoder( + patch_embeds, + embed_ids=None if self.no_decode_pos else visible_ids, + pred_ids=pred_ids, + embed_token_mask=visible_token_mask, + pred_token_mask=pred_token_mask, + packed_output=packed_output, + ) + + def forward_loss( + self, + preds: Float[Tensor, "T P"], + targets_patches: Float[Tensor, "B N P"], + pred_mask_patches: Float[Tensor, "B N P"], + pred_ids: Int[Tensor, "B Q"], + pred_token_mask: Tensor, + ) -> Tensor: + """Average valid-voxel MSE within each scan, then average across scans.""" + batch_ids, slot_ids = pred_token_mask.nonzero(as_tuple=True) + patch_ids = pred_ids[batch_ids, slot_ids] + targets = targets_patches[batch_ids, patch_ids] + voxel_mask = pred_mask_patches[batch_ids, patch_ids] + + patch_errors = ((preds - targets) ** 2 * voxel_mask).sum(dim=1) + patch_voxels = voxel_mask.sum(dim=1).to(dtype=patch_errors.dtype) + batch_size = targets_patches.shape[0] + scan_errors = patch_errors.new_zeros(batch_size).scatter_add_(0, batch_ids, patch_errors) + scan_voxels = patch_voxels.new_zeros(batch_size).scatter_add_(0, batch_ids, patch_voxels) + return (scan_errors / scan_voxels).mean() + + @torch.no_grad() + def forward_pred_images( + self, + preds: Float[Tensor, "B Q P"], + pred_ids: Int[Tensor, "B Q"], + pred_token_mask: Tensor | None = None, + img_mask: Tensor | None = None, + targets_stats: tuple[Tensor, Tensor] | None = None, + ) -> Tensor: + B, _, P = preds.shape + N = self.pred_patchify.num_patches + if pred_token_mask is not None: + preds = preds.masked_fill(~pred_token_mask.unsqueeze(-1), 0) + + preds = torch.zeros((B, N, P), dtype=preds.dtype, device=preds.device).scatter_add_( + 1, pred_ids.unsqueeze(-1).expand(-1, -1, P), preds + ) + + if targets_stats is not None: + targets_mean, targets_std = targets_stats + preds = preds * targets_std + targets_mean + + pred_images = self.pred_patchify.unpatchify(preds) + if img_mask is not None: + pred_images = pred_images.masked_fill(~img_mask, 0) + return pred_images + + def forward( + self, + images: Tensor, + img_mask: Tensor, + mask_ratio: float, + pred_mask_ratio: float | None = None, + pad_to_multiple: int | None = None, + with_state: bool = True, + ) -> Tensor | tuple[Tensor, dict]: + img_mask, visible_mask, pred_mask = self.prepare_masks( + img_mask, + None, + None, + device=images.device, + ) + targets_patches, targets_stats = self.prepare_targets(images, img_mask) + + ( + cls_embeds, + reg_embeds, + patch_embeds, + visible_mask, + visible_ids, + visible_token_mask, + ) = self.encoder( + images, + mask=visible_mask, + mask_ratio=mask_ratio, + pad_to_multiple=pad_to_multiple, + ) + + pred_mask_patches, pred_ids, pred_token_mask = self.prepare_pred_mask( + visible_mask, + pred_mask=pred_mask, + pred_mask_ratio=pred_mask_ratio, + pad_to_multiple=pad_to_multiple, + ) + + preds = self.forward_decoder( + patch_embeds, + visible_ids, + pred_ids, + visible_token_mask=visible_token_mask, + pred_token_mask=pred_token_mask, + packed_output=not with_state, + ) + + loss_preds = preds if not with_state else preds[pred_token_mask] + loss = self.forward_loss( + loss_preds, + targets_patches, + pred_mask_patches, + pred_ids, + pred_token_mask, + ) + + if not with_state: + return loss + + pred_mask = self.pred_patchify.unpatchify(pred_mask_patches) + pred_images = self.forward_pred_images( + preds, + pred_ids, + pred_token_mask=pred_token_mask, + img_mask=img_mask, + targets_stats=targets_stats, + ) + + state = { + "targets_patches": targets_patches, + "targets_stats": targets_stats, + "patch_embeds": patch_embeds, + "cls_embeds": cls_embeds, + "reg_embeds": reg_embeds, + "visible_mask": visible_mask, + "visible_ids": visible_ids, + "visible_token_mask": visible_token_mask, + "pred_mask": pred_mask, + "pred_ids": pred_ids, + "pred_token_mask": pred_token_mask, + "preds": preds, + "pred_images": pred_images, + } + return loss, state + + def forward_embedding( + self, + x: Tensor, + mask: Tensor | None = None, + mask_ratio: float | None = None, + ): + return self.encoder.forward_embedding(x, mask=mask, mask_ratio=mask_ratio) + + +class MaskedViT(MaskedEncoder, PyTorchModelHubMixin): + def __init__( + self, + img_size: int | tuple[int, int, int] = (208, 240, 208), + in_chans: int = 1, + patch_size: int | tuple[int, int, int] = (16, 16, 16), + depth: int = 12, + embed_dim: int = 768, + num_heads: int = 12, + qkv_bias: bool = True, + proj_bias: bool = True, + mlp_ratio: int | float = 4, + class_token: bool = True, + reg_tokens: int = 0, + no_embed_class: bool = False, + final_norm: bool = True, + drop_path_rate: float = 0.0, + mask_drop_scale: bool = False, + pos_embed: Literal["abs", "sep", "sincos"] = "sincos", + ): + img_size = _to_3d_tuple(img_size, "img_size") + patch_size = _to_3d_tuple(patch_size, "patch_size") + + patchify = Patchify3D(img_size, patch_size, in_chans=in_chans) + patch_embed = nn.Linear(patchify.patch_dim, embed_dim) + if pos_embed == "sincos": + pos_embed_layer = SinCosPosEmbed3D + else: + pos_embed_layer = {"abs": AbsolutePosEmbed, "sep": SeparablePosEmbed}[pos_embed] + pos_embed = pos_embed_layer(embed_dim, patchify.grid_size) + + super().__init__( + patchify=patchify, + patch_embed=patch_embed, + pos_embed=pos_embed, + depth=depth, + embed_dim=embed_dim, + num_heads=num_heads, + qkv_bias=qkv_bias, + proj_bias=proj_bias, + mlp_ratio=mlp_ratio, + class_token=class_token, + reg_tokens=reg_tokens, + no_embed_class=no_embed_class, + final_norm=final_norm, + drop_path_rate=drop_path_rate, + mask_drop_scale=mask_drop_scale, + ) + + self.init_weights() + + def init_weights(self): + self.apply(_init_weights) + + +def _to_3d_tuple(value: int | Sequence[int], name: str) -> tuple[int, int, int]: + if isinstance(value, int): + return (value, value, value) + if len(value) != 3: + raise ValueError(f"{name} must have exactly 3 spatial dimensions, got {tuple(value)}") + return tuple(int(item) for item in value) + + +# JAX ViT xavier uniform init +# https://github.com/facebookresearch/capi/blob/main/model.py +def _init_weights(m: nn.Module) -> None: + if isinstance(m, nn.Linear): + nn.init.xavier_uniform_(m.weight) + if m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.LayerNorm) and m.elementwise_affine: + nn.init.constant_(m.weight, 1.0) + if m.bias is not None: + nn.init.constant_(m.bias, 0) + + +def _create_vit(**kwargs): + model = MaskedViT(**kwargs) + return model + + +def _create_mae_vit(**kwargs): + model = MaskedAutoencoderViT(**kwargs) + return model + + +def mae_vit_small(**kwargs): + model_args = dict(embed_dim=384, depth=12, num_heads=6) + return _create_mae_vit(**model_args, **kwargs) + + +def mae_vit_base(**kwargs): + model_args = dict(embed_dim=768, depth=12, num_heads=12) + return _create_mae_vit(**model_args, **kwargs) + + +def mae_vit_large(**kwargs): + model_args = dict(embed_dim=1024, depth=24, num_heads=16) + return _create_mae_vit(**model_args, **kwargs) + + +def mae_vit_huge(**kwargs): + model_args = dict(embed_dim=1280, depth=32, num_heads=16) + return _create_mae_vit(**model_args, **kwargs) + + +# "patch embed" baseline model, depth 0 ViT (hah) +def patch_embed_small(**kwargs): + model_args = dict(embed_dim=384, depth=0) + return _create_vit(**model_args, **kwargs) + + +def patch_embed_base(**kwargs): + model_args = dict(embed_dim=768, depth=0) + return _create_vit(**model_args, **kwargs) diff --git a/fomo_tune_baseline/output/task3/build/smri_mae/modules.py b/fomo_tune_baseline/output/task3/build/smri_mae/modules.py new file mode 100644 index 0000000000000000000000000000000000000000..828aecb1bd57182f1690feb4e943395db1bdd515 --- /dev/null +++ b/fomo_tune_baseline/output/task3/build/smri_mae/modules.py @@ -0,0 +1,453 @@ +# This source code is licensed under the Apache License, Version 2.0 +# +# References: +# capi: https://github.com/facebookresearch/capi/blob/main/model.py +# timm: https://github.com/huggingface/pytorch-image-models/blob/v1.0.20/timm/models/vision_transformer.py +# vjepa2: https://github.com/facebookresearch/vjepa2/blob/main/src/models/utils/pos_embs.py + +import math +from functools import partial +from typing import NamedTuple, Type + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch import Tensor +from einops import rearrange +from jaxtyping import Float, Int +from timm.layers import DropPath, to_3tuple + +Layer = Type[nn.Module] + + +class JaggedBatch(NamedTuple): + """Sequence boundaries and cached launch metadata for jagged attention.""" + + offsets: Tensor + max_seqlen: int + + @classmethod + def from_mask(cls, mask: Tensor) -> "JaggedBatch": + mask = mask.to(dtype=torch.bool) + counts = mask.sum(dim=1) + return cls( + offsets=F.pad(counts.cumsum(dim=0), (1, 0)), + max_seqlen=mask.shape[1], + ) + + def as_nested(self, tokens: Tensor) -> Tensor: + # Cached conservative bounds avoid min/max reductions and GPU-to-CPU + # synchronization when Flash SDPA inspects the jagged sequence lengths. + return torch.nested.nested_tensor_from_jagged( + tokens, + self.offsets, + min_seqlen=1, + max_seqlen=self.max_seqlen, + ).transpose(1, 2) + + +def unpack_tokens(tokens: Tensor, token_mask: Tensor) -> Tensor: + """Restore packed values to a padded batch, filling invalid slots with zero.""" + output = tokens.new_zeros((*token_mask.shape, *tokens.shape[1:])) + return output.index_put((token_mask,), tokens) + + +def jagged_scaled_dot_product_attention( + query: Tensor, + key: Tensor, + value: Tensor, + jagged_batch: JaggedBatch, +) -> Tensor: + """Run SDPA on a packed batch of variable-length sequences.""" + output_jagged = F.scaled_dot_product_attention( + jagged_batch.as_nested(query), + jagged_batch.as_nested(key), + jagged_batch.as_nested(value), + ) + return output_jagged.transpose(1, 2).values() + + +class Attention(nn.Module): + def __init__( + self, + dim: int, + num_heads: int, + qkv_bias: bool = False, + proj_bias: bool = False, + ) -> None: + super().__init__() + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.qkv = nn.Linear(dim, 3 * dim, bias=qkv_bias) + self.proj = nn.Linear(dim, dim, bias=proj_bias) + + def extra_repr(self): + return f"num_heads={self.num_heads}" + + def forward( + self, + x: Float[Tensor, "L D"], + jagged_batch: JaggedBatch, + ) -> Float[Tensor, "L D"]: + L, D = x.shape + h, dh = self.num_heads, self.head_dim + + qkv = self.qkv(x).reshape(L, 3, h, dh) + q, k, v = qkv.unbind(1) + + x = jagged_scaled_dot_product_attention( + q, + k, + v, + jagged_batch=jagged_batch, + ) + x = x.reshape(L, D) + x = self.proj(x) + return x + + +class Mlp(nn.Module): + def __init__( + self, + dim: int, + mlp_ratio: int | float = 4, + bias: bool = False, + ) -> None: + super().__init__() + hidden_features = int(dim * mlp_ratio) + self.fc1 = nn.Linear(dim, hidden_features, bias=bias) + self.act = nn.GELU() + self.fc2 = nn.Linear(hidden_features, dim, bias=bias) + + def forward(self, x: Float[Tensor, "... D"]) -> Float[Tensor, "... D"]: + x = self.fc1(x) + x = self.act(x) + x = self.fc2(x) + return x + + +# timm default eps=1e-6 +LayerNorm = partial(nn.LayerNorm, eps=1e-6) + + +class Block(nn.Module): + def __init__( + self, + dim: int, + num_heads: int, + qkv_bias: bool = False, + proj_bias: bool = False, + mlp_ratio: int | float = 4, + drop_path: float = 0.0, + norm_layer: Layer = LayerNorm, + ) -> None: + super().__init__() + self.norm1 = norm_layer(dim) + self.attn = Attention( + dim=dim, + num_heads=num_heads, + qkv_bias=qkv_bias, + proj_bias=proj_bias, + ) + self.drop_path1 = DropPath(drop_path) if drop_path > 0 else nn.Identity() + + self.norm2 = norm_layer(dim) + self.mlp = Mlp( + dim=dim, + mlp_ratio=mlp_ratio, + bias=proj_bias, + ) + self.drop_path2 = DropPath(drop_path) if drop_path > 0 else nn.Identity() + + def forward( + self, + x: Float[Tensor, "L D"], + jagged_batch: JaggedBatch, + ) -> Float[Tensor, "L D"]: + x = x + self.drop_path1( + self.attn( + self.norm1(x), + jagged_batch=jagged_batch, + ) + ) + x = x + self.drop_path2(self.mlp(self.norm2(x))) + return x + + +# Patching and position embedding modules + + +class Patchify3D(nn.Module): + def __init__( + self, + img_size: int | tuple[int, int, int], + patch_size: int | tuple[int, int, int], + in_chans: int = 3, + ) -> None: + super().__init__() + self.img_size = to_3tuple(img_size) + self.patch_size = to_3tuple(patch_size) + self.in_chans = in_chans + + T, H, W = self.img_size + p_t, p_h, p_w = self.patch_size + if T % p_t or H % p_h or W % p_w: + raise ValueError( + f"img_size {self.img_size} must be divisible by patch_size {self.patch_size}" + ) + self.grid_size = (T // p_t, H // p_h, W // p_w) + self.num_patches = math.prod(self.grid_size) + self.patch_dim = in_chans * math.prod(self.patch_size) + + def forward(self, x: Float[Tensor, "B C T H W"]) -> Float[Tensor, "B N P"]: + x = patchify3d(x, self.patch_size) + return x + + def unpatchify(self, x: Float[Tensor, "B N P"]) -> Float[Tensor, "B C T H W"]: + x = unpatchify3d(x, patch_size=self.patch_size, img_size=self.img_size) + return x + + def extra_repr(self): + return f"{self.img_size}, {self.patch_size}, in_chans={self.in_chans}" + + +def patchify3d(x: Tensor, patch_size: tuple[int, int, int]) -> Tensor: + p_t, p_h, p_w = to_3tuple(patch_size) + B, C, T, H, W = x.shape + x = rearrange(x, "b c (t u) (h p) (w q) -> b (t h w) (c u p q)", u=p_t, p=p_h, q=p_w) + return x + + +def unpatchify3d( + x: Tensor, + patch_size: tuple[int, int, int], + img_size: tuple[int, int, int], +) -> Tensor: + B, N, P = x.shape + p_t, p_h, p_w = to_3tuple(patch_size) + T, H, W = to_3tuple(img_size) + x = rearrange( + x, + "b (t h w) (c u p q) -> b c (t u) (h p) (w q)", + t=T // p_t, + h=H // p_h, + w=W // p_w, + u=p_t, + p=p_h, + q=p_w, + ) + return x + + +class AbsolutePosEmbed(nn.Module): + def __init__(self, embed_dim: int, grid_size: tuple[int, ...]) -> None: + super().__init__() + self.embed_dim = embed_dim + self.grid_size = grid_size + self.num_patches = math.prod(grid_size) + + self.weight = nn.Parameter(torch.empty(self.num_patches, embed_dim)) + self.reset_parameters() + + def reset_parameters(self): + nn.init.trunc_normal_(self.weight, std=0.02) + + def forward( + self, + x: Float[Tensor, "B L D"], + pos_ids: Int[Tensor, "B L"] | None = None, + ) -> Float[Tensor, "B L D"]: + x = apply_pos_embed(x, self.weight, pos_ids=pos_ids) + return x + + def extra_repr(self): + return f"{self.embed_dim}, {self.grid_size}" + + +class SeparablePosEmbed(nn.Module): + def __init__(self, embed_dim: int, grid_size: tuple[int, ...]) -> None: + super().__init__() + self.embed_dim = embed_dim + self.grid_size = grid_size + self.num_patches = math.prod(grid_size) + + N_t, *grid_size_spatial = grid_size + N_s = math.prod(grid_size_spatial) + self.weight_spatial = nn.Parameter(torch.empty(1, N_s, embed_dim)) + self.weight_temporal = nn.Parameter(torch.empty(N_t, 1, embed_dim)) + self.reset_parameters() + + def reset_parameters(self): + nn.init.trunc_normal_(self.weight_spatial, std=0.02) + nn.init.trunc_normal_(self.weight_temporal, std=0.02) + + def forward( + self, + x: Float[Tensor, "B L D"], + pos_ids: Int[Tensor, "B L"] | None = None, + ) -> Float[Tensor, "B L D"]: + B, N, D = x.shape + weight = (self.weight_temporal + self.weight_spatial).flatten(0, 1) # [N, D] + x = apply_pos_embed(x, weight, pos_ids=pos_ids) + return x + + def extra_repr(self): + return f"{self.embed_dim}, {self.grid_size}" + + +class SinCosPosEmbed3D(nn.Module): + def __init__(self, embed_dim: int, grid_size: tuple[int, int, int]) -> None: + super().__init__() + self.embed_dim = embed_dim + self.grid_size = grid_size + self.num_patches = math.prod(grid_size) + + N_t, N_h, N_w = grid_size + weight = get_3d_sincos_pos_embed( + embed_dim=embed_dim, + grid_size=(N_h, N_w), + grid_depth=N_t, + uniform_power=True, + ) + self.weight = nn.Parameter(torch.from_numpy(weight).float(), requires_grad=False) + + def forward( + self, + x: Float[Tensor, "B L D"], + pos_ids: Int[Tensor, "B L"] | None = None, + ) -> Float[Tensor, "B L D"]: + x = apply_pos_embed(x, self.weight, pos_ids=pos_ids) + return x + + def extra_repr(self): + return f"{self.embed_dim}, {self.grid_size}" + + +# sincos pos embed utils from vjepa2, but fixed the confusing meshgrid indexing + + +def get_3d_sincos_pos_embed(embed_dim, grid_size, grid_depth, cls_token=False, uniform_power=False): + """ + grid_size: tuple of int of the grid height and width + grid_depth: int of the grid depth + returns: + pos_embed: [grid_depth*grid_height*grid_width, embed_dim] (w/o cls_token) + or [1+grid_depth*grid_height*grid_width, embed_dim] (w/ cls_token) + """ + grid_d = np.arange(grid_depth, dtype=float) + grid_h = np.arange(grid_size[0], dtype=float) + grid_w = np.arange(grid_size[1], dtype=float) + grid_d, grid_h, grid_w = np.meshgrid(grid_d, grid_h, grid_w, indexing="ij") + + if not uniform_power: + h_embed_dim = embed_dim // 4 + w_embed_dim = embed_dim // 4 + d_embed_dim = embed_dim // 2 + else: + h_embed_dim = w_embed_dim = d_embed_dim = int(np.ceil(embed_dim / 6) * 2) + + emb_h = get_1d_sincos_pos_embed_from_grid(h_embed_dim, grid_h) # (T*H*W, D1) + emb_w = get_1d_sincos_pos_embed_from_grid(w_embed_dim, grid_w) # (T*H*W, D2) + emb_d = get_1d_sincos_pos_embed_from_grid(d_embed_dim, grid_d) # (T*H*W, D3) + pos_embed = np.concatenate([emb_d, emb_h, emb_w], axis=1) + pos_embed = pos_embed[:, :embed_dim] + if cls_token: + pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0) + return pos_embed + + +def get_1d_sincos_pos_embed_from_grid(embed_dim, pos): + """ + embed_dim: output dimension for each position + pos: a list of positions to be encoded: size (M,) + returns: (M, D) + """ + assert embed_dim % 2 == 0 + omega = np.arange(embed_dim // 2, dtype=float) + omega /= embed_dim / 2.0 + omega = 1.0 / 10000**omega # (D/2,) + + pos = pos.reshape(-1) # (M,) + out = np.einsum("m,d->md", pos, omega) # (M, D/2), outer product + + emb_sin = np.sin(out) # (M, D/2) + emb_cos = np.cos(out) # (M, D/2) + + emb = np.concatenate([emb_sin, emb_cos], axis=1) # (M, D) + return emb + + +def apply_pos_embed( + x: Float[Tensor, "B L D"], + weight: Float[Tensor, "N D"], + pos_ids: Int[Tensor, "B L"] | None = None, +) -> Float[Tensor, "B L D"]: + B, L, D = x.shape + weight = weight.expand(B, -1, -1) + if pos_ids is not None: + weight = weight.gather(1, pos_ids.unsqueeze(-1).expand(-1, -1, D)) + x = x + weight + return x + + +# (masked) normalization used for MAE target normalization + + +class Normalize(nn.Module): + def __init__( + self, + grid_size: tuple[int, ...], + dim: int | tuple[int, ...] | None = -1, + eps: float = 1e-6, + ) -> None: + super().__init__() + self.grid_size = grid_size + self.dim = dim + self.eps = eps + + def forward(self, x: Tensor, mask: Tensor | None = None) -> tuple[Tensor, Tensor, Tensor]: + """ + Normalize input sequence along dim(s) after reshaping to grid. + Returns tuple of (x, mean, std). + """ + B, N, D = x.shape + x = x.reshape((B, *self.grid_size, D)) + if mask is not None: + mask = mask.reshape((B, *self.grid_size, D)) + x, mean, std = masked_normalize(x, mask, dim=self.dim, eps=self.eps) + else: + x, mean, std = normalize(x, dim=self.dim, eps=self.eps) + mean = mean.expand_as(x).reshape(B, N, D) + std = std.expand_as(x).reshape(B, N, D) + x = x.reshape(B, N, D) + return x, mean, std + + def extra_repr(self): + return f"{self.grid_size}, dim={self.dim}" + + +def masked_normalize( + x: Tensor, + mask: Tensor, + dim: int | tuple[int, ...] | None = -1, + eps: float = 1e-6, +) -> tuple[Tensor, Tensor, Tensor]: + num_obs = mask.sum(dim=dim, keepdim=True).clamp(min=1) + mean = (mask * x).sum(dim=dim, keepdim=True) / num_obs + var = (mask * (x - mean) ** 2).sum(dim=dim, keepdim=True) / num_obs + std = (var + eps) ** 0.5 + x = mask * (x - mean) / std + return x, mean, std + + +def normalize( + x: Tensor, + dim: int | tuple[int, ...] | None = -1, + eps: float = 1e-6, +) -> tuple[Tensor, Tensor, Tensor]: + mean = x.mean(dim=dim, keepdim=True) + var = torch.var(x, dim=dim, keepdim=True, unbiased=False) + std = (var + eps) ** 0.5 + x = (x - mean) / std + return x, mean, std diff --git a/fomo_tune_baseline/output/task3/build/smri_mae/utils.py b/fomo_tune_baseline/output/task3/build/smri_mae/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..0f6274f813429552c9ddfccfe4033ae677b52342 --- /dev/null +++ b/fomo_tune_baseline/output/task3/build/smri_mae/utils.py @@ -0,0 +1,581 @@ +# Copyright (c) Sophont, Inc +# This source code is licensed under the Apache License, Version 2.0 +# +# References: +# deit: https://github.com/facebookresearch/deit/blob/main/utils.py +# beit3: https://github.com/microsoft/unilm/blob/master/beit3/utils.py +# capi: https://github.com/facebookresearch/capi/blob/main/utils.py +# dinov2: https://github.com/facebookresearch/dinov2/blob/main/dinov2/utils/param_groups.py +# timm: https://github.com/huggingface/pytorch-image-models/blob/main/timm/utils/cuda.py +# dino: https://github.com/facebookresearch/dino/blob/main/utils.py + +import datetime +import inspect +import math +import os +import random +import subprocess +import time +from collections import defaultdict, deque +from omegaconf import OmegaConf +from pathlib import Path + +import numpy as np +import torch +import torch.distributed as dist +import torch.nn as nn +from torch import Tensor +from torch.amp import GradScaler +from torch.optim import Optimizer + + +# these very useful utils copied from deit with only minor changes +# thanks to the original authors, wherever you are + + +def configure_flash_sdpa() -> None: + """Use Flash Attention exclusively for CUDA SDPA.""" + torch.backends.cuda.enable_flash_sdp(True) + torch.backends.cuda.enable_mem_efficient_sdp(False) + torch.backends.cuda.enable_math_sdp(False) + torch.backends.cuda.enable_cudnn_sdp(False) + print("SDPA backend: flash") + + +class SmoothedValue: + """Track a series of values and provide access to smoothed values over a + window or the global series average. + """ + + def __init__(self, window_size=20, fmt=None): + if fmt is None: + fmt = "{median:.4f} ({global_avg:.4f})" + self.deque = deque(maxlen=window_size) + self.total = 0.0 + self.count = 0 + self.fmt = fmt + + def update(self, value, n=1): + value = float(value) + if math.isfinite(value): + self.deque.append(value) + self.count += n + self.total += value * n + + def synchronize_between_processes(self): + """ + Warning: does not synchronize the deque! + """ + if not is_dist_avail_and_initialized(): + return + t = torch.tensor([self.count, self.total], dtype=torch.float64, device="cuda") + dist.barrier() + dist.all_reduce(t) + t = t.tolist() + self.count = int(t[0]) + self.total = t[1] + + @property + def median(self): + if not self.count: + return float("nan") + d = torch.tensor(list(self.deque)) + return d.median().item() + + @property + def avg(self): + if not self.count: + return float("nan") + d = torch.tensor(list(self.deque), dtype=torch.float32) + return d.mean().item() + + @property + def global_avg(self): + if not self.count: + return float("nan") + return self.total / self.count + + @property + def max(self): + if not self.count: + return float("nan") + return max(self.deque) + + @property + def value(self): + if not self.count: + return float("nan") + return self.deque[-1] + + def __str__(self): + return self.fmt.format( + median=self.median, + avg=self.avg, + global_avg=self.global_avg, + max=self.max, + value=self.value, + ) + + +class MetricLogger: + def __init__(self, delimiter="\t"): + self.meters = defaultdict(SmoothedValue) + self.delimiter = delimiter + + def update(self, **kwargs): + for k, v in kwargs.items(): + if v is None: + continue + if isinstance(v, (torch.Tensor, np.generic)): + v = v.item() + assert isinstance(v, (float, int)) + self.meters[k].update(v) + + def __getattr__(self, attr): + if attr in self.meters: + return self.meters[attr] + if attr in self.__dict__: + return self.__dict__[attr] + raise AttributeError("'{}' object has no attribute '{}'".format(type(self).__name__, attr)) + + def __str__(self): + loss_str = [] + for name, meter in self.meters.items(): + loss_str.append("{}: {}".format(name, str(meter))) + return self.delimiter.join(loss_str) + + def synchronize_between_processes(self): + for meter in self.meters.values(): + meter.synchronize_between_processes() + + def add_meter(self, name, meter): + self.meters[name] = meter + + def log_every(self, iterable, print_freq, header=None, total_steps=None): + i = 0 + total_steps = total_steps or len(iterable) + if not header: + header = "" + start_time = time.time() + end = time.time() + iter_time = SmoothedValue(fmt="{avg:.4f}") + data_time = SmoothedValue(fmt="{avg:.4f}") + space_fmt = ":" + str(len(str(total_steps))) + "d" + log_msg = [ + header, + "[{0" + space_fmt + "}/{1}]", + "eta: {eta}", + "{meters}", + "time: {time}", + "data: {data}", + ] + if torch.cuda.is_available(): + log_msg.append("max mem: {memory:.0f}") + log_msg = self.delimiter.join(log_msg) + MB = 1024.0 * 1024.0 + for obj in iterable: + if i >= total_steps: + break + data_time.update(time.time() - end) + yield obj + iter_time.update(time.time() - end) + if i % print_freq == 0 or i == total_steps - 1: + eta_seconds = iter_time.global_avg * (total_steps - i) + eta_string = str(datetime.timedelta(seconds=int(eta_seconds))) + if torch.cuda.is_available(): + print( + log_msg.format( + i, + total_steps, + eta=eta_string, + meters=str(self), + time=str(iter_time), + data=str(data_time), + memory=torch.cuda.max_memory_allocated() / MB, + ) + ) + + else: + print( + log_msg.format( + i, + total_steps, + eta=eta_string, + meters=str(self), + time=str(iter_time), + data=str(data_time), + ) + ) + i += 1 + end = time.time() + total_time = time.time() - start_time + total_time_str = str(datetime.timedelta(seconds=int(total_time))) + print( + "{} Total time: {} ({:.4f} s / it)".format( + header, total_time_str, total_time / total_steps + ) + ) + + +def setup_for_distributed(log_path=None): + """ + This function disables printing when not in master process + """ + import builtins as __builtin__ + + builtin_print = __builtin__.print + + is_master = is_main_process() + + def print(*args, **kwargs): + force = kwargs.pop("force", False) + if is_master or force: + builtin_print(*args, **kwargs) + # tee to log file + if log_path and "file" not in kwargs: + with open(log_path, "a") as f: + builtin_print(*args, file=f, **kwargs) + + __builtin__.print = print + + +def is_dist_avail_and_initialized(): + if not dist.is_available(): + return False + if not dist.is_initialized(): + return False + return True + + +def get_world_size(): + if not is_dist_avail_and_initialized(): + return 1 + return dist.get_world_size() + + +def get_rank(): + if not is_dist_avail_and_initialized(): + return 0 + return dist.get_rank() + + +def is_main_process(): + return get_rank() == 0 + + +def save_on_master(obj, path, *args, **kwargs): + if is_main_process(): + path = Path(path) + tmp_path = path.with_name(f".{path.name}.tmp-{os.getpid()}") + try: + torch.save(obj, tmp_path, *args, **kwargs) + os.replace(tmp_path, path) + except Exception: + tmp_path.unlink(missing_ok=True) + raise + + +def init_distributed_mode(args): + # removed slurm block, can add if we use slurm + if "RANK" in os.environ and "WORLD_SIZE" in os.environ: + args.rank = int(os.environ["RANK"]) + args.world_size = int(os.environ["WORLD_SIZE"]) + args.gpu = int(os.environ["LOCAL_RANK"]) + else: + args.distributed = False + return + + args.distributed = True + + torch.cuda.set_device(args.gpu) + args.dist_backend = "nccl" + print(f"| distributed init (rank {args.rank})") + torch.distributed.init_process_group( + backend=args.dist_backend, + world_size=args.world_size, + rank=args.rank, + device_id=args.gpu, + ) + torch.distributed.barrier() + + +# checkpoint saving utils adapted from beit3 + + +def capture_rng_state() -> dict: + state = {"torch": torch.get_rng_state()} + if torch.cuda.is_available(): + state["cuda"] = torch.cuda.get_rng_state() + return state + + +def restore_rng_state(state: dict) -> None: + torch.set_rng_state(state["torch"]) + if "cuda" in state and torch.cuda.is_available(): + torch.cuda.set_rng_state(state["cuda"]) + + +def _all_rank_rng_states() -> list[dict]: + local_state = capture_rng_state() + if not is_dist_avail_and_initialized(): + return [local_state] + states = [None] * get_world_size() + dist.all_gather_object(states, local_state) + return states + + +def save_model(args, epoch, model_without_ddp, optimizer, loss_scaler): + output_dir = Path(args.output_dir) + checkpoint_path = output_dir / f"checkpoint-{epoch:05d}.pth" + last_checkpoint_path = output_dir / "checkpoint-last.pth" + if epoch % args.checkpoint_period != 0 and epoch != args.epochs - 1: + return + + to_save = { + "model": model_without_ddp.state_dict(), + "optimizer": optimizer.state_dict(), + "epoch": epoch, + "scaler": None if loss_scaler is None else loss_scaler.state_dict(), + "args": OmegaConf.to_container(args), + "rng_states": _all_rank_rng_states(), + } + + print(f"saving checkpoint {last_checkpoint_path}") + save_on_master(to_save, last_checkpoint_path) + print(f"saving checkpoint {checkpoint_path}") + save_on_master(to_save, checkpoint_path) + + if args.max_checkpoints and is_main_process(): + all_checkpoints = sorted(output_dir.glob("checkpoint-[0-9]*.pth")) + del_count = max(0, len(all_checkpoints) - args.max_checkpoints) + for checkpoint_path in all_checkpoints[:del_count]: + print(f"removing checkpoint {checkpoint_path}") + checkpoint_path.unlink() + + +def load_model(args, model_without_ddp, optimizer, loss_scaler): + auto_resume = getattr(args, "auto_resume", True) + output_dir = Path(args.output_dir) + + last_checkpoint_path = output_dir / "checkpoint-last.pth" + if auto_resume and last_checkpoint_path.exists(): + args.ckpt = str(last_checkpoint_path) + args.resume = True + + if args.ckpt: + ckpt = torch.load(args.ckpt, map_location="cpu", weights_only=True) + model_without_ddp.load_state_dict(ckpt["model"]) + print(f"loaded model state from checkpoint {args.ckpt}") + + if args.resume: + optimizer.load_state_dict(ckpt["optimizer"]) + if loss_scaler is not None: + loss_scaler.load_state_dict(ckpt["scaler"]) + args.start_epoch = ckpt["epoch"] + 1 + rng_states = ckpt.get("rng_states") + if rng_states is not None: + if len(rng_states) != get_world_size(): + raise ValueError( + "checkpoint RNG state world size does not match current world size: " + f"{len(rng_states)} != {get_world_size()}" + ) + restore_rng_state(rng_states[get_rank()]) + print(f"restored RNG state for rank {get_rank()}") + print(f"loaded optimizer state, resuming training from {args.start_epoch}") + + +# optimization utils + + +# from capi +class WarmupThenCosine: + def __init__( + self, + base_value: float, + final_value: float, + total_iters: int, + warmup_iters: int = 0, + start_warmup_value: float = 0.0, + freeze_iters: int = 0, + truncate_cos: float = 1.0, + ): + super().__init__() + self.final_value = final_value + self.total_iters = total_iters + + freeze_schedule = np.zeros(freeze_iters) + + warmup_schedule = np.linspace(start_warmup_value, base_value, warmup_iters) + + iters = np.arange(total_iters - warmup_iters - freeze_iters) + schedule = final_value + 0.5 * (base_value - final_value) * ( + 1 + np.cos(np.pi * truncate_cos * iters / len(iters)) + ) + self.schedule = np.concatenate((freeze_schedule, warmup_schedule, schedule)) + assert len(self.schedule) == self.total_iters + + def __getitem__(self, it: int) -> float: + if it >= self.total_iters: + return self.final_value + # cast to float or else it can corrupt the checkpoint + return float(self.schedule[it]) + + +# adapted from timm backward logic +# https://github.com/huggingface/pytorch-image-models/blob/main/timm/utils/cuda.py +def backward_step( + loss: Tensor, + optimizer: Optimizer, + scaler: GradScaler = None, + need_update: bool = True, + max_norm: float | None = None, +) -> Tensor | None: + if scaler is not None: + scaler.scale(loss).backward() + else: + loss.backward() + + if need_update: + if scaler is not None: + scaler.unscale_(optimizer) + + total_norm = clip_grad(optimizer, max_norm) + + if scaler is not None: + scaler.step(optimizer) + scaler.update() + else: + optimizer.step() + optimizer.zero_grad() + else: + total_norm = None + return total_norm + + +def clip_grad(optimizer: Optimizer, max_norm: float | None = None) -> Tensor: + params = [p for group in optimizer.param_groups for p in group["params"]] + if max_norm: + total_norm = nn.utils.clip_grad_norm_(params, max_norm, error_if_nonfinite=False) + else: + grads = [p.grad for p in params if p.grad is not None] + total_norm = nn.utils.get_total_norm(grads, error_if_nonfinite=False) + torch._assert_async(torch.isfinite(total_norm), "non-finite gradient norm") + return total_norm + + +# from dinov2 with some minor changes +def get_param_groups(model, patch_embed_lr_mult=1.0): + # no lr decay, we could add this later if needed + all_params = [] + + for name, param in model.named_parameters(): + if not param.requires_grad: + continue + d = {"param": param, "lr_multiplier": 1.0, "wd_multiplier": 1.0, "name": name} + + if name.endswith(".bias") or "norm" in name or "gamma" in name: + d["wd_multiplier"] = 0.0 + + if "patch_embed" in name: + d["lr_multiplier"] = d["lr_multiplier"] * patch_embed_lr_mult + + all_params.append(d) + + param_groups = _fuse_param_groups(all_params) + return param_groups + + +def _fuse_param_groups(all_param_groups): + fused_param_groups = defaultdict(lambda: {"params": []}) + for d in all_param_groups: + keys = sorted(set(d.keys()) - {"param", "name"}) + identifier = "_".join(f"{k}{d[k]}" for k in keys) + for k in keys: + fused_param_groups[identifier][k] = d[k] + fused_param_groups[identifier]["params"].append(d["param"]) + + param_groups = list(fused_param_groups.values()) + return param_groups + + +def update_lr(param_groups, lr: float): + for group in param_groups: + group["lr"] = lr * group["lr_multiplier"] + + +def update_wd(param_groups, weight_decay: float | None = None): + for group in param_groups: + group["weight_decay"] = weight_decay * group["wd_multiplier"] + + +# moving data to cuda utils copied from capi +# added device argument + + +def send_data(x, device=None, dtype_map=None): + if device is None: + device = torch.device("cuda") + else: + device = torch.device(device) + + if isinstance(x, torch.Tensor): + dtype = dtype_map.get(x.dtype) if dtype_map else None + return x.to(device=device, dtype=dtype, non_blocking=True) + if isinstance(x, dict): + return {k: send_data(v, device=device, dtype_map=dtype_map) for k, v in x.items()} + if isinstance(x, list): + return [send_data(v, device=device, dtype_map=dtype_map) for v in x] + return x + + +def pre_send_to_cuda_wrapper(generator, device=None, dtype_map=None): + """From apex""" + data = None + stream = torch.cuda.Stream(device) + for next_data in generator: + with torch.cuda.stream(stream): + next_data = send_data(next_data, device=device, dtype_map=dtype_map) + if data is not None: + yield data + torch.cuda.current_stream(device).wait_stream(stream) + data = next_data + if data is not None: + yield data + + +# other misc utils + + +# from dino +def get_sha(): + cwd = os.path.dirname(os.path.abspath(__file__)) + + def _run(command): + return subprocess.check_output(command, cwd=cwd).decode("ascii").strip() + + sha = "N/A" + diff = "clean" + branch = "N/A" + try: + sha = _run(["git", "rev-parse", "HEAD"]) + diff = _run(["git", "diff-index", "HEAD"]) + diff = "has uncommitted changes" if diff else "clean" + branch = _run(["git", "rev-parse", "--abbrev-ref", "HEAD"]) + except Exception: + pass + message = f"sha: {sha}, status: {diff}, branch: {branch}" + return message + + +# from timm +def random_seed(seed=42, rank=0): + torch.manual_seed(seed + rank) + np.random.seed(seed + rank) + random.seed(seed + rank) + + +# mine :) +def filter_kwargs(func, kwargs): + sigature = inspect.signature(func) + kwargs = {k: v for k, v in kwargs.items() if k in sigature.parameters} + return kwargs diff --git a/fomo_tune_baseline/output/task3/build/smri_mae/visualization.py b/fomo_tune_baseline/output/task3/build/smri_mae/visualization.py new file mode 100644 index 0000000000000000000000000000000000000000..f721fe06e93dc1a2dfb4a679069029bdf55afc6b --- /dev/null +++ b/fomo_tune_baseline/output/task3/build/smri_mae/visualization.py @@ -0,0 +1,400 @@ +from collections.abc import Mapping +from io import BytesIO + +import torch + +from matplotlib import patches +from matplotlib import pyplot as plt +from PIL import Image +from torch import Tensor + +VIEW_NAMES = { + "sagittal": "Sagittal", + "saggital": "Sagittal", + "axial": "Axial", + "coronal": "Coronal", +} + + +def fig2pil(fig) -> Image.Image: + buffer = BytesIO() + fig.savefig(buffer, format="png", dpi=fig.dpi, facecolor=fig.get_facecolor()) + buffer.seek(0) + image = Image.open(buffer).convert("RGB") + buffer.close() + return image + + +def raw_stats_from_batch(batch: dict) -> tuple[Tensor | None, Tensor | None]: + metas = batch.get("meta") + if not metas: + return None, None + + means = [] + stds = [] + for meta in metas: + try: + mean = meta["raw_mean"] + std = meta["raw_std"] + except (KeyError, TypeError): + return None, None + if mean in ("", None) or std in ("", None): + return None, None + means.append(float(mean)) + stds.append(float(std)) + return torch.tensor(means), torch.tensor(stds) + + +def plot_mask_pred( + target: Tensor, + pred: Tensor, + pred_mask: Tensor | None = None, + img_mask: Tensor | None = None, + sample_idx: int = 0, + channel_idx: int = 0, + slice_idx: int | Mapping[str, int] | None = None, + patch_size: int | tuple[int, int, int] = 16, + views: tuple[str, ...] = ("sagittal", "axial", "coronal"), + cmap: str = "gray", + figsize: tuple[float, float] | None = None, + mask_style: str = "blank", + raw_mean: float | Tensor | None = None, + raw_std: float | Tensor | None = None, +): + target_vol = _select_volume(target, sample_idx=sample_idx, channel_idx=channel_idx) + pred_vol = _select_volume(pred, sample_idx=sample_idx, channel_idx=channel_idx) + if raw_mean is not None and raw_std is not None: + raw_mean = _select_scalar(raw_mean, sample_idx=sample_idx) + raw_std = _select_scalar(raw_std, sample_idx=sample_idx) + target_vol = target_vol * raw_std + raw_mean + pred_vol = pred_vol * raw_std + raw_mean + pred_mask_vol = ( + torch.zeros_like(target_vol) + if pred_mask is None + else _select_volume(pred_mask, sample_idx=sample_idx, channel_idx=channel_idx) > 0 + ) + img_mask_vol = None + if img_mask is not None: + img_mask_vol = _select_volume(img_mask, sample_idx=sample_idx, channel_idx=channel_idx) > 0 + + composite_vol = _prediction_composite(target_vol, pred_vol, pred_mask_vol) + vmin, vmax = _intensity_limits(target_vol, img_mask_vol) + + patch_size = _as_3tuple(patch_size) + view_items = [] + for view in views: + view_key = view.lower() + if view_key not in VIEW_NAMES: + raise ValueError(f"unknown MRI view {view!r}; expected one of {tuple(VIEW_NAMES)}") + target_slice = _extract_view_slice(target_vol, view_key, slice_idx) + composite_slice = _extract_view_slice(composite_vol, view_key, slice_idx) + mask_slice = _extract_view_slice(pred_mask_vol.float(), view_key, slice_idx) > 0 + img_mask_slice = None + if img_mask_vol is not None: + img_mask_slice = _extract_view_slice(img_mask_vol.float(), view_key, slice_idx) > 0 + view_items.append( + { + "key": view_key, + "title": VIEW_NAMES[view_key], + "target": _masked_input_display(target_slice, mask_slice, img_mask_slice, vmin), + "composite": _apply_display_mask(composite_slice, img_mask_slice, vmin), + "actual": _apply_display_mask(target_slice, img_mask_slice, vmin), + "mask": mask_slice, + "img_mask": img_mask_slice, + "patch_rc": _view_patch_size(view_key, patch_size), + } + ) + + _crop_view_items(view_items) + + fig, axes, layout = _make_figure_canvas(view_items, figsize=figsize) + for item, x in zip(view_items, layout["col_centers"]): + fig.text( + x, + layout["title_y"], + item["title"], + ha="center", + va="center", + color="#f8fafc", + fontsize=7, + ) + for label, y in zip(("Masked", "Pred", "Actual"), layout["row_centers"]): + fig.text(layout["label_x"], y, label, ha="right", va="center", color="#cbd5e1", fontsize=6) + + for item, top_ax, middle_ax, bottom_ax in zip(view_items, axes[0], axes[1], axes[2]): + top_ax.imshow( + item["target"], + cmap=cmap, + vmin=vmin, + vmax=vmax, + interpolation="nearest", + origin="upper", + ) + if mask_style == "boxes": + _draw_patch_boxes(top_ax, item["mask"], item["patch_rc"]) + elif mask_style != "blank": + raise ValueError("mask_style must be 'blank' or 'boxes'") + _style_axis(top_ax) + + middle_ax.imshow( + item["composite"], + cmap=cmap, + vmin=vmin, + vmax=vmax, + interpolation="nearest", + origin="upper", + ) + _style_axis(middle_ax) + + bottom_ax.imshow( + item["actual"], + cmap=cmap, + vmin=vmin, + vmax=vmax, + interpolation="nearest", + origin="upper", + ) + _style_axis(bottom_ax) + return fig + + +def _select_volume( + x: Tensor, + sample_idx: int = 0, + channel_idx: int = 0, +) -> Tensor: + x = x.detach().float().cpu() + if x.ndim == 5: + return x[sample_idx, channel_idx] + if x.ndim == 4: + return x[sample_idx] + if x.ndim == 3: + return x + raise ValueError(f"expected a 3D volume tensor, got shape {tuple(x.shape)}") + + +def _select_scalar(value: float | Tensor, sample_idx: int = 0) -> float: + if isinstance(value, Tensor): + value = value.detach().float().cpu() + if value.ndim > 0: + value = value.reshape(-1)[sample_idx] + return float(value) + return float(value) + + +def _prediction_composite(target: Tensor, pred: Tensor, pred_mask: Tensor) -> Tensor: + pred_mask = pred_mask.to(dtype=target.dtype) + return target * (1 - pred_mask) + pred * pred_mask + + +def _extract_view_slice( + volume: Tensor, + view: str, + slice_idx: int | Mapping[str, int] | None = None, +) -> Tensor: + if isinstance(slice_idx, Mapping): + slice_idx = slice_idx.get(view) + + if view in ("sagittal", "saggital"): + idx = _resolve_slice_idx(volume.shape[0], slice_idx) + return volume[idx, :, :].transpose(0, 1).flip(0) + if view == "axial": + idx = _resolve_slice_idx(volume.shape[2], slice_idx) + return volume[:, :, idx].transpose(0, 1).flip(0) + if view == "coronal": + idx = _resolve_slice_idx(volume.shape[1], slice_idx) + return volume[:, idx, :].transpose(0, 1).flip(0) + raise ValueError(f"unknown MRI view {view!r}") + + +def _resolve_slice_idx(size: int, slice_idx: int | None) -> int: + idx = size // 2 if slice_idx is None else int(slice_idx) + if idx < 0: + idx += size + if idx < 0 or idx >= size: + raise IndexError(f"slice index {idx} is out of bounds for axis with size {size}") + return idx + + +def _intensity_limits(volume: Tensor, mask: Tensor | None = None) -> tuple[float, float]: + values = volume[mask] if mask is not None and mask.any() else volume.flatten() + values = values[torch.isfinite(values)] + if values.numel() == 0: + return 0.0, 1.0 + if values.numel() < 32: + vmin = values.min() + vmax = values.max() + else: + vmin, vmax = torch.quantile(values, torch.tensor([0.005, 0.995])) + if torch.isclose(vmin, vmax): + delta = max(abs(float(vmin)) * 0.05, 1.0) + return float(vmin) - delta, float(vmax) + delta + return float(vmin), float(vmax) + + +def _apply_display_mask(image: Tensor, mask: Tensor | None, fill_value: float) -> Tensor: + if mask is None: + return image + return torch.where(mask, image, torch.full_like(image, fill_value)) + + +def _masked_input_display( + image: Tensor, + pred_mask: Tensor, + img_mask: Tensor | None, + fill_value: float, +) -> Tensor: + display = torch.where(pred_mask, torch.full_like(image, fill_value), image) + return _apply_display_mask(display, img_mask, fill_value) + + +def _crop_view_items(view_items: list[dict]) -> None: + for item in view_items: + mask = item["img_mask"] + if mask is None: + mask = item["actual"] != item["actual"].min() + row_slice, col_slice = _content_crop(mask, item["patch_rc"]) + for key in ("target", "composite", "actual", "mask"): + item[key] = item[key][row_slice, col_slice] + if item["img_mask"] is not None: + item["img_mask"] = item["img_mask"][row_slice, col_slice] + + +def _content_crop(mask: Tensor, patch_size: tuple[int, int]) -> tuple[slice, slice]: + mask = mask.detach().cpu() > 0 + if not mask.any(): + return slice(None), slice(None) + + rows, cols = mask.nonzero(as_tuple=True) + patch_h, patch_w = patch_size + height, width = mask.shape + row0 = max((int(rows.min()) // patch_h - 1) * patch_h, 0) + row1 = min((int(rows.max()) // patch_h + 2) * patch_h, height) + col0 = max((int(cols.min()) // patch_w - 1) * patch_w, 0) + col1 = min((int(cols.max()) // patch_w + 2) * patch_w, width) + return slice(row0, row1), slice(col0, col1) + + +def _as_3tuple(value: int | tuple[int, int, int]) -> tuple[int, int, int]: + if isinstance(value, int): + return (value, value, value) + if len(value) != 3: + raise ValueError(f"expected a 3-tuple patch size, got {value!r}") + return tuple(int(v) for v in value) + + +def _view_patch_size(view: str, patch_size: tuple[int, int, int]) -> tuple[int, int]: + p_x, p_y, p_z = patch_size + if view in ("sagittal", "saggital"): + return p_z, p_y + if view == "axial": + return p_y, p_x + if view == "coronal": + return p_z, p_x + raise ValueError(f"unknown MRI view {view!r}") + + +def _make_figure_canvas( + view_items: list[dict[str, Tensor | str | tuple[int, int]]], + figsize: tuple[float, float] | None = None, +): + dpi = 160 + left = 58 + right = 6 + top = 18 + bottom = 8 + row_gap = 14 + col_gap = 8 + widths = [int(item["target"].shape[1]) for item in view_items] + heights = [int(item["target"].shape[0]) for item in view_items] + row_h = max(heights) + num_rows = 3 + fig_w = left + right + sum(widths) + col_gap * (len(widths) - 1) + fig_h = top + bottom + row_h * num_rows + row_gap * (num_rows - 1) + + scale = 1.35 + if figsize is not None: + requested_w = figsize[0] * dpi + requested_h = figsize[1] * dpi + scale = max(requested_w / fig_w, requested_h / fig_h) + figsize = (fig_w * scale / dpi, fig_h * scale / dpi) + fig = plt.figure(figsize=figsize, dpi=dpi, facecolor="#0b0f14") + + axes = [[] for _ in range(num_rows)] + col_centers = [] + x = left + for width, height in zip(widths, heights): + col_centers.append((x + width / 2) / fig_w) + ys = [ + bottom + (num_rows - row - 1) * (row_h + row_gap) + (row_h - height) / 2 + for row in range(num_rows) + ] + for row, y in enumerate(ys): + axes[row].append( + fig.add_axes( + [ + x / fig_w, + y / fig_h, + width / fig_w, + height / fig_h, + ], + facecolor="black", + ) + ) + x += width + col_gap + row_centers = [ + (bottom + (num_rows - row - 1) * (row_h + row_gap) + row_h / 2) / fig_h + for row in range(num_rows) + ] + layout = { + "col_centers": col_centers, + "row_centers": row_centers, + "label_x": (left - 8) / fig_w, + "title_y": (fig_h - top / 2) / fig_h, + } + + return fig, axes, layout + + +def _style_axis(ax) -> None: + ax.set_xticks([]) + ax.set_yticks([]) + for spine in ax.spines.values(): + spine.set_visible(False) + + +def _draw_patch_boxes( + ax, + mask: Tensor, + patch_size: tuple[int, int], + color: str = "#facc15", +) -> None: + for col, row, width, height in _patch_rectangles(mask, patch_size): + ax.add_patch( + patches.Rectangle( + (col - 0.5, row - 0.5), + width, + height, + fill=False, + edgecolor=color, + linewidth=0.75, + alpha=0.95, + ) + ) + + +def _patch_rectangles( + mask: Tensor, + patch_size: tuple[int, int], +) -> list[tuple[int, int, int, int]]: + mask = mask.detach().cpu() > 0 + patch_h, patch_w = patch_size + height, width = mask.shape + rectangles = [] + for row in range(0, height, patch_h): + box_h = min(patch_h, height - row) + for col in range(0, width, patch_w): + box_w = min(patch_w, width - col) + if mask[row : row + box_h, col : col + box_w].any(): + rectangles.append((col, row, box_w, box_h)) + return rectangles diff --git a/fomo_tune_baseline/output/task3/config.yaml b/fomo_tune_baseline/output/task3/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..42af3932e7182587c4a9c2a47db74aeb12d18aa3 --- /dev/null +++ b/fomo_tune_baseline/output/task3/config.yaml @@ -0,0 +1,6 @@ +task: task3 +ckpt_path: /data/mihir-stuff/smri-pretrained/pretrain_full_90_10_h100/checkpoint-last.pth +output_root: experiments/fomo_tune_baseline/output +name: task3 +device: cuda +seed: 4466 diff --git a/fomo_tune_baseline/output/task3/log.txt b/fomo_tune_baseline/output/task3/log.txt new file mode 100644 index 0000000000000000000000000000000000000000..ca1b8422fc1ba17a1d5af221dacb6f610021aebf --- /dev/null +++ b/fomo_tune_baseline/output/task3/log.txt @@ -0,0 +1,30 @@ +16:34:12 run task3 (git 1df2e5d) +16:34:12 config: +task: task3 +ckpt_path: /data/mihir-stuff/smri-pretrained/pretrain_full_90_10_h100/checkpoint-last.pth +output_root: experiments/fomo_tune_baseline/output +name: task3 +device: cuda +seed: 4466 +16:35:16 dataset: 494 subjects, age 19-80 mean 45.2 +16:38:39 fold 1/20 n=25 mae=3.91 (190s) +16:38:51 fold 2/20 n=25 mae=4.50 (202s) +16:38:57 fold 3/20 n=25 mae=3.74 (208s) +16:39:03 fold 4/20 n=25 mae=3.58 (214s) +16:39:09 fold 5/20 n=25 mae=3.57 (220s) +16:39:14 fold 6/20 n=25 mae=3.12 (226s) +16:39:20 fold 7/20 n=25 mae=3.42 (232s) +16:39:26 fold 8/20 n=25 mae=3.12 (238s) +16:39:32 fold 9/20 n=25 mae=3.74 (244s) +16:39:38 fold 10/20 n=25 mae=4.20 (249s) +16:39:44 fold 11/20 n=25 mae=3.22 (255s) +16:39:50 fold 12/20 n=25 mae=3.50 (261s) +16:39:55 fold 13/20 n=25 mae=3.39 (267s) +16:40:01 fold 14/20 n=25 mae=4.17 (273s) +16:40:07 fold 15/20 n=24 mae=4.36 (278s) +16:40:12 fold 16/20 n=24 mae=4.15 (284s) +16:40:18 fold 17/20 n=24 mae=3.66 (289s) +16:40:23 fold 18/20 n=24 mae=3.54 (295s) +16:40:29 fold 19/20 n=24 mae=3.42 (300s) +16:40:34 fold 20/20 n=24 mae=3.53 (306s) +16:40:35 result: pearson_r=0.9631 pearson_r_ci_low=0.9567 pearson_r_ci_high=0.9687 mae=3.6910 mae_ci_low=3.4464 mae_ci_high=3.9466 (306s) diff --git a/fomo_tune_baseline/output/task3/metrics.json b/fomo_tune_baseline/output/task3/metrics.json new file mode 100644 index 0000000000000000000000000000000000000000..e4c32ce134cb5f41aea37219dee1b0f6e9abc615 --- /dev/null +++ b/fomo_tune_baseline/output/task3/metrics.json @@ -0,0 +1 @@ +{"name": "task3", "pearson_r": 0.9630771532969529, "pearson_r_ci_low": 0.9567144985968425, "pearson_r_ci_high": 0.9687044710309238, "mae": 3.6910024558001684, "mae_ci_low": 3.446356968454986, "mae_ci_high": 3.946628273643462, "run_time": 306.1} diff --git a/fomo_tune_baseline/output/task3/model/config.yaml b/fomo_tune_baseline/output/task3/model/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..42af3932e7182587c4a9c2a47db74aeb12d18aa3 --- /dev/null +++ b/fomo_tune_baseline/output/task3/model/config.yaml @@ -0,0 +1,6 @@ +task: task3 +ckpt_path: /data/mihir-stuff/smri-pretrained/pretrain_full_90_10_h100/checkpoint-last.pth +output_root: experiments/fomo_tune_baseline/output +name: task3 +device: cuda +seed: 4466 diff --git a/fomo_tune_baseline/output/task3/model/head.joblib b/fomo_tune_baseline/output/task3/model/head.joblib new file mode 100644 index 0000000000000000000000000000000000000000..7ade23f9dc38a75d598a1a8dfeb4930ff2d4b1a9 --- /dev/null +++ b/fomo_tune_baseline/output/task3/model/head.joblib @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e263a95e347fbaf86b4edc44717e55891711bf37a7a1309c7f04753af46a4dbf +size 29906 diff --git a/fomo_tune_baseline/output/task3/task3.sif b/fomo_tune_baseline/output/task3/task3.sif new file mode 100644 index 0000000000000000000000000000000000000000..45dde18ad211fb8f1bf87576156a131d977f71ea --- /dev/null +++ b/fomo_tune_baseline/output/task3/task3.sif @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a77daed42a7520e4509214c931f4aa082d333fae57944f44bfcdd79428ee16ce +size 5328908288 diff --git a/fomo_tune_baseline/output/task5/build/Apptainer.def b/fomo_tune_baseline/output/task5/build/Apptainer.def new file mode 100644 index 0000000000000000000000000000000000000000..98a8bdfc1fefdbffdaed0dedbc7166a3da28b472 --- /dev/null +++ b/fomo_tune_baseline/output/task5/build/Apptainer.def @@ -0,0 +1,31 @@ +Bootstrap: docker +From: python:3.11-slim + +# NOT buildable where it sits: the %files paths below are relative to the build cwd, which is the +# staging dir `build.py` writes. Build it with `python -m fomo_tune.build `, not by +# pointing apptainer at this file. +# +# Versions are pinned to the training environment: numpy, scikit-learn and joblib because they +# unpickle `head.joblib`, torch because that is what the checkpoint was written by. + +%files + fomo_tune /app/fomo_tune + smri_mae /app/smri_mae + model /app/model + predict.py /app/predict.py + +%post + pip install --no-cache-dir \ + torch==2.8.0 \ + numpy==2.4.6 \ + nibabel==5.4.2 \ + einops==0.8.2 \ + jaxtyping==0.3.10 \ + timm==1.0.27 \ + huggingface-hub==0.36.2 \ + scikit-learn==1.8.0 \ + joblib==1.5.3 \ + omegaconf==2.3.0 + +%runscript + exec python /app/predict.py "$@" diff --git a/fomo_tune_baseline/output/task5/build/fomo_tune/README.md b/fomo_tune_baseline/output/task5/build/fomo_tune/README.md new file mode 100644 index 0000000000000000000000000000000000000000..aa7a6cdadfe314f9d1b0fc2d48ef67b85032d3c5 --- /dev/null +++ b/fomo_tune_baseline/output/task5/build/fomo_tune/README.md @@ -0,0 +1,239 @@ +# fomo_tune + +The five FOMO26 challenge tasks, one script each, tuned independently. + +This is a spinoff of `nanobrain.eval`, which scored every backbone on every task through one fixed +probe. That was the right shape for a benchmark and the wrong shape for a competition: here we care +about one backbone (sMRI MAE) and five scores, and each task wants a different method. **Nothing +here imports `nanobrain.eval`, and it should stay that way** — this package may be shared with +people who won't get the eval suite. + +## Layout + +| File | | +|---|---| +| `datasets.py` | core, **frozen**. One `load_fomo_task()` per task, streaming the challenge zips into an HF dataset. Raw niftis, no resampling — the backbone transform does that. | +| `backbone.py` | core, **frozen**. `load_backbone(ckpt_path) -> (SmriMaeBackbone, SmriMaeTransform)`. Frozen sMRI MAE encoder; the transform canonicalizes to RAS, rescales to 1mm, fits to the pretraining shape, z-scores in a mean-threshold brain mask. | +| `utils.py` | core. `set_seed`, `git_sha`, `setup_logging`. | +| `main_task.py` | shell. One task, end to end. Task 1 is the worked example; copy it. | +| `build.py` + `Apptainer.def` | shell. Package a run dir into the challenge `.sif`. Shared by every task. | + +`datasets.py` and `backbone.py` are settled and their caches are warm. Treat them as read-only: +new work goes in `main_task.py`. If one of them genuinely needs to change, that is a +conversation first, because it invalidates every score already recorded. + +## The pattern + +`main_task1.py` is in three sections, and the split is the point of the whole design. + +**`Task1Method` — the part we tune.** Features, head, hyperparameters, anything that might move +the score. Its interface is: + +```python +method.fit(rows) # rows are dataset records: subject, label, images +method.predict(images) # -> the challenge's output for one subject +method.save(model_dir) # config.yaml + head.joblib +Task1Method.load(model_dir, **overrides) +``` + +**The protocol — fixed.** Pool out-of-fold predictions over all subjects, bootstrap subjects for +the CI. No repeats, no stratification; the bootstrap is the only variance estimate. Splitting is +per-task but fixed within a task — leave-one-out where n is tiny (task 1, n=21), **20-fold** where +it isn't (tasks 3 and 5), which is close enough to LOO without paying for 494 refits. Once a task's +scheme is set, hold it or scores stop being comparable across iterations. That is also why +`cross_validate` seeds its shuffle at 0 rather than from `cfg.seed`: the folds are part of the +protocol, so tuning the run's seed must not silently redraw them. + +**Two entrypoints.** `train` runs the protocol then fits a head on all subjects and saves it; +`predict` is the challenge CLI. Both go through `Method.predict`, which is why every fold +exercises the code the submission will run. + +That last point is the load-bearing one. `predict` is not a wrapper written at packaging time — it +is the same call cross-validation already made once per held-out subject. When you add a task, keep +that property. + +```bash +uv run python -m fomo_tune.main_task1 train modalities=[dwi_b1000,flair] name=task1_dwi_flair +uv run python -m fomo_tune.main_task1 predict --model-dir output/fomo_tune/task1_dwi/model \ + --adc adc.nii.gz --dwi dwi.nii.gz --flair flair.nii.gz --output prob.txt +``` + +`train` takes omegaconf dotlist overrides against the `Config` dataclass at the top of the file. +It writes `config.yaml`, `log.txt`, `metrics.json`, and `model/` into `{output_root}/{name}/`. + +## Status + +Tasks 1, 5 and 3 are drafted and verified. Task 1 is also packaged — its container passes the +challenge validator; 5 and 3 have not been built yet. **Tasks 2 and 4 are tabled** — both are +segmentation, both need `predict` to emit a nifti on the input grid, and neither is worth opening +until the classification and regression tasks are settled. + +All three on `vitl_fomo300`, one H100, wall being the cross-validation loop: + +| run | result | wall | +|---|---|---| +| `task1_dwi`, dwi_b1000, n=21, LOO | AUROC **0.990** [0.944, 1.000] | 25s | +| `task5_t1w`, t1w, n=48, 20-fold | AUROC **0.984** [0.953, 1.000] | 73s | +| `task3_t1w`, t1w, n=494, 20-fold | r **0.962** [0.956, 0.968], MAE **3.71y** [3.45, 3.97] | 260s | + +**Task 3's row is one fold-seed stale.** It was measured before `cross_validate` froze its shuffle +at 0, so it is a 20-fold run with `random_state=4466`. Task 1 (LOO) and task 5 are current. The +re-run is cheap — 260s on a GPU — it just has not been done. Expect a shift of the same order task +5 saw when its folds moved (0.948 → 0.984, i.e. inside the CI but not negligible). + +Task 1's earlier probe sweep got 0.954 [0.861, 1.000] on the same checkpoint +(`experiments/eval_global_0728`), so it roughly reproduces — the gap is LOO vs 5×5 stratified CV, +one interpolation instead of two, and a head selected on AUROC instead of balanced accuracy. + +Two checks worth repeating per task — `.claude/scratch/verify_task1.py` and +`.claude/scratch/verify_task35.py ` do both: +- features are **bit-identical** whether the nifti comes from the HF dataset wrapper or from + `nib.load` off disk, so CV numbers transfer to the container +- the `predict` CLI agrees with the in-process method + +## What changes per task + +Counts and modalities, read from the local zips: + +| Task | n | Inputs | Output | Split | Notes | +|---|---|---|---|---|---| +| 1 infarct | 21 | adc, dwi_b1000, flair (+t2s/swi) | probability | LOO | done | +| 5 polymicrogyria | 48 | t1w | probability | 20-fold | done | +| 3 brain age | 494 | t1w | age in years | 20-fold | done — RidgeCV head, scored by **Pearson r and MAE**, each with its own bootstrap CI | +| 2 meningioma | 23 | dwi_b1000, flair (+t2s/swi) | mask, input grid | — | tabled | +| 4 trigeminal | 40 | t2w | mask, labels 1=nerve 2=vessel | — | tabled | + +Tasks 5 and 3 diverge from task 1 only where that table says. `cross_validate` over a shuffled +`KFold` replaces `leave_one_out`; both take one modality, so `features` loses the +concat-over-modalities loop and `Config` loses `modalities`; the challenge CLI flag is `--t1` for +both, and it is `--t1` for task 3 too even though the file in the zip is `t1w.nii.gz`. + +Task 3 is the first regression, so its `score` loops over the two metrics rather than returning +one, and it drops task 1's guard against bootstrap resamples with fewer than two distinct labels. +The analogous degenerate case for regression is a resample with no spread in `y`, where Pearson r +is undefined rather than merely unstable — at n=494 it does not happen. + +When tasks 2 and 4 come back: `predict` must write a nifti on the input's grid, and the method +needs localized features rather than a pooled vector — `backbone.forward` returns `patch_coords` +in world mm for exactly that. Task 4's label order (1=nerve, 2=vessel) is still a guess and needs +confirming against the challenge data before per-class numbers mean anything. + +## Gotchas + +**Raw niftis are on disk** at `data/fomo_eval/Task_/preprocessed//ses-01/`, which is the +easy way to exercise `predict` on a real file rather than one written out of the dataset: + +```bash +uv run python -m fomo_tune.main_task1 predict \ + --model-dir output/fomo_tune/task1_dwi/model \ + --adc data/fomo_eval/Task_1/preprocessed/sub-20/ses-01/adc.nii.gz \ + --dwi data/fomo_eval/Task_1/preprocessed/sub-20/ses-01/dwi_b1000.nii.gz \ + --flair data/fomo_eval/Task_1/preprocessed/sub-20/ses-01/flair.nii.gz \ + --output /tmp/prob.txt +``` + +Task 5 breaks the naming: `Task_5/preprocessed/sub_01/ses_01/t1.nii.gz` — underscores throughout, +and `t1` not `t1w`. `datasets.py` already handles it; anything you write by hand won't. + +```bash +uv run python -m fomo_tune.main_task5 predict --model-dir output/fomo_tune/task5_t1w/model \ + --t1 data/fomo_eval/Task_5/preprocessed/sub_01/ses_01/t1.nii.gz --output /tmp/prob.txt +``` + +**Volumes are wildly anisotropic.** Task 1's DWI is 0.46×0.46×**5.6**mm, so the transform +upsamples z by 5.6× to reach 1mm iso. Nothing is wrong, but don't read the 1mm grid as real +resolution. + +**The backbone never saw skull or neck.** Pretraining used a SynthSeg brain mask; the transform +substitutes a mean-intensity threshold, which keeps both. Known fidelity gap — see +`.claude/memory/smri-mae-preprocessing-gap.md`. + +**Probabilities are not calibrated.** `LogisticRegressionCV` on ~20 samples × 1024 features shrinks +hard; task 1's out-of-fold probabilities all land in 0.48–0.52 with near-perfect ranking. Fine for +AUROC, which is what the challenge scores, but don't read them as probabilities. Task 5's do span +0–1, which is n=48 rather than n=21 and not evidence of calibration. + +**n is tiny, so the CI is the result.** Task 1's is ~0.06 wide at the top of the range. Most tuning +deltas you chase will be inside it. `.claude/NOTES.md` thread 1 has the longer argument. + +**GPUs need an allocation** — the login node has no driver. See the `gpu-session` skill. + +## Submission + +`build.py` packages a run dir into the `.sif` the challenge wants. One command, taking the run dir +the shipped head was saved into: + +```bash +uv run python -m fomo_tune.build output/fomo_tune/task1_dwi +``` + +It stages `/app`, then builds from there: + +``` +/app/predict.py # shim: calls fomo_tune.main_task predict +/app/model/config.yaml # from the run dir +/app/model/head.joblib # from the run dir +/app/model/backbone.pth # stripped checkpoint, --ckpt-path points here +``` + +**Both `build.py` and `Apptainer.def` are shared across tasks**, which they can be because nothing +in staging or in the dependency list is task-specific. The one thing that does vary is the module +the shim imports, and that comes from `task` in the run's saved config — so a run dir knows which +task it belongs to, and `build.py` never needs telling. + +`predict.py` is **generated at build time** rather than checked in. It is eight lines whose whole +meaning is the container layout staged around it, so there is nowhere outside a container to run +it. This does not weaken the point above about `predict` not being written at packaging time: the +logic still lives in `main_task.py`, exercised once per fold, and the shim only picks the +subcommand and two paths. + +**`Apptainer.def` is not buildable where it sits.** Its `%files` paths are relative to the build +cwd, which is the staging dir. Pointing `apptainer build` at it in the repo fails confusingly; go +through `build.py`. + +The run dir deliberately does *not* carry backbone weights — that checkpoint is 3.9G and would be +copied on every run. `--ckpt-path` overrides what `config.yaml` recorded, so the saved config stays +a faithful record of what trained rather than being rewritten at build time. + +**The staged checkpoint is stripped to `model` and `args`**, which is 3.9G → 1.3G because the rest +is optimizer state inference never reads. `load_backbone` needs no change for this, and `predict` +gives a bit-identical probability either way (0.524739 on `sub-20`, checked on GPU). + +**The base image is `python:3.11-slim`, not a CUDA image.** The PyPI torch wheel *is* the cu128 +build and vendors the whole CUDA userspace as `nvidia-*` packages, so all the container needs from +the host is the driver, which `--nv`/`--nvccli` binds in. That keeps the SIF at 5.3G (4.0G of +image, 1.3G of weights) against roughly double for `pytorch/pytorch` and far more for NGC. +Versions are pinned to the training environment +mostly so `head.joblib` unpickles against the numpy/sklearn that wrote it. + +Apptainer caches the bootstrap layers but **always re-runs `%post`**, so every build re-downloads +~3G of wheels. If that gets annoying, bake a deps-only base SIF and `Bootstrap: localimage` off it. + +### Validating + +`third_party/container-validator` is the challenge's own validator, test niftis included: + +```bash +python third_party/container-validator/container_validator/validate.py \ + --task task1 --sif output/fomo_tune/task1_dwi/task1.sif +``` + +It runs `python /app/predict.py --flair /input/… --adc … --dwi … --swi … --output /output/.txt` +inside an `apptainer instance` with `/input`, `/output` and `/tmp` bound — which is exactly the +shim's contract, so nothing in `predict.py` is guessing at the interface. + +One thing it does that is easy to miss: it takes GPU via `--nvccli` rather than `--nv`, and one of +its tests runs `nvidia-smi -L` **inside** the container. `python:3.11-slim` ships no `nvidia-smi`, +so that test passes only because `--nvccli` injects the host one — a CUDA base image would hide +that dependency rather than remove it. + +**The `task1_dwi` container passes all 20 validator tests**, and `predict` inside it returns +0.524739 on `sub-20`, identical to the same call outside the container. So the packaging is +verified end to end, not just built. + +**Run it on a compute node with apptainer, which as of 2026-08-11 means `n-6`** — `salloc +--nodelist=n-6`. The other nodes fail the validator's preflight. The login node has apptainer but +no driver, and +`predict` there dies inside `can_use_cudnn_attention` — the jagged-SDPA path reaches into CUDA even +when the tensors are on CPU, so a driver-less host fails at the forward pass rather than falling +back. That is the CPU gap worth remembering; it is not a container problem. diff --git a/fomo_tune_baseline/output/task5/build/fomo_tune/backbone.py b/fomo_tune_baseline/output/task5/build/fomo_tune/backbone.py new file mode 100644 index 0000000000000000000000000000000000000000..4a1275a32e0f9da0963043225e8b3aefcbe8edc6 --- /dev/null +++ b/fomo_tune_baseline/output/task5/build/fomo_tune/backbone.py @@ -0,0 +1,153 @@ +import inspect + +import nibabel as nib +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange +from torch import Tensor + +import smri_mae.model_mae as models_mae + + +class SmriMaeBackbone(nn.Module): + grid_coords: Tensor + + def __init__(self, encoder: models_mae.MaskedEncoder): + super().__init__() + self.encoder = encoder + self.img_size = self.encoder.patchify.img_size + + grid_size = self.encoder.patchify.grid_size + patch_size = np.array(self.encoder.patchify.patch_size) + grid_coords = rearrange(np.indices(grid_size), "c x y z -> (x y z) c") + grid_coords = grid_coords * patch_size + (patch_size - 1) / 2 + grid_coords = torch.as_tensor(grid_coords, dtype=torch.float32) + self.register_buffer("grid_coords", grid_coords) + + def forward(self, batch: dict[str, Tensor]) -> dict[str, Tensor]: + images = batch["image"] + mask = batch["mask"] + affine = batch["affine"] + + B, C, X, Y, Z = images.shape + assert (X, Y, Z) == self.img_size, f"expected {self.img_size}, got {(X, Y, Z)}" + + _, _, patch_embeds, _, patch_ids, token_mask = self.encoder(images, mask=mask) + + # [B, L, 3] world xyz coords of embeddings + patch_coords = self.grid_coords[patch_ids, :] + rot = affine[:, :3, :3] + trans = affine[:, :3, 3] + patch_coords = patch_coords @ rot.transpose(1, 2) + trans[:, None, :] + + return { + "patch_embeds": patch_embeds, + "patch_ids": patch_ids, + "token_mask": token_mask, + "patch_coords": patch_coords, + } + + +class SmriMaeTransform: + def __init__( + self, + img_size: tuple[int, int, int] = (208, 240, 208), + spacing: tuple[float, float, float] = (1.0, 1.0, 1.0), + ): + self.img_size = img_size + self.spacing = spacing + + def __call__(self, img: nib.Nifti1Image) -> dict[str, Tensor]: + # repack image to handle incomplete hf Nifti interface + img = nib.Nifti1Image(img.dataobj, img.affine, img.header) + img = nib.as_closest_canonical(img) + + data = torch.from_numpy(np.ascontiguousarray(img.get_fdata(dtype=np.float32))) + affine = np.asarray(img.affine) + + spacing = img.header.get_zooms() + if max(abs(s - s_) for s, s_ in zip(spacing, self.spacing)) > 0.05: + data, affine = rescale(data, affine, spacing, self.spacing) + + data, affine = fit_to_shape(data, affine, target_shape=self.img_size) + + # mean threshold, not the SynthSeg mask used in pretraining, so skull and neck stay in + mask = data > data.mean() + brain = data[mask] + mean = brain.mean() + # population std (correction=0) to match the pretraining normalization + std = brain.std(correction=0).clamp_min(1e-6) + data = torch.where(mask, (data - mean) / std, 0.0) + + return { + "image": data.unsqueeze(0), + "mask": mask.unsqueeze(0), + "affine": torch.as_tensor(affine, dtype=torch.float32), + } + + +def rescale( + x: torch.Tensor, + affine: np.ndarray, + spacing: tuple[float, ...], + target_spacing: tuple[float, ...] = (1.0, 1.0, 1.0), +) -> tuple[torch.Tensor, np.ndarray]: + scales = tuple([current / target for current, target in zip(spacing, target_spacing)]) + resampled = F.interpolate(x[None, None], scale_factor=scales, mode="trilinear").squeeze(0, 1) + + # align_corners=False reads output voxel j from input voxel (j + 0.5) / scale - 0.5 + scale = np.asarray(scales, dtype=float) + step = np.diag([*(1 / scale), 1.0]) + step[:3, 3] = 0.5 / scale - 0.5 + return resampled, affine @ step + + +def fit_to_shape( + x: torch.Tensor, affine: np.ndarray, target_shape: tuple[int, ...] +) -> tuple[torch.Tensor, np.ndarray]: + """Centre the volume in `target_shape`, padding the short axes and cropping the long ones.""" + pads = [target - size for size, target in zip(x.shape, target_shape)] + padding = [side for pad in reversed(pads) for side in (pad // 2, pad - pad // 2)] + + # a crop is a negative pad, so output voxel k came from input voxel k - pad // 2 either way + step = np.eye(4) + step[:3, 3] = [-(pad // 2) for pad in pads] + return F.pad(x, padding), affine @ step + + +def resolve_ckpt(ckpt_path: str) -> str: + """A local path for a checkpoint, downloading it if it is an hf://// URI.""" + from huggingface_hub import hf_hub_download + + if ckpt_path.startswith("hf://"): + org, repo, *rest = ckpt_path.removeprefix("hf://").split("/") + return hf_hub_download(f"{org}/{repo}", "/".join(rest)) + + return ckpt_path + + +def load_backbone(ckpt_path: str) -> tuple[SmriMaeBackbone, SmriMaeTransform]: + path = resolve_ckpt(ckpt_path) + ckpt = torch.load(path, map_location="cpu", weights_only=True, mmap=True) + args = ckpt["args"] + + model_fn = models_mae.__dict__[args["model"]] + model: models_mae.MaskedAutoencoderViT = model_fn( + img_size=args["img_size"], + in_chans=args.get("in_chans", 1), + patch_size=args["patch_size"], + # older checkpoints carry training flags the current model_mae no longer takes + **filter_kwargs(models_mae.MaskedAutoencoderViT, args.get("model_kwargs") or {}), + ) + model.load_state_dict(ckpt["model"]) + backbone = SmriMaeBackbone(model.encoder) + transform = SmriMaeTransform(img_size=args["img_size"]) + return backbone, transform + + +def filter_kwargs(func, kwargs): + signature = inspect.signature(func) + kwargs = {k: v for k, v in kwargs.items() if k in signature.parameters} + return kwargs diff --git a/fomo_tune_baseline/output/task5/build/fomo_tune/datasets.py b/fomo_tune_baseline/output/task5/build/fomo_tune/datasets.py new file mode 100644 index 0000000000000000000000000000000000000000..64c2ac2da798086e88cbb335601ae850c2dad249 --- /dev/null +++ b/fomo_tune_baseline/output/task5/build/fomo_tune/datasets.py @@ -0,0 +1,205 @@ +import os +import shutil +import tempfile +import zipfile +from collections.abc import Generator +from contextlib import contextmanager +from pathlib import Path + +import fsspec +from datasets import Dataset, Features, Nifti, Value + +FOMO_EVAL_BASE_URL = os.getenv( + "FOMO_EVAL_BASE_URL", + "https://sid.erda.dk/share_redirect/fmeuvo1EdF", +) +FOMO_EVAL_TASK5_URL = os.getenv( + "FOMO_EVAL_TASK5_URL", + "https://huggingface.co/datasets/medarc/smri-fm/resolve/main/fomo_eval/Task_5.zip", +) + + +@contextmanager +def open_zip(url: str) -> Generator[zipfile.ZipFile, None, None]: + """Open a task zip, copying a remote url to a temp file first.""" + with tempfile.TemporaryDirectory() as tmp: + local = Path(url) + if not local.exists(): + local = Path(tmp) / "task.zip" + with fsspec.open(url) as src, local.open("wb") as dst: + shutil.copyfileobj(src, dst) + with zipfile.ZipFile(local) as zf: + yield zf + + +def subject_ids(zf: zipfile.ZipFile) -> list[str]: + return sorted({name.split("/")[2] for name in zf.namelist() if name.endswith(".nii.gz")}) + + +# ---- Task 1: acute infarct (classification; positives also carry a lesion mask) -------- + + +def load_fomo_task1() -> Dataset: + # No 4th modality: it is swi on 16 subjects and t2s on the other 5. + suffixes = ("adc", "dwi_b1000", "flair") + features = Features( + { + "subject": Value("string"), + "label": Value("int32"), + **{suffix: Nifti() for suffix in suffixes}, + } + ) + dataset = Dataset.from_generator( + _fomo_task1_generator, + features=features, + gen_kwargs={"suffixes": suffixes}, + writer_batch_size=16, + ) + return dataset + + +def _fomo_task1_generator(suffixes: tuple[str, ...]): + url = f"{FOMO_EVAL_BASE_URL}/Task_1.zip" + with open_zip(url) as zf: + for sub in subject_ids(zf): + label = int(zf.read(f"Task_1/labels/{sub}/ses-01/label.txt").strip()) + sample = {"subject": sub, "label": label} + for suffix in suffixes: + name = f"Task_1/preprocessed/{sub}/ses-01/{suffix}.nii.gz" + sample[suffix] = {"path": None, "bytes": zf.read(name)} + yield sample + + +# ---- Task 2: meningioma segmentation --------------------------------------------------- + + +def load_fomo_task2() -> Dataset: + # No 4th modality: it is t2s on 15 subjects and swi on the other 8. + suffixes = ("dwi_b1000", "flair") + features = Features( + { + "subject": Value("string"), + **{suffix: Nifti() for suffix in suffixes}, + "seg": Nifti(), + } + ) + dataset = Dataset.from_generator( + _fomo_task2_generator, + features=features, + gen_kwargs={"suffixes": suffixes}, + writer_batch_size=16, + ) + return dataset + + +def _fomo_task2_generator(suffixes: tuple[str, ...]): + url = f"{FOMO_EVAL_BASE_URL}/Task_2.zip" + with open_zip(url) as zf: + for sub in subject_ids(zf): + sample = {"subject": sub} + for suffix in suffixes: + name = f"Task_2/preprocessed/{sub}/ses-01/{suffix}.nii.gz" + sample[suffix] = {"path": None, "bytes": zf.read(name)} + # Seg is on the image grid (shapes match) but its affine differs by up to 0.03mm. + name = f"Task_2/labels/{sub}/ses-01/seg.nii.gz" + sample["seg"] = {"path": None, "bytes": zf.read(name)} + yield sample + + +# ---- Task 3: brain age regression ------------------------------------------------------ + + +def load_fomo_task3() -> Dataset: + features = Features( + { + "subject": Value("string"), + "age": Value("int32"), + "t1w": Nifti(), + } + ) + dataset = Dataset.from_generator( + _fomo_task3_generator, + features=features, + writer_batch_size=16, + ) + return dataset + + +def _fomo_task3_generator(): + url = f"{FOMO_EVAL_BASE_URL}/Task_3.zip" + with open_zip(url) as zf: + for sub in subject_ids(zf): + age = int(zf.read(f"Task_3/labels/{sub}/ses-01/labels.txt").strip()) + image_gz = zf.read(f"Task_3/preprocessed/{sub}/ses-01/t1w.nii.gz") + sample = { + "subject": sub, + "age": age, + "t1w": {"path": None, "bytes": image_gz}, + } + yield sample + + +# ---- Task 4: trigeminal nerve/vessel segmentation -------------------------------------- + + +def load_fomo_task4() -> Dataset: + # Volumes are uncropped 0.5mm near-iso, ~360x512x512; crop before feeding a model. + features = Features( + { + "subject": Value("string"), + "t2w": Nifti(), + "seg": Nifti(), + } + ) + dataset = Dataset.from_generator( + _fomo_task4_generator, + features=features, + writer_batch_size=16, + ) + return dataset + + +def _fomo_task4_generator(): + url = f"{FOMO_EVAL_BASE_URL}/Task_4.zip" + with open_zip(url) as zf: + for sub in subject_ids(zf): + image_gz = zf.read(f"Task_4/preprocessed/{sub}/ses-01/t2w.nii.gz") + seg_gz = zf.read(f"Task_4/labels/{sub}/ses-01/seg.nii.gz") + sample = { + "subject": sub, + "t2w": {"path": None, "bytes": image_gz}, + "seg": {"path": None, "bytes": seg_gz}, + } + yield sample + + +# ---- Task 5: polymicrogyria classification --------------------------------------------- + + +def load_fomo_task5() -> Dataset: + features = Features( + { + "subject": Value("string"), + "label": Value("int32"), + "t1w": Nifti(), + } + ) + dataset = Dataset.from_generator( + _fomo_task5_generator, + features=features, + writer_batch_size=16, + ) + return dataset + + +def _fomo_task5_generator(): + with open_zip(FOMO_EVAL_TASK5_URL) as zf: + for sub in subject_ids(zf): + label = int(zf.read(f"Task_5/labels/{sub}/ses_01/labels.txt").strip()) + image_gz = zf.read(f"Task_5/preprocessed/{sub}/ses_01/t1.nii.gz") + sample = { + "subject": sub, + "label": label, + "t1w": {"path": None, "bytes": image_gz}, + } + yield sample diff --git a/fomo_tune_baseline/output/task5/build/fomo_tune/main_task1.py b/fomo_tune_baseline/output/task5/build/fomo_tune/main_task1.py new file mode 100644 index 0000000000000000000000000000000000000000..bfbc47d49ceeed6fc43dc52a0a007d933ff86d4b --- /dev/null +++ b/fomo_tune_baseline/output/task5/build/fomo_tune/main_task1.py @@ -0,0 +1,253 @@ +"""FOMO task 1: acute infarct classification, scored by AUROC as the challenge scores it. + +`Task1Method` is the part we tune -- features, head, hyperparameters. The protocol below it is +fixed so scores stay comparable across iterations: leave one subject out, pool the out-of-fold +predictions, bootstrap subjects for the CI. + +`train` runs that protocol then fits and saves a head; `predict` is the challenge contract, +modality paths in and one probability out. Both go through `Task1Method.predict`, so every fold +exercises the path the submission will run. +""" + +import argparse +import json +import logging +import time +from dataclasses import dataclass, field +from pathlib import Path + +import joblib +import nibabel as nib +import numpy as np +import torch +from omegaconf import OmegaConf +from sklearn.linear_model import LogisticRegressionCV +from sklearn.metrics import roc_auc_score +from sklearn.pipeline import make_pipeline +from sklearn.preprocessing import StandardScaler + +from fomo_tune.backbone import load_backbone +from fomo_tune.utils import git_sha, set_seed, setup_logging + +logger = logging.getLogger("fomo_tune") + +Images = dict[str, nib.Nifti1Image] + + +@dataclass +class Config: + task: str = "task1" + ckpt_path: str = ( + "/data/mihir-stuff/smri-pretrained/pretrain_full_90_10_h100/checkpoint-last.pth" + ) + modalities: list[str] = field(default_factory=lambda: ["dwi_b1000"]) + output_root: str = "output/fomo_tune" + name: str = "task1" + device: str = "cuda" + seed: int = 4466 + + +# ---- method: the part we tune ----------------------------------------------------------- + + +class Task1Method: + """Frozen sMRI MAE, mean-pooled tokens per modality concatenated, logistic head.""" + + def __init__(self, cfg: Config): + self.cfg = cfg + self.backbone, self.transform = load_backbone(cfg.ckpt_path) + self.device = torch.device(cfg.device) + self.backbone.to(self.device).eval().requires_grad_(False) + self.modalities = list(cfg.modalities) + self.cache: dict[str, np.ndarray] = {} + self.head = None + + @torch.inference_mode() + def features(self, images: Images) -> np.ndarray: + """(D,) per subject. A pure function of the images, so training and inference agree.""" + pooled = [] + for modality in self.modalities: + sample = self.transform(images[modality]) + batch = {key: value[None].to(self.device) for key, value in sample.items()} + + with torch.autocast("cuda", torch.bfloat16, enabled=self.device.type == "cuda"): + out = self.backbone(batch) + + patch_embeds = out["patch_embeds"] + token_mask = out["token_mask"].bool().unsqueeze(-1) + embed = (patch_embeds * token_mask).sum(dim=1) / token_mask.sum(dim=1) + pooled.append(embed[0].float().cpu()) + + return torch.cat(pooled).numpy() + + def cached_features(self, row: dict) -> np.ndarray: + if row["subject"] not in self.cache: + self.cache[row["subject"]] = self.features(row) + return self.cache[row["subject"]] + + def fit(self, rows: list[dict]) -> None: + X = np.stack([self.cached_features(row) for row in rows]) + y = np.array([row["label"] for row in rows]) + + clf = LogisticRegressionCV( + Cs=10, + class_weight="balanced", + scoring="roc_auc", + max_iter=1000, + l1_ratios=(0,), + use_legacy_attributes=False, + ) + self.head = make_pipeline(StandardScaler(), clf) + self.head.fit(X, y) + self.positive = list(self.head.classes_).index(1) + + def predict(self, images: Images) -> float: + """Positive-class probability. Indexes `classes_` rather than assuming column 1, which + would silently score the wrong class if the label order differed.""" + X = self.features(images)[None] + probs = self.head.predict_proba(X)[0] + return float(probs[self.positive]) + + def save(self, model_dir: Path) -> None: + """Everything `load` needs but the backbone weights, which stay wherever `ckpt_path` + points -- a few hundred KB, so a run saves one without copying a 3.7G checkpoint.""" + model_dir.mkdir(parents=True, exist_ok=True) + OmegaConf.save(self.cfg, model_dir / "config.yaml") + joblib.dump({"head": self.head, "positive": self.positive}, model_dir / "head.joblib") + + @classmethod + def load(cls, model_dir: Path, **overrides) -> "Task1Method": + """Rebuild a fitted method from `save`. Overrides are Config fields, for what differs + between here and the container -- the backbone path, the device.""" + cfg = OmegaConf.merge( + OmegaConf.structured(Config), OmegaConf.load(model_dir / "config.yaml"), overrides + ) + method = cls(cfg) + state = joblib.load(model_dir / "head.joblib") + method.head, method.positive = state["head"], state["positive"] + return method + + +# ---- protocol: the part we hold fixed --------------------------------------------------- + +# Every image the task ships. The method picks which of them it wants, as at inference, where +# the challenge hands over all four modalities whether or not a model uses them. +IMAGE_COLS = ("adc", "dwi_b1000", "flair") + + +def leave_one_out(rows: list[dict], method: Task1Method) -> tuple[np.ndarray, np.ndarray]: + """Out-of-fold score for every subject, each predicted by a head fit on the other n-1.""" + y = np.array([row["label"] for row in rows]) + oof = np.zeros(len(rows), dtype=float) + start = time.perf_counter() + for held_out, row in enumerate(rows): + method.fit([r for r in rows if r["subject"] != row["subject"]]) + oof[held_out] = method.predict({key: row[key] for key in IMAGE_COLS}) + logger.info( + f"fold {held_out + 1}/{len(rows)} {row['subject']} " + f"y={y[held_out]} p={oof[held_out]:.3f} ({time.perf_counter() - start:.0f}s)" + ) + return y, oof + + +def score( + y: np.ndarray, oof: np.ndarray, seed: int = 0, n_boot: int = 2000, alpha: float = 0.05 +) -> dict: + """AUROC, the challenge metric, plus a percentile CI resampling subjects with replacement.""" + rng = np.random.default_rng(seed) + samples = [] + for _ in range(n_boot): + rows = rng.integers(0, len(y), size=len(y)) + if len(np.unique(y[rows])) < 2: + continue + samples.append(roc_auc_score(y[rows], oof[rows])) + + low, high = np.percentile(samples, [100 * alpha / 2, 100 * (1 - alpha / 2)]) + return { + "auroc": float(roc_auc_score(y, oof)), + "auroc_ci_low": float(low), + "auroc_ci_high": float(high), + } + + +# ---- entrypoints ------------------------------------------------------------------------ + + +def train(args: argparse.Namespace) -> None: + # imported here, not at the top, so the container needs no dataset stack to run `predict` + from fomo_tune.datasets import load_fomo_task1 + + cfg = OmegaConf.merge(OmegaConf.structured(Config), OmegaConf.from_dotlist(args.overrides)) + run_dir = Path(cfg.output_root) / cfg.name + run_dir.mkdir(parents=True, exist_ok=True) + + setup_logging(run_dir) + set_seed(cfg.seed) + logger.info(f"run {cfg.name} (git {git_sha()})") + logger.info(f"config:\n{OmegaConf.to_yaml(cfg).rstrip()}") + OmegaConf.save(cfg, run_dir / "config.yaml") + + # decoded once: leave-one-out revisits every subject n times, and the niftis are small + rows = list(load_fomo_task1()) + logger.info(f"dataset: {len(rows)} subjects, {sum(r['label'] for r in rows)} positive") + + method = Task1Method(cfg) + start = time.perf_counter() + y, oof = leave_one_out(rows, method) + run_time = time.perf_counter() - start + summary = score(y, oof) + + # the shipped head sees all n subjects, so it is not any of the models scored above + method.fit(rows) + method.save(run_dir / "model") + + record = {"name": cfg.name, **summary, "run_time": round(run_time, 1)} + (run_dir / "metrics.json").write_text(json.dumps(record) + "\n") + scores = " ".join(f"{k}={v:.4f}" for k, v in summary.items()) + logger.info(f"result: {scores} ({run_time:.0f}s)") + + +def predict(args: argparse.Namespace) -> None: + """The challenge contract: modality paths in, one probability written to `--output`. + + `/app/predict.py` in the container is a shim over this, so what scores the submission is the + code leave-one-out already ran, not something generated at build time. + """ + overrides = {"device": args.device} + if args.ckpt_path: + overrides["ckpt_path"] = args.ckpt_path + method = Task1Method.load(args.model_dir, **overrides) + + # every image the challenge hands over, as in `leave_one_out`; the method takes what it uses + paths = {"adc": args.adc, "dwi_b1000": args.dwi, "flair": args.flair} + probability = method.predict({key: nib.load(path) for key, path in paths.items()}) + + args.output.write_text(f"{probability:.6f}\n") + + +def main() -> None: + parser = argparse.ArgumentParser() + modes = parser.add_subparsers(required=True) + + train_parser = modes.add_parser("train", help="leave-one-out over the task, then fit and save") + train_parser.add_argument("overrides", nargs="*", help="config overrides, e.g. device=cpu") + train_parser.set_defaults(run=train) + + predict_parser = modes.add_parser("predict", help="one subject, one probability") + for flag in ("--flair", "--adc", "--dwi"): + predict_parser.add_argument(flag, type=Path, required=True) + # accepted and ignored: the 4th modality is swi on some subjects and t2s on others + for flag in ("--t2s", "--swi"): + predict_parser.add_argument(flag, type=Path) + predict_parser.add_argument("--output", type=Path, required=True) + predict_parser.add_argument("--model-dir", type=Path, default=Path("/app/model")) + predict_parser.add_argument("--ckpt-path", help="overrides the trained config's backbone path") + predict_parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") + predict_parser.set_defaults(run=predict) + + args = parser.parse_args() + args.run(args) + + +if __name__ == "__main__": + main() diff --git a/fomo_tune_baseline/output/task5/build/fomo_tune/main_task3.py b/fomo_tune_baseline/output/task5/build/fomo_tune/main_task3.py new file mode 100644 index 0000000000000000000000000000000000000000..02db5413c10a0d44eb248e7f2739a436c45806c2 --- /dev/null +++ b/fomo_tune_baseline/output/task5/build/fomo_tune/main_task3.py @@ -0,0 +1,241 @@ +"""FOMO task 3: brain age regression, scored by Pearson r and MAE as the challenge scores it. + +`Task3Method` is the part we tune -- features, head, hyperparameters. The protocol below it is +fixed so scores stay comparable across iterations: 20-fold over the 494 subjects, pool the +out-of-fold predictions, bootstrap subjects for the CI. + +`train` runs that protocol then fits and saves a head; `predict` is the challenge contract, one t1 +path in and one age out. Both go through `Task3Method.predict`, so every fold exercises the path +the submission will run. +""" + +import argparse +import json +import logging +import time +from dataclasses import dataclass +from pathlib import Path + +import joblib +import nibabel as nib +import numpy as np +import torch +from omegaconf import OmegaConf +from sklearn.linear_model import RidgeCV +from sklearn.model_selection import KFold +from sklearn.pipeline import make_pipeline +from sklearn.preprocessing import StandardScaler + +from fomo_tune.backbone import load_backbone +from fomo_tune.utils import git_sha, set_seed, setup_logging + +logger = logging.getLogger("fomo_tune") + +Images = dict[str, nib.Nifti1Image] + + +@dataclass +class Config: + task: str = "task3" + ckpt_path: str = ( + "/data/mihir-stuff/smri-pretrained/pretrain_full_90_10_h100/checkpoint-last.pth" + ) + output_root: str = "output/fomo_tune" + name: str = "task3" + device: str = "cuda" + seed: int = 4466 + + +# ---- method: the part we tune ----------------------------------------------------------- + + +class Task3Method: + """Frozen sMRI MAE, mean-pooled tokens over the t1w, ridge head.""" + + def __init__(self, cfg: Config): + self.cfg = cfg + self.backbone, self.transform = load_backbone(cfg.ckpt_path) + self.device = torch.device(cfg.device) + self.backbone.to(self.device).eval().requires_grad_(False) + self.cache: dict[str, np.ndarray] = {} + self.head = None + + @torch.inference_mode() + def features(self, images: Images) -> np.ndarray: + """(D,) per subject. A pure function of the images, so training and inference agree.""" + sample = self.transform(images["t1w"]) + batch = {key: value[None].to(self.device) for key, value in sample.items()} + + with torch.autocast("cuda", torch.bfloat16, enabled=self.device.type == "cuda"): + out = self.backbone(batch) + + patch_embeds = out["patch_embeds"] + token_mask = out["token_mask"].bool().unsqueeze(-1) + embed = (patch_embeds * token_mask).sum(dim=1) / token_mask.sum(dim=1) + return embed[0].float().cpu().numpy() + + def cached_features(self, row: dict) -> np.ndarray: + if row["subject"] not in self.cache: + self.cache[row["subject"]] = self.features(row) + return self.cache[row["subject"]] + + def fit(self, rows: list[dict]) -> None: + X = np.stack([self.cached_features(row) for row in rows]) + y = np.array([row["age"] for row in rows], dtype=float) + + # RidgeCV picks alpha by its own efficient leave-one-out, so the fold's own split is + # never touched by model selection + self.head = make_pipeline(StandardScaler(), RidgeCV(alphas=np.logspace(-3, 6, 19))) + self.head.fit(X, y) + + def predict(self, images: Images) -> float: + """Age in years.""" + X = self.features(images)[None] + return float(self.head.predict(X)[0]) + + def save(self, model_dir: Path) -> None: + """Everything `load` needs but the backbone weights, which stay wherever `ckpt_path` + points -- a few hundred KB, so a run saves one without copying a 3.7G checkpoint.""" + model_dir.mkdir(parents=True, exist_ok=True) + OmegaConf.save(self.cfg, model_dir / "config.yaml") + joblib.dump(self.head, model_dir / "head.joblib") + + @classmethod + def load(cls, model_dir: Path, **overrides) -> "Task3Method": + """Rebuild a fitted method from `save`. Overrides are Config fields, for what differs + between here and the container -- the backbone path, the device.""" + cfg = OmegaConf.merge( + OmegaConf.structured(Config), OmegaConf.load(model_dir / "config.yaml"), overrides + ) + method = cls(cfg) + method.head = joblib.load(model_dir / "head.joblib") + return method + + +# ---- protocol: the part we hold fixed --------------------------------------------------- + +# Every image the task ships. The method picks which of them it wants, as at inference, where the +# challenge hands over the modalities whether or not a model uses them. +IMAGE_COLS = ("t1w",) + + +def cross_validate( + rows: list[dict], method: Task3Method, seed: int = 0, n_folds: int = 20 +) -> tuple[np.ndarray, np.ndarray]: + """Out-of-fold age for every subject, each predicted by a head fit on the other folds.""" + y = np.array([row["age"] for row in rows], dtype=float) + oof = np.zeros(len(rows), dtype=float) + folds = KFold(n_splits=n_folds, shuffle=True, random_state=seed) + start = time.perf_counter() + for fold, (train, test) in enumerate(folds.split(rows)): + method.fit([rows[i] for i in train]) + for i in test: + oof[i] = method.predict({key: rows[i][key] for key in IMAGE_COLS}) + logger.info( + f"fold {fold + 1}/{n_folds} n={len(test)} mae={np.abs(y[test] - oof[test]).mean():.2f} " + f"({time.perf_counter() - start:.0f}s)" + ) + return y, oof + + +def metrics(y: np.ndarray, oof: np.ndarray) -> dict: + return { + "pearson_r": float(np.corrcoef(y, oof)[0, 1]), + "mae": float(np.abs(y - oof).mean()), + } + + +def score( + y: np.ndarray, oof: np.ndarray, seed: int = 0, n_boot: int = 2000, alpha: float = 0.05 +) -> dict: + """Both challenge metrics, each with a percentile CI resampling subjects with replacement.""" + rng = np.random.default_rng(seed) + resamples = rng.integers(0, len(y), size=(n_boot, len(y))) + + summary = {} + for name, point in metrics(y, oof).items(): + samples = [metrics(y[rows], oof[rows])[name] for rows in resamples] + low, high = np.percentile(samples, [100 * alpha / 2, 100 * (1 - alpha / 2)]) + summary[name] = point + summary[f"{name}_ci_low"] = float(low) + summary[f"{name}_ci_high"] = float(high) + return summary + + +# ---- entrypoints ------------------------------------------------------------------------ + + +def train(args: argparse.Namespace) -> None: + # imported here, not at the top, so the container needs no dataset stack to run `predict` + from fomo_tune.datasets import load_fomo_task3 + + cfg = OmegaConf.merge(OmegaConf.structured(Config), OmegaConf.from_dotlist(args.overrides)) + run_dir = Path(cfg.output_root) / cfg.name + run_dir.mkdir(parents=True, exist_ok=True) + + setup_logging(run_dir) + set_seed(cfg.seed) + logger.info(f"run {cfg.name} (git {git_sha()})") + logger.info(f"config:\n{OmegaConf.to_yaml(cfg).rstrip()}") + OmegaConf.save(cfg, run_dir / "config.yaml") + + rows = list(load_fomo_task3()) + ages = np.array([row["age"] for row in rows]) + logger.info( + f"dataset: {len(rows)} subjects, age {ages.min()}-{ages.max()} mean {ages.mean():.1f}" + ) + + method = Task3Method(cfg) + start = time.perf_counter() + y, oof = cross_validate(rows, method) + run_time = time.perf_counter() - start + summary = score(y, oof) + + # the shipped head sees all n subjects, so it is not any of the models scored above + method.fit(rows) + method.save(run_dir / "model") + + record = {"name": cfg.name, **summary, "run_time": round(run_time, 1)} + (run_dir / "metrics.json").write_text(json.dumps(record) + "\n") + scores = " ".join(f"{k}={v:.4f}" for k, v in summary.items()) + logger.info(f"result: {scores} ({run_time:.0f}s)") + + +def predict(args: argparse.Namespace) -> None: + """The challenge contract: a t1 path in, one age written to `--output`. + + `/app/predict.py` in the container is a shim over this, so what scores the submission is the + code cross-validation already ran, not something generated at build time. + """ + overrides = {"device": args.device} + if args.ckpt_path: + overrides["ckpt_path"] = args.ckpt_path + method = Task3Method.load(args.model_dir, **overrides) + + age = method.predict({"t1w": nib.load(args.t1)}) + + args.output.write_text(f"{age:.6f}\n") + + +def main() -> None: + parser = argparse.ArgumentParser() + modes = parser.add_subparsers(required=True) + + train_parser = modes.add_parser("train", help="cross-validate over the task, then fit and save") + train_parser.add_argument("overrides", nargs="*", help="config overrides, e.g. device=cpu") + train_parser.set_defaults(run=train) + + predict_parser = modes.add_parser("predict", help="one subject, one age in years") + predict_parser.add_argument("--t1", type=Path, required=True) + predict_parser.add_argument("--output", type=Path, required=True) + predict_parser.add_argument("--model-dir", type=Path, default=Path("/app/model")) + predict_parser.add_argument("--ckpt-path", help="overrides the trained config's backbone path") + predict_parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") + predict_parser.set_defaults(run=predict) + + args = parser.parse_args() + args.run(args) + + +if __name__ == "__main__": + main() diff --git a/fomo_tune_baseline/output/task5/build/fomo_tune/main_task5.py b/fomo_tune_baseline/output/task5/build/fomo_tune/main_task5.py new file mode 100644 index 0000000000000000000000000000000000000000..d8cccbecf130497c483b3a4b7665e835ca428873 --- /dev/null +++ b/fomo_tune_baseline/output/task5/build/fomo_tune/main_task5.py @@ -0,0 +1,245 @@ +"""FOMO task 5: polymicrogyria classification, scored by AUROC as the challenge scores it. + +`Task5Method` is the part we tune -- features, head, hyperparameters. The protocol below it is +fixed so scores stay comparable across iterations: 20-fold over the 48 subjects, pool the +out-of-fold predictions, bootstrap subjects for the CI. + +`train` runs that protocol then fits and saves a head; `predict` is the challenge contract, one t1 +path in and one probability out. Both go through `Task5Method.predict`, so every fold exercises +the path the submission will run. +""" + +import argparse +import json +import logging +import time +from dataclasses import dataclass +from pathlib import Path + +import joblib +import nibabel as nib +import numpy as np +import torch +from omegaconf import OmegaConf +from sklearn.linear_model import LogisticRegressionCV +from sklearn.metrics import roc_auc_score +from sklearn.model_selection import KFold +from sklearn.pipeline import make_pipeline +from sklearn.preprocessing import StandardScaler + +from fomo_tune.backbone import load_backbone +from fomo_tune.utils import git_sha, set_seed, setup_logging + +logger = logging.getLogger("fomo_tune") + +Images = dict[str, nib.Nifti1Image] + + +@dataclass +class Config: + task: str = "task5" + ckpt_path: str = ( + "/data/mihir-stuff/smri-pretrained/pretrain_full_90_10_h100/checkpoint-last.pth" + ) + output_root: str = "output/fomo_tune" + name: str = "task5" + device: str = "cuda" + seed: int = 4466 + + +# ---- method: the part we tune ----------------------------------------------------------- + + +class Task5Method: + """Frozen sMRI MAE, mean-pooled tokens over the t1w, logistic head.""" + + def __init__(self, cfg: Config): + self.cfg = cfg + self.backbone, self.transform = load_backbone(cfg.ckpt_path) + self.device = torch.device(cfg.device) + self.backbone.to(self.device).eval().requires_grad_(False) + self.cache: dict[str, np.ndarray] = {} + self.head = None + + @torch.inference_mode() + def features(self, images: Images) -> np.ndarray: + """(D,) per subject. A pure function of the images, so training and inference agree.""" + sample = self.transform(images["t1w"]) + batch = {key: value[None].to(self.device) for key, value in sample.items()} + + with torch.autocast("cuda", torch.bfloat16, enabled=self.device.type == "cuda"): + out = self.backbone(batch) + + patch_embeds = out["patch_embeds"] + token_mask = out["token_mask"].bool().unsqueeze(-1) + embed = (patch_embeds * token_mask).sum(dim=1) / token_mask.sum(dim=1) + return embed[0].float().cpu().numpy() + + def cached_features(self, row: dict) -> np.ndarray: + if row["subject"] not in self.cache: + self.cache[row["subject"]] = self.features(row) + return self.cache[row["subject"]] + + def fit(self, rows: list[dict]) -> None: + X = np.stack([self.cached_features(row) for row in rows]) + y = np.array([row["label"] for row in rows]) + + clf = LogisticRegressionCV( + Cs=10, + class_weight="balanced", + scoring="roc_auc", + max_iter=1000, + l1_ratios=(0,), + use_legacy_attributes=False, + ) + self.head = make_pipeline(StandardScaler(), clf) + self.head.fit(X, y) + self.positive = list(self.head.classes_).index(1) + + def predict(self, images: Images) -> float: + """Positive-class probability. Indexes `classes_` rather than assuming column 1, which + would silently score the wrong class if the label order differed.""" + X = self.features(images)[None] + probs = self.head.predict_proba(X)[0] + return float(probs[self.positive]) + + def save(self, model_dir: Path) -> None: + """Everything `load` needs but the backbone weights, which stay wherever `ckpt_path` + points -- a few hundred KB, so a run saves one without copying a 3.7G checkpoint.""" + model_dir.mkdir(parents=True, exist_ok=True) + OmegaConf.save(self.cfg, model_dir / "config.yaml") + joblib.dump({"head": self.head, "positive": self.positive}, model_dir / "head.joblib") + + @classmethod + def load(cls, model_dir: Path, **overrides) -> "Task5Method": + """Rebuild a fitted method from `save`. Overrides are Config fields, for what differs + between here and the container -- the backbone path, the device.""" + cfg = OmegaConf.merge( + OmegaConf.structured(Config), OmegaConf.load(model_dir / "config.yaml"), overrides + ) + method = cls(cfg) + state = joblib.load(model_dir / "head.joblib") + method.head, method.positive = state["head"], state["positive"] + return method + + +# ---- protocol: the part we hold fixed --------------------------------------------------- + +# Every image the task ships. The method picks which of them it wants, as at inference, where the +# challenge hands over the modalities whether or not a model uses them. +IMAGE_COLS = ("t1w",) + + +def cross_validate( + rows: list[dict], method: Task5Method, seed: int = 0, n_folds: int = 20 +) -> tuple[np.ndarray, np.ndarray]: + """Out-of-fold score for every subject, each predicted by a head fit on the other folds.""" + y = np.array([row["label"] for row in rows]) + oof = np.zeros(len(rows), dtype=float) + folds = KFold(n_splits=n_folds, shuffle=True, random_state=seed) + start = time.perf_counter() + for fold, (train, test) in enumerate(folds.split(rows)): + method.fit([rows[i] for i in train]) + for i in test: + oof[i] = method.predict({key: rows[i][key] for key in IMAGE_COLS}) + logger.info( + f"fold {fold + 1}/{n_folds} n={len(test)} y={y[test]} " + f"p={np.round(oof[test], 3)} ({time.perf_counter() - start:.0f}s)" + ) + return y, oof + + +def score( + y: np.ndarray, oof: np.ndarray, seed: int = 0, n_boot: int = 2000, alpha: float = 0.05 +) -> dict: + """AUROC, the challenge metric, plus a percentile CI resampling subjects with replacement.""" + rng = np.random.default_rng(seed) + samples = [] + for _ in range(n_boot): + rows = rng.integers(0, len(y), size=len(y)) + if len(np.unique(y[rows])) < 2: + continue + samples.append(roc_auc_score(y[rows], oof[rows])) + + low, high = np.percentile(samples, [100 * alpha / 2, 100 * (1 - alpha / 2)]) + return { + "auroc": float(roc_auc_score(y, oof)), + "auroc_ci_low": float(low), + "auroc_ci_high": float(high), + } + + +# ---- entrypoints ------------------------------------------------------------------------ + + +def train(args: argparse.Namespace) -> None: + # imported here, not at the top, so the container needs no dataset stack to run `predict` + from fomo_tune.datasets import load_fomo_task5 + + cfg = OmegaConf.merge(OmegaConf.structured(Config), OmegaConf.from_dotlist(args.overrides)) + run_dir = Path(cfg.output_root) / cfg.name + run_dir.mkdir(parents=True, exist_ok=True) + + setup_logging(run_dir) + set_seed(cfg.seed) + logger.info(f"run {cfg.name} (git {git_sha()})") + logger.info(f"config:\n{OmegaConf.to_yaml(cfg).rstrip()}") + OmegaConf.save(cfg, run_dir / "config.yaml") + + rows = list(load_fomo_task5()) + logger.info(f"dataset: {len(rows)} subjects, {sum(r['label'] for r in rows)} positive") + + method = Task5Method(cfg) + start = time.perf_counter() + y, oof = cross_validate(rows, method) + run_time = time.perf_counter() - start + summary = score(y, oof) + + # the shipped head sees all n subjects, so it is not any of the models scored above + method.fit(rows) + method.save(run_dir / "model") + + record = {"name": cfg.name, **summary, "run_time": round(run_time, 1)} + (run_dir / "metrics.json").write_text(json.dumps(record) + "\n") + scores = " ".join(f"{k}={v:.4f}" for k, v in summary.items()) + logger.info(f"result: {scores} ({run_time:.0f}s)") + + +def predict(args: argparse.Namespace) -> None: + """The challenge contract: a t1 path in, one probability written to `--output`. + + `/app/predict.py` in the container is a shim over this, so what scores the submission is the + code cross-validation already ran, not something generated at build time. + """ + overrides = {"device": args.device} + if args.ckpt_path: + overrides["ckpt_path"] = args.ckpt_path + method = Task5Method.load(args.model_dir, **overrides) + + probability = method.predict({"t1w": nib.load(args.t1)}) + + args.output.write_text(f"{probability:.6f}\n") + + +def main() -> None: + parser = argparse.ArgumentParser() + modes = parser.add_subparsers(required=True) + + train_parser = modes.add_parser("train", help="cross-validate over the task, then fit and save") + train_parser.add_argument("overrides", nargs="*", help="config overrides, e.g. device=cpu") + train_parser.set_defaults(run=train) + + predict_parser = modes.add_parser("predict", help="one subject, one probability") + predict_parser.add_argument("--t1", type=Path, required=True) + predict_parser.add_argument("--output", type=Path, required=True) + predict_parser.add_argument("--model-dir", type=Path, default=Path("/app/model")) + predict_parser.add_argument("--ckpt-path", help="overrides the trained config's backbone path") + predict_parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") + predict_parser.set_defaults(run=predict) + + args = parser.parse_args() + args.run(args) + + +if __name__ == "__main__": + main() diff --git a/fomo_tune_baseline/output/task5/build/fomo_tune/utils.py b/fomo_tune_baseline/output/task5/build/fomo_tune/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..57a5c6d0396d8811e763254a987b331cf41ee1f1 --- /dev/null +++ b/fomo_tune_baseline/output/task5/build/fomo_tune/utils.py @@ -0,0 +1,33 @@ +import logging +import random +import subprocess +import sys +from pathlib import Path + +import numpy as np +import torch + +logger = logging.getLogger("fomo_tune") + + +def set_seed(seed: int) -> None: + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + + +def git_sha() -> str: + kwargs = dict(cwd=Path(__file__).parent, capture_output=True, text=True, check=True) + sha = subprocess.run(["git", "rev-parse", "--short", "HEAD"], **kwargs).stdout.strip() + dirty = subprocess.run(["git", "status", "--porcelain", "-uno"], **kwargs).stdout.strip() + return f"{sha}-dirty" if dirty else sha + + +def setup_logging(run_dir: Path) -> None: + handlers = [logging.StreamHandler(sys.stdout), logging.FileHandler(run_dir / "log.txt")] + logger.setLevel(logging.INFO) + logger.handlers.clear() + for handler in handlers: + handler.setFormatter(logging.Formatter("%(asctime)s %(message)s", datefmt="%H:%M:%S")) + logger.addHandler(handler) + logger.propagate = False diff --git a/fomo_tune_baseline/output/task5/build/model/backbone.pth b/fomo_tune_baseline/output/task5/build/model/backbone.pth new file mode 100644 index 0000000000000000000000000000000000000000..793b42f5cea0b055534c29cd18c21f3200873219 --- /dev/null +++ b/fomo_tune_baseline/output/task5/build/model/backbone.pth @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:aacf2582e9464e7fd4bd3d5ed41a5b9aa770e5b04a24b2a82c1c4ba5fb5f6e7a +size 1389681124 diff --git a/fomo_tune_baseline/output/task5/build/model/config.yaml b/fomo_tune_baseline/output/task5/build/model/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..95ebb19790bae19d7927e9621aab1d6520030460 --- /dev/null +++ b/fomo_tune_baseline/output/task5/build/model/config.yaml @@ -0,0 +1,6 @@ +task: task5 +ckpt_path: /data/mihir-stuff/smri-pretrained/pretrain_full_90_10_h100/checkpoint-last.pth +output_root: experiments/fomo_tune_baseline/output +name: task5 +device: cuda +seed: 4466 diff --git a/fomo_tune_baseline/output/task5/build/model/head.joblib b/fomo_tune_baseline/output/task5/build/model/head.joblib new file mode 100644 index 0000000000000000000000000000000000000000..fb645916a1c0161712332085d845e2a4bd8f5014 --- /dev/null +++ b/fomo_tune_baseline/output/task5/build/model/head.joblib @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:dc3c55c52bd9306b8385abe2232837d7ff0181320f93c487db2eaea27a895c1c +size 445215 diff --git a/fomo_tune_baseline/output/task5/build/predict.py b/fomo_tune_baseline/output/task5/build/predict.py new file mode 100644 index 0000000000000000000000000000000000000000..053ab311ea60f8d459a3a68349a2483bde1bfe35 --- /dev/null +++ b/fomo_tune_baseline/output/task5/build/predict.py @@ -0,0 +1,16 @@ + +import sys + +from fomo_tune.main_task5 import main + +sys.argv = [ + sys.argv[0], + "predict", + *sys.argv[1:], + "--model-dir", + "/app/model", + "--ckpt-path", + "/app/model/backbone.pth", +] + +main() diff --git a/fomo_tune_baseline/output/task5/build/smri_mae/config/default_pretrain.yaml b/fomo_tune_baseline/output/task5/build/smri_mae/config/default_pretrain.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c16290b1e2c24e4d206f7e32a7b9c2ee5266067e --- /dev/null +++ b/fomo_tune_baseline/output/task5/build/smri_mae/config/default_pretrain.yaml @@ -0,0 +1,98 @@ +# Name of the run. Used for output directory suffix and wandb. +name: pretrain + +# Description of the run. Goes in wandb notes. +notes: null + +# Root output directory. +# The run writes to checkpoints/ when name is set. +output_dir: checkpoints + +# Standard 3D structural MRI volume size. +img_size: [208, 240, 208] +patch_size: 8 + +# Masking. +mask_ratio: 0.80 +pred_mask_ratio: null +pad_to_multiple: 32 + +# Model. +model: mae_vit_large +model_kwargs: + # target normalization: null/none, global, slice, or patch. + target_norm: none + + no_decode_pos: false + mask_drop_scale: false + + class_token: true + reg_tokens: 0 + no_embed_class: false + + decoder_depth: 4 + drop_path_rate: 0.0 + +# Datasets. +datasets: + fomo_train: + url: datasets/FOMO_with_dwi/shard.{000000..001620}.tar + samples_per_epoch: 243200 + shuffle: true + buffer_size: 8000 + drop_last: true + + fomo_val: + url: datasets/FOMO_with_dwi/shard.{001621..001800}.tar + samples_per_epoch: 26880 + shuffle: false + buffer_size: 1000 + drop_last: true + +train_dataset: fomo_train +eval_datasets: + - fomo_val + +# Data loader. +num_workers: 4 +prefetch_factor: 2 +presend_cuda: true + +# Optimization. +epochs: 100 +batch_size: 64 +accum_iter: 1 + +base_lr: 0.001 +min_lr: 1e-6 +warmup_epochs: 10 +weight_decay: 0.05 +betas: [0.9, 0.95] +clip_grad: 1.0 + +amp: true +amp_dtype: bfloat16 + +# Checkpointing. +ckpt: null +resume: false +auto_resume: true +start_epoch: 0 +checkpoint_period: 10 +max_checkpoints: 5 + +# Evaluation. +eval_period: 10 + +# Sync checkpoints to an R2 bucket using the AWS CLI. Set to a URL to enable. +r2_sync: null + +device: cuda +distributed: false +seed: 7338 +eval_seed: 7338 +debug: false + +wandb: false +wandb_entity: null +wandb_project: smri-fm diff --git a/fomo_tune_baseline/output/task5/build/smri_mae/main_pretrain.py b/fomo_tune_baseline/output/task5/build/smri_mae/main_pretrain.py new file mode 100644 index 0000000000000000000000000000000000000000..5470551d5b56bf4653fad4c2a23888c417c0c6e8 --- /dev/null +++ b/fomo_tune_baseline/output/task5/build/smri_mae/main_pretrain.py @@ -0,0 +1,486 @@ +# Copyright (c) Sophont, Inc +# This source code is licensed under the Apache License, Version 2.0 +# +# References: +# deit: https://github.com/facebookresearch/deit/blob/main/main.py +# capi: https://github.com/facebookresearch/capi/blob/main/train_capi.py + +import argparse +import datetime +import json +import math +import random +import subprocess +import time +from contextlib import nullcontext +from functools import partial +from itertools import islice +from pathlib import Path +from typing import Iterable, Sequence + +import torch +import torch.nn as nn +import wandb +import webdataset as wds +from omegaconf import DictConfig, OmegaConf +from PIL import Image + +from matplotlib import pyplot as plt +from torch import Tensor + +import data.mri_data as mri_data +import smri_mae.model_mae as models_mae +import smri_mae.utils as ut +import smri_mae.visualization as vis + +DEFAULT_CONFIG = Path(__file__).parent / "config/default_pretrain.yaml" + +MODELS_DICT = models_mae.__dict__ + + +def main(args: DictConfig): + # setup + ut.init_distributed_mode(args) + global_rank = ut.get_rank() + is_master = global_rank == 0 + world_size = ut.get_world_size() + device = torch.device(args.device) + ut.configure_flash_sdpa() + ut.random_seed(args.seed, rank=global_rank) + + if args.name and not args.output_dir.endswith(args.name): + args.output_dir = f"{args.output_dir}/{args.name}" + output_dir = Path(args.output_dir) + + if is_master: + output_dir.mkdir(parents=True, exist_ok=True) + out_cfg_path = output_dir / "config.yaml" + if out_cfg_path.exists(): + prev_cfg = OmegaConf.load(out_cfg_path) + assert args == prev_cfg, "current config doesn't match previous config" + else: + OmegaConf.save(args, out_cfg_path) + + if args.wandb: + wandb.init( + entity=args.wandb_entity, + project=args.wandb_project, + name=args.name, + notes=args.notes, + config=OmegaConf.to_container(args), + ) + + ut.setup_for_distributed(log_path=output_dir / "log.txt") + + print("pretraining 3D ViTMAE") + print(f"start: {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}") + print(f"cwd: {Path.cwd()}") + print(ut.get_sha()) + print("config:", OmegaConf.to_yaml(args), sep="\n") + + # data loaders + train_loader, eval_loaders = create_data_loaders(args) + + # model + model = MODELS_DICT[args.model]( + img_size=args.img_size, + in_chans=args.get("in_chans", 1), + patch_size=args.patch_size, + **(args.get("model_kwargs") or {}), + ) + model.to(device) + print("model:", model, sep="\n") + num_params = sum(p.numel() for p in model.parameters() if p.requires_grad) + print(f"num params: {num_params / 1e6:.1f}M") + + model_without_ddp = model + if args.distributed: + model = torch.nn.parallel.DistributedDataParallel( + model, + device_ids=[args.gpu], + gradient_as_bucket_view=True, + ) + model_without_ddp = model.module + + # optimizer + total_batch_size = args.batch_size * args.accum_iter * world_size + print( + f"total batch size: {total_batch_size} = " + f"{args.batch_size} bs per gpu x {args.accum_iter} accum x {world_size} gpus" + ) + + if not args.get("lr"): + args.lr = args.base_lr * total_batch_size / 256 + print(f"lr: {args.lr:.2e} = {args.base_lr:.2e} x {total_batch_size} / 256") + else: + print(f"lr: {args.lr:.2e}") + + param_groups = ut.get_param_groups(model) + ut.update_lr(param_groups, args.lr) + ut.update_wd(param_groups, args.weight_decay) + # cast or else it corrupts the checkpoint + betas = tuple(args.betas) if args.betas is not None else None + optimizer = torch.optim.AdamW(param_groups, betas=betas, fused=True) + + epoch_num_batches = len(train_loader) + steps_per_epoch = math.ceil(epoch_num_batches / args.accum_iter) + total_steps = args.epochs * steps_per_epoch + warmup_steps = args.warmup_epochs * steps_per_epoch + lr_schedule = ut.WarmupThenCosine( + base_value=args.lr, + final_value=args.min_lr, + total_iters=total_steps, + warmup_iters=warmup_steps, + ) + print(f"full schedule: epochs = {args.epochs} (steps = {total_steps})") + print(f"warmup: epochs = {args.warmup_epochs} (steps = {warmup_steps})") + + # loss scaling not needed for bfloat16 (according to timm) + if args.amp and args.amp_dtype != "bfloat16": + loss_scaler = torch.GradScaler(device.type) + else: + loss_scaler = None + + # load checkpoint/resume training + ut.load_model(args, model_without_ddp, optimizer, loss_scaler) + + print(f"start training for {args.epochs} epochs") + start_time = time.monotonic() + for epoch in range(args.start_epoch, args.epochs): + train_stats = train_one_epoch( + args, + model, + train_loader, + optimizer, + loss_scaler, + lr_schedule, + epoch, + device, + ) + eval_stats = {} + eval_plots = {} + eval_period = args.get("eval_period", 1) + if eval_period and (epoch % eval_period == 0 or epoch == args.epochs - 1): + for name, loader in eval_loaders.items(): + stats, plots = evaluate( + args, + model, + loader, + epoch, + device, + eval_name=name, + ) + eval_stats.update(stats) + eval_plots.update(plots) + + merged_stats = {"epoch": epoch, **train_stats, **eval_stats} + if is_master: + with (output_dir / "log.json").open("a") as f: + print(json.dumps(merged_stats), file=f) + + for plot_name, img in eval_plots.items(): + plot_name = plot_name.replace("/", "__") + img.save(output_dir / f"{plot_name}__{epoch:05d}.png") + + ut.save_model(args, epoch, model_without_ddp, optimizer, loss_scaler) + sync_checkpoints_to_r2(args, output_dir) + + if args.distributed: + torch.distributed.destroy_process_group() + + total_time = time.monotonic() - start_time + print(f"done! training time: {datetime.timedelta(seconds=int(total_time))}") + + +def create_data_loaders(args: DictConfig): + data_loaders = {} + dataset_names = [args.train_dataset] + args.eval_datasets + + for dataset_name in dataset_names: + dataset_config = args.datasets[dataset_name].copy() + drop_last = dataset_config.pop("drop_last") + is_train = dataset_name == args.train_dataset + + print(f"loading dataset: {dataset_name}\n\n{OmegaConf.to_yaml(dataset_config)}") + shuffle = dataset_config["shuffle"] + samples_per_epoch = dataset_config.pop("samples_per_epoch") + dataset = mri_data.make_sparse_wds_dataset( + dataset_config["url"], + shuffle=shuffle, + buffer_size=dataset_config["buffer_size"], + ) + num_workers = int(args.num_workers) + loader_kwargs = { + "batch_size": args.batch_size, + "collate_fn": partial(mri_data.collate, include_meta=not is_train), + "shuffle": False, + "num_workers": num_workers, + "persistent_workers": num_workers > 0, + "pin_memory": True, + "drop_last": drop_last, + "prefetch_factor": args.prefetch_factor, + } + loader = wds.WebLoader(dataset, **loader_kwargs) + num_batches = samples_per_epoch // (ut.get_world_size() * args.batch_size) + loader = loader.with_epoch(num_batches) + loader = loader.with_length(num_batches, silent=True) + + data_loaders[dataset_name] = loader + + train_loader = data_loaders.pop(args.train_dataset) + return train_loader, data_loaders + + +def sync_checkpoints_to_r2(args: DictConfig, output_dir: Path) -> None: + r2_sync_url = args.get("r2_sync") + if not r2_sync_url or not ut.is_main_process(): + return + + cmd = ["aws", "s3", "sync", str(output_dir), str(r2_sync_url), "--profile", "r2"] + print(f"syncing checkpoints to R2: {output_dir} -> {r2_sync_url}") + subprocess.run(cmd, check=True) + + +def train_one_epoch( + args: DictConfig, + model: nn.Module, + data_loader: Iterable, + optimizer: torch.optim.Optimizer, + loss_scaler: torch.GradScaler | None, + lr_schedule: Sequence[float], + epoch: int, + device: torch.device, +): + model.train() + + metric_logger = ut.MetricLogger(delimiter=" ") + metric_logger.add_meter("lr", ut.SmoothedValue(window_size=1, fmt="{value:.6f}")) + metric_logger.add_meter("grad", ut.SmoothedValue()) + header = f"Train: [{epoch}]" + log_wandb = args.wandb and ut.is_main_process() + + epoch_num_batches = len(data_loader) + steps_per_epoch = math.ceil(epoch_num_batches / args.accum_iter) + + print_freq = args.get("print_freq", 100) if not args.debug else 1 + num_batches = epoch_num_batches if not args.debug else 10 + amp_dtype = getattr(torch, args.amp_dtype) + use_cuda = device.type == "cuda" + if use_cuda and args.presend_cuda: + data_loader = ut.pre_send_to_cuda_wrapper( + data_loader, device, dtype_map={torch.float16: amp_dtype} + ) + + optimizer.zero_grad() + + for batch_idx, batch in enumerate( + metric_logger.log_every(data_loader, print_freq, header, total_steps=num_batches) + ): + if use_cuda and not args.presend_cuda: + batch = ut.send_data(batch, device, dtype_map={torch.float16: amp_dtype}) + + batch_step = batch_idx + 1 + log_step = batch_step % print_freq == 0 or batch_step == num_batches + update_in_epoch = batch_idx // args.accum_iter + group_size = min(args.accum_iter, num_batches - update_in_epoch * args.accum_iter) + need_update = batch_step % args.accum_iter == 0 or batch_step == num_batches + global_step = epoch * steps_per_epoch + update_in_epoch + lr = lr_schedule[global_step] + if need_update: + ut.update_lr(optimizer.param_groups, lr) + + images, img_mask = mri_data.densify_sparse_image_batch( + batch["image_values"], + batch["img_mask"], + (int(args.get("in_chans", 1)), *args.img_size), + dtype=amp_dtype, + ) + + sync_context = model.no_sync() if args.distributed and not need_update else nullcontext() + with sync_context: + with torch.autocast(device_type=device.type, dtype=amp_dtype, enabled=args.amp): + loss = model( + images, + img_mask=img_mask, + mask_ratio=args.mask_ratio, + pred_mask_ratio=args.pred_mask_ratio, + pad_to_multiple=args.pad_to_multiple, + with_state=False, + ) + + loss_for_log = loss.detach() + torch._assert_async(torch.isfinite(loss_for_log), "non-finite loss") + + grad_norm = ut.backward_step( + loss / group_size, + optimizer, + scaler=loss_scaler, + need_update=need_update, + max_norm=args.clip_grad, + ) + + if need_update and log_step: + loss_value = loss_for_log.item() + grad_norm_value = grad_norm.item() + metric_logger.update(loss=loss_value, lr=lr, grad=grad_norm_value) + if log_wandb: + wandb.log( + { + "train/loss": loss_value, + "train/lr": lr, + "train/grad": grad_norm_value, + }, + step=int(1000 * (epoch + batch_step / epoch_num_batches)), + ) + + # gather the stats from all processes + metric_logger.synchronize_between_processes() + print("Averaged stats:", metric_logger) + return {f"train/{k}": meter.global_avg for k, meter in metric_logger.meters.items()} + + +@torch.inference_mode() +def evaluate( + args: DictConfig, + model: nn.Module, + data_loader: Iterable, + epoch: int, + device: torch.device, + eval_name: str, +): + model.eval() + + metric_logger = ut.MetricLogger(delimiter=" ") + header = f"Eval ({eval_name}): [{epoch}]" + is_master = ut.is_main_process() + log_wandb = args.wandb and is_master + + epoch_num_batches = len(data_loader) + if epoch_num_batches <= 0: + raise ValueError(f"eval loader {eval_name!r} has zero batches") + + print_freq = args.get("print_freq", 100) if not args.debug else 1 + num_batches = epoch_num_batches if not args.debug else 10 + num_batches = min(num_batches, epoch_num_batches) + eval_seed = int(args.get("eval_seed", args.seed)) + ut.get_rank() + example_step = random.Random(eval_seed).randint(1, num_batches) + amp_dtype = getattr(torch, args.amp_dtype) + use_cuda = device.type == "cuda" + rng_state = ut.capture_rng_state() + torch.set_rng_state(torch.Generator().manual_seed(eval_seed).get_state()) + if use_cuda: + torch.cuda.manual_seed(eval_seed) + if use_cuda and args.presend_cuda: + data_loader = ut.pre_send_to_cuda_wrapper( + data_loader, device, dtype_map={torch.float16: amp_dtype} + ) + + eval_batches = islice(data_loader, num_batches) + for batch_idx, batch in enumerate( + metric_logger.log_every(eval_batches, print_freq, header, total_steps=num_batches) + ): + if use_cuda and not args.presend_cuda: + batch = ut.send_data(batch, device, dtype_map={torch.float16: amp_dtype}) + + batch_step = batch_idx + 1 + + images, img_mask = mri_data.densify_sparse_image_batch( + batch["image_values"], + batch["img_mask"], + (int(args.get("in_chans", 1)), *args.img_size), + dtype=amp_dtype, + ) + + with torch.autocast(device_type=device.type, dtype=amp_dtype, enabled=args.amp): + loss, state = model( + images, + img_mask=img_mask, + mask_ratio=args.mask_ratio, + pred_mask_ratio=args.pred_mask_ratio, + pad_to_multiple=args.pad_to_multiple, + ) + + loss_value = loss.detach().float().item() + finite = torch.tensor(int(math.isfinite(loss_value)), dtype=torch.int32, device=device) + if args.distributed: + torch.distributed.all_reduce(finite, op=torch.distributed.ReduceOp.MIN) + if not finite.item(): + raise RuntimeError("non-finite validation loss detected") + metric_logger.meters["loss"].update(loss_value, n=int(batch["img_mask"].shape[0])) + + if is_master and batch_step == example_step: + example_batch = {"image": images[:1], "img_mask": img_mask[:1]} + if "meta" in batch: + example_batch["meta"] = batch["meta"][:1] + example_state = { + "pred_images": state["pred_images"][:1], + "pred_mask": state["pred_mask"][:1], + } + example_data = { + "batch": ut.send_data(example_batch, "cpu"), + "state": ut.send_data(example_state, "cpu"), + } + + # gather the stats from all processes + metric_logger.synchronize_between_processes() + print(f"Averaged stats ({eval_name}):", metric_logger) + stats = {f"eval/{eval_name}/{k}": meter.global_avg for k, meter in metric_logger.meters.items()} + + plots = {} + if is_master: + print(f"Making plots ({eval_name}): example={example_step}") + plots = make_plots(args, **example_data) + plots = {f"eval/{eval_name}/{k}": img for k, img in plots.items()} + + if log_wandb: + wandb.log(stats, step=1000 * (epoch + 1)) + wandb.log( + {k: wandb.Image(img, caption=f"example={example_step}") for k, img in plots.items()}, + step=1000 * (epoch + 1), + ) + ut.restore_rng_state(rng_state) + return stats, plots + + +def make_plots( + args: DictConfig, + batch: dict[str, Tensor], + state: dict[str, Tensor], +) -> dict[str, Image.Image]: + fig_kwargs = args.get("fig_kwargs", {}) + + images = batch["image"] + img_mask = batch.get("img_mask") + if img_mask is not None: + img_mask = img_mask.expand_as(images) + + raw_mean, raw_std = vis.raw_stats_from_batch(batch) + + plots = {} + mask_pred_fig = vis.plot_mask_pred( + target=images, + pred=state["pred_images"], + pred_mask=state["pred_mask"], + img_mask=img_mask, + patch_size=args.patch_size, + raw_mean=raw_mean, + raw_std=raw_std, + **ut.filter_kwargs(vis.plot_mask_pred, fig_kwargs), + ) + plots["mask_pred"] = vis.fig2pil(mask_pred_fig) + plt.close(mask_pred_fig) + + return plots + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--cfg-path", type=str, default=None) + parser.add_argument("--overrides", type=str, default=None, nargs="+") + args = parser.parse_args() + cfg = OmegaConf.load(DEFAULT_CONFIG) + if args.cfg_path: + cfg = OmegaConf.unsafe_merge(cfg, OmegaConf.load(args.cfg_path)) + if args.overrides: + cfg = OmegaConf.unsafe_merge(cfg, OmegaConf.from_dotlist(args.overrides)) + main(cfg) diff --git a/fomo_tune_baseline/output/task5/build/smri_mae/masking.py b/fomo_tune_baseline/output/task5/build/smri_mae/masking.py new file mode 100644 index 0000000000000000000000000000000000000000..28f9d53cbbd8d32a3923f9b1a6655c92cd526bf7 --- /dev/null +++ b/fomo_tune_baseline/output/task5/build/smri_mae/masking.py @@ -0,0 +1,80 @@ +import torch +from jaxtyping import Float, Int +from torch import Tensor + + +def pad_patch_mask( + patch_mask: Float[Tensor, "B N"], + mask_ratio: float, + shuffle: bool = False, + generator: torch.Generator | None = None, + pad_to_multiple: int | None = None, +) -> tuple[Float[Tensor, "B N"], Int[Tensor, "B L"], Tensor]: + """ + Select each row's own mask-ratio count, then pad ids to the batch max length. + + Returns: + - selected patch mask [B, N] + - padded selected patch ids [B, Lpad] + - token mask [B, Lpad], true for real ids and false for padding + """ + if not 0.0 <= mask_ratio <= 1.0: + raise ValueError(f"mask_ratio must be in [0, 1], got {mask_ratio}") + + B, N = patch_mask.shape + device = patch_mask.device + patch_mask = patch_mask.to(dtype=torch.bool) + + valid_counts = patch_mask.sum(dim=1) + num_keep = torch.floor(valid_counts.to(torch.float64) * (1.0 - mask_ratio)).to(torch.long) + if not shuffle: + selected = patch_mask & (patch_mask.cumsum(dim=1) <= num_keep.unsqueeze(1)) + padded_ids, token_mask = patch_ids_from_mask( + selected, + pad_to_multiple=pad_to_multiple, + ) + return selected, padded_ids, token_mask + + # One masked sort directly produces random valid IDs. The previous + # shuffle/select/inverse-shuffle path required two full argsorts plus a + # dynamic nonzero/scatter solely to recover the same selected set. + noise = torch.rand(B, N, generator=generator, device=device) + noise.masked_fill_(~patch_mask, torch.inf) + shuffled_ids = torch.argsort(noise, dim=1) + + max_count = int(num_keep.max().item()) + if pad_to_multiple is not None: + if pad_to_multiple <= 0: + raise ValueError(f"pad_to_multiple must be positive, got {pad_to_multiple}") + max_count = (max_count + pad_to_multiple - 1) // pad_to_multiple * pad_to_multiple + padded_ids = shuffled_ids[:, :max_count] + token_mask = torch.arange(max_count, device=device).unsqueeze(0) < num_keep.unsqueeze(1) + selected = torch.zeros_like(patch_mask).scatter_(1, padded_ids, token_mask) + return selected, padded_ids, token_mask + + +def patch_ids_from_mask( + patch_mask: Tensor, + pad_to_multiple: int | None = None, +) -> tuple[Int[Tensor, "B L"], Tensor]: + """Return optionally rounded patch IDs and their token-validity mask.""" + if pad_to_multiple is not None and pad_to_multiple <= 0: + raise ValueError(f"pad_to_multiple must be positive, got {pad_to_multiple}") + + patch_mask = patch_mask.to(dtype=torch.bool) + B, N = patch_mask.shape + device = patch_mask.device + counts = patch_mask.sum(dim=1) + max_count = int(counts.max().item()) + if pad_to_multiple is not None: + max_count = (max_count + pad_to_multiple - 1) // pad_to_multiple * pad_to_multiple + + patch_ids = torch.zeros((B, max_count), dtype=torch.long, device=device) + token_mask = torch.arange(max_count, device=device).unsqueeze(0) < counts.unsqueeze(1) + if max_count == 0: + return patch_ids, token_mask + + batch_ids, selected_ids = patch_mask.nonzero(as_tuple=True) + slot_ids = patch_mask.cumsum(dim=1)[batch_ids, selected_ids].to(torch.long) - 1 + patch_ids[batch_ids, slot_ids] = selected_ids + return patch_ids, token_mask diff --git a/fomo_tune_baseline/output/task5/build/smri_mae/model_mae.py b/fomo_tune_baseline/output/task5/build/smri_mae/model_mae.py new file mode 100644 index 0000000000000000000000000000000000000000..f9a4d884784e53f55ae557c46dd4afa2eb70e7d3 --- /dev/null +++ b/fomo_tune_baseline/output/task5/build/smri_mae/model_mae.py @@ -0,0 +1,916 @@ +# Copyright (c) Sophont, Inc +# This source code is licensed under the Apache License, Version 2.0 +# +# References: +# capi: https://github.com/facebookresearch/capi/blob/main/model.py +# timm: https://github.com/huggingface/pytorch-image-models/blob/v1.0.20/timm/models/vision_transformer.py + +""" +From-scratch re-implementation of the original MAE model. + +MaskedEncoder: standard ViT with masking +MaskedDecoder: standard self-attention MAE decoder +MaskedAutoEncoderViT: full MAE model for 3D structural MRI volumes +""" + +from collections.abc import Sequence +from typing import Literal + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch import Tensor +from torch.utils.checkpoint import checkpoint +from huggingface_hub import PyTorchModelHubMixin +from jaxtyping import Float, Int + +from .modules import ( + AbsolutePosEmbed, + Block, + LayerNorm, + Normalize, + JaggedBatch, + Patchify3D, + SeparablePosEmbed, + SinCosPosEmbed3D, + unpack_tokens, +) +from .masking import pad_patch_mask + + +class MaskedEncoder(nn.Module): + """ + Masked transformer encoder. + """ + + def __init__( + self, + patchify: nn.Module, + patch_embed: nn.Module, + pos_embed: nn.Module, + depth: int = 12, + embed_dim: int = 768, + num_heads: int = 12, + qkv_bias: bool = True, + proj_bias: bool = True, + mlp_ratio: int | float = 4, + class_token: bool = True, + reg_tokens: int = 0, + no_embed_class: bool = False, + final_norm: bool = True, + drop_path_rate: float = 0.0, + mask_drop_scale: bool = False, + ): + super().__init__() + self.num_prefix_tokens = int(class_token) + reg_tokens + self.num_reg_tokens = reg_tokens + self.has_class_token = class_token + self.no_embed_class = no_embed_class + + # scale inputs by 1 / observed rate (like dropout) + self.mask_drop_scale = mask_drop_scale + + # inject tokenization modules, so that the encoder doesn't specifically need to + # know how the data are tokenized, while still implementing a complete + # self-contained model. + self.patchify = patchify + self.patch_embed = patch_embed + self.pos_embed = pos_embed + + R = reg_tokens + self.cls_token = nn.Parameter(torch.empty(1, 1, embed_dim)) if class_token else None + self.reg_token = nn.Parameter(torch.empty(1, R, embed_dim)) if reg_tokens else None + + if not no_embed_class: + self.cls_token_pos = nn.Parameter(torch.empty(1, 1, embed_dim)) if class_token else None + self.reg_token_pos = nn.Parameter(torch.empty(1, R, embed_dim)) if reg_tokens else None + else: + self.cls_token_pos = self.reg_token_pos = None + + # stochastic depth decay rule + dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)] + + self.blocks = nn.ModuleList( + [ + Block( + dim=embed_dim, + num_heads=num_heads, + qkv_bias=qkv_bias, + proj_bias=proj_bias, + mlp_ratio=mlp_ratio, + drop_path=dpr[ii], + ) + for ii in range(depth) + ] + ) + + self.norm = LayerNorm(embed_dim) if final_norm else nn.Identity() + + self.reset_parameters() + + def extra_repr(self): + return ( + f"class_token={self.has_class_token}, reg_tokens={self.num_reg_tokens}, " + f"no_embed_class={self.no_embed_class}, mask_drop_scale={self.mask_drop_scale}" + ) + + def reset_parameters(self) -> None: + for p in [self.cls_token, self.cls_token_pos, self.reg_token, self.reg_token_pos]: + if p is not None: + nn.init.trunc_normal_(p, std=0.02) + + def cat_tokens(self, x: Tensor) -> Tensor: + # prepend cls and reg tokens with optional learned position embedding + # the cls and reg pos embedding is ofc redundant, but included in many other + # implementations. + B, _, _ = x.shape + + to_cat = [] + if self.has_class_token: + cls_token = self.cls_token + if not self.no_embed_class: + cls_token = cls_token + self.cls_token_pos + to_cat.append(cls_token.expand(B, -1, -1)) + + if self.num_reg_tokens: + reg_token = self.reg_token + if not self.no_embed_class: + reg_token = reg_token + self.reg_token_pos + to_cat.append(reg_token.expand(B, -1, -1)) + + if to_cat: + x = torch.cat(to_cat + [x], dim=1) + return x + + def cat_token_mask(self, token_mask: Tensor, batch_size: int) -> Tensor: + if self.num_prefix_tokens: + prefix_mask = torch.ones( + (batch_size, self.num_prefix_tokens), + dtype=torch.bool, + device=token_mask.device, + ) + token_mask = torch.cat([prefix_mask, token_mask], dim=1) + return token_mask + + def chunk_tokens(self, x: Tensor) -> tuple[Tensor | None, Tensor | None, Tensor]: + cls_offset = int(self.has_class_token) + cls = x[:, :cls_offset] if self.has_class_token else None + if self.num_reg_tokens: + reg = x[:, cls_offset : self.num_prefix_tokens, :] + else: + reg = None + patch = x[:, self.num_prefix_tokens :, :] + return cls, reg, patch + + def forward( + self, + x: Tensor, + mask: Tensor | None = None, + mask_ratio: float | None = None, + pad_to_multiple: int | None = None, + ) -> tuple[ + Float[Tensor, "B 1 D"] | None, + Float[Tensor, "B R D"] | None, + Float[Tensor, "B L D"], + Tensor | None, + Int[Tensor, "B L"] | None, + Tensor | None, + ]: + """ + x: input data shape [B, C, D, H, W] + mask: visible mask, 1 = visible, 0 = invisible. broadcastable shape + mask_ratio: mask ratio for uniform random masking + + returns: + - cls_embeds: [B, 1, D] + - reg_embeds: [B, R, D] + - patch_embeds: [B, L, D], where L is the number of visible patches + - mask: observed mask, 1 = observed, 0 = unobserved. same shape as input + - mask_ids: indices of visible patches [B L] + - token_mask: valid token mask for padded per-sample masking [B L] + """ + # apply mask to the input + if mask is not None: + mask = mask.to(device=x.device, dtype=torch.bool).expand_as(x) + x = x.masked_fill(~mask, 0) + + # patchify input + x = self.patchify(x) + B, N, P = x.shape + + # patchify mask and apply dropout style scaling + if mask is not None: + mask_patches = self.patchify(mask) + patch_num_obs = mask_patches.sum(dim=-1) + patch_mask = patch_num_obs > 0 + if self.mask_drop_scale: + patch_num_obs = patch_num_obs.to(x.dtype) + x = x * (P / patch_num_obs.unsqueeze(-1).clamp(min=1.0)) + elif mask_ratio is not None: + patch_mask = torch.ones((B, N), dtype=torch.bool, device=x.device) + mask_patches = patch_mask.unsqueeze(-1).expand(-1, -1, P) + else: + patch_mask = mask_patches = None + + # patch and position embed + x = self.patch_embed(x) + x = self.pos_embed(x) + + if mask is not None or mask_ratio is not None: + mask_ratio = 0.0 if mask_ratio is None else mask_ratio + patch_mask, mask_ids, token_mask = pad_patch_mask( + patch_mask, + mask_ratio=mask_ratio, + shuffle=mask_ratio > 0, + pad_to_multiple=pad_to_multiple, + ) + + mask_patches = mask_patches & patch_mask.unsqueeze(-1) + mask = self.patchify.unpatchify(mask_patches) + x = x.gather(1, mask_ids.unsqueeze(-1).expand(-1, -1, x.shape[-1])) + else: + mask_ids = None + token_mask = None + + cls_embeds, reg_embeds, patch_embeds = self.forward_patch_embeds( + x, + token_mask=token_mask, + ) + return cls_embeds, reg_embeds, patch_embeds, mask, mask_ids, token_mask + + def forward_patch_embeds( + self, + x: Float[Tensor, "B L D"], + token_mask: Tensor | None = None, + ) -> tuple[ + Float[Tensor, "B 1 D"] | None, + Float[Tensor, "B R D"] | None, + Float[Tensor, "B L D"], + ]: + B = x.shape[0] + if token_mask is None: + token_mask = torch.ones(x.shape[:2], dtype=torch.bool, device=x.device) + x = self.cat_tokens(x) + token_mask = self.cat_token_mask(token_mask, B) + jagged_batch = JaggedBatch.from_mask(token_mask) + x = x[token_mask] + for block in self.blocks: + x = block(x, jagged_batch=jagged_batch) + x = self.norm(x) + x = unpack_tokens(x, token_mask) + + cls_embeds, reg_embeds, patch_embeds = self.chunk_tokens(x) + return cls_embeds, reg_embeds, patch_embeds + + def forward_visible_ids( + self, + x: Tensor, + visible_ids: Int[Tensor, "B L"], + img_mask: Tensor | None = None, + ) -> tuple[ + Float[Tensor, "B 1 D"] | None, + Float[Tensor, "B R D"] | None, + Float[Tensor, "B L D"], + ]: + if img_mask is not None: + img_mask = img_mask.to(device=x.device, dtype=torch.bool).expand_as(x) + x = x.masked_fill(~img_mask, 0) + + x = self.patchify(x) + if self.mask_drop_scale and img_mask is not None: + mask_patches = self.patchify(img_mask) + patch_num_obs = mask_patches.sum(dim=-1).to(x.dtype) + x = x * (self.patchify.patch_dim / patch_num_obs.unsqueeze(-1).clamp(min=1.0)) + x = self.patch_embed(x) + x = self.pos_embed(x) + visible_ids = visible_ids.to(device=x.device) + x = x.gather(1, visible_ids.unsqueeze(-1).expand(-1, -1, x.shape[-1])) + return self.forward_patch_embeds(x) + + def forward_embedding( + self, + x: Tensor, + mask: Tensor | None = None, + mask_ratio: float | None = None, + ): + cls_embeds, reg_embeds, patch_embeds, *_ = self.forward( + x, + mask=mask, + mask_ratio=mask_ratio, + ) + return cls_embeds, reg_embeds, patch_embeds + + +class MaskedDecoder(nn.Module): + """Self-attention MAE decoder supporting sparse subset decoding via pred_ids.""" + + def __init__( + self, + pos_embed: nn.Module, + head: nn.Module | None = None, + input_dim: int | None = None, + depth: int = 12, + embed_dim: int = 768, + num_heads: int = 12, + qkv_bias: bool = True, + proj_bias: bool = True, + mlp_ratio: int | float = 4, + class_token: bool = True, + no_embed_class: bool = False, + final_norm: bool = True, + ): + super().__init__() + input_dim = embed_dim if input_dim is None else input_dim + self.has_class_token = class_token + self.no_embed_class = no_embed_class + + self.cls_token = nn.Parameter(torch.empty(1, 1, embed_dim)) if class_token else None + self.cls_token_pos = ( + nn.Parameter(torch.empty(1, 1, embed_dim)) + if class_token and not no_embed_class + else None + ) + self.mask_token = nn.Parameter(torch.empty(1, 1, embed_dim)) + + # decoder position embedding, encodes query position information into masks + self.pos_embed = pos_embed + + self.proj = nn.Identity() if input_dim == embed_dim else nn.Linear(input_dim, embed_dim) + + self.blocks = nn.ModuleList( + [ + Block( + dim=embed_dim, + num_heads=num_heads, + qkv_bias=qkv_bias, + proj_bias=proj_bias, + mlp_ratio=mlp_ratio, + ) + for _ in range(depth) + ] + ) + + self.norm = LayerNorm(embed_dim) if final_norm else nn.Identity() + + # optional injected prediction head + self.head = nn.Identity() if head is None else head + + self.reset_parameters() + + def extra_repr(self): + return f"class_token={self.has_class_token}, no_embed_class={self.no_embed_class}" + + def reset_parameters(self) -> None: + # official mae initializes decoder cls token to zeros + # although perhaps this was an oversight + if self.cls_token is not None: + nn.init.zeros_(self.cls_token) + if self.cls_token_pos is not None: + nn.init.trunc_normal_(self.cls_token_pos, std=0.02) + nn.init.trunc_normal_(self.mask_token, std=0.02) + + def cat_tokens(self, x: Tensor) -> Tensor: + if not self.has_class_token: + return x + cls_token = self.cls_token + if not self.no_embed_class: + cls_token = cls_token + self.cls_token_pos + return torch.cat([cls_token.expand(x.shape[0], -1, -1), x], dim=1) + + def cat_token_mask(self, token_mask: Tensor, batch_size: int) -> Tensor: + if self.has_class_token: + cls_mask = torch.ones( + (batch_size, 1), + dtype=torch.bool, + device=token_mask.device, + ) + token_mask = torch.cat([cls_mask, token_mask], dim=1) + return token_mask + + def chunk_tokens(self, x: Tensor) -> tuple[Tensor | None, Tensor]: + cls_offset = int(self.has_class_token) + cls = x[:, :cls_offset] if self.has_class_token else None + patch = x[:, cls_offset:, :] + return cls, patch + + def forward( + self, + embeds: Float[Tensor, "B L D"], + embed_ids: Int[Tensor, "B L"] | None = None, + pred_ids: Int[Tensor, "B Q"] | None = None, + embed_token_mask: Tensor | None = None, + pred_token_mask: Tensor | None = None, + packed_output: bool = False, + ) -> Float[Tensor, "B Q P"] | Float[Tensor, "T P"]: + """ + embeds: input patch embeddings. + embed_ids: optional patch indices for input embeddings. If not provided, no + position will be added to the embeddings. + pred_ids: patch indices of query mask positions. If None, decode *all* patches. + + returns: + - pred [B, Q, P] where Q is the number of prediction patches and P is the output + dimension + """ + B, L, _ = embeds.shape + + Q = self.pos_embed.num_patches if pred_ids is None else pred_ids.shape[1] + mask = self.mask_token.expand(B, Q, -1) + mask = self.pos_embed(mask, pos_ids=pred_ids) + + embeds = self.proj(embeds) + + if embed_ids is not None: + embeds = self.pos_embed(embeds, pos_ids=embed_ids) + if embed_token_mask is None: + embed_token_mask = torch.ones((B, L), dtype=torch.bool, device=embeds.device) + if pred_token_mask is None: + pred_token_mask = torch.ones((B, Q), dtype=torch.bool, device=embeds.device) + x = torch.cat([embeds, mask], dim=1) + token_mask = torch.cat([embed_token_mask, pred_token_mask], dim=1) + + x = self.cat_tokens(x) + token_mask = self.cat_token_mask(token_mask, B) + jagged_batch = JaggedBatch.from_mask(token_mask) + x = x[token_mask] + # Keep headroom for rare maximum-length PSP batches. + checkpoint_start = max(0, len(self.blocks) - 2) + for block_index, block in enumerate(self.blocks): + if self.training and torch.is_grad_enabled() and block_index >= checkpoint_start: + x = checkpoint(block, x, jagged_batch, use_reentrant=False) + else: + x = block(x, jagged_batch=jagged_batch) + + x = self.norm(x) + if packed_output: + pred_offset = int(self.has_class_token) + L + prediction_mask = F.pad(pred_token_mask, (pred_offset, 0)) + return self.head(x[prediction_mask[token_mask]]) + + x = unpack_tokens(x, token_mask) + _, x = self.chunk_tokens(x) + + pred = x[:, L:] + pred = pred.masked_fill(~pred_token_mask.unsqueeze(-1), 0) + pred = self.head(pred) + return pred + + +class MaskedAutoencoderViT(nn.Module, PyTorchModelHubMixin): + def __init__( + self, + img_size: int | tuple[int, int, int] = (208, 240, 208), + patch_size: int | tuple[int, int, int] = (16, 16, 16), + in_chans: int = 1, + depth: int = 12, + embed_dim: int = 768, + num_heads: int = 12, + decoder_depth: int = 4, + decoder_embed_dim: int | None = 512, + decoder_num_heads: int | None = 16, # default from mae, head dim = 32 + qkv_bias: bool = True, + proj_bias: bool = True, + mlp_ratio: int | float = 4, + class_token: bool = True, + reg_tokens: int = 0, + no_embed_class: bool = False, + drop_path_rate: float = 0.0, + mask_drop_scale: bool = False, + no_decode_pos: bool = False, + pos_embed: Literal["abs", "sep", "sincos"] = "sincos", + target_norm: Literal["none", "global", "slice", "patch"] | None = None, + ): + super().__init__() + img_size = _to_3d_tuple(img_size, "img_size") + patch_size = _to_3d_tuple(patch_size, "patch_size") + + self.no_decode_pos = no_decode_pos # don't pos encode embeddings in decoder + + # patchify reshapes input into sequence of flattened patches, shape [B, N, P] + ndim = 3 + patchify = Patchify3D(img_size, patch_size, in_chans=in_chans) + + # linear patch embedding P -> D + patch_embed = nn.Linear(patchify.patch_dim, embed_dim) + + # position embedding + # separable position embedding decouples the first spatial axis from the + # others. Fixed sin/cos embeddings are the default for sMRI volumes. + if pos_embed == "sincos": + pos_embed_layer = SinCosPosEmbed3D + else: + pos_embed_layer = {"abs": AbsolutePosEmbed, "sep": SeparablePosEmbed}[pos_embed] + pos_embed = pos_embed_layer(embed_dim, patchify.grid_size) + + # encoder. for inference, this model can be extracted and used like a regular vit + self.encoder = MaskedEncoder( + patchify=patchify, + patch_embed=patch_embed, + pos_embed=pos_embed, + depth=depth, + embed_dim=embed_dim, + num_heads=num_heads, + qkv_bias=qkv_bias, + proj_bias=proj_bias, + mlp_ratio=mlp_ratio, + class_token=class_token, + reg_tokens=reg_tokens, + no_embed_class=no_embed_class, + drop_path_rate=drop_path_rate, + mask_drop_scale=mask_drop_scale, + ) + + self.pred_patchify = patchify + + # fall back to encoder architecture width + decoder_embed_dim = decoder_embed_dim or embed_dim + decoder_num_heads = decoder_num_heads or num_heads + + decoder_pos_embed = pos_embed_layer(decoder_embed_dim, self.pred_patchify.grid_size) + # we might want to try tying the weights of the prediction head to the patch + # embedding at some point. + decoder_head = nn.Linear(decoder_embed_dim, self.pred_patchify.patch_dim) + + self.decoder = MaskedDecoder( + pos_embed=decoder_pos_embed, + head=decoder_head, + input_dim=embed_dim, + depth=decoder_depth, + embed_dim=decoder_embed_dim, + num_heads=decoder_num_heads, + qkv_bias=qkv_bias, + proj_bias=proj_bias, + mlp_ratio=mlp_ratio, + class_token=class_token, + no_embed_class=no_embed_class, + ) + + # mae style target normalization + # dim is relative to an unflattened embedding tensor of shape [B, *grid_size, D] + if target_norm not in {"none", None}: + norm_dim = { + "global": tuple(range(1, ndim + 2)), # full sequence + "slice": tuple(range(2, ndim + 2)), # each depth slice along first dim + "patch": -1, # normalize each patch independently (mae pix norm loss) + }[target_norm] + self.target_norm = Normalize(self.pred_patchify.grid_size, dim=norm_dim) + else: + self.target_norm = None + + self.init_weights() + + def extra_repr(self): + return f"no_decode_pos={self.no_decode_pos}" + + def init_weights(self): + self.apply(_init_weights) + + def prepare_targets(self, images: Tensor, img_mask: Tensor | None): + """ + images: [B, C, D, H, W] + img_mask: mask of valid data. only used for computing correct normalization + stats. same shape as images. + """ + targets_patches = self.pred_patchify(images) # [B, N, P] + + # target normalization + if self.target_norm is not None: + # full image data mask used for normalization stats only + if img_mask is not None: + img_mask_patches = self.pred_patchify(img_mask) + else: + img_mask_patches = None + targets_patches, *targets_stats = self.target_norm(targets_patches, img_mask_patches) + else: + targets_stats = None + + return targets_patches, targets_stats + + def prepare_masks( + self, + img_mask: Tensor, + visible_mask: Tensor | None, + pred_mask: Tensor | None, + device: torch.device, + ): + img_mask = img_mask.to(device=device, dtype=torch.bool) + + if visible_mask is None: + visible_mask = img_mask + else: + visible_mask = img_mask & visible_mask.to(device=device, dtype=torch.bool) + + if pred_mask is None: + pred_mask = img_mask + else: + pred_mask = img_mask & pred_mask.to(device=device, dtype=torch.bool) + + return img_mask, visible_mask, pred_mask + + def prepare_pred_mask( + self, + visible_mask: Tensor, + pred_mask: Tensor | None = None, + pred_mask_ratio: float | None = None, + pad_to_multiple: int | None = None, + ): + """ + prepare prediction mask by removing visible content + visible_mask: [B, C, D, H, W], 1 = visible, 0 = invisible + pred_mask: same shape, 1 = predict, 0 = don't predict + """ + if pred_mask is None: + pred_mask = torch.ones_like(visible_mask) + + pred_mask = pred_mask & ~visible_mask + + pred_mask_patches = self.pred_patchify(pred_mask) + pred_patch_mask = pred_mask_patches.any(dim=-1) + # Optionally subsample the prediction candidates. + mask_ratio = 0.0 if pred_mask_ratio is None else pred_mask_ratio + pred_patch_mask, pred_ids, pred_token_mask = pad_patch_mask( + pred_patch_mask, + mask_ratio=mask_ratio, + # With per-sample padding every candidate is retained when the ratio + # is zero, so randomizing their order is pure overhead. + shuffle=mask_ratio > 0, + pad_to_multiple=pad_to_multiple, + ) + pred_mask_patches = pred_mask_patches & pred_patch_mask.unsqueeze(-1) + return pred_mask_patches, pred_ids, pred_token_mask + + def forward_decoder( + self, + patch_embeds: Float[Tensor, "B L D"], + visible_ids: Int[Tensor, "B L"], + pred_ids: Int[Tensor, "B Q"] | None, + visible_token_mask: Tensor | None = None, + pred_token_mask: Tensor | None = None, + packed_output: bool = False, + ) -> Float[Tensor, "B Q P"] | Float[Tensor, "T P"]: + return self.decoder( + patch_embeds, + embed_ids=None if self.no_decode_pos else visible_ids, + pred_ids=pred_ids, + embed_token_mask=visible_token_mask, + pred_token_mask=pred_token_mask, + packed_output=packed_output, + ) + + def forward_loss( + self, + preds: Float[Tensor, "T P"], + targets_patches: Float[Tensor, "B N P"], + pred_mask_patches: Float[Tensor, "B N P"], + pred_ids: Int[Tensor, "B Q"], + pred_token_mask: Tensor, + ) -> Tensor: + """Average valid-voxel MSE within each scan, then average across scans.""" + batch_ids, slot_ids = pred_token_mask.nonzero(as_tuple=True) + patch_ids = pred_ids[batch_ids, slot_ids] + targets = targets_patches[batch_ids, patch_ids] + voxel_mask = pred_mask_patches[batch_ids, patch_ids] + + patch_errors = ((preds - targets) ** 2 * voxel_mask).sum(dim=1) + patch_voxels = voxel_mask.sum(dim=1).to(dtype=patch_errors.dtype) + batch_size = targets_patches.shape[0] + scan_errors = patch_errors.new_zeros(batch_size).scatter_add_(0, batch_ids, patch_errors) + scan_voxels = patch_voxels.new_zeros(batch_size).scatter_add_(0, batch_ids, patch_voxels) + return (scan_errors / scan_voxels).mean() + + @torch.no_grad() + def forward_pred_images( + self, + preds: Float[Tensor, "B Q P"], + pred_ids: Int[Tensor, "B Q"], + pred_token_mask: Tensor | None = None, + img_mask: Tensor | None = None, + targets_stats: tuple[Tensor, Tensor] | None = None, + ) -> Tensor: + B, _, P = preds.shape + N = self.pred_patchify.num_patches + if pred_token_mask is not None: + preds = preds.masked_fill(~pred_token_mask.unsqueeze(-1), 0) + + preds = torch.zeros((B, N, P), dtype=preds.dtype, device=preds.device).scatter_add_( + 1, pred_ids.unsqueeze(-1).expand(-1, -1, P), preds + ) + + if targets_stats is not None: + targets_mean, targets_std = targets_stats + preds = preds * targets_std + targets_mean + + pred_images = self.pred_patchify.unpatchify(preds) + if img_mask is not None: + pred_images = pred_images.masked_fill(~img_mask, 0) + return pred_images + + def forward( + self, + images: Tensor, + img_mask: Tensor, + mask_ratio: float, + pred_mask_ratio: float | None = None, + pad_to_multiple: int | None = None, + with_state: bool = True, + ) -> Tensor | tuple[Tensor, dict]: + img_mask, visible_mask, pred_mask = self.prepare_masks( + img_mask, + None, + None, + device=images.device, + ) + targets_patches, targets_stats = self.prepare_targets(images, img_mask) + + ( + cls_embeds, + reg_embeds, + patch_embeds, + visible_mask, + visible_ids, + visible_token_mask, + ) = self.encoder( + images, + mask=visible_mask, + mask_ratio=mask_ratio, + pad_to_multiple=pad_to_multiple, + ) + + pred_mask_patches, pred_ids, pred_token_mask = self.prepare_pred_mask( + visible_mask, + pred_mask=pred_mask, + pred_mask_ratio=pred_mask_ratio, + pad_to_multiple=pad_to_multiple, + ) + + preds = self.forward_decoder( + patch_embeds, + visible_ids, + pred_ids, + visible_token_mask=visible_token_mask, + pred_token_mask=pred_token_mask, + packed_output=not with_state, + ) + + loss_preds = preds if not with_state else preds[pred_token_mask] + loss = self.forward_loss( + loss_preds, + targets_patches, + pred_mask_patches, + pred_ids, + pred_token_mask, + ) + + if not with_state: + return loss + + pred_mask = self.pred_patchify.unpatchify(pred_mask_patches) + pred_images = self.forward_pred_images( + preds, + pred_ids, + pred_token_mask=pred_token_mask, + img_mask=img_mask, + targets_stats=targets_stats, + ) + + state = { + "targets_patches": targets_patches, + "targets_stats": targets_stats, + "patch_embeds": patch_embeds, + "cls_embeds": cls_embeds, + "reg_embeds": reg_embeds, + "visible_mask": visible_mask, + "visible_ids": visible_ids, + "visible_token_mask": visible_token_mask, + "pred_mask": pred_mask, + "pred_ids": pred_ids, + "pred_token_mask": pred_token_mask, + "preds": preds, + "pred_images": pred_images, + } + return loss, state + + def forward_embedding( + self, + x: Tensor, + mask: Tensor | None = None, + mask_ratio: float | None = None, + ): + return self.encoder.forward_embedding(x, mask=mask, mask_ratio=mask_ratio) + + +class MaskedViT(MaskedEncoder, PyTorchModelHubMixin): + def __init__( + self, + img_size: int | tuple[int, int, int] = (208, 240, 208), + in_chans: int = 1, + patch_size: int | tuple[int, int, int] = (16, 16, 16), + depth: int = 12, + embed_dim: int = 768, + num_heads: int = 12, + qkv_bias: bool = True, + proj_bias: bool = True, + mlp_ratio: int | float = 4, + class_token: bool = True, + reg_tokens: int = 0, + no_embed_class: bool = False, + final_norm: bool = True, + drop_path_rate: float = 0.0, + mask_drop_scale: bool = False, + pos_embed: Literal["abs", "sep", "sincos"] = "sincos", + ): + img_size = _to_3d_tuple(img_size, "img_size") + patch_size = _to_3d_tuple(patch_size, "patch_size") + + patchify = Patchify3D(img_size, patch_size, in_chans=in_chans) + patch_embed = nn.Linear(patchify.patch_dim, embed_dim) + if pos_embed == "sincos": + pos_embed_layer = SinCosPosEmbed3D + else: + pos_embed_layer = {"abs": AbsolutePosEmbed, "sep": SeparablePosEmbed}[pos_embed] + pos_embed = pos_embed_layer(embed_dim, patchify.grid_size) + + super().__init__( + patchify=patchify, + patch_embed=patch_embed, + pos_embed=pos_embed, + depth=depth, + embed_dim=embed_dim, + num_heads=num_heads, + qkv_bias=qkv_bias, + proj_bias=proj_bias, + mlp_ratio=mlp_ratio, + class_token=class_token, + reg_tokens=reg_tokens, + no_embed_class=no_embed_class, + final_norm=final_norm, + drop_path_rate=drop_path_rate, + mask_drop_scale=mask_drop_scale, + ) + + self.init_weights() + + def init_weights(self): + self.apply(_init_weights) + + +def _to_3d_tuple(value: int | Sequence[int], name: str) -> tuple[int, int, int]: + if isinstance(value, int): + return (value, value, value) + if len(value) != 3: + raise ValueError(f"{name} must have exactly 3 spatial dimensions, got {tuple(value)}") + return tuple(int(item) for item in value) + + +# JAX ViT xavier uniform init +# https://github.com/facebookresearch/capi/blob/main/model.py +def _init_weights(m: nn.Module) -> None: + if isinstance(m, nn.Linear): + nn.init.xavier_uniform_(m.weight) + if m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.LayerNorm) and m.elementwise_affine: + nn.init.constant_(m.weight, 1.0) + if m.bias is not None: + nn.init.constant_(m.bias, 0) + + +def _create_vit(**kwargs): + model = MaskedViT(**kwargs) + return model + + +def _create_mae_vit(**kwargs): + model = MaskedAutoencoderViT(**kwargs) + return model + + +def mae_vit_small(**kwargs): + model_args = dict(embed_dim=384, depth=12, num_heads=6) + return _create_mae_vit(**model_args, **kwargs) + + +def mae_vit_base(**kwargs): + model_args = dict(embed_dim=768, depth=12, num_heads=12) + return _create_mae_vit(**model_args, **kwargs) + + +def mae_vit_large(**kwargs): + model_args = dict(embed_dim=1024, depth=24, num_heads=16) + return _create_mae_vit(**model_args, **kwargs) + + +def mae_vit_huge(**kwargs): + model_args = dict(embed_dim=1280, depth=32, num_heads=16) + return _create_mae_vit(**model_args, **kwargs) + + +# "patch embed" baseline model, depth 0 ViT (hah) +def patch_embed_small(**kwargs): + model_args = dict(embed_dim=384, depth=0) + return _create_vit(**model_args, **kwargs) + + +def patch_embed_base(**kwargs): + model_args = dict(embed_dim=768, depth=0) + return _create_vit(**model_args, **kwargs) diff --git a/fomo_tune_baseline/output/task5/build/smri_mae/modules.py b/fomo_tune_baseline/output/task5/build/smri_mae/modules.py new file mode 100644 index 0000000000000000000000000000000000000000..828aecb1bd57182f1690feb4e943395db1bdd515 --- /dev/null +++ b/fomo_tune_baseline/output/task5/build/smri_mae/modules.py @@ -0,0 +1,453 @@ +# This source code is licensed under the Apache License, Version 2.0 +# +# References: +# capi: https://github.com/facebookresearch/capi/blob/main/model.py +# timm: https://github.com/huggingface/pytorch-image-models/blob/v1.0.20/timm/models/vision_transformer.py +# vjepa2: https://github.com/facebookresearch/vjepa2/blob/main/src/models/utils/pos_embs.py + +import math +from functools import partial +from typing import NamedTuple, Type + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch import Tensor +from einops import rearrange +from jaxtyping import Float, Int +from timm.layers import DropPath, to_3tuple + +Layer = Type[nn.Module] + + +class JaggedBatch(NamedTuple): + """Sequence boundaries and cached launch metadata for jagged attention.""" + + offsets: Tensor + max_seqlen: int + + @classmethod + def from_mask(cls, mask: Tensor) -> "JaggedBatch": + mask = mask.to(dtype=torch.bool) + counts = mask.sum(dim=1) + return cls( + offsets=F.pad(counts.cumsum(dim=0), (1, 0)), + max_seqlen=mask.shape[1], + ) + + def as_nested(self, tokens: Tensor) -> Tensor: + # Cached conservative bounds avoid min/max reductions and GPU-to-CPU + # synchronization when Flash SDPA inspects the jagged sequence lengths. + return torch.nested.nested_tensor_from_jagged( + tokens, + self.offsets, + min_seqlen=1, + max_seqlen=self.max_seqlen, + ).transpose(1, 2) + + +def unpack_tokens(tokens: Tensor, token_mask: Tensor) -> Tensor: + """Restore packed values to a padded batch, filling invalid slots with zero.""" + output = tokens.new_zeros((*token_mask.shape, *tokens.shape[1:])) + return output.index_put((token_mask,), tokens) + + +def jagged_scaled_dot_product_attention( + query: Tensor, + key: Tensor, + value: Tensor, + jagged_batch: JaggedBatch, +) -> Tensor: + """Run SDPA on a packed batch of variable-length sequences.""" + output_jagged = F.scaled_dot_product_attention( + jagged_batch.as_nested(query), + jagged_batch.as_nested(key), + jagged_batch.as_nested(value), + ) + return output_jagged.transpose(1, 2).values() + + +class Attention(nn.Module): + def __init__( + self, + dim: int, + num_heads: int, + qkv_bias: bool = False, + proj_bias: bool = False, + ) -> None: + super().__init__() + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.qkv = nn.Linear(dim, 3 * dim, bias=qkv_bias) + self.proj = nn.Linear(dim, dim, bias=proj_bias) + + def extra_repr(self): + return f"num_heads={self.num_heads}" + + def forward( + self, + x: Float[Tensor, "L D"], + jagged_batch: JaggedBatch, + ) -> Float[Tensor, "L D"]: + L, D = x.shape + h, dh = self.num_heads, self.head_dim + + qkv = self.qkv(x).reshape(L, 3, h, dh) + q, k, v = qkv.unbind(1) + + x = jagged_scaled_dot_product_attention( + q, + k, + v, + jagged_batch=jagged_batch, + ) + x = x.reshape(L, D) + x = self.proj(x) + return x + + +class Mlp(nn.Module): + def __init__( + self, + dim: int, + mlp_ratio: int | float = 4, + bias: bool = False, + ) -> None: + super().__init__() + hidden_features = int(dim * mlp_ratio) + self.fc1 = nn.Linear(dim, hidden_features, bias=bias) + self.act = nn.GELU() + self.fc2 = nn.Linear(hidden_features, dim, bias=bias) + + def forward(self, x: Float[Tensor, "... D"]) -> Float[Tensor, "... D"]: + x = self.fc1(x) + x = self.act(x) + x = self.fc2(x) + return x + + +# timm default eps=1e-6 +LayerNorm = partial(nn.LayerNorm, eps=1e-6) + + +class Block(nn.Module): + def __init__( + self, + dim: int, + num_heads: int, + qkv_bias: bool = False, + proj_bias: bool = False, + mlp_ratio: int | float = 4, + drop_path: float = 0.0, + norm_layer: Layer = LayerNorm, + ) -> None: + super().__init__() + self.norm1 = norm_layer(dim) + self.attn = Attention( + dim=dim, + num_heads=num_heads, + qkv_bias=qkv_bias, + proj_bias=proj_bias, + ) + self.drop_path1 = DropPath(drop_path) if drop_path > 0 else nn.Identity() + + self.norm2 = norm_layer(dim) + self.mlp = Mlp( + dim=dim, + mlp_ratio=mlp_ratio, + bias=proj_bias, + ) + self.drop_path2 = DropPath(drop_path) if drop_path > 0 else nn.Identity() + + def forward( + self, + x: Float[Tensor, "L D"], + jagged_batch: JaggedBatch, + ) -> Float[Tensor, "L D"]: + x = x + self.drop_path1( + self.attn( + self.norm1(x), + jagged_batch=jagged_batch, + ) + ) + x = x + self.drop_path2(self.mlp(self.norm2(x))) + return x + + +# Patching and position embedding modules + + +class Patchify3D(nn.Module): + def __init__( + self, + img_size: int | tuple[int, int, int], + patch_size: int | tuple[int, int, int], + in_chans: int = 3, + ) -> None: + super().__init__() + self.img_size = to_3tuple(img_size) + self.patch_size = to_3tuple(patch_size) + self.in_chans = in_chans + + T, H, W = self.img_size + p_t, p_h, p_w = self.patch_size + if T % p_t or H % p_h or W % p_w: + raise ValueError( + f"img_size {self.img_size} must be divisible by patch_size {self.patch_size}" + ) + self.grid_size = (T // p_t, H // p_h, W // p_w) + self.num_patches = math.prod(self.grid_size) + self.patch_dim = in_chans * math.prod(self.patch_size) + + def forward(self, x: Float[Tensor, "B C T H W"]) -> Float[Tensor, "B N P"]: + x = patchify3d(x, self.patch_size) + return x + + def unpatchify(self, x: Float[Tensor, "B N P"]) -> Float[Tensor, "B C T H W"]: + x = unpatchify3d(x, patch_size=self.patch_size, img_size=self.img_size) + return x + + def extra_repr(self): + return f"{self.img_size}, {self.patch_size}, in_chans={self.in_chans}" + + +def patchify3d(x: Tensor, patch_size: tuple[int, int, int]) -> Tensor: + p_t, p_h, p_w = to_3tuple(patch_size) + B, C, T, H, W = x.shape + x = rearrange(x, "b c (t u) (h p) (w q) -> b (t h w) (c u p q)", u=p_t, p=p_h, q=p_w) + return x + + +def unpatchify3d( + x: Tensor, + patch_size: tuple[int, int, int], + img_size: tuple[int, int, int], +) -> Tensor: + B, N, P = x.shape + p_t, p_h, p_w = to_3tuple(patch_size) + T, H, W = to_3tuple(img_size) + x = rearrange( + x, + "b (t h w) (c u p q) -> b c (t u) (h p) (w q)", + t=T // p_t, + h=H // p_h, + w=W // p_w, + u=p_t, + p=p_h, + q=p_w, + ) + return x + + +class AbsolutePosEmbed(nn.Module): + def __init__(self, embed_dim: int, grid_size: tuple[int, ...]) -> None: + super().__init__() + self.embed_dim = embed_dim + self.grid_size = grid_size + self.num_patches = math.prod(grid_size) + + self.weight = nn.Parameter(torch.empty(self.num_patches, embed_dim)) + self.reset_parameters() + + def reset_parameters(self): + nn.init.trunc_normal_(self.weight, std=0.02) + + def forward( + self, + x: Float[Tensor, "B L D"], + pos_ids: Int[Tensor, "B L"] | None = None, + ) -> Float[Tensor, "B L D"]: + x = apply_pos_embed(x, self.weight, pos_ids=pos_ids) + return x + + def extra_repr(self): + return f"{self.embed_dim}, {self.grid_size}" + + +class SeparablePosEmbed(nn.Module): + def __init__(self, embed_dim: int, grid_size: tuple[int, ...]) -> None: + super().__init__() + self.embed_dim = embed_dim + self.grid_size = grid_size + self.num_patches = math.prod(grid_size) + + N_t, *grid_size_spatial = grid_size + N_s = math.prod(grid_size_spatial) + self.weight_spatial = nn.Parameter(torch.empty(1, N_s, embed_dim)) + self.weight_temporal = nn.Parameter(torch.empty(N_t, 1, embed_dim)) + self.reset_parameters() + + def reset_parameters(self): + nn.init.trunc_normal_(self.weight_spatial, std=0.02) + nn.init.trunc_normal_(self.weight_temporal, std=0.02) + + def forward( + self, + x: Float[Tensor, "B L D"], + pos_ids: Int[Tensor, "B L"] | None = None, + ) -> Float[Tensor, "B L D"]: + B, N, D = x.shape + weight = (self.weight_temporal + self.weight_spatial).flatten(0, 1) # [N, D] + x = apply_pos_embed(x, weight, pos_ids=pos_ids) + return x + + def extra_repr(self): + return f"{self.embed_dim}, {self.grid_size}" + + +class SinCosPosEmbed3D(nn.Module): + def __init__(self, embed_dim: int, grid_size: tuple[int, int, int]) -> None: + super().__init__() + self.embed_dim = embed_dim + self.grid_size = grid_size + self.num_patches = math.prod(grid_size) + + N_t, N_h, N_w = grid_size + weight = get_3d_sincos_pos_embed( + embed_dim=embed_dim, + grid_size=(N_h, N_w), + grid_depth=N_t, + uniform_power=True, + ) + self.weight = nn.Parameter(torch.from_numpy(weight).float(), requires_grad=False) + + def forward( + self, + x: Float[Tensor, "B L D"], + pos_ids: Int[Tensor, "B L"] | None = None, + ) -> Float[Tensor, "B L D"]: + x = apply_pos_embed(x, self.weight, pos_ids=pos_ids) + return x + + def extra_repr(self): + return f"{self.embed_dim}, {self.grid_size}" + + +# sincos pos embed utils from vjepa2, but fixed the confusing meshgrid indexing + + +def get_3d_sincos_pos_embed(embed_dim, grid_size, grid_depth, cls_token=False, uniform_power=False): + """ + grid_size: tuple of int of the grid height and width + grid_depth: int of the grid depth + returns: + pos_embed: [grid_depth*grid_height*grid_width, embed_dim] (w/o cls_token) + or [1+grid_depth*grid_height*grid_width, embed_dim] (w/ cls_token) + """ + grid_d = np.arange(grid_depth, dtype=float) + grid_h = np.arange(grid_size[0], dtype=float) + grid_w = np.arange(grid_size[1], dtype=float) + grid_d, grid_h, grid_w = np.meshgrid(grid_d, grid_h, grid_w, indexing="ij") + + if not uniform_power: + h_embed_dim = embed_dim // 4 + w_embed_dim = embed_dim // 4 + d_embed_dim = embed_dim // 2 + else: + h_embed_dim = w_embed_dim = d_embed_dim = int(np.ceil(embed_dim / 6) * 2) + + emb_h = get_1d_sincos_pos_embed_from_grid(h_embed_dim, grid_h) # (T*H*W, D1) + emb_w = get_1d_sincos_pos_embed_from_grid(w_embed_dim, grid_w) # (T*H*W, D2) + emb_d = get_1d_sincos_pos_embed_from_grid(d_embed_dim, grid_d) # (T*H*W, D3) + pos_embed = np.concatenate([emb_d, emb_h, emb_w], axis=1) + pos_embed = pos_embed[:, :embed_dim] + if cls_token: + pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0) + return pos_embed + + +def get_1d_sincos_pos_embed_from_grid(embed_dim, pos): + """ + embed_dim: output dimension for each position + pos: a list of positions to be encoded: size (M,) + returns: (M, D) + """ + assert embed_dim % 2 == 0 + omega = np.arange(embed_dim // 2, dtype=float) + omega /= embed_dim / 2.0 + omega = 1.0 / 10000**omega # (D/2,) + + pos = pos.reshape(-1) # (M,) + out = np.einsum("m,d->md", pos, omega) # (M, D/2), outer product + + emb_sin = np.sin(out) # (M, D/2) + emb_cos = np.cos(out) # (M, D/2) + + emb = np.concatenate([emb_sin, emb_cos], axis=1) # (M, D) + return emb + + +def apply_pos_embed( + x: Float[Tensor, "B L D"], + weight: Float[Tensor, "N D"], + pos_ids: Int[Tensor, "B L"] | None = None, +) -> Float[Tensor, "B L D"]: + B, L, D = x.shape + weight = weight.expand(B, -1, -1) + if pos_ids is not None: + weight = weight.gather(1, pos_ids.unsqueeze(-1).expand(-1, -1, D)) + x = x + weight + return x + + +# (masked) normalization used for MAE target normalization + + +class Normalize(nn.Module): + def __init__( + self, + grid_size: tuple[int, ...], + dim: int | tuple[int, ...] | None = -1, + eps: float = 1e-6, + ) -> None: + super().__init__() + self.grid_size = grid_size + self.dim = dim + self.eps = eps + + def forward(self, x: Tensor, mask: Tensor | None = None) -> tuple[Tensor, Tensor, Tensor]: + """ + Normalize input sequence along dim(s) after reshaping to grid. + Returns tuple of (x, mean, std). + """ + B, N, D = x.shape + x = x.reshape((B, *self.grid_size, D)) + if mask is not None: + mask = mask.reshape((B, *self.grid_size, D)) + x, mean, std = masked_normalize(x, mask, dim=self.dim, eps=self.eps) + else: + x, mean, std = normalize(x, dim=self.dim, eps=self.eps) + mean = mean.expand_as(x).reshape(B, N, D) + std = std.expand_as(x).reshape(B, N, D) + x = x.reshape(B, N, D) + return x, mean, std + + def extra_repr(self): + return f"{self.grid_size}, dim={self.dim}" + + +def masked_normalize( + x: Tensor, + mask: Tensor, + dim: int | tuple[int, ...] | None = -1, + eps: float = 1e-6, +) -> tuple[Tensor, Tensor, Tensor]: + num_obs = mask.sum(dim=dim, keepdim=True).clamp(min=1) + mean = (mask * x).sum(dim=dim, keepdim=True) / num_obs + var = (mask * (x - mean) ** 2).sum(dim=dim, keepdim=True) / num_obs + std = (var + eps) ** 0.5 + x = mask * (x - mean) / std + return x, mean, std + + +def normalize( + x: Tensor, + dim: int | tuple[int, ...] | None = -1, + eps: float = 1e-6, +) -> tuple[Tensor, Tensor, Tensor]: + mean = x.mean(dim=dim, keepdim=True) + var = torch.var(x, dim=dim, keepdim=True, unbiased=False) + std = (var + eps) ** 0.5 + x = (x - mean) / std + return x, mean, std diff --git a/fomo_tune_baseline/output/task5/build/smri_mae/utils.py b/fomo_tune_baseline/output/task5/build/smri_mae/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..0f6274f813429552c9ddfccfe4033ae677b52342 --- /dev/null +++ b/fomo_tune_baseline/output/task5/build/smri_mae/utils.py @@ -0,0 +1,581 @@ +# Copyright (c) Sophont, Inc +# This source code is licensed under the Apache License, Version 2.0 +# +# References: +# deit: https://github.com/facebookresearch/deit/blob/main/utils.py +# beit3: https://github.com/microsoft/unilm/blob/master/beit3/utils.py +# capi: https://github.com/facebookresearch/capi/blob/main/utils.py +# dinov2: https://github.com/facebookresearch/dinov2/blob/main/dinov2/utils/param_groups.py +# timm: https://github.com/huggingface/pytorch-image-models/blob/main/timm/utils/cuda.py +# dino: https://github.com/facebookresearch/dino/blob/main/utils.py + +import datetime +import inspect +import math +import os +import random +import subprocess +import time +from collections import defaultdict, deque +from omegaconf import OmegaConf +from pathlib import Path + +import numpy as np +import torch +import torch.distributed as dist +import torch.nn as nn +from torch import Tensor +from torch.amp import GradScaler +from torch.optim import Optimizer + + +# these very useful utils copied from deit with only minor changes +# thanks to the original authors, wherever you are + + +def configure_flash_sdpa() -> None: + """Use Flash Attention exclusively for CUDA SDPA.""" + torch.backends.cuda.enable_flash_sdp(True) + torch.backends.cuda.enable_mem_efficient_sdp(False) + torch.backends.cuda.enable_math_sdp(False) + torch.backends.cuda.enable_cudnn_sdp(False) + print("SDPA backend: flash") + + +class SmoothedValue: + """Track a series of values and provide access to smoothed values over a + window or the global series average. + """ + + def __init__(self, window_size=20, fmt=None): + if fmt is None: + fmt = "{median:.4f} ({global_avg:.4f})" + self.deque = deque(maxlen=window_size) + self.total = 0.0 + self.count = 0 + self.fmt = fmt + + def update(self, value, n=1): + value = float(value) + if math.isfinite(value): + self.deque.append(value) + self.count += n + self.total += value * n + + def synchronize_between_processes(self): + """ + Warning: does not synchronize the deque! + """ + if not is_dist_avail_and_initialized(): + return + t = torch.tensor([self.count, self.total], dtype=torch.float64, device="cuda") + dist.barrier() + dist.all_reduce(t) + t = t.tolist() + self.count = int(t[0]) + self.total = t[1] + + @property + def median(self): + if not self.count: + return float("nan") + d = torch.tensor(list(self.deque)) + return d.median().item() + + @property + def avg(self): + if not self.count: + return float("nan") + d = torch.tensor(list(self.deque), dtype=torch.float32) + return d.mean().item() + + @property + def global_avg(self): + if not self.count: + return float("nan") + return self.total / self.count + + @property + def max(self): + if not self.count: + return float("nan") + return max(self.deque) + + @property + def value(self): + if not self.count: + return float("nan") + return self.deque[-1] + + def __str__(self): + return self.fmt.format( + median=self.median, + avg=self.avg, + global_avg=self.global_avg, + max=self.max, + value=self.value, + ) + + +class MetricLogger: + def __init__(self, delimiter="\t"): + self.meters = defaultdict(SmoothedValue) + self.delimiter = delimiter + + def update(self, **kwargs): + for k, v in kwargs.items(): + if v is None: + continue + if isinstance(v, (torch.Tensor, np.generic)): + v = v.item() + assert isinstance(v, (float, int)) + self.meters[k].update(v) + + def __getattr__(self, attr): + if attr in self.meters: + return self.meters[attr] + if attr in self.__dict__: + return self.__dict__[attr] + raise AttributeError("'{}' object has no attribute '{}'".format(type(self).__name__, attr)) + + def __str__(self): + loss_str = [] + for name, meter in self.meters.items(): + loss_str.append("{}: {}".format(name, str(meter))) + return self.delimiter.join(loss_str) + + def synchronize_between_processes(self): + for meter in self.meters.values(): + meter.synchronize_between_processes() + + def add_meter(self, name, meter): + self.meters[name] = meter + + def log_every(self, iterable, print_freq, header=None, total_steps=None): + i = 0 + total_steps = total_steps or len(iterable) + if not header: + header = "" + start_time = time.time() + end = time.time() + iter_time = SmoothedValue(fmt="{avg:.4f}") + data_time = SmoothedValue(fmt="{avg:.4f}") + space_fmt = ":" + str(len(str(total_steps))) + "d" + log_msg = [ + header, + "[{0" + space_fmt + "}/{1}]", + "eta: {eta}", + "{meters}", + "time: {time}", + "data: {data}", + ] + if torch.cuda.is_available(): + log_msg.append("max mem: {memory:.0f}") + log_msg = self.delimiter.join(log_msg) + MB = 1024.0 * 1024.0 + for obj in iterable: + if i >= total_steps: + break + data_time.update(time.time() - end) + yield obj + iter_time.update(time.time() - end) + if i % print_freq == 0 or i == total_steps - 1: + eta_seconds = iter_time.global_avg * (total_steps - i) + eta_string = str(datetime.timedelta(seconds=int(eta_seconds))) + if torch.cuda.is_available(): + print( + log_msg.format( + i, + total_steps, + eta=eta_string, + meters=str(self), + time=str(iter_time), + data=str(data_time), + memory=torch.cuda.max_memory_allocated() / MB, + ) + ) + + else: + print( + log_msg.format( + i, + total_steps, + eta=eta_string, + meters=str(self), + time=str(iter_time), + data=str(data_time), + ) + ) + i += 1 + end = time.time() + total_time = time.time() - start_time + total_time_str = str(datetime.timedelta(seconds=int(total_time))) + print( + "{} Total time: {} ({:.4f} s / it)".format( + header, total_time_str, total_time / total_steps + ) + ) + + +def setup_for_distributed(log_path=None): + """ + This function disables printing when not in master process + """ + import builtins as __builtin__ + + builtin_print = __builtin__.print + + is_master = is_main_process() + + def print(*args, **kwargs): + force = kwargs.pop("force", False) + if is_master or force: + builtin_print(*args, **kwargs) + # tee to log file + if log_path and "file" not in kwargs: + with open(log_path, "a") as f: + builtin_print(*args, file=f, **kwargs) + + __builtin__.print = print + + +def is_dist_avail_and_initialized(): + if not dist.is_available(): + return False + if not dist.is_initialized(): + return False + return True + + +def get_world_size(): + if not is_dist_avail_and_initialized(): + return 1 + return dist.get_world_size() + + +def get_rank(): + if not is_dist_avail_and_initialized(): + return 0 + return dist.get_rank() + + +def is_main_process(): + return get_rank() == 0 + + +def save_on_master(obj, path, *args, **kwargs): + if is_main_process(): + path = Path(path) + tmp_path = path.with_name(f".{path.name}.tmp-{os.getpid()}") + try: + torch.save(obj, tmp_path, *args, **kwargs) + os.replace(tmp_path, path) + except Exception: + tmp_path.unlink(missing_ok=True) + raise + + +def init_distributed_mode(args): + # removed slurm block, can add if we use slurm + if "RANK" in os.environ and "WORLD_SIZE" in os.environ: + args.rank = int(os.environ["RANK"]) + args.world_size = int(os.environ["WORLD_SIZE"]) + args.gpu = int(os.environ["LOCAL_RANK"]) + else: + args.distributed = False + return + + args.distributed = True + + torch.cuda.set_device(args.gpu) + args.dist_backend = "nccl" + print(f"| distributed init (rank {args.rank})") + torch.distributed.init_process_group( + backend=args.dist_backend, + world_size=args.world_size, + rank=args.rank, + device_id=args.gpu, + ) + torch.distributed.barrier() + + +# checkpoint saving utils adapted from beit3 + + +def capture_rng_state() -> dict: + state = {"torch": torch.get_rng_state()} + if torch.cuda.is_available(): + state["cuda"] = torch.cuda.get_rng_state() + return state + + +def restore_rng_state(state: dict) -> None: + torch.set_rng_state(state["torch"]) + if "cuda" in state and torch.cuda.is_available(): + torch.cuda.set_rng_state(state["cuda"]) + + +def _all_rank_rng_states() -> list[dict]: + local_state = capture_rng_state() + if not is_dist_avail_and_initialized(): + return [local_state] + states = [None] * get_world_size() + dist.all_gather_object(states, local_state) + return states + + +def save_model(args, epoch, model_without_ddp, optimizer, loss_scaler): + output_dir = Path(args.output_dir) + checkpoint_path = output_dir / f"checkpoint-{epoch:05d}.pth" + last_checkpoint_path = output_dir / "checkpoint-last.pth" + if epoch % args.checkpoint_period != 0 and epoch != args.epochs - 1: + return + + to_save = { + "model": model_without_ddp.state_dict(), + "optimizer": optimizer.state_dict(), + "epoch": epoch, + "scaler": None if loss_scaler is None else loss_scaler.state_dict(), + "args": OmegaConf.to_container(args), + "rng_states": _all_rank_rng_states(), + } + + print(f"saving checkpoint {last_checkpoint_path}") + save_on_master(to_save, last_checkpoint_path) + print(f"saving checkpoint {checkpoint_path}") + save_on_master(to_save, checkpoint_path) + + if args.max_checkpoints and is_main_process(): + all_checkpoints = sorted(output_dir.glob("checkpoint-[0-9]*.pth")) + del_count = max(0, len(all_checkpoints) - args.max_checkpoints) + for checkpoint_path in all_checkpoints[:del_count]: + print(f"removing checkpoint {checkpoint_path}") + checkpoint_path.unlink() + + +def load_model(args, model_without_ddp, optimizer, loss_scaler): + auto_resume = getattr(args, "auto_resume", True) + output_dir = Path(args.output_dir) + + last_checkpoint_path = output_dir / "checkpoint-last.pth" + if auto_resume and last_checkpoint_path.exists(): + args.ckpt = str(last_checkpoint_path) + args.resume = True + + if args.ckpt: + ckpt = torch.load(args.ckpt, map_location="cpu", weights_only=True) + model_without_ddp.load_state_dict(ckpt["model"]) + print(f"loaded model state from checkpoint {args.ckpt}") + + if args.resume: + optimizer.load_state_dict(ckpt["optimizer"]) + if loss_scaler is not None: + loss_scaler.load_state_dict(ckpt["scaler"]) + args.start_epoch = ckpt["epoch"] + 1 + rng_states = ckpt.get("rng_states") + if rng_states is not None: + if len(rng_states) != get_world_size(): + raise ValueError( + "checkpoint RNG state world size does not match current world size: " + f"{len(rng_states)} != {get_world_size()}" + ) + restore_rng_state(rng_states[get_rank()]) + print(f"restored RNG state for rank {get_rank()}") + print(f"loaded optimizer state, resuming training from {args.start_epoch}") + + +# optimization utils + + +# from capi +class WarmupThenCosine: + def __init__( + self, + base_value: float, + final_value: float, + total_iters: int, + warmup_iters: int = 0, + start_warmup_value: float = 0.0, + freeze_iters: int = 0, + truncate_cos: float = 1.0, + ): + super().__init__() + self.final_value = final_value + self.total_iters = total_iters + + freeze_schedule = np.zeros(freeze_iters) + + warmup_schedule = np.linspace(start_warmup_value, base_value, warmup_iters) + + iters = np.arange(total_iters - warmup_iters - freeze_iters) + schedule = final_value + 0.5 * (base_value - final_value) * ( + 1 + np.cos(np.pi * truncate_cos * iters / len(iters)) + ) + self.schedule = np.concatenate((freeze_schedule, warmup_schedule, schedule)) + assert len(self.schedule) == self.total_iters + + def __getitem__(self, it: int) -> float: + if it >= self.total_iters: + return self.final_value + # cast to float or else it can corrupt the checkpoint + return float(self.schedule[it]) + + +# adapted from timm backward logic +# https://github.com/huggingface/pytorch-image-models/blob/main/timm/utils/cuda.py +def backward_step( + loss: Tensor, + optimizer: Optimizer, + scaler: GradScaler = None, + need_update: bool = True, + max_norm: float | None = None, +) -> Tensor | None: + if scaler is not None: + scaler.scale(loss).backward() + else: + loss.backward() + + if need_update: + if scaler is not None: + scaler.unscale_(optimizer) + + total_norm = clip_grad(optimizer, max_norm) + + if scaler is not None: + scaler.step(optimizer) + scaler.update() + else: + optimizer.step() + optimizer.zero_grad() + else: + total_norm = None + return total_norm + + +def clip_grad(optimizer: Optimizer, max_norm: float | None = None) -> Tensor: + params = [p for group in optimizer.param_groups for p in group["params"]] + if max_norm: + total_norm = nn.utils.clip_grad_norm_(params, max_norm, error_if_nonfinite=False) + else: + grads = [p.grad for p in params if p.grad is not None] + total_norm = nn.utils.get_total_norm(grads, error_if_nonfinite=False) + torch._assert_async(torch.isfinite(total_norm), "non-finite gradient norm") + return total_norm + + +# from dinov2 with some minor changes +def get_param_groups(model, patch_embed_lr_mult=1.0): + # no lr decay, we could add this later if needed + all_params = [] + + for name, param in model.named_parameters(): + if not param.requires_grad: + continue + d = {"param": param, "lr_multiplier": 1.0, "wd_multiplier": 1.0, "name": name} + + if name.endswith(".bias") or "norm" in name or "gamma" in name: + d["wd_multiplier"] = 0.0 + + if "patch_embed" in name: + d["lr_multiplier"] = d["lr_multiplier"] * patch_embed_lr_mult + + all_params.append(d) + + param_groups = _fuse_param_groups(all_params) + return param_groups + + +def _fuse_param_groups(all_param_groups): + fused_param_groups = defaultdict(lambda: {"params": []}) + for d in all_param_groups: + keys = sorted(set(d.keys()) - {"param", "name"}) + identifier = "_".join(f"{k}{d[k]}" for k in keys) + for k in keys: + fused_param_groups[identifier][k] = d[k] + fused_param_groups[identifier]["params"].append(d["param"]) + + param_groups = list(fused_param_groups.values()) + return param_groups + + +def update_lr(param_groups, lr: float): + for group in param_groups: + group["lr"] = lr * group["lr_multiplier"] + + +def update_wd(param_groups, weight_decay: float | None = None): + for group in param_groups: + group["weight_decay"] = weight_decay * group["wd_multiplier"] + + +# moving data to cuda utils copied from capi +# added device argument + + +def send_data(x, device=None, dtype_map=None): + if device is None: + device = torch.device("cuda") + else: + device = torch.device(device) + + if isinstance(x, torch.Tensor): + dtype = dtype_map.get(x.dtype) if dtype_map else None + return x.to(device=device, dtype=dtype, non_blocking=True) + if isinstance(x, dict): + return {k: send_data(v, device=device, dtype_map=dtype_map) for k, v in x.items()} + if isinstance(x, list): + return [send_data(v, device=device, dtype_map=dtype_map) for v in x] + return x + + +def pre_send_to_cuda_wrapper(generator, device=None, dtype_map=None): + """From apex""" + data = None + stream = torch.cuda.Stream(device) + for next_data in generator: + with torch.cuda.stream(stream): + next_data = send_data(next_data, device=device, dtype_map=dtype_map) + if data is not None: + yield data + torch.cuda.current_stream(device).wait_stream(stream) + data = next_data + if data is not None: + yield data + + +# other misc utils + + +# from dino +def get_sha(): + cwd = os.path.dirname(os.path.abspath(__file__)) + + def _run(command): + return subprocess.check_output(command, cwd=cwd).decode("ascii").strip() + + sha = "N/A" + diff = "clean" + branch = "N/A" + try: + sha = _run(["git", "rev-parse", "HEAD"]) + diff = _run(["git", "diff-index", "HEAD"]) + diff = "has uncommitted changes" if diff else "clean" + branch = _run(["git", "rev-parse", "--abbrev-ref", "HEAD"]) + except Exception: + pass + message = f"sha: {sha}, status: {diff}, branch: {branch}" + return message + + +# from timm +def random_seed(seed=42, rank=0): + torch.manual_seed(seed + rank) + np.random.seed(seed + rank) + random.seed(seed + rank) + + +# mine :) +def filter_kwargs(func, kwargs): + sigature = inspect.signature(func) + kwargs = {k: v for k, v in kwargs.items() if k in sigature.parameters} + return kwargs diff --git a/fomo_tune_baseline/output/task5/build/smri_mae/visualization.py b/fomo_tune_baseline/output/task5/build/smri_mae/visualization.py new file mode 100644 index 0000000000000000000000000000000000000000..f721fe06e93dc1a2dfb4a679069029bdf55afc6b --- /dev/null +++ b/fomo_tune_baseline/output/task5/build/smri_mae/visualization.py @@ -0,0 +1,400 @@ +from collections.abc import Mapping +from io import BytesIO + +import torch + +from matplotlib import patches +from matplotlib import pyplot as plt +from PIL import Image +from torch import Tensor + +VIEW_NAMES = { + "sagittal": "Sagittal", + "saggital": "Sagittal", + "axial": "Axial", + "coronal": "Coronal", +} + + +def fig2pil(fig) -> Image.Image: + buffer = BytesIO() + fig.savefig(buffer, format="png", dpi=fig.dpi, facecolor=fig.get_facecolor()) + buffer.seek(0) + image = Image.open(buffer).convert("RGB") + buffer.close() + return image + + +def raw_stats_from_batch(batch: dict) -> tuple[Tensor | None, Tensor | None]: + metas = batch.get("meta") + if not metas: + return None, None + + means = [] + stds = [] + for meta in metas: + try: + mean = meta["raw_mean"] + std = meta["raw_std"] + except (KeyError, TypeError): + return None, None + if mean in ("", None) or std in ("", None): + return None, None + means.append(float(mean)) + stds.append(float(std)) + return torch.tensor(means), torch.tensor(stds) + + +def plot_mask_pred( + target: Tensor, + pred: Tensor, + pred_mask: Tensor | None = None, + img_mask: Tensor | None = None, + sample_idx: int = 0, + channel_idx: int = 0, + slice_idx: int | Mapping[str, int] | None = None, + patch_size: int | tuple[int, int, int] = 16, + views: tuple[str, ...] = ("sagittal", "axial", "coronal"), + cmap: str = "gray", + figsize: tuple[float, float] | None = None, + mask_style: str = "blank", + raw_mean: float | Tensor | None = None, + raw_std: float | Tensor | None = None, +): + target_vol = _select_volume(target, sample_idx=sample_idx, channel_idx=channel_idx) + pred_vol = _select_volume(pred, sample_idx=sample_idx, channel_idx=channel_idx) + if raw_mean is not None and raw_std is not None: + raw_mean = _select_scalar(raw_mean, sample_idx=sample_idx) + raw_std = _select_scalar(raw_std, sample_idx=sample_idx) + target_vol = target_vol * raw_std + raw_mean + pred_vol = pred_vol * raw_std + raw_mean + pred_mask_vol = ( + torch.zeros_like(target_vol) + if pred_mask is None + else _select_volume(pred_mask, sample_idx=sample_idx, channel_idx=channel_idx) > 0 + ) + img_mask_vol = None + if img_mask is not None: + img_mask_vol = _select_volume(img_mask, sample_idx=sample_idx, channel_idx=channel_idx) > 0 + + composite_vol = _prediction_composite(target_vol, pred_vol, pred_mask_vol) + vmin, vmax = _intensity_limits(target_vol, img_mask_vol) + + patch_size = _as_3tuple(patch_size) + view_items = [] + for view in views: + view_key = view.lower() + if view_key not in VIEW_NAMES: + raise ValueError(f"unknown MRI view {view!r}; expected one of {tuple(VIEW_NAMES)}") + target_slice = _extract_view_slice(target_vol, view_key, slice_idx) + composite_slice = _extract_view_slice(composite_vol, view_key, slice_idx) + mask_slice = _extract_view_slice(pred_mask_vol.float(), view_key, slice_idx) > 0 + img_mask_slice = None + if img_mask_vol is not None: + img_mask_slice = _extract_view_slice(img_mask_vol.float(), view_key, slice_idx) > 0 + view_items.append( + { + "key": view_key, + "title": VIEW_NAMES[view_key], + "target": _masked_input_display(target_slice, mask_slice, img_mask_slice, vmin), + "composite": _apply_display_mask(composite_slice, img_mask_slice, vmin), + "actual": _apply_display_mask(target_slice, img_mask_slice, vmin), + "mask": mask_slice, + "img_mask": img_mask_slice, + "patch_rc": _view_patch_size(view_key, patch_size), + } + ) + + _crop_view_items(view_items) + + fig, axes, layout = _make_figure_canvas(view_items, figsize=figsize) + for item, x in zip(view_items, layout["col_centers"]): + fig.text( + x, + layout["title_y"], + item["title"], + ha="center", + va="center", + color="#f8fafc", + fontsize=7, + ) + for label, y in zip(("Masked", "Pred", "Actual"), layout["row_centers"]): + fig.text(layout["label_x"], y, label, ha="right", va="center", color="#cbd5e1", fontsize=6) + + for item, top_ax, middle_ax, bottom_ax in zip(view_items, axes[0], axes[1], axes[2]): + top_ax.imshow( + item["target"], + cmap=cmap, + vmin=vmin, + vmax=vmax, + interpolation="nearest", + origin="upper", + ) + if mask_style == "boxes": + _draw_patch_boxes(top_ax, item["mask"], item["patch_rc"]) + elif mask_style != "blank": + raise ValueError("mask_style must be 'blank' or 'boxes'") + _style_axis(top_ax) + + middle_ax.imshow( + item["composite"], + cmap=cmap, + vmin=vmin, + vmax=vmax, + interpolation="nearest", + origin="upper", + ) + _style_axis(middle_ax) + + bottom_ax.imshow( + item["actual"], + cmap=cmap, + vmin=vmin, + vmax=vmax, + interpolation="nearest", + origin="upper", + ) + _style_axis(bottom_ax) + return fig + + +def _select_volume( + x: Tensor, + sample_idx: int = 0, + channel_idx: int = 0, +) -> Tensor: + x = x.detach().float().cpu() + if x.ndim == 5: + return x[sample_idx, channel_idx] + if x.ndim == 4: + return x[sample_idx] + if x.ndim == 3: + return x + raise ValueError(f"expected a 3D volume tensor, got shape {tuple(x.shape)}") + + +def _select_scalar(value: float | Tensor, sample_idx: int = 0) -> float: + if isinstance(value, Tensor): + value = value.detach().float().cpu() + if value.ndim > 0: + value = value.reshape(-1)[sample_idx] + return float(value) + return float(value) + + +def _prediction_composite(target: Tensor, pred: Tensor, pred_mask: Tensor) -> Tensor: + pred_mask = pred_mask.to(dtype=target.dtype) + return target * (1 - pred_mask) + pred * pred_mask + + +def _extract_view_slice( + volume: Tensor, + view: str, + slice_idx: int | Mapping[str, int] | None = None, +) -> Tensor: + if isinstance(slice_idx, Mapping): + slice_idx = slice_idx.get(view) + + if view in ("sagittal", "saggital"): + idx = _resolve_slice_idx(volume.shape[0], slice_idx) + return volume[idx, :, :].transpose(0, 1).flip(0) + if view == "axial": + idx = _resolve_slice_idx(volume.shape[2], slice_idx) + return volume[:, :, idx].transpose(0, 1).flip(0) + if view == "coronal": + idx = _resolve_slice_idx(volume.shape[1], slice_idx) + return volume[:, idx, :].transpose(0, 1).flip(0) + raise ValueError(f"unknown MRI view {view!r}") + + +def _resolve_slice_idx(size: int, slice_idx: int | None) -> int: + idx = size // 2 if slice_idx is None else int(slice_idx) + if idx < 0: + idx += size + if idx < 0 or idx >= size: + raise IndexError(f"slice index {idx} is out of bounds for axis with size {size}") + return idx + + +def _intensity_limits(volume: Tensor, mask: Tensor | None = None) -> tuple[float, float]: + values = volume[mask] if mask is not None and mask.any() else volume.flatten() + values = values[torch.isfinite(values)] + if values.numel() == 0: + return 0.0, 1.0 + if values.numel() < 32: + vmin = values.min() + vmax = values.max() + else: + vmin, vmax = torch.quantile(values, torch.tensor([0.005, 0.995])) + if torch.isclose(vmin, vmax): + delta = max(abs(float(vmin)) * 0.05, 1.0) + return float(vmin) - delta, float(vmax) + delta + return float(vmin), float(vmax) + + +def _apply_display_mask(image: Tensor, mask: Tensor | None, fill_value: float) -> Tensor: + if mask is None: + return image + return torch.where(mask, image, torch.full_like(image, fill_value)) + + +def _masked_input_display( + image: Tensor, + pred_mask: Tensor, + img_mask: Tensor | None, + fill_value: float, +) -> Tensor: + display = torch.where(pred_mask, torch.full_like(image, fill_value), image) + return _apply_display_mask(display, img_mask, fill_value) + + +def _crop_view_items(view_items: list[dict]) -> None: + for item in view_items: + mask = item["img_mask"] + if mask is None: + mask = item["actual"] != item["actual"].min() + row_slice, col_slice = _content_crop(mask, item["patch_rc"]) + for key in ("target", "composite", "actual", "mask"): + item[key] = item[key][row_slice, col_slice] + if item["img_mask"] is not None: + item["img_mask"] = item["img_mask"][row_slice, col_slice] + + +def _content_crop(mask: Tensor, patch_size: tuple[int, int]) -> tuple[slice, slice]: + mask = mask.detach().cpu() > 0 + if not mask.any(): + return slice(None), slice(None) + + rows, cols = mask.nonzero(as_tuple=True) + patch_h, patch_w = patch_size + height, width = mask.shape + row0 = max((int(rows.min()) // patch_h - 1) * patch_h, 0) + row1 = min((int(rows.max()) // patch_h + 2) * patch_h, height) + col0 = max((int(cols.min()) // patch_w - 1) * patch_w, 0) + col1 = min((int(cols.max()) // patch_w + 2) * patch_w, width) + return slice(row0, row1), slice(col0, col1) + + +def _as_3tuple(value: int | tuple[int, int, int]) -> tuple[int, int, int]: + if isinstance(value, int): + return (value, value, value) + if len(value) != 3: + raise ValueError(f"expected a 3-tuple patch size, got {value!r}") + return tuple(int(v) for v in value) + + +def _view_patch_size(view: str, patch_size: tuple[int, int, int]) -> tuple[int, int]: + p_x, p_y, p_z = patch_size + if view in ("sagittal", "saggital"): + return p_z, p_y + if view == "axial": + return p_y, p_x + if view == "coronal": + return p_z, p_x + raise ValueError(f"unknown MRI view {view!r}") + + +def _make_figure_canvas( + view_items: list[dict[str, Tensor | str | tuple[int, int]]], + figsize: tuple[float, float] | None = None, +): + dpi = 160 + left = 58 + right = 6 + top = 18 + bottom = 8 + row_gap = 14 + col_gap = 8 + widths = [int(item["target"].shape[1]) for item in view_items] + heights = [int(item["target"].shape[0]) for item in view_items] + row_h = max(heights) + num_rows = 3 + fig_w = left + right + sum(widths) + col_gap * (len(widths) - 1) + fig_h = top + bottom + row_h * num_rows + row_gap * (num_rows - 1) + + scale = 1.35 + if figsize is not None: + requested_w = figsize[0] * dpi + requested_h = figsize[1] * dpi + scale = max(requested_w / fig_w, requested_h / fig_h) + figsize = (fig_w * scale / dpi, fig_h * scale / dpi) + fig = plt.figure(figsize=figsize, dpi=dpi, facecolor="#0b0f14") + + axes = [[] for _ in range(num_rows)] + col_centers = [] + x = left + for width, height in zip(widths, heights): + col_centers.append((x + width / 2) / fig_w) + ys = [ + bottom + (num_rows - row - 1) * (row_h + row_gap) + (row_h - height) / 2 + for row in range(num_rows) + ] + for row, y in enumerate(ys): + axes[row].append( + fig.add_axes( + [ + x / fig_w, + y / fig_h, + width / fig_w, + height / fig_h, + ], + facecolor="black", + ) + ) + x += width + col_gap + row_centers = [ + (bottom + (num_rows - row - 1) * (row_h + row_gap) + row_h / 2) / fig_h + for row in range(num_rows) + ] + layout = { + "col_centers": col_centers, + "row_centers": row_centers, + "label_x": (left - 8) / fig_w, + "title_y": (fig_h - top / 2) / fig_h, + } + + return fig, axes, layout + + +def _style_axis(ax) -> None: + ax.set_xticks([]) + ax.set_yticks([]) + for spine in ax.spines.values(): + spine.set_visible(False) + + +def _draw_patch_boxes( + ax, + mask: Tensor, + patch_size: tuple[int, int], + color: str = "#facc15", +) -> None: + for col, row, width, height in _patch_rectangles(mask, patch_size): + ax.add_patch( + patches.Rectangle( + (col - 0.5, row - 0.5), + width, + height, + fill=False, + edgecolor=color, + linewidth=0.75, + alpha=0.95, + ) + ) + + +def _patch_rectangles( + mask: Tensor, + patch_size: tuple[int, int], +) -> list[tuple[int, int, int, int]]: + mask = mask.detach().cpu() > 0 + patch_h, patch_w = patch_size + height, width = mask.shape + rectangles = [] + for row in range(0, height, patch_h): + box_h = min(patch_h, height - row) + for col in range(0, width, patch_w): + box_w = min(patch_w, width - col) + if mask[row : row + box_h, col : col + box_w].any(): + rectangles.append((col, row, box_w, box_h)) + return rectangles diff --git a/fomo_tune_baseline/output/task5/config.yaml b/fomo_tune_baseline/output/task5/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..95ebb19790bae19d7927e9621aab1d6520030460 --- /dev/null +++ b/fomo_tune_baseline/output/task5/config.yaml @@ -0,0 +1,6 @@ +task: task5 +ckpt_path: /data/mihir-stuff/smri-pretrained/pretrain_full_90_10_h100/checkpoint-last.pth +output_root: experiments/fomo_tune_baseline/output +name: task5 +device: cuda +seed: 4466 diff --git a/fomo_tune_baseline/output/task5/log.txt b/fomo_tune_baseline/output/task5/log.txt new file mode 100644 index 0000000000000000000000000000000000000000..0d9c585afa99f273da7263aa378f008b739b9589 --- /dev/null +++ b/fomo_tune_baseline/output/task5/log.txt @@ -0,0 +1,30 @@ +16:32:19 run task5 (git 1df2e5d) +16:32:19 config: +task: task5 +ckpt_path: /data/mihir-stuff/smri-pretrained/pretrain_full_90_10_h100/checkpoint-last.pth +output_root: experiments/fomo_tune_baseline/output +name: task5 +device: cuda +seed: 4466 +16:32:42 dataset: 48 subjects, 24 positive +16:33:21 fold 1/20 n=3 y=[0 1 1] p=[0.013 0.998 0.974] (35s) +16:33:24 fold 2/20 n=3 y=[1 1 1] p=[0.796 0.719 0.999] (37s) +16:33:26 fold 3/20 n=3 y=[0 1 1] p=[0.155 0.998 0.905] (39s) +16:33:28 fold 4/20 n=3 y=[0 0 1] p=[0.024 0.032 0.811] (42s) +16:33:29 fold 5/20 n=3 y=[1 1 1] p=[0.988 0.652 0.552] (43s) +16:33:31 fold 6/20 n=3 y=[0 0 1] p=[0.46 0.744 0.994] (44s) +16:33:33 fold 7/20 n=3 y=[0 0 1] p=[0.134 0.148 0.996] (47s) +16:33:36 fold 8/20 n=3 y=[0 0 1] p=[0.037 0.012 0.915] (50s) +16:33:37 fold 9/20 n=2 y=[0 1] p=[0.137 0.997] (51s) +16:33:38 fold 10/20 n=2 y=[0 1] p=[0.223 0.993] (51s) +16:33:39 fold 11/20 n=2 y=[0 0] p=[0.118 0.32 ] (53s) +16:33:42 fold 12/20 n=2 y=[0 1] p=[0.122 0.849] (55s) +16:33:43 fold 13/20 n=2 y=[0 1] p=[0.207 0.282] (57s) +16:33:45 fold 14/20 n=2 y=[0 1] p=[0.177 0.319] (58s) +16:33:47 fold 15/20 n=2 y=[0 1] p=[0.112 0.943] (60s) +16:33:48 fold 16/20 n=2 y=[0 1] p=[0.034 0.998] (61s) +16:33:50 fold 17/20 n=2 y=[0 0] p=[0.018 0.021] (63s) +16:33:51 fold 18/20 n=2 y=[0 1] p=[0.02 0.937] (64s) +16:33:52 fold 19/20 n=2 y=[0 1] p=[0.004 0.96 ] (65s) +16:33:54 fold 20/20 n=2 y=[0 1] p=[0.075 0.929] (68s) +16:33:56 result: auroc=0.9844 auroc_ci_low=0.9528 auroc_ci_high=1.0000 (68s) diff --git a/fomo_tune_baseline/output/task5/metrics.json b/fomo_tune_baseline/output/task5/metrics.json new file mode 100644 index 0000000000000000000000000000000000000000..d98cd3d599914b86e5592ea9dd6d9cf2b575fa5a --- /dev/null +++ b/fomo_tune_baseline/output/task5/metrics.json @@ -0,0 +1 @@ +{"name": "task5", "auroc": 0.984375, "auroc_ci_low": 0.9527867965367964, "auroc_ci_high": 1.0, "run_time": 67.5} diff --git a/fomo_tune_baseline/output/task5/model/config.yaml b/fomo_tune_baseline/output/task5/model/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..95ebb19790bae19d7927e9621aab1d6520030460 --- /dev/null +++ b/fomo_tune_baseline/output/task5/model/config.yaml @@ -0,0 +1,6 @@ +task: task5 +ckpt_path: /data/mihir-stuff/smri-pretrained/pretrain_full_90_10_h100/checkpoint-last.pth +output_root: experiments/fomo_tune_baseline/output +name: task5 +device: cuda +seed: 4466 diff --git a/fomo_tune_baseline/output/task5/model/head.joblib b/fomo_tune_baseline/output/task5/model/head.joblib new file mode 100644 index 0000000000000000000000000000000000000000..fb645916a1c0161712332085d845e2a4bd8f5014 --- /dev/null +++ b/fomo_tune_baseline/output/task5/model/head.joblib @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:dc3c55c52bd9306b8385abe2232837d7ff0181320f93c487db2eaea27a895c1c +size 445215 diff --git a/fomo_tune_baseline/output/task5/task5.sif b/fomo_tune_baseline/output/task5/task5.sif new file mode 100644 index 0000000000000000000000000000000000000000..4c7601c15a40e08a21d7e5c355f5081c4c76dac9 --- /dev/null +++ b/fomo_tune_baseline/output/task5/task5.sif @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a54f2e32e4a955d3b46f1a87abeacf09bd097ae476aa1557283502ccb4a2abf4 +size 5329235968 diff --git a/fomo_tune_baseline/slurms/build.log b/fomo_tune_baseline/slurms/build.log new file mode 100644 index 0000000000000000000000000000000000000000..0ae71fedfab66912d8cf9feaeece62b35cc4cad2 --- /dev/null +++ b/fomo_tune_baseline/slurms/build.log @@ -0,0 +1,718 @@ +=== task1 === +INFO: /etc/singularity/ exists; cleanup by system administrator is not complete (see https://apptainer.org/docs/admin/latest/singularity_migration.html) +INFO: /etc/singularity/ exists; cleanup by system administrator is not complete (see https://apptainer.org/docs/admin/latest/singularity_migration.html) +INFO: Starting build... +INFO: Fetching OCI image... +INFO: Extracting OCI image... +INFO: Inserting Apptainer configuration... +INFO: Copying fomo_tune to /app/fomo_tune +INFO: Copying smri_mae to /app/smri_mae +INFO: Copying model to /app/model +INFO: Copying predict.py to /app/predict.py +INFO: Running post scriptlet ++ pip install --no-cache-dir torch==2.8.0 numpy==2.4.6 nibabel==5.4.2 einops==0.8.2 jaxtyping==0.3.10 timm==1.0.27 huggingface-hub==0.36.2 scikit-learn==1.8.0 joblib==1.5.3 omegaconf==2.3.0 +Collecting torch==2.8.0 + Downloading torch-2.8.0-cp311-cp311-manylinux_2_28_x86_64.whl.metadata (30 kB) +Collecting numpy==2.4.6 + Downloading numpy-2.4.6-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl.metadata (6.6 kB) +Collecting nibabel==5.4.2 + Downloading nibabel-5.4.2-py3-none-any.whl.metadata (8.9 kB) +Collecting einops==0.8.2 + Downloading einops-0.8.2-py3-none-any.whl.metadata (13 kB) +Collecting jaxtyping==0.3.10 + Downloading jaxtyping-0.3.10-py3-none-any.whl.metadata (6.2 kB) +Collecting timm==1.0.27 + Downloading timm-1.0.27-py3-none-any.whl.metadata (40 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 40.2/40.2 kB 115.4 MB/s eta 0:00:00 +Collecting huggingface-hub==0.36.2 + Downloading huggingface_hub-0.36.2-py3-none-any.whl.metadata (15 kB) +Collecting scikit-learn==1.8.0 + Downloading scikit_learn-1.8.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl.metadata (11 kB) +Collecting joblib==1.5.3 + Downloading joblib-1.5.3-py3-none-any.whl.metadata (5.5 kB) +Collecting omegaconf==2.3.0 + Downloading omegaconf-2.3.0-py3-none-any.whl.metadata (3.9 kB) +Collecting filelock (from torch==2.8.0) + Downloading filelock-3.32.2-py3-none-any.whl.metadata (2.0 kB) +Collecting typing-extensions>=4.10.0 (from torch==2.8.0) + Downloading typing_extensions-4.16.0-py3-none-any.whl.metadata (3.3 kB) +Collecting sympy>=1.13.3 (from torch==2.8.0) + Downloading sympy-1.14.0-py3-none-any.whl.metadata (12 kB) +Collecting networkx (from torch==2.8.0) + Downloading networkx-3.6.1-py3-none-any.whl.metadata (6.8 kB) +Collecting jinja2 (from torch==2.8.0) + Downloading jinja2-3.1.6-py3-none-any.whl.metadata (2.9 kB) +Collecting fsspec (from torch==2.8.0) + Downloading fsspec-2026.7.0-py3-none-any.whl.metadata (10 kB) +Collecting nvidia-cuda-nvrtc-cu12==12.8.93 (from torch==2.8.0) + Downloading nvidia_cuda_nvrtc_cu12-12.8.93-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl.metadata (1.7 kB) +Collecting nvidia-cuda-runtime-cu12==12.8.90 (from torch==2.8.0) + Downloading nvidia_cuda_runtime_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl.metadata (1.7 kB) +Collecting nvidia-cuda-cupti-cu12==12.8.90 (from torch==2.8.0) + Downloading nvidia_cuda_cupti_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl.metadata (1.7 kB) +Collecting nvidia-cudnn-cu12==9.10.2.21 (from torch==2.8.0) + Downloading nvidia_cudnn_cu12-9.10.2.21-py3-none-manylinux_2_27_x86_64.whl.metadata (1.8 kB) +Collecting nvidia-cublas-cu12==12.8.4.1 (from torch==2.8.0) + Downloading nvidia_cublas_cu12-12.8.4.1-py3-none-manylinux_2_27_x86_64.whl.metadata (1.7 kB) +Collecting nvidia-cufft-cu12==11.3.3.83 (from torch==2.8.0) + Downloading nvidia_cufft_cu12-11.3.3.83-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl.metadata (1.7 kB) +Collecting nvidia-curand-cu12==10.3.9.90 (from torch==2.8.0) + Downloading nvidia_curand_cu12-10.3.9.90-py3-none-manylinux_2_27_x86_64.whl.metadata (1.7 kB) +Collecting nvidia-cusolver-cu12==11.7.3.90 (from torch==2.8.0) + Downloading nvidia_cusolver_cu12-11.7.3.90-py3-none-manylinux_2_27_x86_64.whl.metadata (1.8 kB) +Collecting nvidia-cusparse-cu12==12.5.8.93 (from torch==2.8.0) + Downloading nvidia_cusparse_cu12-12.5.8.93-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl.metadata (1.8 kB) +Collecting nvidia-cusparselt-cu12==0.7.1 (from torch==2.8.0) + Downloading nvidia_cusparselt_cu12-0.7.1-py3-none-manylinux2014_x86_64.whl.metadata (7.0 kB) +Collecting nvidia-nccl-cu12==2.27.3 (from torch==2.8.0) + Downloading nvidia_nccl_cu12-2.27.3-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl.metadata (2.0 kB) +Collecting nvidia-nvtx-cu12==12.8.90 (from torch==2.8.0) + Downloading nvidia_nvtx_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl.metadata (1.8 kB) +Collecting nvidia-nvjitlink-cu12==12.8.93 (from torch==2.8.0) + Downloading nvidia_nvjitlink_cu12-12.8.93-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl.metadata (1.7 kB) +Collecting nvidia-cufile-cu12==1.13.1.3 (from torch==2.8.0) + Downloading nvidia_cufile_cu12-1.13.1.3-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl.metadata (1.7 kB) +Collecting triton==3.4.0 (from torch==2.8.0) + Downloading triton-3.4.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl.metadata (1.7 kB) +Collecting importlib-resources>=5.12 (from nibabel==5.4.2) + Downloading importlib_resources-7.1.0-py3-none-any.whl.metadata (4.0 kB) +Requirement already satisfied: packaging>=20 in /usr/local/lib/python3.11/site-packages (from nibabel==5.4.2) (26.3) +Collecting wadler-lindig>=0.1.3 (from jaxtyping==0.3.10) + Downloading wadler_lindig-0.1.7-py3-none-any.whl.metadata (17 kB) +Collecting torchvision (from timm==1.0.27) + Downloading torchvision-0.28.0-cp311-cp311-manylinux_2_28_x86_64.whl.metadata (5.6 kB) +Collecting pyyaml (from timm==1.0.27) + Downloading pyyaml-6.0.3-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl.metadata (2.4 kB) +Collecting safetensors (from timm==1.0.27) + Downloading safetensors-0.8.0-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (4.2 kB) +Collecting hf-xet<2.0.0,>=1.1.3 (from huggingface-hub==0.36.2) + Downloading hf_xet-1.6.0-cp38-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl.metadata (4.9 kB) +Collecting requests (from huggingface-hub==0.36.2) + Downloading requests-2.34.2-py3-none-any.whl.metadata (4.8 kB) +Collecting tqdm>=4.42.1 (from huggingface-hub==0.36.2) + Downloading tqdm-4.70.0-py3-none-any.whl.metadata (57 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 57.3/57.3 kB 49.5 MB/s eta 0:00:00 +Collecting scipy>=1.10.0 (from scikit-learn==1.8.0) + Downloading scipy-1.17.1-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl.metadata (62 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 62.1/62.1 kB 59.8 MB/s eta 0:00:00 +Collecting threadpoolctl>=3.2.0 (from scikit-learn==1.8.0) + Downloading threadpoolctl-3.6.0-py3-none-any.whl.metadata (13 kB) +Collecting antlr4-python3-runtime==4.9.* (from omegaconf==2.3.0) + Downloading antlr4-python3-runtime-4.9.3.tar.gz (117 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 117.0/117.0 kB 53.3 MB/s eta 0:00:00 + Preparing metadata (setup.py): started + Preparing metadata (setup.py): finished with status 'done' +Requirement already satisfied: setuptools>=40.8.0 in /usr/local/lib/python3.11/site-packages (from triton==3.4.0->torch==2.8.0) (79.0.1) +Collecting mpmath<1.4,>=1.1.0 (from sympy>=1.13.3->torch==2.8.0) + Downloading mpmath-1.3.0-py3-none-any.whl.metadata (8.6 kB) +Collecting MarkupSafe>=2.0 (from jinja2->torch==2.8.0) + Downloading markupsafe-3.0.3-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl.metadata (2.7 kB) +Collecting charset_normalizer<4,>=2 (from requests->huggingface-hub==0.36.2) + Downloading charset_normalizer-3.4.9-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl.metadata (41 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 41.7/41.7 kB 322.5 MB/s eta 0:00:00 +Collecting idna<4,>=2.5 (from requests->huggingface-hub==0.36.2) + Downloading idna-3.18-py3-none-any.whl.metadata (6.1 kB) +Collecting urllib3<3,>=1.26 (from requests->huggingface-hub==0.36.2) + Downloading urllib3-2.7.0-py3-none-any.whl.metadata (6.9 kB) +Collecting certifi>=2023.5.7 (from requests->huggingface-hub==0.36.2) + Downloading certifi-2026.7.22-py3-none-any.whl.metadata (2.5 kB) +INFO: pip is looking at multiple versions of torchvision to determine which version is compatible with other requirements. This could take a while. +Collecting torchvision (from timm==1.0.27) + Downloading torchvision-0.27.1-cp311-cp311-manylinux_2_28_x86_64.whl.metadata (5.5 kB) + Downloading torchvision-0.27.0-cp311-cp311-manylinux_2_28_x86_64.whl.metadata (5.5 kB) + Downloading torchvision-0.26.0-cp311-cp311-manylinux_2_28_x86_64.whl.metadata (5.5 kB) + Downloading torchvision-0.25.0-cp311-cp311-manylinux_2_28_x86_64.whl.metadata (5.4 kB) + Downloading torchvision-0.24.1-cp311-cp311-manylinux_2_28_x86_64.whl.metadata (5.9 kB) + Downloading torchvision-0.24.0-cp311-cp311-manylinux_2_28_x86_64.whl.metadata (5.9 kB) + Downloading torchvision-0.23.0-cp311-cp311-manylinux_2_28_x86_64.whl.metadata (6.1 kB) +Collecting pillow!=8.3.*,>=5.3.0 (from torchvision->timm==1.0.27) + Downloading pillow-12.3.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl.metadata (9.1 kB) +Downloading torch-2.8.0-cp311-cp311-manylinux_2_28_x86_64.whl (888.1 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 888.1/888.1 MB 227.6 MB/s eta 0:00:00 +Downloading numpy-2.4.6-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl (16.9 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 16.9/16.9 MB 286.0 MB/s eta 0:00:00 +Downloading nibabel-5.4.2-py3-none-any.whl (3.3 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 3.3/3.3 MB 288.9 MB/s eta 0:00:00 +Downloading einops-0.8.2-py3-none-any.whl (65 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 65.6/65.6 kB 290.5 MB/s eta 0:00:00 +Downloading jaxtyping-0.3.10-py3-none-any.whl (56 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 56.4/56.4 kB 356.7 MB/s eta 0:00:00 +Downloading timm-1.0.27-py3-none-any.whl (2.6 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 2.6/2.6 MB 215.7 MB/s eta 0:00:00 +Downloading huggingface_hub-0.36.2-py3-none-any.whl (566 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 566.4/566.4 kB 221.4 MB/s eta 0:00:00 +Downloading scikit_learn-1.8.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl (9.1 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 9.1/9.1 MB 169.5 MB/s eta 0:00:00 +Downloading joblib-1.5.3-py3-none-any.whl (309 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 309.1/309.1 kB 112.6 MB/s eta 0:00:00 +Downloading omegaconf-2.3.0-py3-none-any.whl (79 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 79.5/79.5 kB 406.4 MB/s eta 0:00:00 +Downloading nvidia_cublas_cu12-12.8.4.1-py3-none-manylinux_2_27_x86_64.whl (594.3 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 594.3/594.3 MB 213.7 MB/s eta 0:00:00 +Downloading nvidia_cuda_cupti_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl (10.2 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 10.2/10.2 MB 218.2 MB/s eta 0:00:00 +Downloading nvidia_cuda_nvrtc_cu12-12.8.93-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl (88.0 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 88.0/88.0 MB 301.5 MB/s eta 0:00:00 +Downloading nvidia_cuda_runtime_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl (954 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 954.8/954.8 kB 223.9 MB/s eta 0:00:00 +Downloading nvidia_cudnn_cu12-9.10.2.21-py3-none-manylinux_2_27_x86_64.whl (706.8 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 706.8/706.8 MB 302.6 MB/s eta 0:00:00 +Downloading nvidia_cufft_cu12-11.3.3.83-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl (193.1 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 193.1/193.1 MB 305.5 MB/s eta 0:00:00 +Downloading nvidia_cufile_cu12-1.13.1.3-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl (1.2 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 1.2/1.2 MB 228.2 MB/s eta 0:00:00 +Downloading nvidia_curand_cu12-10.3.9.90-py3-none-manylinux_2_27_x86_64.whl (63.6 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 63.6/63.6 MB 312.1 MB/s eta 0:00:00 +Downloading nvidia_cusolver_cu12-11.7.3.90-py3-none-manylinux_2_27_x86_64.whl (267.5 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 267.5/267.5 MB 301.3 MB/s eta 0:00:00 +Downloading nvidia_cusparse_cu12-12.5.8.93-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl (288.2 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 288.2/288.2 MB 299.6 MB/s eta 0:00:00 +Downloading nvidia_cusparselt_cu12-0.7.1-py3-none-manylinux2014_x86_64.whl (287.2 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 287.2/287.2 MB 304.3 MB/s eta 0:00:00 +Downloading nvidia_nccl_cu12-2.27.3-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl (322.4 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 322.4/322.4 MB 275.0 MB/s eta 0:00:00 +Downloading nvidia_nvjitlink_cu12-12.8.93-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl (39.3 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 39.3/39.3 MB 255.2 MB/s eta 0:00:00 +Downloading nvidia_nvtx_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl (89 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 90.0/90.0 kB 393.8 MB/s eta 0:00:00 +Downloading triton-3.4.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl (155.5 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 155.5/155.5 MB 302.6 MB/s eta 0:00:00 +Downloading fsspec-2026.7.0-py3-none-any.whl (206 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 206.6/206.6 kB 464.9 MB/s eta 0:00:00 +Downloading hf_xet-1.6.0-cp38-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl (4.5 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 4.5/4.5 MB 338.8 MB/s eta 0:00:00 +Downloading importlib_resources-7.1.0-py3-none-any.whl (37 kB) +Downloading pyyaml-6.0.3-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl (806 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 806.6/806.6 kB 427.8 MB/s eta 0:00:00 +Downloading scipy-1.17.1-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl (35.3 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 35.3/35.3 MB 268.4 MB/s eta 0:00:00 +Downloading sympy-1.14.0-py3-none-any.whl (6.3 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 6.3/6.3 MB 331.0 MB/s eta 0:00:00 +Downloading threadpoolctl-3.6.0-py3-none-any.whl (18 kB) +Downloading tqdm-4.70.0-py3-none-any.whl (80 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 80.2/80.2 kB 319.2 MB/s eta 0:00:00 +Downloading typing_extensions-4.16.0-py3-none-any.whl (45 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 45.6/45.6 kB 289.4 MB/s eta 0:00:00 +Downloading wadler_lindig-0.1.7-py3-none-any.whl (20 kB) +Downloading filelock-3.32.2-py3-none-any.whl (98 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 98.8/98.8 kB 317.4 MB/s eta 0:00:00 +Downloading jinja2-3.1.6-py3-none-any.whl (134 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 134.9/134.9 kB 406.2 MB/s eta 0:00:00 +Downloading networkx-3.6.1-py3-none-any.whl (2.1 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 2.1/2.1 MB 220.8 MB/s eta 0:00:00 +Downloading requests-2.34.2-py3-none-any.whl (73 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 73.1/73.1 kB 367.4 MB/s eta 0:00:00 +Downloading safetensors-0.8.0-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (516 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 516.0/516.0 kB 227.7 MB/s eta 0:00:00 +Downloading torchvision-0.23.0-cp311-cp311-manylinux_2_28_x86_64.whl (8.6 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 8.6/8.6 MB 168.8 MB/s eta 0:00:00 +Downloading certifi-2026.7.22-py3-none-any.whl (136 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 137.0/137.0 kB 412.0 MB/s eta 0:00:00 +Downloading charset_normalizer-3.4.9-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl (221 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 221.3/221.3 kB 455.8 MB/s eta 0:00:00 +Downloading idna-3.18-py3-none-any.whl (65 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 65.5/65.5 kB 258.9 MB/s eta 0:00:00 +Downloading markupsafe-3.0.3-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl (22 kB) +Downloading mpmath-1.3.0-py3-none-any.whl (536 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 536.2/536.2 kB 241.6 MB/s eta 0:00:00 +Downloading pillow-12.3.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl (6.9 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 6.9/6.9 MB 176.0 MB/s eta 0:00:00 +Downloading urllib3-2.7.0-py3-none-any.whl (131 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 131.1/131.1 kB 349.6 MB/s eta 0:00:00 +Building wheels for collected packages: antlr4-python3-runtime + Building wheel for antlr4-python3-runtime (setup.py): started + Building wheel for antlr4-python3-runtime (setup.py): finished with status 'done' + Created wheel for antlr4-python3-runtime: filename=antlr4_python3_runtime-4.9.3-py3-none-any.whl size=144590 sha256=7049f12da321fd7083d3bdaa9af713223bfc3702787beff8fbc0980b4dfdc5a8 + Stored in directory: /tmp/pip-ephem-wheel-cache-1binwsnn/wheels/1a/97/32/461f837398029ad76911109f07047fde1d7b661a147c7c56d1 +Successfully built antlr4-python3-runtime +Installing collected packages: nvidia-cusparselt-cu12, mpmath, antlr4-python3-runtime, wadler-lindig, urllib3, typing-extensions, triton, tqdm, threadpoolctl, sympy, safetensors, pyyaml, pillow, nvidia-nvtx-cu12, nvidia-nvjitlink-cu12, nvidia-nccl-cu12, nvidia-curand-cu12, nvidia-cufile-cu12, nvidia-cuda-runtime-cu12, nvidia-cuda-nvrtc-cu12, nvidia-cuda-cupti-cu12, nvidia-cublas-cu12, numpy, networkx, MarkupSafe, joblib, importlib-resources, idna, hf-xet, fsspec, filelock, einops, charset_normalizer, certifi, scipy, requests, omegaconf, nvidia-cusparse-cu12, nvidia-cufft-cu12, nvidia-cudnn-cu12, nibabel, jinja2, jaxtyping, scikit-learn, nvidia-cusolver-cu12, huggingface-hub, torch, torchvision, timm +Successfully installed MarkupSafe-3.0.3 antlr4-python3-runtime-4.9.3 certifi-2026.7.22 charset_normalizer-3.4.9 einops-0.8.2 filelock-3.32.2 fsspec-2026.7.0 hf-xet-1.6.0 huggingface-hub-0.36.2 idna-3.18 importlib-resources-7.1.0 jaxtyping-0.3.10 jinja2-3.1.6 joblib-1.5.3 mpmath-1.3.0 networkx-3.6.1 nibabel-5.4.2 numpy-2.4.6 nvidia-cublas-cu12-12.8.4.1 nvidia-cuda-cupti-cu12-12.8.90 nvidia-cuda-nvrtc-cu12-12.8.93 nvidia-cuda-runtime-cu12-12.8.90 nvidia-cudnn-cu12-9.10.2.21 nvidia-cufft-cu12-11.3.3.83 nvidia-cufile-cu12-1.13.1.3 nvidia-curand-cu12-10.3.9.90 nvidia-cusolver-cu12-11.7.3.90 nvidia-cusparse-cu12-12.5.8.93 nvidia-cusparselt-cu12-0.7.1 nvidia-nccl-cu12-2.27.3 nvidia-nvjitlink-cu12-12.8.93 nvidia-nvtx-cu12-12.8.90 omegaconf-2.3.0 pillow-12.3.0 pyyaml-6.0.3 requests-2.34.2 safetensors-0.8.0 scikit-learn-1.8.0 scipy-1.17.1 sympy-1.14.0 threadpoolctl-3.6.0 timm-1.0.27 torch-2.8.0 torchvision-0.23.0 tqdm-4.70.0 triton-3.4.0 typing-extensions-4.16.0 urllib3-2.7.0 wadler-lindig-0.1.7 +WARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv + +[notice] A new release of pip is available: 24.0 -> 26.2.1 +[notice] To update, run: pip install --upgrade pip +INFO: Adding runscript +INFO: /etc/singularity/ exists; cleanup by system administrator is not complete (see https://apptainer.org/docs/admin/latest/singularity_migration.html) +INFO: Creating SIF file... +INFO: Build complete: /data/connor/nanobrain.1/experiments/fomo_tune_baseline/output/task1/task1.sif +built experiments/fomo_tune_baseline/output/task1/task1.sif +=== task5 === +INFO: /etc/singularity/ exists; cleanup by system administrator is not complete (see https://apptainer.org/docs/admin/latest/singularity_migration.html) +INFO: /etc/singularity/ exists; cleanup by system administrator is not complete (see https://apptainer.org/docs/admin/latest/singularity_migration.html) +INFO: Starting build... +INFO: Fetching OCI image... +INFO: Extracting OCI image... +INFO: Inserting Apptainer configuration... +INFO: Copying fomo_tune to /app/fomo_tune +INFO: Copying smri_mae to /app/smri_mae +INFO: Copying model to /app/model +INFO: Copying predict.py to /app/predict.py +INFO: Running post scriptlet ++ pip install --no-cache-dir torch==2.8.0 numpy==2.4.6 nibabel==5.4.2 einops==0.8.2 jaxtyping==0.3.10 timm==1.0.27 huggingface-hub==0.36.2 scikit-learn==1.8.0 joblib==1.5.3 omegaconf==2.3.0 +Collecting torch==2.8.0 + Downloading torch-2.8.0-cp311-cp311-manylinux_2_28_x86_64.whl.metadata (30 kB) +Collecting numpy==2.4.6 + Downloading numpy-2.4.6-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl.metadata (6.6 kB) +Collecting nibabel==5.4.2 + Downloading nibabel-5.4.2-py3-none-any.whl.metadata (8.9 kB) +Collecting einops==0.8.2 + Downloading einops-0.8.2-py3-none-any.whl.metadata (13 kB) +Collecting jaxtyping==0.3.10 + Downloading jaxtyping-0.3.10-py3-none-any.whl.metadata (6.2 kB) +Collecting timm==1.0.27 + Downloading timm-1.0.27-py3-none-any.whl.metadata (40 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 40.2/40.2 kB 20.8 MB/s eta 0:00:00 +Collecting huggingface-hub==0.36.2 + Downloading huggingface_hub-0.36.2-py3-none-any.whl.metadata (15 kB) +Collecting scikit-learn==1.8.0 + Downloading scikit_learn-1.8.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl.metadata (11 kB) +Collecting joblib==1.5.3 + Downloading joblib-1.5.3-py3-none-any.whl.metadata (5.5 kB) +Collecting omegaconf==2.3.0 + Downloading omegaconf-2.3.0-py3-none-any.whl.metadata (3.9 kB) +Collecting filelock (from torch==2.8.0) + Downloading filelock-3.32.2-py3-none-any.whl.metadata (2.0 kB) +Collecting typing-extensions>=4.10.0 (from torch==2.8.0) + Downloading typing_extensions-4.16.0-py3-none-any.whl.metadata (3.3 kB) +Collecting sympy>=1.13.3 (from torch==2.8.0) + Downloading sympy-1.14.0-py3-none-any.whl.metadata (12 kB) +Collecting networkx (from torch==2.8.0) + Downloading networkx-3.6.1-py3-none-any.whl.metadata (6.8 kB) +Collecting jinja2 (from torch==2.8.0) + Downloading jinja2-3.1.6-py3-none-any.whl.metadata (2.9 kB) +Collecting fsspec (from torch==2.8.0) + Downloading fsspec-2026.7.0-py3-none-any.whl.metadata (10 kB) +Collecting nvidia-cuda-nvrtc-cu12==12.8.93 (from torch==2.8.0) + Downloading nvidia_cuda_nvrtc_cu12-12.8.93-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl.metadata (1.7 kB) +Collecting nvidia-cuda-runtime-cu12==12.8.90 (from torch==2.8.0) + Downloading nvidia_cuda_runtime_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl.metadata (1.7 kB) +Collecting nvidia-cuda-cupti-cu12==12.8.90 (from torch==2.8.0) + Downloading nvidia_cuda_cupti_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl.metadata (1.7 kB) +Collecting nvidia-cudnn-cu12==9.10.2.21 (from torch==2.8.0) + Downloading nvidia_cudnn_cu12-9.10.2.21-py3-none-manylinux_2_27_x86_64.whl.metadata (1.8 kB) +Collecting nvidia-cublas-cu12==12.8.4.1 (from torch==2.8.0) + Downloading nvidia_cublas_cu12-12.8.4.1-py3-none-manylinux_2_27_x86_64.whl.metadata (1.7 kB) +Collecting nvidia-cufft-cu12==11.3.3.83 (from torch==2.8.0) + Downloading nvidia_cufft_cu12-11.3.3.83-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl.metadata (1.7 kB) +Collecting nvidia-curand-cu12==10.3.9.90 (from torch==2.8.0) + Downloading nvidia_curand_cu12-10.3.9.90-py3-none-manylinux_2_27_x86_64.whl.metadata (1.7 kB) +Collecting nvidia-cusolver-cu12==11.7.3.90 (from torch==2.8.0) + Downloading nvidia_cusolver_cu12-11.7.3.90-py3-none-manylinux_2_27_x86_64.whl.metadata (1.8 kB) +Collecting nvidia-cusparse-cu12==12.5.8.93 (from torch==2.8.0) + Downloading nvidia_cusparse_cu12-12.5.8.93-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl.metadata (1.8 kB) +Collecting nvidia-cusparselt-cu12==0.7.1 (from torch==2.8.0) + Downloading nvidia_cusparselt_cu12-0.7.1-py3-none-manylinux2014_x86_64.whl.metadata (7.0 kB) +Collecting nvidia-nccl-cu12==2.27.3 (from torch==2.8.0) + Downloading nvidia_nccl_cu12-2.27.3-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl.metadata (2.0 kB) +Collecting nvidia-nvtx-cu12==12.8.90 (from torch==2.8.0) + Downloading nvidia_nvtx_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl.metadata (1.8 kB) +Collecting nvidia-nvjitlink-cu12==12.8.93 (from torch==2.8.0) + Downloading nvidia_nvjitlink_cu12-12.8.93-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl.metadata (1.7 kB) +Collecting nvidia-cufile-cu12==1.13.1.3 (from torch==2.8.0) + Downloading nvidia_cufile_cu12-1.13.1.3-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl.metadata (1.7 kB) +Collecting triton==3.4.0 (from torch==2.8.0) + Downloading triton-3.4.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl.metadata (1.7 kB) +Collecting importlib-resources>=5.12 (from nibabel==5.4.2) + Downloading importlib_resources-7.1.0-py3-none-any.whl.metadata (4.0 kB) +Requirement already satisfied: packaging>=20 in /usr/local/lib/python3.11/site-packages (from nibabel==5.4.2) (26.3) +Collecting wadler-lindig>=0.1.3 (from jaxtyping==0.3.10) + Downloading wadler_lindig-0.1.7-py3-none-any.whl.metadata (17 kB) +Collecting torchvision (from timm==1.0.27) + Downloading torchvision-0.28.0-cp311-cp311-manylinux_2_28_x86_64.whl.metadata (5.6 kB) +Collecting pyyaml (from timm==1.0.27) + Downloading pyyaml-6.0.3-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl.metadata (2.4 kB) +Collecting safetensors (from timm==1.0.27) + Downloading safetensors-0.8.0-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (4.2 kB) +Collecting hf-xet<2.0.0,>=1.1.3 (from huggingface-hub==0.36.2) + Downloading hf_xet-1.6.0-cp38-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl.metadata (4.9 kB) +Collecting requests (from huggingface-hub==0.36.2) + Downloading requests-2.34.2-py3-none-any.whl.metadata (4.8 kB) +Collecting tqdm>=4.42.1 (from huggingface-hub==0.36.2) + Downloading tqdm-4.70.0-py3-none-any.whl.metadata (57 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 57.3/57.3 kB 39.7 MB/s eta 0:00:00 +Collecting scipy>=1.10.0 (from scikit-learn==1.8.0) + Downloading scipy-1.17.1-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl.metadata (62 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 62.1/62.1 kB 67.7 MB/s eta 0:00:00 +Collecting threadpoolctl>=3.2.0 (from scikit-learn==1.8.0) + Downloading threadpoolctl-3.6.0-py3-none-any.whl.metadata (13 kB) +Collecting antlr4-python3-runtime==4.9.* (from omegaconf==2.3.0) + Downloading antlr4-python3-runtime-4.9.3.tar.gz (117 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 117.0/117.0 kB 86.0 MB/s eta 0:00:00 + Preparing metadata (setup.py): started + Preparing metadata (setup.py): finished with status 'done' +Requirement already satisfied: setuptools>=40.8.0 in /usr/local/lib/python3.11/site-packages (from triton==3.4.0->torch==2.8.0) (79.0.1) +Collecting mpmath<1.4,>=1.1.0 (from sympy>=1.13.3->torch==2.8.0) + Downloading mpmath-1.3.0-py3-none-any.whl.metadata (8.6 kB) +Collecting MarkupSafe>=2.0 (from jinja2->torch==2.8.0) + Downloading markupsafe-3.0.3-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl.metadata (2.7 kB) +Collecting charset_normalizer<4,>=2 (from requests->huggingface-hub==0.36.2) + Downloading charset_normalizer-3.4.9-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl.metadata (41 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 41.7/41.7 kB 295.5 MB/s eta 0:00:00 +Collecting idna<4,>=2.5 (from requests->huggingface-hub==0.36.2) + Downloading idna-3.18-py3-none-any.whl.metadata (6.1 kB) +Collecting urllib3<3,>=1.26 (from requests->huggingface-hub==0.36.2) + Downloading urllib3-2.7.0-py3-none-any.whl.metadata (6.9 kB) +Collecting certifi>=2023.5.7 (from requests->huggingface-hub==0.36.2) + Downloading certifi-2026.7.22-py3-none-any.whl.metadata (2.5 kB) +INFO: pip is looking at multiple versions of torchvision to determine which version is compatible with other requirements. This could take a while. +Collecting torchvision (from timm==1.0.27) + Downloading torchvision-0.27.1-cp311-cp311-manylinux_2_28_x86_64.whl.metadata (5.5 kB) + Downloading torchvision-0.27.0-cp311-cp311-manylinux_2_28_x86_64.whl.metadata (5.5 kB) + Downloading torchvision-0.26.0-cp311-cp311-manylinux_2_28_x86_64.whl.metadata (5.5 kB) + Downloading torchvision-0.25.0-cp311-cp311-manylinux_2_28_x86_64.whl.metadata (5.4 kB) + Downloading torchvision-0.24.1-cp311-cp311-manylinux_2_28_x86_64.whl.metadata (5.9 kB) + Downloading torchvision-0.24.0-cp311-cp311-manylinux_2_28_x86_64.whl.metadata (5.9 kB) + Downloading torchvision-0.23.0-cp311-cp311-manylinux_2_28_x86_64.whl.metadata (6.1 kB) +Collecting pillow!=8.3.*,>=5.3.0 (from torchvision->timm==1.0.27) + Downloading pillow-12.3.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl.metadata (9.1 kB) +Downloading torch-2.8.0-cp311-cp311-manylinux_2_28_x86_64.whl (888.1 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 888.1/888.1 MB 212.0 MB/s eta 0:00:00 +Downloading numpy-2.4.6-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl (16.9 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 16.9/16.9 MB 311.3 MB/s eta 0:00:00 +Downloading nibabel-5.4.2-py3-none-any.whl (3.3 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 3.3/3.3 MB 210.5 MB/s eta 0:00:00 +Downloading einops-0.8.2-py3-none-any.whl (65 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 65.6/65.6 kB 354.1 MB/s eta 0:00:00 +Downloading jaxtyping-0.3.10-py3-none-any.whl (56 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 56.4/56.4 kB 336.3 MB/s eta 0:00:00 +Downloading timm-1.0.27-py3-none-any.whl (2.6 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 2.6/2.6 MB 175.2 MB/s eta 0:00:00 +Downloading huggingface_hub-0.36.2-py3-none-any.whl (566 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 566.4/566.4 kB 129.1 MB/s eta 0:00:00 +Downloading scikit_learn-1.8.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl (9.1 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 9.1/9.1 MB 308.6 MB/s eta 0:00:00 +Downloading joblib-1.5.3-py3-none-any.whl (309 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 309.1/309.1 kB 431.8 MB/s eta 0:00:00 +Downloading omegaconf-2.3.0-py3-none-any.whl (79 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 79.5/79.5 kB 388.1 MB/s eta 0:00:00 +Downloading nvidia_cublas_cu12-12.8.4.1-py3-none-manylinux_2_27_x86_64.whl (594.3 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 594.3/594.3 MB 139.4 MB/s eta 0:00:00 +Downloading nvidia_cuda_cupti_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl (10.2 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 10.2/10.2 MB 174.8 MB/s eta 0:00:00 +Downloading nvidia_cuda_nvrtc_cu12-12.8.93-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl (88.0 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 88.0/88.0 MB 244.3 MB/s eta 0:00:00 +Downloading nvidia_cuda_runtime_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl (954 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 954.8/954.8 kB 289.0 MB/s eta 0:00:00 +Downloading nvidia_cudnn_cu12-9.10.2.21-py3-none-manylinux_2_27_x86_64.whl (706.8 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 706.8/706.8 MB 232.5 MB/s eta 0:00:00 +Downloading nvidia_cufft_cu12-11.3.3.83-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl (193.1 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 193.1/193.1 MB 155.1 MB/s eta 0:00:00 +Downloading nvidia_cufile_cu12-1.13.1.3-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl (1.2 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 1.2/1.2 MB 173.3 MB/s eta 0:00:00 +Downloading nvidia_curand_cu12-10.3.9.90-py3-none-manylinux_2_27_x86_64.whl (63.6 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 63.6/63.6 MB 280.8 MB/s eta 0:00:00 +Downloading nvidia_cusolver_cu12-11.7.3.90-py3-none-manylinux_2_27_x86_64.whl (267.5 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 267.5/267.5 MB 206.9 MB/s eta 0:00:00 +Downloading nvidia_cusparse_cu12-12.5.8.93-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl (288.2 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 288.2/288.2 MB 178.9 MB/s eta 0:00:00 +Downloading nvidia_cusparselt_cu12-0.7.1-py3-none-manylinux2014_x86_64.whl (287.2 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 287.2/287.2 MB 219.2 MB/s eta 0:00:00 +Downloading nvidia_nccl_cu12-2.27.3-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl (322.4 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 322.4/322.4 MB 296.8 MB/s eta 0:00:00 +Downloading nvidia_nvjitlink_cu12-12.8.93-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl (39.3 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 39.3/39.3 MB 207.8 MB/s eta 0:00:00 +Downloading nvidia_nvtx_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl (89 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 90.0/90.0 kB 401.1 MB/s eta 0:00:00 +Downloading triton-3.4.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl (155.5 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 155.5/155.5 MB 299.9 MB/s eta 0:00:00 +Downloading fsspec-2026.7.0-py3-none-any.whl (206 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 206.6/206.6 kB 461.6 MB/s eta 0:00:00 +Downloading hf_xet-1.6.0-cp38-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl (4.5 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 4.5/4.5 MB 305.4 MB/s eta 0:00:00 +Downloading importlib_resources-7.1.0-py3-none-any.whl (37 kB) +Downloading pyyaml-6.0.3-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl (806 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 806.6/806.6 kB 461.9 MB/s eta 0:00:00 +Downloading scipy-1.17.1-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl (35.3 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 35.3/35.3 MB 157.6 MB/s eta 0:00:00 +Downloading sympy-1.14.0-py3-none-any.whl (6.3 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 6.3/6.3 MB 254.2 MB/s eta 0:00:00 +Downloading threadpoolctl-3.6.0-py3-none-any.whl (18 kB) +Downloading tqdm-4.70.0-py3-none-any.whl (80 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 80.2/80.2 kB 384.1 MB/s eta 0:00:00 +Downloading typing_extensions-4.16.0-py3-none-any.whl (45 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 45.6/45.6 kB 314.4 MB/s eta 0:00:00 +Downloading wadler_lindig-0.1.7-py3-none-any.whl (20 kB) +Downloading filelock-3.32.2-py3-none-any.whl (98 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 98.8/98.8 kB 308.0 MB/s eta 0:00:00 +Downloading jinja2-3.1.6-py3-none-any.whl (134 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 134.9/134.9 kB 357.5 MB/s eta 0:00:00 +Downloading networkx-3.6.1-py3-none-any.whl (2.1 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 2.1/2.1 MB 327.2 MB/s eta 0:00:00 +Downloading requests-2.34.2-py3-none-any.whl (73 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 73.1/73.1 kB 379.7 MB/s eta 0:00:00 +Downloading safetensors-0.8.0-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (516 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 516.0/516.0 kB 201.6 MB/s eta 0:00:00 +Downloading torchvision-0.23.0-cp311-cp311-manylinux_2_28_x86_64.whl (8.6 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 8.6/8.6 MB 189.9 MB/s eta 0:00:00 +Downloading certifi-2026.7.22-py3-none-any.whl (136 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 137.0/137.0 kB 390.2 MB/s eta 0:00:00 +Downloading charset_normalizer-3.4.9-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl (221 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 221.3/221.3 kB 440.2 MB/s eta 0:00:00 +Downloading idna-3.18-py3-none-any.whl (65 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 65.5/65.5 kB 278.9 MB/s eta 0:00:00 +Downloading markupsafe-3.0.3-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl (22 kB) +Downloading mpmath-1.3.0-py3-none-any.whl (536 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 536.2/536.2 kB 470.6 MB/s eta 0:00:00 +Downloading pillow-12.3.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl (6.9 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 6.9/6.9 MB 295.5 MB/s eta 0:00:00 +Downloading urllib3-2.7.0-py3-none-any.whl (131 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 131.1/131.1 kB 434.9 MB/s eta 0:00:00 +Building wheels for collected packages: antlr4-python3-runtime + Building wheel for antlr4-python3-runtime (setup.py): started + Building wheel for antlr4-python3-runtime (setup.py): finished with status 'done' + Created wheel for antlr4-python3-runtime: filename=antlr4_python3_runtime-4.9.3-py3-none-any.whl size=144590 sha256=9153c9732e66786ceb0134b28a0118f6e9fedc66584880fabf6533f1989fac2e + Stored in directory: /tmp/pip-ephem-wheel-cache-5dznrpve/wheels/1a/97/32/461f837398029ad76911109f07047fde1d7b661a147c7c56d1 +Successfully built antlr4-python3-runtime +Installing collected packages: nvidia-cusparselt-cu12, mpmath, antlr4-python3-runtime, wadler-lindig, urllib3, typing-extensions, triton, tqdm, threadpoolctl, sympy, safetensors, pyyaml, pillow, nvidia-nvtx-cu12, nvidia-nvjitlink-cu12, nvidia-nccl-cu12, nvidia-curand-cu12, nvidia-cufile-cu12, nvidia-cuda-runtime-cu12, nvidia-cuda-nvrtc-cu12, nvidia-cuda-cupti-cu12, nvidia-cublas-cu12, numpy, networkx, MarkupSafe, joblib, importlib-resources, idna, hf-xet, fsspec, filelock, einops, charset_normalizer, certifi, scipy, requests, omegaconf, nvidia-cusparse-cu12, nvidia-cufft-cu12, nvidia-cudnn-cu12, nibabel, jinja2, jaxtyping, scikit-learn, nvidia-cusolver-cu12, huggingface-hub, torch, torchvision, timm +Successfully installed MarkupSafe-3.0.3 antlr4-python3-runtime-4.9.3 certifi-2026.7.22 charset_normalizer-3.4.9 einops-0.8.2 filelock-3.32.2 fsspec-2026.7.0 hf-xet-1.6.0 huggingface-hub-0.36.2 idna-3.18 importlib-resources-7.1.0 jaxtyping-0.3.10 jinja2-3.1.6 joblib-1.5.3 mpmath-1.3.0 networkx-3.6.1 nibabel-5.4.2 numpy-2.4.6 nvidia-cublas-cu12-12.8.4.1 nvidia-cuda-cupti-cu12-12.8.90 nvidia-cuda-nvrtc-cu12-12.8.93 nvidia-cuda-runtime-cu12-12.8.90 nvidia-cudnn-cu12-9.10.2.21 nvidia-cufft-cu12-11.3.3.83 nvidia-cufile-cu12-1.13.1.3 nvidia-curand-cu12-10.3.9.90 nvidia-cusolver-cu12-11.7.3.90 nvidia-cusparse-cu12-12.5.8.93 nvidia-cusparselt-cu12-0.7.1 nvidia-nccl-cu12-2.27.3 nvidia-nvjitlink-cu12-12.8.93 nvidia-nvtx-cu12-12.8.90 omegaconf-2.3.0 pillow-12.3.0 pyyaml-6.0.3 requests-2.34.2 safetensors-0.8.0 scikit-learn-1.8.0 scipy-1.17.1 sympy-1.14.0 threadpoolctl-3.6.0 timm-1.0.27 torch-2.8.0 torchvision-0.23.0 tqdm-4.70.0 triton-3.4.0 typing-extensions-4.16.0 urllib3-2.7.0 wadler-lindig-0.1.7 +WARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv + +[notice] A new release of pip is available: 24.0 -> 26.2.1 +[notice] To update, run: pip install --upgrade pip +INFO: Adding runscript +INFO: /etc/singularity/ exists; cleanup by system administrator is not complete (see https://apptainer.org/docs/admin/latest/singularity_migration.html) +INFO: Creating SIF file... +INFO: Build complete: /data/connor/nanobrain.1/experiments/fomo_tune_baseline/output/task5/task5.sif +built experiments/fomo_tune_baseline/output/task5/task5.sif +=== task3 === +INFO: /etc/singularity/ exists; cleanup by system administrator is not complete (see https://apptainer.org/docs/admin/latest/singularity_migration.html) +INFO: /etc/singularity/ exists; cleanup by system administrator is not complete (see https://apptainer.org/docs/admin/latest/singularity_migration.html) +INFO: Starting build... +INFO: Fetching OCI image... +INFO: Extracting OCI image... +INFO: Inserting Apptainer configuration... +INFO: Copying fomo_tune to /app/fomo_tune +INFO: Copying smri_mae to /app/smri_mae +INFO: Copying model to /app/model +INFO: Copying predict.py to /app/predict.py +INFO: Running post scriptlet ++ pip install --no-cache-dir torch==2.8.0 numpy==2.4.6 nibabel==5.4.2 einops==0.8.2 jaxtyping==0.3.10 timm==1.0.27 huggingface-hub==0.36.2 scikit-learn==1.8.0 joblib==1.5.3 omegaconf==2.3.0 +Collecting torch==2.8.0 + Downloading torch-2.8.0-cp311-cp311-manylinux_2_28_x86_64.whl.metadata (30 kB) +Collecting numpy==2.4.6 + Downloading numpy-2.4.6-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl.metadata (6.6 kB) +Collecting nibabel==5.4.2 + Downloading nibabel-5.4.2-py3-none-any.whl.metadata (8.9 kB) +Collecting einops==0.8.2 + Downloading einops-0.8.2-py3-none-any.whl.metadata (13 kB) +Collecting jaxtyping==0.3.10 + Downloading jaxtyping-0.3.10-py3-none-any.whl.metadata (6.2 kB) +Collecting timm==1.0.27 + Downloading timm-1.0.27-py3-none-any.whl.metadata (40 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 40.2/40.2 kB 330.8 MB/s eta 0:00:00 +Collecting huggingface-hub==0.36.2 + Downloading huggingface_hub-0.36.2-py3-none-any.whl.metadata (15 kB) +Collecting scikit-learn==1.8.0 + Downloading scikit_learn-1.8.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl.metadata (11 kB) +Collecting joblib==1.5.3 + Downloading joblib-1.5.3-py3-none-any.whl.metadata (5.5 kB) +Collecting omegaconf==2.3.0 + Downloading omegaconf-2.3.0-py3-none-any.whl.metadata (3.9 kB) +Collecting filelock (from torch==2.8.0) + Downloading filelock-3.32.2-py3-none-any.whl.metadata (2.0 kB) +Collecting typing-extensions>=4.10.0 (from torch==2.8.0) + Downloading typing_extensions-4.16.0-py3-none-any.whl.metadata (3.3 kB) +Collecting sympy>=1.13.3 (from torch==2.8.0) + Downloading sympy-1.14.0-py3-none-any.whl.metadata (12 kB) +Collecting networkx (from torch==2.8.0) + Downloading networkx-3.6.1-py3-none-any.whl.metadata (6.8 kB) +Collecting jinja2 (from torch==2.8.0) + Downloading jinja2-3.1.6-py3-none-any.whl.metadata (2.9 kB) +Collecting fsspec (from torch==2.8.0) + Downloading fsspec-2026.7.0-py3-none-any.whl.metadata (10 kB) +Collecting nvidia-cuda-nvrtc-cu12==12.8.93 (from torch==2.8.0) + Downloading nvidia_cuda_nvrtc_cu12-12.8.93-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl.metadata (1.7 kB) +Collecting nvidia-cuda-runtime-cu12==12.8.90 (from torch==2.8.0) + Downloading nvidia_cuda_runtime_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl.metadata (1.7 kB) +Collecting nvidia-cuda-cupti-cu12==12.8.90 (from torch==2.8.0) + Downloading nvidia_cuda_cupti_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl.metadata (1.7 kB) +Collecting nvidia-cudnn-cu12==9.10.2.21 (from torch==2.8.0) + Downloading nvidia_cudnn_cu12-9.10.2.21-py3-none-manylinux_2_27_x86_64.whl.metadata (1.8 kB) +Collecting nvidia-cublas-cu12==12.8.4.1 (from torch==2.8.0) + Downloading nvidia_cublas_cu12-12.8.4.1-py3-none-manylinux_2_27_x86_64.whl.metadata (1.7 kB) +Collecting nvidia-cufft-cu12==11.3.3.83 (from torch==2.8.0) + Downloading nvidia_cufft_cu12-11.3.3.83-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl.metadata (1.7 kB) +Collecting nvidia-curand-cu12==10.3.9.90 (from torch==2.8.0) + Downloading nvidia_curand_cu12-10.3.9.90-py3-none-manylinux_2_27_x86_64.whl.metadata (1.7 kB) +Collecting nvidia-cusolver-cu12==11.7.3.90 (from torch==2.8.0) + Downloading nvidia_cusolver_cu12-11.7.3.90-py3-none-manylinux_2_27_x86_64.whl.metadata (1.8 kB) +Collecting nvidia-cusparse-cu12==12.5.8.93 (from torch==2.8.0) + Downloading nvidia_cusparse_cu12-12.5.8.93-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl.metadata (1.8 kB) +Collecting nvidia-cusparselt-cu12==0.7.1 (from torch==2.8.0) + Downloading nvidia_cusparselt_cu12-0.7.1-py3-none-manylinux2014_x86_64.whl.metadata (7.0 kB) +Collecting nvidia-nccl-cu12==2.27.3 (from torch==2.8.0) + Downloading nvidia_nccl_cu12-2.27.3-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl.metadata (2.0 kB) +Collecting nvidia-nvtx-cu12==12.8.90 (from torch==2.8.0) + Downloading nvidia_nvtx_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl.metadata (1.8 kB) +Collecting nvidia-nvjitlink-cu12==12.8.93 (from torch==2.8.0) + Downloading nvidia_nvjitlink_cu12-12.8.93-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl.metadata (1.7 kB) +Collecting nvidia-cufile-cu12==1.13.1.3 (from torch==2.8.0) + Downloading nvidia_cufile_cu12-1.13.1.3-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl.metadata (1.7 kB) +Collecting triton==3.4.0 (from torch==2.8.0) + Downloading triton-3.4.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl.metadata (1.7 kB) +Collecting importlib-resources>=5.12 (from nibabel==5.4.2) + Downloading importlib_resources-7.1.0-py3-none-any.whl.metadata (4.0 kB) +Requirement already satisfied: packaging>=20 in /usr/local/lib/python3.11/site-packages (from nibabel==5.4.2) (26.3) +Collecting wadler-lindig>=0.1.3 (from jaxtyping==0.3.10) + Downloading wadler_lindig-0.1.7-py3-none-any.whl.metadata (17 kB) +Collecting torchvision (from timm==1.0.27) + Downloading torchvision-0.28.0-cp311-cp311-manylinux_2_28_x86_64.whl.metadata (5.6 kB) +Collecting pyyaml (from timm==1.0.27) + Downloading pyyaml-6.0.3-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl.metadata (2.4 kB) +Collecting safetensors (from timm==1.0.27) + Downloading safetensors-0.8.0-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (4.2 kB) +Collecting hf-xet<2.0.0,>=1.1.3 (from huggingface-hub==0.36.2) + Downloading hf_xet-1.6.0-cp38-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl.metadata (4.9 kB) +Collecting requests (from huggingface-hub==0.36.2) + Downloading requests-2.34.2-py3-none-any.whl.metadata (4.8 kB) +Collecting tqdm>=4.42.1 (from huggingface-hub==0.36.2) + Downloading tqdm-4.70.0-py3-none-any.whl.metadata (57 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 57.3/57.3 kB 362.2 MB/s eta 0:00:00 +Collecting scipy>=1.10.0 (from scikit-learn==1.8.0) + Downloading scipy-1.17.1-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl.metadata (62 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 62.1/62.1 kB 394.1 MB/s eta 0:00:00 +Collecting threadpoolctl>=3.2.0 (from scikit-learn==1.8.0) + Downloading threadpoolctl-3.6.0-py3-none-any.whl.metadata (13 kB) +Collecting antlr4-python3-runtime==4.9.* (from omegaconf==2.3.0) + Downloading antlr4-python3-runtime-4.9.3.tar.gz (117 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 117.0/117.0 kB 424.4 MB/s eta 0:00:00 + Preparing metadata (setup.py): started + Preparing metadata (setup.py): finished with status 'done' +Requirement already satisfied: setuptools>=40.8.0 in /usr/local/lib/python3.11/site-packages (from triton==3.4.0->torch==2.8.0) (79.0.1) +Collecting mpmath<1.4,>=1.1.0 (from sympy>=1.13.3->torch==2.8.0) + Downloading mpmath-1.3.0-py3-none-any.whl.metadata (8.6 kB) +Collecting MarkupSafe>=2.0 (from jinja2->torch==2.8.0) + Downloading markupsafe-3.0.3-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl.metadata (2.7 kB) +Collecting charset_normalizer<4,>=2 (from requests->huggingface-hub==0.36.2) + Downloading charset_normalizer-3.4.9-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl.metadata (41 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 41.7/41.7 kB 330.0 MB/s eta 0:00:00 +Collecting idna<4,>=2.5 (from requests->huggingface-hub==0.36.2) + Downloading idna-3.18-py3-none-any.whl.metadata (6.1 kB) +Collecting urllib3<3,>=1.26 (from requests->huggingface-hub==0.36.2) + Downloading urllib3-2.7.0-py3-none-any.whl.metadata (6.9 kB) +Collecting certifi>=2023.5.7 (from requests->huggingface-hub==0.36.2) + Downloading certifi-2026.7.22-py3-none-any.whl.metadata (2.5 kB) +INFO: pip is looking at multiple versions of torchvision to determine which version is compatible with other requirements. This could take a while. +Collecting torchvision (from timm==1.0.27) + Downloading torchvision-0.27.1-cp311-cp311-manylinux_2_28_x86_64.whl.metadata (5.5 kB) + Downloading torchvision-0.27.0-cp311-cp311-manylinux_2_28_x86_64.whl.metadata (5.5 kB) + Downloading torchvision-0.26.0-cp311-cp311-manylinux_2_28_x86_64.whl.metadata (5.5 kB) + Downloading torchvision-0.25.0-cp311-cp311-manylinux_2_28_x86_64.whl.metadata (5.4 kB) + Downloading torchvision-0.24.1-cp311-cp311-manylinux_2_28_x86_64.whl.metadata (5.9 kB) + Downloading torchvision-0.24.0-cp311-cp311-manylinux_2_28_x86_64.whl.metadata (5.9 kB) + Downloading torchvision-0.23.0-cp311-cp311-manylinux_2_28_x86_64.whl.metadata (6.1 kB) +Collecting pillow!=8.3.*,>=5.3.0 (from torchvision->timm==1.0.27) + Downloading pillow-12.3.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl.metadata (9.1 kB) +Downloading torch-2.8.0-cp311-cp311-manylinux_2_28_x86_64.whl (888.1 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 888.1/888.1 MB 183.5 MB/s eta 0:00:00 +Downloading numpy-2.4.6-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl (16.9 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 16.9/16.9 MB 258.7 MB/s eta 0:00:00 +Downloading nibabel-5.4.2-py3-none-any.whl (3.3 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 3.3/3.3 MB 227.2 MB/s eta 0:00:00 +Downloading einops-0.8.2-py3-none-any.whl (65 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 65.6/65.6 kB 281.2 MB/s eta 0:00:00 +Downloading jaxtyping-0.3.10-py3-none-any.whl (56 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 56.4/56.4 kB 341.7 MB/s eta 0:00:00 +Downloading timm-1.0.27-py3-none-any.whl (2.6 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 2.6/2.6 MB 220.8 MB/s eta 0:00:00 +Downloading huggingface_hub-0.36.2-py3-none-any.whl (566 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 566.4/566.4 kB 451.6 MB/s eta 0:00:00 +Downloading scikit_learn-1.8.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl (9.1 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 9.1/9.1 MB 264.0 MB/s eta 0:00:00 +Downloading joblib-1.5.3-py3-none-any.whl (309 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 309.1/309.1 kB 480.7 MB/s eta 0:00:00 +Downloading omegaconf-2.3.0-py3-none-any.whl (79 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 79.5/79.5 kB 377.4 MB/s eta 0:00:00 +Downloading nvidia_cublas_cu12-12.8.4.1-py3-none-manylinux_2_27_x86_64.whl (594.3 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 594.3/594.3 MB 261.8 MB/s eta 0:00:00 +Downloading nvidia_cuda_cupti_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl (10.2 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 10.2/10.2 MB 212.3 MB/s eta 0:00:00 +Downloading nvidia_cuda_nvrtc_cu12-12.8.93-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl (88.0 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 88.0/88.0 MB 224.5 MB/s eta 0:00:00 +Downloading nvidia_cuda_runtime_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl (954 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 954.8/954.8 kB 309.2 MB/s eta 0:00:00 +Downloading nvidia_cudnn_cu12-9.10.2.21-py3-none-manylinux_2_27_x86_64.whl (706.8 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 706.8/706.8 MB 361.1 MB/s eta 0:00:00 +Downloading nvidia_cufft_cu12-11.3.3.83-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl (193.1 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 193.1/193.1 MB 210.5 MB/s eta 0:00:00 +Downloading nvidia_cufile_cu12-1.13.1.3-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl (1.2 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 1.2/1.2 MB 406.4 MB/s eta 0:00:00 +Downloading nvidia_curand_cu12-10.3.9.90-py3-none-manylinux_2_27_x86_64.whl (63.6 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 63.6/63.6 MB 339.4 MB/s eta 0:00:00 +Downloading nvidia_cusolver_cu12-11.7.3.90-py3-none-manylinux_2_27_x86_64.whl (267.5 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 267.5/267.5 MB 292.8 MB/s eta 0:00:00 +Downloading nvidia_cusparse_cu12-12.5.8.93-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl (288.2 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 288.2/288.2 MB 291.6 MB/s eta 0:00:00 +Downloading nvidia_cusparselt_cu12-0.7.1-py3-none-manylinux2014_x86_64.whl (287.2 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 287.2/287.2 MB 301.2 MB/s eta 0:00:00 +Downloading nvidia_nccl_cu12-2.27.3-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl (322.4 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 322.4/322.4 MB 221.2 MB/s eta 0:00:00 +Downloading nvidia_nvjitlink_cu12-12.8.93-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl (39.3 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 39.3/39.3 MB 343.4 MB/s eta 0:00:00 +Downloading nvidia_nvtx_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl (89 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 90.0/90.0 kB 389.6 MB/s eta 0:00:00 +Downloading triton-3.4.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl (155.5 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 155.5/155.5 MB 206.4 MB/s eta 0:00:00 +Downloading fsspec-2026.7.0-py3-none-any.whl (206 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 206.6/206.6 kB 369.1 MB/s eta 0:00:00 +Downloading hf_xet-1.6.0-cp38-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl (4.5 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 4.5/4.5 MB 376.2 MB/s eta 0:00:00 +Downloading importlib_resources-7.1.0-py3-none-any.whl (37 kB) +Downloading pyyaml-6.0.3-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl (806 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 806.6/806.6 kB 457.9 MB/s eta 0:00:00 +Downloading scipy-1.17.1-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl (35.3 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 35.3/35.3 MB 210.6 MB/s eta 0:00:00 +Downloading sympy-1.14.0-py3-none-any.whl (6.3 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 6.3/6.3 MB 365.8 MB/s eta 0:00:00 +Downloading threadpoolctl-3.6.0-py3-none-any.whl (18 kB) +Downloading tqdm-4.70.0-py3-none-any.whl (80 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 80.2/80.2 kB 406.0 MB/s eta 0:00:00 +Downloading typing_extensions-4.16.0-py3-none-any.whl (45 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 45.6/45.6 kB 218.5 MB/s eta 0:00:00 +Downloading wadler_lindig-0.1.7-py3-none-any.whl (20 kB) +Downloading filelock-3.32.2-py3-none-any.whl (98 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 98.8/98.8 kB 321.0 MB/s eta 0:00:00 +Downloading jinja2-3.1.6-py3-none-any.whl (134 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 134.9/134.9 kB 418.1 MB/s eta 0:00:00 +Downloading networkx-3.6.1-py3-none-any.whl (2.1 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 2.1/2.1 MB 377.2 MB/s eta 0:00:00 +Downloading requests-2.34.2-py3-none-any.whl (73 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 73.1/73.1 kB 298.4 MB/s eta 0:00:00 +Downloading safetensors-0.8.0-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (516 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 516.0/516.0 kB 453.2 MB/s eta 0:00:00 +Downloading torchvision-0.23.0-cp311-cp311-manylinux_2_28_x86_64.whl (8.6 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 8.6/8.6 MB 168.6 MB/s eta 0:00:00 +Downloading certifi-2026.7.22-py3-none-any.whl (136 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 137.0/137.0 kB 368.6 MB/s eta 0:00:00 +Downloading charset_normalizer-3.4.9-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl (221 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 221.3/221.3 kB 471.0 MB/s eta 0:00:00 +Downloading idna-3.18-py3-none-any.whl (65 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 65.5/65.5 kB 392.2 MB/s eta 0:00:00 +Downloading markupsafe-3.0.3-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl (22 kB) +Downloading mpmath-1.3.0-py3-none-any.whl (536 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 536.2/536.2 kB 478.8 MB/s eta 0:00:00 +Downloading pillow-12.3.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl (6.9 MB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 6.9/6.9 MB 369.2 MB/s eta 0:00:00 +Downloading urllib3-2.7.0-py3-none-any.whl (131 kB) + ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 131.1/131.1 kB 438.3 MB/s eta 0:00:00 +Building wheels for collected packages: antlr4-python3-runtime + Building wheel for antlr4-python3-runtime (setup.py): started + Building wheel for antlr4-python3-runtime (setup.py): finished with status 'done' + Created wheel for antlr4-python3-runtime: filename=antlr4_python3_runtime-4.9.3-py3-none-any.whl size=144590 sha256=5aa799b6efbe9e55190ab92e61a9ecd9b458a6b893b508bf688fa3754a709a14 + Stored in directory: /tmp/pip-ephem-wheel-cache-pem62r8t/wheels/1a/97/32/461f837398029ad76911109f07047fde1d7b661a147c7c56d1 +Successfully built antlr4-python3-runtime +Installing collected packages: nvidia-cusparselt-cu12, mpmath, antlr4-python3-runtime, wadler-lindig, urllib3, typing-extensions, triton, tqdm, threadpoolctl, sympy, safetensors, pyyaml, pillow, nvidia-nvtx-cu12, nvidia-nvjitlink-cu12, nvidia-nccl-cu12, nvidia-curand-cu12, nvidia-cufile-cu12, nvidia-cuda-runtime-cu12, nvidia-cuda-nvrtc-cu12, nvidia-cuda-cupti-cu12, nvidia-cublas-cu12, numpy, networkx, MarkupSafe, joblib, importlib-resources, idna, hf-xet, fsspec, filelock, einops, charset_normalizer, certifi, scipy, requests, omegaconf, nvidia-cusparse-cu12, nvidia-cufft-cu12, nvidia-cudnn-cu12, nibabel, jinja2, jaxtyping, scikit-learn, nvidia-cusolver-cu12, huggingface-hub, torch, torchvision, timm +Successfully installed MarkupSafe-3.0.3 antlr4-python3-runtime-4.9.3 certifi-2026.7.22 charset_normalizer-3.4.9 einops-0.8.2 filelock-3.32.2 fsspec-2026.7.0 hf-xet-1.6.0 huggingface-hub-0.36.2 idna-3.18 importlib-resources-7.1.0 jaxtyping-0.3.10 jinja2-3.1.6 joblib-1.5.3 mpmath-1.3.0 networkx-3.6.1 nibabel-5.4.2 numpy-2.4.6 nvidia-cublas-cu12-12.8.4.1 nvidia-cuda-cupti-cu12-12.8.90 nvidia-cuda-nvrtc-cu12-12.8.93 nvidia-cuda-runtime-cu12-12.8.90 nvidia-cudnn-cu12-9.10.2.21 nvidia-cufft-cu12-11.3.3.83 nvidia-cufile-cu12-1.13.1.3 nvidia-curand-cu12-10.3.9.90 nvidia-cusolver-cu12-11.7.3.90 nvidia-cusparse-cu12-12.5.8.93 nvidia-cusparselt-cu12-0.7.1 nvidia-nccl-cu12-2.27.3 nvidia-nvjitlink-cu12-12.8.93 nvidia-nvtx-cu12-12.8.90 omegaconf-2.3.0 pillow-12.3.0 pyyaml-6.0.3 requests-2.34.2 safetensors-0.8.0 scikit-learn-1.8.0 scipy-1.17.1 sympy-1.14.0 threadpoolctl-3.6.0 timm-1.0.27 torch-2.8.0 torchvision-0.23.0 tqdm-4.70.0 triton-3.4.0 typing-extensions-4.16.0 urllib3-2.7.0 wadler-lindig-0.1.7 +WARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv + +[notice] A new release of pip is available: 24.0 -> 26.2.1 +[notice] To update, run: pip install --upgrade pip +INFO: Adding runscript +INFO: /etc/singularity/ exists; cleanup by system administrator is not complete (see https://apptainer.org/docs/admin/latest/singularity_migration.html) +INFO: Creating SIF file... +INFO: Build complete: /data/connor/nanobrain.1/experiments/fomo_tune_baseline/output/task3/task3.sif +built experiments/fomo_tune_baseline/output/task3/task3.sif +=== sifs === +-rwxr-xr-x 1 connor hpcroot 5.0G Aug 11 16:48 experiments/fomo_tune_baseline/output/task1/task1.sif +-rwxr-xr-x 1 connor hpcroot 5.0G Aug 11 16:58 experiments/fomo_tune_baseline/output/task3/task3.sif +-rwxr-xr-x 1 connor hpcroot 5.0G Aug 11 16:53 experiments/fomo_tune_baseline/output/task5/task5.sif diff --git a/fomo_tune_baseline/slurms/slurm-378508.out b/fomo_tune_baseline/slurms/slurm-378508.out new file mode 100644 index 0000000000000000000000000000000000000000..83dad8e8c95b64c9a11f211a0edbbe72f10914f9 --- /dev/null +++ b/fomo_tune_baseline/slurms/slurm-378508.out @@ -0,0 +1,100 @@ +=== task1 === +16:31:38 run task1 (git 1df2e5d) +16:31:38 config: +task: task1 +ckpt_path: /data/mihir-stuff/smri-pretrained/pretrain_full_90_10_h100/checkpoint-last.pth +modalities: +- dwi_b1000 +output_root: experiments/fomo_tune_baseline/output +name: task1 +device: cuda +seed: 4466 +16:31:40 dataset: 21 subjects, 13 positive +16:31:48 fold 1/21 sub-01 y=1 p=0.512 (4s) +16:31:49 fold 2/21 sub-02 y=1 p=0.514 (4s) +16:31:49 fold 3/21 sub-03 y=0 p=0.486 (5s) +16:31:49 fold 4/21 sub-04 y=1 p=0.499 (5s) +16:31:50 fold 5/21 sub-05 y=1 p=0.507 (5s) +16:31:50 fold 6/21 sub-06 y=1 p=0.493 (6s) +16:31:50 fold 7/21 sub-07 y=0 p=0.489 (6s) +16:31:51 fold 8/21 sub-08 y=0 p=0.490 (6s) +16:31:51 fold 9/21 sub-09 y=1 p=0.523 (7s) +16:31:52 fold 10/21 sub-10 y=1 p=0.525 (7s) +16:31:52 fold 11/21 sub-11 y=0 p=0.479 (7s) +16:31:52 fold 12/21 sub-12 y=0 p=0.474 (8s) +16:31:53 fold 13/21 sub-13 y=1 p=0.516 (8s) +16:31:53 fold 14/21 sub-14 y=1 p=0.536 (9s) +16:31:53 fold 15/21 sub-15 y=0 p=0.486 (9s) +16:31:54 fold 16/21 sub-16 y=0 p=0.482 (9s) +16:31:54 fold 17/21 sub-17 y=1 p=0.520 (10s) +16:31:54 fold 18/21 sub-18 y=1 p=0.518 (10s) +16:31:55 fold 19/21 sub-19 y=1 p=0.520 (10s) +16:31:55 fold 20/21 sub-20 y=1 p=0.519 (11s) +16:31:56 fold 21/21 sub-21 y=0 p=0.498 (11s) +16:31:57 result: auroc=0.9904 auroc_ci_low=0.9444 auroc_ci_high=1.0000 (11s) +=== task5 === +16:32:19 run task5 (git 1df2e5d) +16:32:19 config: +task: task5 +ckpt_path: /data/mihir-stuff/smri-pretrained/pretrain_full_90_10_h100/checkpoint-last.pth +output_root: experiments/fomo_tune_baseline/output +name: task5 +device: cuda +seed: 4466 +16:32:42 dataset: 48 subjects, 24 positive +16:33:21 fold 1/20 n=3 y=[0 1 1] p=[0.013 0.998 0.974] (35s) +16:33:24 fold 2/20 n=3 y=[1 1 1] p=[0.796 0.719 0.999] (37s) +16:33:26 fold 3/20 n=3 y=[0 1 1] p=[0.155 0.998 0.905] (39s) +16:33:28 fold 4/20 n=3 y=[0 0 1] p=[0.024 0.032 0.811] (42s) +16:33:29 fold 5/20 n=3 y=[1 1 1] p=[0.988 0.652 0.552] (43s) +16:33:31 fold 6/20 n=3 y=[0 0 1] p=[0.46 0.744 0.994] (44s) +16:33:33 fold 7/20 n=3 y=[0 0 1] p=[0.134 0.148 0.996] (47s) +16:33:36 fold 8/20 n=3 y=[0 0 1] p=[0.037 0.012 0.915] (50s) +16:33:37 fold 9/20 n=2 y=[0 1] p=[0.137 0.997] (51s) +16:33:38 fold 10/20 n=2 y=[0 1] p=[0.223 0.993] (51s) +16:33:39 fold 11/20 n=2 y=[0 0] p=[0.118 0.32 ] (53s) +16:33:42 fold 12/20 n=2 y=[0 1] p=[0.122 0.849] (55s) +16:33:43 fold 13/20 n=2 y=[0 1] p=[0.207 0.282] (57s) +16:33:45 fold 14/20 n=2 y=[0 1] p=[0.177 0.319] (58s) +16:33:47 fold 15/20 n=2 y=[0 1] p=[0.112 0.943] (60s) +16:33:48 fold 16/20 n=2 y=[0 1] p=[0.034 0.998] (61s) +16:33:50 fold 17/20 n=2 y=[0 0] p=[0.018 0.021] (63s) +16:33:51 fold 18/20 n=2 y=[0 1] p=[0.02 0.937] (64s) +16:33:52 fold 19/20 n=2 y=[0 1] p=[0.004 0.96 ] (65s) +16:33:54 fold 20/20 n=2 y=[0 1] p=[0.075 0.929] (68s) +16:33:56 result: auroc=0.9844 auroc_ci_low=0.9528 auroc_ci_high=1.0000 (68s) +=== task3 === +16:34:12 run task3 (git 1df2e5d) +16:34:12 config: +task: task3 +ckpt_path: /data/mihir-stuff/smri-pretrained/pretrain_full_90_10_h100/checkpoint-last.pth +output_root: experiments/fomo_tune_baseline/output +name: task3 +device: cuda +seed: 4466 +16:35:16 dataset: 494 subjects, age 19-80 mean 45.2 +16:38:39 fold 1/20 n=25 mae=3.91 (190s) +16:38:51 fold 2/20 n=25 mae=4.50 (202s) +16:38:57 fold 3/20 n=25 mae=3.74 (208s) +16:39:03 fold 4/20 n=25 mae=3.58 (214s) +16:39:09 fold 5/20 n=25 mae=3.57 (220s) +16:39:14 fold 6/20 n=25 mae=3.12 (226s) +16:39:20 fold 7/20 n=25 mae=3.42 (232s) +16:39:26 fold 8/20 n=25 mae=3.12 (238s) +16:39:32 fold 9/20 n=25 mae=3.74 (244s) +16:39:38 fold 10/20 n=25 mae=4.20 (249s) +16:39:44 fold 11/20 n=25 mae=3.22 (255s) +16:39:50 fold 12/20 n=25 mae=3.50 (261s) +16:39:55 fold 13/20 n=25 mae=3.39 (267s) +16:40:01 fold 14/20 n=25 mae=4.17 (273s) +16:40:07 fold 15/20 n=24 mae=4.36 (278s) +16:40:12 fold 16/20 n=24 mae=4.15 (284s) +16:40:18 fold 17/20 n=24 mae=3.66 (289s) +16:40:23 fold 18/20 n=24 mae=3.54 (295s) +16:40:29 fold 19/20 n=24 mae=3.42 (300s) +16:40:34 fold 20/20 n=24 mae=3.53 (306s) +16:40:35 result: pearson_r=0.9631 pearson_r_ci_low=0.9567 pearson_r_ci_high=0.9687 mae=3.6910 mae_ci_low=3.4464 mae_ci_high=3.9466 (306s) +=== results === +{"name": "task1", "auroc": 0.9903846153846154, "auroc_ci_low": 0.9444444444444445, "auroc_ci_high": 1.0, "run_time": 11.2} +{"name": "task3", "pearson_r": 0.9630771532969529, "pearson_r_ci_low": 0.9567144985968425, "pearson_r_ci_high": 0.9687044710309238, "mae": 3.6910024558001684, "mae_ci_low": 3.446356968454986, "mae_ci_high": 3.946628273643462, "run_time": 306.1} +{"name": "task5", "auroc": 0.984375, "auroc_ci_low": 0.9527867965367964, "auroc_ci_high": 1.0, "run_time": 67.5}