Apppublic

joannadzialo/msc-pub-vae-smoke-space-v2

sourceHugging Faceupdated 3mo agoView on Hugging Face

0likes

App README

VAE Training Pipeline

Train convolutional VAEs on medical image data from either:

—a local DICOM manifest
—a Hugging Face imagefolder dataset

The repo also covers:

—spatial and baseline VAE variants
—Hub checkpoint upload / resume
—latent export
—downstream age / sex prediction from latents
—Hugging Face Jobs runs through a Docker Space image

What Is In The Repo

—src/train/train_vae.py: main VAE training entrypoint
—src/train/export_latents.py: export latent vectors from a checkpoint
—src/train/train_age_predictor.py: downstream regression / classification on latent CSVs
—src/tools/run_latent_classification_pipeline.py: export latents and train downstream model in one shot
—src/tools/precompute_hf_foreground_bboxes.py: precompute patch foreground boxes for HF datasets
—src/models/vae.py: baseline + spatial VAE implementations
—src/data/: manifest loaders, HF dataset loaders, datamodules, preprocessing

Environment

Local Python used in this repo:

—3.13.x

Install dependencies:

bash

pip install -r requirements.txt

Core libraries:

—lightning
—torch
—torchvision
—mlflow
—datasets
—huggingface_hub

Supported VAE Configurations

Architectures:

—baseline
—spatial

Encoders:

—resnet18
—resnet34
—efficientnet_b0
—convnext_tiny

Current encoder support by architecture:

—baseline: resnet18, resnet34, efficientnet_b0, convnext_tiny
—spatial: resnet18, resnet34, convnext_tiny

Useful training features already implemented:

—--resume-from repo::path.ckpt
—--resume-weights-only
—--optimizer adam|adamw
—--decoder-lr for decoder-side parameter groups
—--lr-warmup-values
—--beta-warmup-epochs
—--decoder-skip-mode none|shallow|deep
—--hf-foreground-bbox-cache
—Hub checkpoint upload through --hub-model-repo

Local Training

Train from a DICOM manifest:

bash

python -m src.train.train_vae ^
  --manifest data/dicom_metadata_unique_sop.csv ^
  --path-col filename_path ^
  --vae-architecture baseline ^
  --encoder resnet18 ^
  --recon-loss l1 ^
  --image-size 256 ^
  --resize-mode stretch ^
  --output-dir runs/lightning_vae

Train from a Hugging Face dataset:

bash

python -m src.train.train_vae ^
  --data-source hf_dataset ^
  --hf-dataset your-org/your-imagefolder-dataset ^
  --hf-split train ^
  --hf-image-column image ^
  --vae-architecture spatial ^
  --encoder convnext_tiny ^
  --pretrained-encoder ^
  --latent-dim 64 ^
  --patch-size 256 ^
  --resize-mode none ^
  --recon-loss l1 ^
  --beta 0.01 ^
  --beta-warmup-epochs 30 ^
  --decoder-skip-mode none ^
  --output-dir runs/vae_hf

Foreground Bounding Box Cache

If you train on native patches with --hf-foreground-aware-patches, precomputing bounding boxes once is much faster than recomputing them in every run.

Precompute on a local imagefolder dataset:

bash

python -m src.tools.precompute_hf_foreground_bboxes ^
  --hf-dataset data/hf_imagefolder_AGE_PRED_export_full ^
  --hf-split train ^
  --output-csv runs/hf_foreground_bboxes/foreground_bboxes.csv

Use that cache in training:

bash

python -m src.train.train_vae ^
  --data-source hf_dataset ^
  --hf-dataset your-org/your-imagefolder-dataset ^
  --hf-split train ^
  --hf-image-column image ^
  --hf-foreground-aware-patches ^
  --no-hf-precompute-foreground-bboxes ^
  --hf-foreground-bbox-cache runs/hf_foreground_bboxes/foreground_bboxes.csv

You can also point the cache to a dataset file on the Hub:

bash

--hf-foreground-bbox-cache your-org/your-imagefolder-dataset::foreground_bboxes.csv

Checkpoint Resume

Resume full training state:

bash

python -m src.train.train_vae ^
  --resume-from your-org/vae-checkpoints::stage1/last.ckpt

Resume weights only and start a fresh optimizer:

bash

python -m src.train.train_vae ^
  --resume-from your-org/vae-checkpoints::stage1/last.ckpt ^
  --resume-weights-only ^
  --lr 1e-5

