galilai-group/stable-pretraining

stable‑pretraining – 自监督学习的灵活框架

是什么

• 基于 PyTorch Lightning（附带实验性 JAX/Flax‑NNX 后端）构建的开源 Python 库，使研究人员能够以最少的样板代码快速原型设计、训练和评估自监督学习（SSL）方法。
• 核心思想是仅实现一个 forward 函数，该函数返回一个包含 loss（用于训练）和任何你想记录的其他张量的字典。该库会处理常规的 Lightning 逻辑——数据加载、检查点、恢复、日志记录以及本地实验注册表。

主要组件

组件	作用
数据	处理字典格式的数据集。支持 HuggingFace 数据集、标准 PyTorch 数据集的转换，以及自定义返回字典的数据集。提供 gpu_transform 钩子，通过 Kornia 将繁重的多视图增强移动到 GPU 上，以实现更高的吞吐量。
模块	对骨干网络、投影器和用户提供的 forward 可调用对象（例如 forward.simclr）的轻量封装。无需编写 training_step/validation_step；封装器会从返回的字典中提取 loss。
回调	提供即用型工具用于在线评估（线性探测、K‑NN、RankMe、LiDAR、CLIP 零样本、图像检索、潜在可视化）、教师-学生 EMA、权重衰减调度、检查点处理（WandB、HuggingFace、sklearn 模型）以及系统日志（SLURM 信息、环境转储等）。
训练器 / 管理器	对 Lightning 的 Trainer 的轻量封装。Manager 将训练器、模块和数据模块捆绑在一起，运行训练循环，写入日志，并将运行注册到本地 cache_dir 注册表中。
JAX 后端（实验性）	在 stable_pretraining.jax 下镜像 PyTorch API。使用 Flax‑NNX、Optax 和 JAX 原生原语，同时保持相同的配置文件、回调和日志格式。

内置 SSL 方法

• SimCLR、BYOL、VICReg、Barlow Twins、SimSiam、SwAV‑Sinkhorn、NT‑Xent、InfoNCE，以及较新的 LeJEPA 和 Jet 可逆编码器。通过提供自定义 forward 函数即可添加新方法。

骨干网络与优化器

• MLP、ResNet‑{9,18,34,50,101,152}、ConvMixer、ViT‑{tiny,small,base,large}。
• 优化器包括 Adam、LARS 等；支持 CosineAnnealingLR 等调度器。

快速开始

# 安装最新版本
pip install stable-pretraining

# 或安装开发版本
git clone https://github.com/galilai-group/stable-pretraining.git
cd stable-pretraining
pip install -e .
python -m stable_pretraining.quickstart   # 在合成数据上运行 SimCLR 的一个 CPU 时代

该命令启动一个小型训练运行，并打开本地 Web 查看器（spt web ./spt-quickstart/runs），显示损失曲线和在线线性探测的准确率。

典型工作流程

1. 定义数据 – 使用 spt.data.HFDataset 或包装 PyTorch 数据集；可选地添加 gpu_transform 以实现快速多视图增强。
2. 选择方法 – 选择预实现的 forward（如 spt.forward.simclr）或编写自己的。
3. 创建 Module – 传入骨干网络、投影器、损失、优化器配置和 forward 函数。
4. 添加回调 – 例如 OnlineProbe 以监控下游性能。
5. 运行 – spt.Manager(trainer=trainer, module=module, data=dm)()。

为何重要

• 稳定性：所有训练循环都通过 Lightning，因此可获得稳健的检查点、多 GPU（DDP）支持和自动恢复。
• 灵活性：更改 SSL 目标只需交换 forward 可调用对象；其余管道（数据处理、日志记录、评估）保持不变。
• 速度：GPU 侧增强可显著提高现代硬件的吞吐量，已在仓库的基准表中得到验证。
• 跨框架一致性：实验性 JAX 路径允许使用相同配置和日志格式并排比较 Torch 与 JAX 实现。

文档与社区

• 完整文档：https://galilai-group.github.io/stable-pretraining/
• 基准测试：https://github.com/galilai-group/stable-pretraining/tree/main/benchmarks
• 支持用 Discord：https://discord.gg/adzpqWKM25

谁应该使用它

• 构建新 SSL 目标的研究人员，希望拥有可靠的训练框架。
• 需要可复现、检查点友好型管道进行大规模预训练的工程师。
• 任何对比较 Torch 与 JAX 实现同一方法的人。

---

以上所有细节均直接取自仓库的 README；未推断任何额外功能。