消息

[2024.10.18]发布Inf-CLIP培训和评估代码。

要求和安装

基本依赖关系：

Python >= 3.8
火炬 >= 2.0.0
CUDA 版本 >= 11.8

[远程]安装Inf-CL：

 # 远程安装pip install inf_cl -i https://pypi.org/simple

[本地]安装Inf-CL：

 pip install -e 。

安装所需的包：

 git 克隆 https://github.com/DAMO-NLP-SG/Inf-CLIPcd Inf-CLIP
pip install -r 要求.txt

特征

inf_cl是 Inf-CL 损失的 triton 实现：

环-CL (inf_cl/ring.py#L238)
Inf-CL (inf_cl/ring.py#L251)

inf_clip是具有 Inf-CL 损失和其他训练功能的 CLIP 训练代码库：

梯度累积 (inf_clip/train/train.py#L180)
梯度缓存 (inf_clip/train/train.py#L292)

用法

一个关于如何采用 Inf-CL 损失进行对比学习的简单示例。使用这样的命令来尝试：

torchrun --nproc_per_node 2 tests/example.py

 import torchimport torch.nn.function as Fimport torch.distributed as distimport numpy as npfrom inf_cl import cal_inf_lossdef create_cl_tensors(rank, world_size):# 参数 dtype = torch.float32num_heads = 3 # 注意力头数量 seq_length_q = 32768 # 序列长度 seq_length_k = 32768d_model = 256 #每个头的维度（必须是16、32、64或128）#随机初始化inputsq = torch.rand((seq_length_q // world_size, num_heads * d_model), dtype=dtype, device=f"cuda:{rank}" )k = torch.rand((seq_length_k // world_size, num_heads * d_model), dtype=dtype, device=f"cuda:{rank}")l = torch.ones([], dtype=dtype, device=f" cuda:{rank}") * np.log(1 / 0.07)q = F.normalize(q, p=2, dim=-1).requires_grad_() # Queryk = F.normalize(k, p=2, dim=-1).requires_grad_() # Keyl = l.requires_grad_() # Logit scalereturn q, k, lif __name__ == "__main__":# 假设分布式环境已初始化ddist.init_process_group("nccl")rank = dist.get_rank()world_size = dist.get_world_size()torch.cuda.set_device(rank)#以图文对比学习为例，q是全局图像特征，# k是文本特征，l是logit尺度。 q, k, l = create_cl_tensors(rank, world_size)# 标签默认是对角线元素。 # labels = torch.arange(q.shape[0])loss = cal_inf_loss(q, k, scale=l.exp())print(loss)

主要结果

内存成本

* 表示采用“数据卸载”策略。

支持的最大批量大小

速度

批量大小缩放

更大数据规模的训练需要更大的批量大小。

培训与评估

快速入门

为了促进在我们的代码库之上的进一步开发，我们提供了如何使用 Inf-CLIP 训练定制 CLIP 并在主流剪辑基准上评估训练后的模型的快速入门指南。

训练数据结构：

信息片段
├── 数据集
│ ├── cc3m/ # https://github.com/rom1504/img2dataset/blob/main/dataset_examples/cc3m.md|   |   ├── 0000.tar|   |   ├── 0001.tar|   |   ├── ...|   |   └── 0301.tar
│ ├── cc12m/ # https://github.com/rom1504/img2dataset/blob/main/dataset_examples/cc12m.md|   |   ├── 0000.tar|   |   ├── 0001.tar|   |   ├── ...|   |   └── 1044.tar
│ ├── laion400m/ # https://github.com/rom1504/img2dataset/blob/main/dataset_examples/laion400m.md|   |   ├── 00000.tar|   |   ├── 00001.tar|   |   ├── ...|   |   └── 41407.tar

命令：

 bash 脚本/cc3m/lit_vit-b-32_bs16k.sh
bash 脚本/cc12m/lit_vit-b-32_bs32k.sh
bash脚本/laion400m/lit_vit-b-32_bs256k.sh

评估数据结构：

信息片段
├── 数据集
│ ├── imagenet-1k/ # 下载imagenet的val_images.tar.gz|   |   └── 值/|   |   |   ├── n01440764|   |   |   ├── n01443537|   |   |   ├── ...|   |   |   └── n15075141
│ ├── Clip-benchmark/ # bash 数据集/benchmarks_download.sh|   |   ├── wds_mscoco_captions|   |   ├── wds_flickr8k|   |   ├── wds_flickr30k|   |   ├── wds_imagenet1k|   |   ├── wds_imagenetv2|   |   ├── wds_imagenet_sketch|   |   ├── wds_imagenet-a|   |   ├── wds_imagenet-r|   |   ├── wds_imagenet-o|   |   └── wds_objectnet

命令：

 # imagenet评估bash脚本/imagenet_eval.sh#总体评估bash脚本/benchmarks_eval.sh

引文

如果您发现 Inf-CLIP 对您的研究和应用有用，请使用此 BibTeX 进行引用：

 @article{damovl2024infcl, title={打破内存障碍：对比损失的近乎无限批量大小缩放}，作者={Zesen Cheng, Hang Zhang, Kehan Li, Sicong Leng,zhiqiang Hu, Fei Wu, Deli Zhao, Xin Li,李东冰}，期刊={arXiv预印本arXiv:2410.17243}，年份={2024}，url={https://arxiv.org/abs/2410.12787}}