MLflow

Both train_vae.py and train_age_predictor.py read:

—--tracking-uri
—or MLFLOW_TRACKING_URI

Example:

bash

python -m src.train.train_vae ^
  --tracking-uri http://127.0.0.1:5000 ^
  --output-dir runs/lightning_vae

Hugging Face Jobs

This repo currently launches Jobs through a Docker Space image, not through a raw Docker Hub image name.

Use:

bash

hf.co/spaces/joannadzialo/msc-pub-vae-space-control

not:

bash

joannadzialo/msc-pub-vae-space-control

Example smoke run:

bash

hf jobs run -d ^
  --secrets HF_TOKEN ^
  --flavor t4-small ^
  -e MLFLOW_TRACKING_URI=https://your-mlflow-host ^
  hf.co/spaces/joannadzialo/msc-pub-vae-space-control ^
  python /home/user/app/run_repo_command.py --repo-id joannadzialo/msc-pub-vae-smoke-space-v2 --repo-type space --local-dir /tmp/remote-repo ^
    python -m src.train.train_vae ^
      --tracking-uri https://your-mlflow-host ^
      --data-source hf_dataset ^
      --hf-dataset your-org/your-dataset ^
      --hf-split train ^
      --hf-image-column image ^
      --hf-foreground-aware-patches ^
      --no-hf-precompute-foreground-bboxes ^
      --hf-foreground-bbox-cache your-org/your-dataset::foreground_bboxes.csv ^
      --vae-architecture spatial ^
      --encoder convnext_tiny ^
      --pretrained-encoder ^
      --latent-dim 64 ^
      --patch-size 256 ^
      --resize-mode none ^
      --beta 0.01 ^
      --beta-warmup-epochs 30 ^
      --recon-loss l1 ^
      --optimizer adam ^
      --lr 1e-4 ^
      --lr-warmup-values 1e-5,3e-5,5e-5,7e-5,1e-4 ^
      --decoder-skip-mode none ^
      --epochs 1 ^
      --early-stopping-patience 0 ^
      --batch-size 8 ^
      --workers 0 ^
      --val-split 0.1 ^
      --max-samples 64 ^
      --precision 16-mixed ^
      --accelerator gpu ^
      --devices 1 ^
      --output-dir /tmp/vae_smoke

Checkpoint upload from Jobs:

bash

--hub-model-repo your-org/vae-checkpoints --hub-checkpoint-prefix run_name

Export Latents

Export from a checkpoint:

bash

python -m src.train.export_latents ^
  --checkpoint your-org/vae-checkpoints::some_run/last.ckpt ^
  --data-source hf_dataset ^
  --hf-dataset your-org/your-dataset ^
  --hf-split train ^
  --hf-image-column image ^
  --hf-foreground-aware-patches ^
  --no-hf-precompute-foreground-bboxes ^
  --hf-foreground-bbox-cache your-org/your-dataset::foreground_bboxes.csv ^
  --output-csv runs/lightning_vae/latents.csv

Downstream Prediction From Latents

Regression:

bash

python -m src.train.train_age_predictor ^
  --latent-csv runs/lightning_vae/latents.csv ^
  --target-col age ^
  --task regression

Classification:

bash

python -m src.train.train_age_predictor ^
  --latent-csv runs/lightning_vae/latents.csv ^
  --target-col sex ^
  --task classification ^
  --class-balance both

One-shot pipeline:

bash

python -m src.tools.run_latent_classification_pipeline ^
  --checkpoint your-org/vae-checkpoints::some_run/last.ckpt ^
  --data-source hf_dataset ^
  --hf-dataset your-org/your-dataset ^
  --hf-split train ^
  --hf-image-column image ^
  --hf-foreground-aware-patches ^
  --no-hf-precompute-foreground-bboxes ^
  --hf-foreground-bbox-cache your-org/your-dataset::foreground_bboxes.csv ^
  --target-col age ^
  --task regression ^
  --tracking-uri http://127.0.0.1:5000

Notes

—The VAE loss now uses mean reduction, not sum reduction.
—Spatial ConvNeXt runs support decoder_skip_mode=none|shallow|deep.
—If you want a new LR after checkpoint resume, use --resume-weights-only.
—For HF image datasets, the bbox cache keying supports both current path-based ids and older legacy ids.