How mdivyanshu97/discovr works
1. Overall shape
- This repo is a research training stack for DISCOVR, a self-supervised echocardiography video representation learner.
- The main architecture is a pretraining pipeline plus a downstream fine-tuning/evaluation pipeline, both centered on video clips and checkpointed PyTorch models.
2. Top-level execution entry points
- scripts/run_mae_pretraining.py is the main pretraining launcher.
- scripts/run_class_finetuning.py is the downstream supervised training entry point.
- run_discovr_encoder.py is a lightweight inference wrapper for loading a pretrained checkpoint and extracting embeddings.
- config/run_pretraining.sh appears to be a convenience shell launcher for training jobs.
3. Core model layer
- models/modeling_pretrain.py and models/modeling_finetune.py are the central model definitions for the two training stages.
- Supporting model components include models/masking_generator.py, models/attn_mask.py, and models/info_nce.py, which suggest the pretraining objective combines masked video modeling, attention/masking logic, and contrastive/distillation-style losses.
- The README describes a dual-branch design: a video branch for temporal dynamics and an online image branch for spatial semantics, with cluster distillation between them.
4. Training orchestration
- engine/engine_for_pretraining.py contains the epoch/step logic for self-supervised learning, likely handling forward passes, loss computation, optimization, and checkpointing.
- engine/engine_for_finetuning.py handles supervised downstream training and validation.
- These engines form the boundary between reusable model code and task-specific training behavior.
5. Data ingestion boundary
- data/datasets/datasets.py and data/datasets/kinetics.py define dataset wrappers and likely CSV/file-based sample indexing.
- The README and scripts indicate training is driven by paths to videos plus CSV splits for train/val/test.
- The data layer is where raw echocardiography videos enter the system before any masking or augmentation is applied.
6. Video preprocessing and augmentation
- data/transforms/ contains the augmentation stack for spatiotemporal inputs: video_transforms.py, volume_transforms.py, rand_augment.py, mixup.py, random_erasing.py, and helper functional.py.
- This is the main “input shaping” layer, converting decoded clips into model-ready tensors and applying regularization augmentations.
- utils/visualize_masks.py and scripts/visualise_dataset.py support debugging/inspection of what the model sees.
7. Optimization, metrics, and utility layer
- utils/optim_factory.py centralizes optimizer construction and parameter grouping.
- utils/metrics_utils.py and utils/utils.py provide common training utilities, metric tracking, logging helpers, and general runtime support.
- utils/utils_datasets.py likely bridges dataset metadata and training code.
8. Dataflow through pretraining
- The pretraining path is: video files + CSV splits → dataset loaders → transforms/masking → pretraining engine → modeling_pretrain → losses such as InfoNCE/distillation/masked prediction → checkpoints.
- The pretrained checkpoint is then reused either for inference (run_discovr_encoder.py) or as initialization for downstream evaluation.
9. Dataflow through fine-tuning
- The fine-tuning path is: labeled video data → dataset/transforms → engine_for_finetuning.py → modeling_finetune.py → task-specific predictions and metrics.
- This separation keeps the self-supervised representation learning code distinct from supervised evaluation code.
10. Runtime and infrastructure considerations
- The project is designed for distributed GPU training via torch.distributed.launch and PyTorch.
- The README makes FFmpeg/TorchCodec part of the runtime because video decoding is a first-class dependency of the pipeline.
- Hugging Face is used for checkpoint distribution, but not as a core runtime service.
11. Standalone TubeViT subtree
- models/TubeViT/ is a self-contained subproject with its own setup.py, scripts, dataset/model code, and requirements.
- It includes training, evaluation, inference, visualization, and ViT weight conversion utilities.
- Architecturally, this looks like a bundled external model implementation or experimental dependency rather than the main DISCOVR training path.
12. What an architecture diagram should emphasize
- Show three boundaries clearly: data decoding/augmentation, training engines, and model/loss definitions.
- Show two primary flows: self-supervised pretraining and supervised fine-tuning/inference.
- Include the optional/secondary TubeViT submodule as a separate model stack, and note that pretrained checkpoints are externalized through Hugging Face.
Main entry points: engine_for_pretraining.py Pretrain engine · transforms.py Video transforms · engine_for_finetuning.py Finetune engine · datasets.py Datasets
Components 19
DISCOVR stack
Primary pretraining, fine-tuning, and inference pipeline for self-supervised echocardiography video representation learning.
- Pretrain run · Main launcher for masked self-supervised pretraining.
scripts/run_mae_pretraining.py - Finetune run · Supervised downstream training entrypoint.
scripts/run_class_finetuning.py - Encoder infer · Loads a pretrained checkpoint and extracts embeddings.
run_discovr_encoder.py - Pretrain engine · Epoch and step orchestration for self-supervised training.
engine/engine_for_pretraining.py - Finetune engine · Supervised training and validation orchestration.
engine/engine_for_finetuning.py - Pretrain model · Masked video representation learner with online image/video branches.
models/modeling_pretrain.py - Finetune model · Task-specific downstream prediction model.
models/modeling_finetune.py - Masking · Mask generation and attention masking logic for pretraining.
models/masking_generator.py - Attn mask · Attention mask construction for masked sequence modeling.
models/attn_mask.py - InfoNCE · Contrastive/distillation-style objective used in training.
models/info_nce.py - Datasets · Dataset wrappers and sample indexing for video splits.
data/datasets/datasets.py - Kinetics IO · Dataset-specific loader and video sample interface.
data/datasets/kinetics.py - Video transforms · Spatiotemporal decoding, augmentation, mixup, and erasing pipeline.
data/transforms/transforms.py - Video aug · Clip-level video transformations.
data/transforms/video_transforms.py - Volume aug · 3D volume and temporal tensor transforms.
data/transforms/volume_transforms.py - Optimizer · Optimizer construction and parameter grouping.
utils/optim_factory.py - Metrics · Metric tracking, logging, and runtime helpers.
utils/metrics_utils.py
TubeViT submodule
Bundled standalone vision transformer subtree with its own scripts, model, and dataset code.
- TubeViT model · Standalone TubeViT implementation and positional encoding code.
models/TubeViT/tubevit/model.py - TubeViT scripts · Separate training, evaluation, inference, and visualization entrypoints.
models/TubeViT/scripts/train.py
Connections 20
- Pretrain run → Pretrain engine: launches · Starts self-supervised training.
- Finetune run → Finetune engine: launches · Starts downstream supervised training.
- Encoder infer → Pretrain model: loads · Uses pretrained checkpoint for embedding extraction.
- Pretrain engine → Pretrain model: trains · Runs masked representation learning forward/backward passes.
- Finetune engine → Finetune model: trains · Optimizes downstream predictions.
- Pretrain engine → Masking: applies · Generates masked inputs for pretraining.
- Pretrain engine → Attn mask: uses · Builds attention masks for masked modeling.
- Pretrain engine → InfoNCE: optimizes · Computes representation and distillation objectives.
- Datasets → Kinetics IO: specializes · Dataset-specific video loading adapter.
- Pretrain run → Datasets: reads · Creates training and validation loaders from split metadata.
- Finetune run → Datasets: reads · Builds labeled downstream data loaders.
- Datasets → Video transforms: feeds · Sends decoded clips into augmentation pipeline.
- Video transforms → Video aug: includes · Clip-level augmentation stage.
- Video transforms → Volume aug: includes · Tensor/volume transformation stage.
- Video transforms → Pretrain engine: prepares · Produces model-ready masked clips.
- Video transforms → Finetune engine: prepares · Produces model-ready labeled clips.
- Pretrain engine → Optimizer: configures · Constructs optimizer and parameter groups.
- Finetune engine → Metrics: reports · Tracks losses and evaluation metrics.
- TubeViT scripts → TubeViT model: uses · Standalone TubeViT training and inference scripts.
- TubeViT model → TubeViT scripts: supports · Shared model code for the submodule tools.