contrastors - 对比学习工具包

项目简介

contrastors 是一款专为研究人员和工程师打造的对比学习工具包，旨在帮助他们高效地训练和评估对比模型。这个工具包集成了多种先进技术和灵活功能，支持多GPU训练、受限内存环境下的大批量训练等，为用户提供了广泛的模型训练支持。

主要特性

快速高效的训练：基于 Flash Attention 构建，支持快速高效的模型训练。
多GPU支持：能够在多GPU环境下运行，提升训练速度。
大批量训练：通过 GradCache 技术，实现在内存受限的环境下使用大批量数据进行训练。
Huggingface 支持：可轻松加载常见模型，如 Pythia/GPTNeoX、BERT 等。
预训练与表示学习：支持 Masked Language Modeling (MLM) 预训练，并可灵活调整 Matryoshka Representation Learning 嵌入尺寸。
对比学习模型：提供 CLIP 和 LiT 风格的对比学习。
ViT 模型支持：可加载流行的 ViT 模型，如 timm。

研究成果

contrastors 不仅是一个工具包，还伴随着研究成果的支持，包括：

Zach Nussbaum 等人在论文 Nomic Embed: Training a Reproducible Long Context Text Embedder 中所做的工作。
以及他们在 Nomic Embed Vision: Expanding the Latent Space 中的研究贡献。

入门指南与环境要求

要设置 contrastors 需要依赖 Flash Attention 所提供的自定义内核。在安装环境前，请确保已安装 Cuda 11.8 或更高版本。

使用以下步骤来设置 Python 虚拟环境并安装必需的软件包：

python3 -m venv env
source env/bin/activate
pip3 install torch torchvision torchaudio
pip install wheel packaging ninja setuptools
pip install --no-cache-dir flash-attn --no-build-isolation ...
pip install -e .

数据访问

提供 nomic-embed-text-v1 数据集的访问权限。用户需创建账户并通过 nomic 包进行登录配置，具体操作指引可通过 nomic 的官方渠道获取。

训练与预训练

contrastors 提供了一系列训练脚本，用户可根据需要自行调整：

MLM 预训练：通过命令运行 BERT 的预训练

deepspeed --num_gpus=8 train.py --config=configs/train/mlm.yaml ...

对比学习预训练和微调：可以进行对比学习实验及微调已有模型

torchrun --nproc-per-node=8 train.py --config=configs/train/contrastive_pretrain.yaml ...

预训练模型

我们在以下平台提供了多个预训练模型，供用户下载使用：

加入 Nomic 社区

欢迎访问 Nomic 的官方网站 nomic.ai，并加入我们的 Discord 和 Twitter 社区，与更多研究者和开发者进行交流。

许可协议

此工具包基于 Apache 2.0 许可证开放，具体的模型许可信息请参考各个模型的 cards。

致谢

感谢 Tri Dao 对 Flash Attention 和自定义内核的贡献，以及 OpenCLIP 和 Huggingface 团队在相关开源库方面的支持和工作。