How leaplabthu/chexworld works
1. What this repo is
- CheXWorld is a PyTorch codebase for self-supervised radiograph pretraining plus downstream fine-tuning on classification and segmentation tasks.
- The architecture is centered on a ViT/JEPA-style backbone with task-specific heads and masking strategies.
2. Top-level workflow split
- There are three main execution paths:
- train_jepa.py for pretraining
- train_finetune.py for classification fine-tuning
- train_finetune_seg.py for segmentation fine-tuning
- opts.py acts as the shared configuration/argument layer across these paths.
3. Data subsystem
- data_utils/ contains dataset adapters for multiple chest X-ray sources: CheXpert, MIMIC, NIH, RSNA, SIIM, JSRT, VinDr, plus generic chest dataset plumbing.
- data_utils/data_path.py centralizes filesystem paths, so datasets are kept as an external boundary rather than embedded in the model code.
- The repo expects datasets and splits to be prepared separately, as described in DATA.md.
4. Preprocessing and augmentation boundary
- data_utils/xray_transform.py, xray_transform_add.py, world_model_transform.py, and randaugment.py define image augmentation and normalization pipelines.
- These transforms sit between dataset loading and model input, and are likely different for pretraining vs fine-tuning.
- This layer is where image-specific assumptions live: grayscale radiographs, resizing/cropping, and domain variation handling.
5. Masking subsystem for self-supervision
- The masks/ package is a distinct architectural block for patch masking / block masking during JEPA pretraining.
- Variants such as multiblock, random, data2vec, multi_multiblock_*, and blockwise indicate several masking policies, including 2D and 3D support.
- This subsystem generates the prediction targets/contexts that define the self-supervised learning objective.
6. Core model family
- models/ contains the main representation-learning stack:
- jepa.py, jepa_vit.py, jepa_vit_3d.py, and jepa_add.py for the self-supervised world-model variants
- finetune.py for downstream classifiers
- jepa_vit_add.py and models_vit.py for architectural variants/backbones
- The backbone is ViT-based, with extensions for 3D inputs, additional prediction heads, and medical-image-specific adaptations.
7. Downstream heads and adapters
- models/attentive_pooler.py, roi.py, unet_adapter.py, unet_adapter_conv.py, upernet.py, and unigrad.py provide task heads and feature adapters.
- These components bridge the pretrained representation to:
- image-level prediction for classification
- dense prediction for segmentation
- models/sscale.py and models/utils.py support feature scaling and shared model utilities.
8. Training-time data flow
- The runtime pipeline is:
- dataset loader in data_utils/
- augmentation / image transform
- optional mask generation from masks/
- model forward pass in models/
- optimization / logging / metrics in util/
- This separation makes pretraining and fine-tuning share the same data and model foundations while differing in objective and heads.
9. Optimization and training utilities
- util/ provides training infrastructure: logging, metrics, learning-rate schedules, LARS optimizer support, mixed precision / misc helpers, and mixup/crop utilities.
- util/lr_sched.py, util/lr_decay.py, and util/lars.py are the key optimization-support modules.
- These are cross-cutting concerns used by both pretraining and fine-tuning loops.
10. Repository boundary and external dependencies
- The repo is primarily a research training framework, not a serving or inference system.
- The main external dependency is PyTorch, with medical datasets and pretrained checkpoints handled outside the codebase.
- README.md, PRETRAIN.md, FINETUNE.md, and DATA.md form the operational boundary: code here defines the models and pipelines, while the docs define how to prepare data and run experiments.
11. How to diagram it
- The clearest architecture diagram would show:
- Datasets feeding into dataset-specific loaders
- Transforms + masks producing model-ready batches
- A shared ViT/JEPA backbone
- Two downstream branches: self-supervised pretraining and task fine-tuning/segmentation
- A shared utility layer for optimization, logging, and metrics
- The major architectural distinction is between representation learning and task-specific adaptation.
Main entry points: train_jepa.py JEPA train · chest_dset.py Chest loader · train_finetune_seg.py Seg train · train_finetune.py Cls train
Components 19
Execution
Training entrypoints and shared configuration
- JEPA train · Self-supervised pretraining script
train_jepa.py - Cls train · Classification fine-tuning script
train_finetune.py - Seg train · Segmentation fine-tuning script
train_finetune_seg.py - Options · Shared argument and experiment configuration
opts.py
Data
Dataset adapters and image preprocessing
- Chest loader · Common chest X-ray dataset plumbing
data_utils/chest_dset.py - Paths · Filesystem path and dataset location boundary
data_utils/data_path.py - Xray tfm · Core radiograph preprocessing and augmentation
data_utils/xray_transform.py - Add tfm · Additional transform variants
data_utils/xray_transform_add.py - World tfm · Transform pipeline for self-supervised world-model training
data_utils/world_model_transform.py - RandAugment · Image augmentation policy support
data_utils/randaugment.py
Self-supervision
JEPA pretraining model and masking
Fine-tuning
Downstream classification and segmentation heads
- Fine-tune · Downstream classification model
models/finetune.py - Pooler · Attention pooling for image-level prediction
models/attentive_pooler.py - UNet adapter · Dense prediction adapter for segmentation
models/unet_adapter.py - UPerNet · Segmentation decoder / head
models/upernet.py
Infrastructure
Optimization, logging, and training utilities
- Training utils · Logging, metrics, schedulers, optimizer helpers
util
Connections 22
- JEPA train → Options: config · Uses shared experiment settings
- Cls train → Options: config · Uses shared experiment settings
- Seg train → Options: config · Uses shared experiment settings
- JEPA train → Chest loader: data · Builds training dataset pipeline
- Cls train → Chest loader: data · Builds classification dataset pipeline
- Seg train → Chest loader: data · Builds segmentation dataset pipeline
- Chest loader → Paths: paths · Resolves dataset locations
- Chest loader → Xray tfm: preprocess · Applies radiograph transforms
- Chest loader → Add tfm: preprocess · Applies extra transform variants
- JEPA train → World tfm: augment · Uses pretraining-specific transforms
- JEPA train → Masking: masking · Generates self-supervised targets and context masks
- Masking → JEPA: objective · Feeds masked prediction objective
- JEPA train → ViT JEPA: model · Instantiates ViT-based pretraining backbone
- JEPA train → 3D JEPA: model · Optional 3D backbone path
- Cls train → Fine-tune: model · Instantiates classification model
- Cls train → Pooler: head · Uses pooled image-level representation
- Seg train → UNet adapter: head · Uses dense prediction adapter
- Seg train → UPerNet: head · Uses segmentation decoder
- ViT JEPA → Training utils: train infra · Uses optimization and logging helpers
- Fine-tune → Training utils: train infra · Uses optimization and logging helpers
- Seg train → Training utils: train infra · Uses optimization and logging helpers
- JEPA train → Training utils: train infra · Uses optimization and logging helpers