DOCSEARCHbyAlgolia
Submit Docs

stable-pretraining

repository·main·Indexed 43 days ago

https://github.com/galilai-group/stable-pretraining

A PyTorch Lightning-based framework for foundation-model research, designed for scalable self-supervised learning (SSL), supervised, and multi-modal pretraining. It features a modular dict-based architecture, GPU-side augmentation via kornia, and SLURM-grade orchestration. The library supports multiple SSL methods including SimCLR, BYOL, VICReg, and Barlow Twins, and includes an experimental JAX/Flax-NNX backend.

Tokens
41.2K
Snippets
101
Records
164
Agent score
55%

What's inside stable-pretraining

  1. Overview of stable_pretraining.data components

    main

    The stable_pretraining.data module provides tools for dataset handling, transforms, sampling, and data loading specifically designed for self-supervised learning workflows. It is organized into several functional categories:

    • Core Components: Fundamental building blocks like DataModule, Collator, and Dataset.
    • Real Data Wrappers: Tools to interface with existing data formats, including FromTorchDataset, HFDataset (Hugging Face), and Subset.
    • Synthetic Data Generators: Classes and functions for creating artificial datasets (e.g., GMM, MinariStepsDataset, swiss_roll, perlin_noise).
    • Noise Models: Models for injecting noise into data, such as Categorical or ExponentialMixtureNoiseModel.
    • Samplers: Specialized sampling strategies like RepeatedRandomSampler or SupervisedBatchSampler.
    • Utility Functions: Helpers for data manipulation like fold_views, random_split, and download functions.

  3. Use optimizers and learning rate schedulers from stable_pretraining.optim

    main

    The stable_pretraining.optim module provides specialized optimization tools designed for self-supervised learning workflows. It includes high-performance optimizers like LARS and a variety of learning rate schedulers to manage training dynamics.

    Available Optimizers

    • LARS: Layer-wise Adaptive Rate Scaling, commonly used for large-batch training in self-supervised learning.

    Available Learning Rate Schedulers

    • CosineDecayer: Decays the learning rate following a cosine curve.
    • LinearWarmup: Starts with a low learning rate and increases it linearly to a target value.
    • LinearWarmupCosineAnnealing: Combines a linear warmup phase with a cosine decay phase.
    • LinearWarmupCyclicAnnealing: Combines a linear warmup with a cyclic annealing schedule.
    • LinearWarmupThreeStepsAnnealing: A scheduler that follows a three-step annealing process after a linear warmup.

  4. Understand the ImageNet-10 benchmark results

    main

    The benchmarks/imagenet10/RESULTS.md file contains a performance sweep of SSL methods using a ViT-S/16 backbone trained for 200 epochs.

    Top performing methods (Linear Probe):

    1. VISReg: 90.2%
    2. SwAV: 89.7%
    3. LeJEPA: 87.1%
    4. DINO: 86.1%

    Common failure modes noted in benchmarks:

    • MIM family (e.g., SimMIM, MaskFeat): Reconstruction loss converges, but linear features require much longer schedules (800+ epochs) to form.
    • Whitening/Decorrelation (e.g., W-MSE, TiCO): Small batch sizes (128-256) may not provide stable enough covariance estimates for these methods.
    • Architecture Mismatch (e.g., PIRL): Methods designed for CNNs may fail on ViTs due to how transforms (like jigsaw) interact with positional embeddings.

  5. Compare MultiViewTransform vs RoundRobin data loading approaches

    main

    The stable-pretraining framework provides benchmarks comparing two primary data loading approaches for multi-view Self-Supervised Learning (SSL): MultiViewTransform and RoundRobin.

    Based on the benchmark results, MultiViewTransform generally provides higher throughput (images per second) compared to RoundRobin across different batch sizes (e.g., ~321-324 img/s vs ~295-300 img/s).

  6. Understand GPU augmentation performance bottlenecks

    main

    When using GPU augmentation, the following operations are the primary contributors to latency:

    • GPURandomResizedCrop: High cost due to grid_sample requirements.
    • GPUGaussianBlur: Significant cost (approx. 17ms) regardless of kernel size, as kornia uses an FFT-like path.

    Optimization Levers:

    • Downgrade/Drop GaussianBlur: Reducing the probability p of GaussianBlur or removing it entirely can save significant time.
    • Cheaper Crops: Using uniform random crops instead of scale-resampling crops can reduce the cost of RandomResizedCrop.
    • Model Scale: As models get larger (e.g., moving from ViT-S to ViT-L), the relative cost of augmentation decreases because the forward/backward pass becomes the dominant bottleneck.

  8. Understand probe convergence and hyperparameter tuning

    main

    Probe Convergence

    Because the backbone features are continuously shifting during pretraining, parametric probes must constantly adapt to the changing representations. While this can slow down probe learning, probe training typically converges once the backbone representations themselves converge.

    Hyperparameter Tuning

    To evaluate different probe hyperparameters (such as hidden dimension or number of layers), create a separate probe instance for each configuration. Since probes are lightweight, you can optimize multiple probe configurations jointly alongside the backbone model.

  9. Monitor representations with Callbacks

    main

    Callbacks in stable-pretraining provide real-time insights into representation quality.

    Key Monitoring Callbacks:

    • OnlineProbe: Trains a lightweight linear probe on frozen representations to track downstream accuracy.
    • OnlineKNN: Non-parametric k-nearest neighbors evaluator using a rolling queue of cached embeddings.
    • RankMe: Tracks effective rank of feature representations to detect dimensional collapse.
    • LatentViz: Online 2D visualization of the latent space.

    Training & Checkpointing:

    • TeacherStudentCallback: Performs EMA teacher updates.
    • WeightDecayUpdater: Updates weight decay on a per-batch schedule.
    • HuggingFaceCheckpointCallback: Exports HuggingFace-compatible checkpoints.
    • EnvironmentDumpCallback: Dumps system info (CUDA, packages, etc.) for reproducibility.
    # Monitor SSL representations with a linear classifier
