# leaplabthu/chexworld architecture

> Architecture diagram of the GitHub repository leaplabthu/chexworld, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-06-18). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/leaplabthu/chexworld
- Repository: https://github.com/leaplabthu/chexworld

## Overview

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`](https://github.com/leaplabthu/chexworld/blob/main/train_jepa.py) (JEPA train), [`data_utils/chest_dset.py`](https://github.com/leaplabthu/chexworld/blob/main/data_utils/chest_dset.py) (Chest loader), [`train_finetune_seg.py`](https://github.com/leaplabthu/chexworld/blob/main/train_finetune_seg.py) (Seg train), [`train_finetune.py`](https://github.com/leaplabthu/chexworld/blob/main/train_finetune.py) (Cls train)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_execution["Execution"]
  node_train_jepa["JEPA train<br/>entrypoint<br/>[train_jepa.py]"]
  node_train_finetune["Cls train<br/>entrypoint<br/>[train_finetune.py]"]
  node_train_finetune_seg["Seg train<br/>entrypoint"]
  node_opts["Options<br/>config<br/>[opts.py]"]
end

subgraph group_data["Data"]
  node_chest_dset[("Chest loader<br/>dataset base<br/>[chest_dset.py]")]
  node_data_path["Paths<br/>path config<br/>[data_path.py]"]
  node_xray_transform["Xray tfm<br/>augment pipeline<br/>[xray_transform.py]"]
  node_xray_transform_add["Add tfm<br/>augment pipeline"]
  node_world_model_transform["World tfm<br/>augment pipeline"]
  node_randaugment["RandAugment<br/>augmentation<br/>[randaugment.py]"]
end

subgraph group_selfsup["Self-supervision"]
  node_masks{{"Masking<br/>self-supervision"}}
  node_jepa["JEPA<br/>selfsup model<br/>[jepa.py]"]
  node_jepa_vit["ViT JEPA<br/>backbone<br/>[jepa_vit.py]"]
  node_jepa_vit_3d["3D JEPA<br/>backbone<br/>[jepa_vit_3d.py]"]
end

subgraph group_finetune["Fine-tuning"]
  node_finetune_model["Fine-tune<br/>task model<br/>[finetune.py]"]
  node_attentive_pooler(("Pooler<br/>head"))
  node_unet_adapter["UNet adapter<br/>segmentation head<br/>[unet_adapter.py]"]
  node_upernet["UPerNet<br/>segmentation head<br/>[upernet.py]"]
end

subgraph group_infra["Infrastructure"]
  node_util["Training utils<br/>infra"]
end

node_train_jepa -->|"config"| node_opts
node_train_finetune -->|"config"| node_opts
node_train_finetune_seg -->|"config"| node_opts
node_train_jepa -->|"data"| node_chest_dset
node_train_finetune -->|"data"| node_chest_dset
node_train_finetune_seg -->|"data"| node_chest_dset
node_chest_dset -->|"paths"| node_data_path
node_chest_dset -->|"preprocess"| node_xray_transform
node_chest_dset -->|"preprocess"| node_xray_transform_add
node_train_jepa -->|"augment"| node_world_model_transform
node_train_jepa -->|"masking"| node_masks
node_masks -->|"objective"| node_jepa
node_train_jepa -->|"model"| node_jepa_vit
node_train_jepa -->|"model"| node_jepa_vit_3d
node_train_finetune -->|"model"| node_finetune_model
node_train_finetune -->|"head"| node_attentive_pooler
node_train_finetune_seg -->|"head"| node_unet_adapter
node_train_finetune_seg -->|"head"| node_upernet
node_jepa_vit -->|"train infra"| node_util
node_finetune_model -->|"train infra"| node_util
node_train_finetune_seg -->|"train infra"| node_util
node_train_jepa -->|"train infra"| node_util

click node_train_jepa "https://github.com/leaplabthu/chexworld/blob/main/train_jepa.py"
click node_train_finetune "https://github.com/leaplabthu/chexworld/blob/main/train_finetune.py"
click node_train_finetune_seg "https://github.com/leaplabthu/chexworld/blob/main/train_finetune_seg.py"
click node_opts "https://github.com/leaplabthu/chexworld/blob/main/opts.py"
click node_chest_dset "https://github.com/leaplabthu/chexworld/blob/main/data_utils/chest_dset.py"
click node_data_path "https://github.com/leaplabthu/chexworld/blob/main/data_utils/data_path.py"
click node_xray_transform "https://github.com/leaplabthu/chexworld/blob/main/data_utils/xray_transform.py"
click node_xray_transform_add "https://github.com/leaplabthu/chexworld/blob/main/data_utils/xray_transform_add.py"
click node_world_model_transform "https://github.com/leaplabthu/chexworld/blob/main/data_utils/world_model_transform.py"
click node_randaugment "https://github.com/leaplabthu/chexworld/blob/main/data_utils/randaugment.py"
click node_masks "https://github.com/leaplabthu/chexworld/tree/main/masks"
click node_jepa "https://github.com/leaplabthu/chexworld/blob/main/models/jepa.py"
click node_jepa_vit "https://github.com/leaplabthu/chexworld/blob/main/models/jepa_vit.py"
click node_jepa_vit_3d "https://github.com/leaplabthu/chexworld/blob/main/models/jepa_vit_3d.py"
click node_finetune_model "https://github.com/leaplabthu/chexworld/blob/main/models/finetune.py"
click node_attentive_pooler "https://github.com/leaplabthu/chexworld/blob/main/models/attentive_pooler.py"
click node_unet_adapter "https://github.com/leaplabthu/chexworld/blob/main/models/unet_adapter.py"
click node_upernet "https://github.com/leaplabthu/chexworld/blob/main/models/upernet.py"
click node_util "https://github.com/leaplabthu/chexworld/tree/main/util"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_train_jepa,node_train_finetune,node_train_finetune_seg,node_opts toneBlue
class node_chest_dset,node_data_path,node_xray_transform,node_xray_transform_add,node_world_model_transform,node_randaugment toneAmber
class node_masks,node_jepa,node_jepa_vit,node_jepa_vit_3d toneMint
class node_finetune_model,node_attentive_pooler,node_unet_adapter,node_upernet toneRose
class node_util toneIndigo
```

