contrastors简介
contrastors是由Nomic AI团队开发的对比学习工具包,旨在帮助研究人员和工程师高效地训练和评估对比学习模型。该项目基于PyTorch构建,集成了多项先进技术,如Flash Attention、GradCache等,可以显著提升大规模对比学习模型的训练效率。
主要特性
contrastors具有以下核心特性:
- 基于Flash Attention实现快速高效的训练
- 支持多GPU训练
- 集成GradCache,支持在内存受限环境下使用大batch size
- 支持通过Hugging Face轻松加载常见预训练模型
- 支持掩码语言建模(MLM)预训练
- 实现Matryoshka表示学习,支持灵活的嵌入维度
- 支持CLIP和LiT风格的对比学习
- 可加载流行的ViT模型(如timm)
快速上手
- 克隆项目仓库:
git clone https://github.com/nomic-ai/contrastors.git
cd contrastors
- 安装依赖:
pip install -r requirements.txt
-
按照README中的说明配置环境,包括安装CUDA、PyTorch、Flash Attention等。
-
运行示例训练脚本:
cd src/contrastors
deepspeed --num_gpus=8 train.py --config=configs/train/mlm.yaml --deepspeed_config=configs/deepspeed/ds_config.json --dtype=bf16
相关资源
相关论文
- Nomic Embed: Training a Reproducible Long Context Text Embedder
- Nomic Embed Vision: Expanding the Latent Space
contrastors为对比学习研究提供了一个强大的工具,欢迎感兴趣的读者访问项目主页深入了解和使用。如有问题,可以在GitHub仓库中提issue或加入Nomic AI的社区讨论。