linear_probe = spt.callbacks.OnlineProbe(
    module,
    name="linear_probe",
    input="embedding",
    target="label",
    probe=torch.nn.Linear(512, 10),
    loss_fn=torch.nn.CrossEntropyLoss(),
    metrics={
        "top1": torchmetrics.classification.MulticlassAccuracy(10),
        "top5": torchmetrics.classification.MulticlassAccuracy(10, top_k=5),
    },
)

# Track representation quality with KNN evaluation
knn_probe = spt.callbacks.OnlineKNN(
    name="knn_probe",
    input="embedding",
    target="label",
    queue_length=20000,
    k=10,
)

  11. Logging and metrics in the JAX backend

    main

    Logging in the JAX backend is identical to the PyTorch path. The JAX Trainer drives the same logger classes (RegistryLogger, WandB, Trackio, SwanLab) via the logger= argument.

    • Metric Keys: Use the same keys as the torch path (e.g., fit/loss, train/<p>_loss, eval/<p>_acc, rankme, lidar).
    • Module Logging: Module.log(name, value) works the same way.
    • Registry Integration: spj.Manager automatically attaches the RegistryLogger, so JAX runs write the same sidecar.json and are visible via spt registry ls/best/... just like torch runs.

  12. How `spt.Module` handles different training modes

    main

    The spt.Module is designed to be versatile across different learning paradigms by using the stage parameter and the presence of a loss key in the returned dictionary:

    • Self-supervised learning: Implement contrastive losses within the forward method using the stage or self.training flag.
    • Supervised learning: Implement standard classification or regression logic.
    • Multi-task learning: Return multiple losses or outputs within the batch dictionary.
    • Evaluation-only/Inference: Omit the "loss" key from the returned dictionary. You can also set optim=False during module creation if no optimization is required.