## Components

### Execution

Training entrypoints and shared configuration

- [JEPA train](https://github.com/leaplabthu/chexworld/blob/main/train_jepa.py): entrypoint · Self-supervised pretraining script · `train_jepa.py`
- [Cls train](https://github.com/leaplabthu/chexworld/blob/main/train_finetune.py): entrypoint · Classification fine-tuning script · `train_finetune.py`
- [Seg train](https://github.com/leaplabthu/chexworld/blob/main/train_finetune_seg.py): entrypoint · Segmentation fine-tuning script · `train_finetune_seg.py`
- [Options](https://github.com/leaplabthu/chexworld/blob/main/opts.py): config · Shared argument and experiment configuration · `opts.py`

### Data

Dataset adapters and image preprocessing

- [Chest loader](https://github.com/leaplabthu/chexworld/blob/main/data_utils/chest_dset.py): dataset base · Common chest X-ray dataset plumbing · `data_utils/chest_dset.py`
- [Paths](https://github.com/leaplabthu/chexworld/blob/main/data_utils/data_path.py): path config · Filesystem path and dataset location boundary · `data_utils/data_path.py`
- [Xray tfm](https://github.com/leaplabthu/chexworld/blob/main/data_utils/xray_transform.py): augment pipeline · Core radiograph preprocessing and augmentation · `data_utils/xray_transform.py`
- [Add tfm](https://github.com/leaplabthu/chexworld/blob/main/data_utils/xray_transform_add.py): augment pipeline · Additional transform variants · `data_utils/xray_transform_add.py`
- [World tfm](https://github.com/leaplabthu/chexworld/blob/main/data_utils/world_model_transform.py): augment pipeline · Transform pipeline for self-supervised world-model training · `data_utils/world_model_transform.py`
- [RandAugment](https://github.com/leaplabthu/chexworld/blob/main/data_utils/randaugment.py): augmentation · Image augmentation policy support · `data_utils/randaugment.py`

### Self-supervision

JEPA pretraining model and masking

- [Masking](https://github.com/leaplabthu/chexworld/tree/main/masks): self-supervision · Patch and block masking policies for JEPA · `masks`
- [JEPA](https://github.com/leaplabthu/chexworld/blob/main/models/jepa.py): selfsup model · Core self-supervised training wrapper · `models/jepa.py`
- [ViT JEPA](https://github.com/leaplabthu/chexworld/blob/main/models/jepa_vit.py): backbone · Vision Transformer JEPA backbone · `models/jepa_vit.py`
- [3D JEPA](https://github.com/leaplabthu/chexworld/blob/main/models/jepa_vit_3d.py): backbone · 3D Vision Transformer variant · `models/jepa_vit_3d.py`

### Fine-tuning

Downstream classification and segmentation heads

- [Fine-tune](https://github.com/leaplabthu/chexworld/blob/main/models/finetune.py): task model · Downstream classification model · `models/finetune.py`
- [Pooler](https://github.com/leaplabthu/chexworld/blob/main/models/attentive_pooler.py): head · Attention pooling for image-level prediction · `models/attentive_pooler.py`
- [UNet adapter](https://github.com/leaplabthu/chexworld/blob/main/models/unet_adapter.py): segmentation head · Dense prediction adapter for segmentation · `models/unet_adapter.py`
- [UPerNet](https://github.com/leaplabthu/chexworld/blob/main/models/upernet.py): segmentation head · Segmentation decoder / head · `models/upernet.py`

### Infrastructure

Optimization, logging, and training utilities

- [Training utils](https://github.com/leaplabthu/chexworld/tree/main/util): infra · Logging, metrics, schedulers, optimizer helpers · `util`

## Connections

- 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