DINO-v2笔记

DINO-v2笔记

DINO-v2一种无监督学习的预训练方法,可以生成具有强大泛化能力的视觉特征,适用于各种图像分布和任务,而无需进行微调。这篇论文重点介绍了数据和模型规模方面的技术贡献,包括自动构建一个多样化和精心筛选的图像数据集、在多个层级上进行训练、使用Sinkhorn-Knopp居中方法和KoLeo正则化等。实验结果表明,该方法在多个图像理解任务上的表现超过了目前公开的最佳无监督和半监督方法。

作者实际上花了大量的篇幅减少了数据如何创建,如何进行预训练和如何优化训练过程。

项目主页,项目开源在GitHub

数据集准备

作者通过从一个大型未筛选数据池中检索与几个精选数据集中的图像接近的图像来组装他们的LVD-142M数据集。作者描述了数据处理流程的主要组成部分,包括精选/未筛选数据源,图像去重步骤和检索系统。整个流程不需要任何元数据或文本,直接使用图像。

整个处理分布在一个由20个节点组成的计算集群上,该集群配备了8个V100-32GB GPU,生成LVD-142M数据集不到两天。

数据来源

作者在包含 ImageNet-22k,ImageNet-1k、Google Landmarks 和几个细粒度数据集的的数据集进行选择。对于不安全的数据源,爬取公开可用的网络数据存储库中收集了原始未过滤的图像数据集。从存储库中的每个网页中,作者从标签中提取图像的 URL 链接。作者在构建数据集过程中丢弃了不受域限制或限制的 URL,并对下载的图像(PCA 哈希重复数据删除、NSFW 过滤和模糊可识别人脸)进行后处理。这导致 1.2B 个独特的图像。

消除重复数据

作者将使用了A Self-Supervised Descriptor for Image Copy Detection 这篇论文中的方法来处理未经处理的数据,并去除接近重复的图像。这减少了冗余并增加了图像之间的多样性。此外还删除了这个工作中使用的任何基准测试或验证集中包含的几乎重复的图像。

自监督的图像检索

首先使用在ImageNet-22k上预训练的自监督ViT-H/16网络来计算图像嵌入,并使用余弦相似性作为图像之间的距离度量。接下来对未分级的数据进行k-means聚类。给定要检索的查询数据集,如果它足够大,那么就为每个查询图像检索N个(通常是4个)最近的邻居。如果它很小,就从与每个查询图像相对应的聚类中采样M个图像。可以通过目视检查检索结果来调整N和M。

判别式自监督的预培训

图像级目标

同一图像的不同裁剪中获得不同的部分,使用ViT进行编码,用过去迭代的指数移动平均值构建教师模型,从学生和教师网络中提取的特征之间的交叉熵损失学习学生模型的参数

patch级目标

随即屏蔽给学生的一些输入补丁,但不屏蔽给老师的。然后,我们在每个屏蔽补丁上的两个网络的补丁特征之间添加交叉熵损失。这种损失与图像级别的损失相结合。

解绑两个目标的权重联系

将上面两个目标相关的权重捆绑在一起会使模型在patch上欠拟合,而在图像级别上过拟合。解开这些权重可以解决这个问题,并提高两个目标的性能。

Sinkhorn-Knopp centering

这是一种替代DINO和iBot模型中的teacher softmax-centering步骤的方法,即使用SwAV模型的Sinkhorn-Knopp(SK)批量归一化。作者在这个方法中运行了3次Sinkhorn-Knopp算法步骤,并对学生应用softmax归一化。这个方法的目的是提高自监督学习模型的性能。

KoLeo regularizer

KoLeo正则化器源自Kozachenko-Leonenko差分熵估计器,它鼓励批处理中特征的均匀跨度。给定一组n个向量(x1, . . . , xn),它被定义为\(\mathcal{L}_\text{koleo}=-\frac1n\sum_{i=1}^n\log(d_{n,i})\) ,其中\(d_{n,i}=\min_{j\neq i}\left\|x_i-x_j\right\|\)\(x_i\)和批处理中任何其他点之间的最小距离。在计算这个正则化器之前,我们还要对特征进行L2-归一化。

Adapting the resolution

在像素级别的下游任务中,如分割或检测,提高图像分辨率是非常重要的,因为低分辨率下小物体容易消失。但是高分辨率的训练需要更多的时间和内存,所以作者提出了一种方法,在预训练的最后一段时间内将图像的分辨率提高到518×518。

有效的实施

作者对于训练大规模模型的几个改进措施,包括使用A100 GPU和PyTorch 2.0进行训练,提供代码和预训练模型,并在附录的Table 17中详细描述了模型的细节。

另外,与iBOT实现相比,DINOv2的代码在相同硬件条件下,运行速度提高了2倍,内存使用量减少了三分之一。

快速高效的注意力

作者自己实现了一个fastattention,需要注意的是作者的ViT-g架构略有不同,采用1536的嵌入维度和24个头(每个头64维),而不是1408的嵌入维度和16个头(每个头88维),以最大化计算效率

自注意中的嵌套张量

作者使用了一种新的技术,可以在同一个正向传递中运行全局裁剪和局部裁剪(具有不同数量的补丁标记),与之前的实现相比,可以获得显着的计算效率提升。此外,作者提到他们使用的基础组件已经在xFormers库中提供。

有效的随机深度

作者使用了一种改进的随机深度(stochastic depth)方法,相比于传统的掩码方法,该方法跳过了被丢弃的残差计算,从而在一定程度上节省了内存和计算资源。在本次实验中,使用高丢弃率(d=40%)时,这种方法使计算效率和内存使用效率得到了显著提高。具体实现方法是通过在批处理维度上随机重新排列B个样本,并在块计算中仅对前\((1-d)×B\)个样本进行计算。

完全共享数据并行(FSDP)

通过将模型副本分配到多个GPU中,可以将模型大小限制在GPU节点总内存的范围内。此外,FSDP的实现方式可以将权重片段存储为float32,但在传播权重和梯度时使用float16,从而降低跨GPU通信成本。相较于DistributedDataParallel(DDP)中使用的float32梯度all-reduce操作,使用Pytorch-FSDP混合精度训练的通信成本减少了约50%,在扩展GPU节点数量时训练过程更加高效。总的来说,Pytorch-FSDP混合精度训练在几乎所有情况下都优于使用autocast的DDP。

模型蒸馏

作者发现即使对于一个规模较大的ViT-L模型,他们的预训练方法也能够取得比从头开始训练更好的性能。此外,他们还提出了一种知识蒸馏方法,与A simple recipe for competitive low-compute self supervised vision models. arXiv preprint arXiv:2301.09451 所描述的方法相似,但没有修改蒸馏的损失项,并评估了学生模型的指数移动平均值。


DINO-v2笔记
https://studyinglover.com/2023/07/27/DINO-v2笔记/
作者
StudyingLover
发布于
2023年7月27日
许可协议