galilai-group/stable-pretraining

stable‑pretraining – 자기지도학습을 위한 유연한 프레임워크

무엇인가요

• PyTorch Lightning 위에 구축된 오픈소스 파이썬 라이브러리이며, 실험적 JAX/Flax‑NNX 백엔드를 포함합니다. 연구자들이 자기지도학습(SSL) 방법을 최소한의 부가 작업으로 프로토타이핑, 훈련, 평가할 수 있게 해줍니다.
• 핵심 아이디어는 loss (훈련용)와 로깅하고 싶은 임의의 텐서를 포함하는 사전을 반환하는 forward 함수만 구현하는 것입니다. 라이브러리는 일반적인 Lightning의 처리(데이터 로드, 체크포인트, 재시작, 로깅, 로컬 실험 레지스트리)를 모두 처리합니다.

주요 구성 요소

구성 요소	역할
데이터	사전 형식의 데이터셋을 처리합니다. HuggingFace 데이터셋, 표준 PyTorch 데이터셋으로부터의 변환, 사용자 정의 사전 반환 데이터셋을 지원합니다. Kornia를 통해 GPU에 무거운 다중 뷰 증강을 이동시켜 더 높은 처리량을 제공하는 gpu_transform 훅을 제공합니다.
모듈	백본, 프로젝터, 사용자 제공 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의 1CPU 에포크 실행

이 명령어는 작고 작은 훈련 실행을 시작하고, 손실 곡선과 온라인 선형 프로브의 정확도를 표시하는 로컬 웹 뷰어(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
• 지원용 디스코드: https://discord.gg/adzpqWKM25

누가 사용해야 하나요

• 새로운 SSL 목적을 구축하고 신뢰할 수 있는 훈련 기반을 원하는 연구자.
• 대규모 사전 훈련을 위한 재현 가능하고 체크포인트 친화적인 파이프라인을 필요로 하는 엔지니어.
• 동일한 방법의 Torch와 JAX 구현을 비교하고 싶은 누구나.

---

위의 모든 세부 정보는 리포지토리의 README에서 직접 가져왔으며, 추가 기능은 추측되지 않았습니다.