[论文解读] Solo-learn: A Library of Self-supervised Methods for Visual Representation Learning

Victor G. Turrisi da Costa, Enrico Fini+3 more|arXiv (Cornell University)|Aug 3, 2021

Domain Adaptation and Few-Shot Learning被引用 61

一句话总结

solo-learn 是一个基于 PyTorch 的开源库，实现了大量最前沿的自监督学习（SSL）方法用于视觉表征学习，支持分布式训练、快速数据加载、在线线性评估，以及易于扩展。

ABSTRACT

This paper presents solo-learn, a library of self-supervised methods for visual representation learning. Implemented in Python, using Pytorch and Pytorch lightning, the library fits both research and industry needs by featuring distributed training pipelines with mixed-precision, faster data loading via Nvidia DALI, online linear evaluation for better prototyping, and many additional training tricks. Our goal is to provide an easy-to-use library comprising a large amount of Self-supervised Learning (SSL) methods, that can be easily extended and fine-tuned by the community. solo-learn opens up avenues for exploiting large-budget SSL solutions on inexpensive smaller infrastructures and seeks to democratize SSL by making it accessible to all. The source code is available at https://github.com/vturrisi/solo-learn.

研究动机与目标

• 提供一个可重复使用、模块化的库，实现广泛的最先进 SSL 方法用于视觉表征学习。
• 通过提供标准化实现、快速数据加载和在线线性评估来促进实验和可重复性。
• 通过使预算较小的研究人员能够有效地训练和比较方法来实现对 SSL 的民主化。
• 整合下游任务和预训练模型使用的工具，以加速原型设计和部署。

提出的方法

• 在 PyTorch 中实现 13 种 SSL 方法：Barlow Twins、BYOL、DeepCluster V2、DINO、MoCo V2+、NNCLR、ReSSL、SimCLR、Supervised Contrastive Learning、SimSiam、SwAV、VICReg 和 W-MSE。
• 设计了一个模块化管道，单独的组件包括：solo.args 用于参数处理、solo.methods 用于 SSL 方法、solo.losses 用于损失函数、solo.utils 用于数据处理和增强管道，以及用于分布式、混合精度训练的 PyTorch Lightning 训练器。
• 利用 Nvidia DALI 提高数据加载速度，并支持混合精度和在线线性评估，以实现快速原型开发。
• 提供预训练模型并可与 DetectronV2 轻松集成用于下游目标检测任务。
• 在 CIFAR-10、CIFAR-100 和 ImageNet-100 上对 SSL 方法进行了基准测试，并对若干方法进行超参数调整以获得具有竞争力的结果。
• 与 VISSL 和 Lightly 进行比较，突出对更现代的方法的支持、低资源训练，以及诸如自动 UMAP 可视化等额外工具。

实验结果

研究问题

• RQ1 一个统一、可扩展的库是否能够在广泛的 SSL 方法之间实现公平、可重复的比较？
• RQ2 使用现代数据加载器（Nvidia DALI）以及混合精度/分布式训练在适度硬件上进行 SSL 工作负载的实际好处（速度、内存、可用性）是什么？
• RQ3 在标准基准（CIFAR-10/100、ImageNet-100）上在线线性评估下，当代 SSL 方法的表现如何，并且它们与先前的库相比如何？
• RQ4 solo-learn 的预训练 SSL 模型是否能够用于下游任务，如通过 DetectronV2 进行对象检测？

主要发现

• 在 PyTorch 中为 13 种 SSL 方法提供标准化实现，便于更易重复和比较。
• 该库提供分布式训练、混合精度支持、使用 Nvidia DALI 的快速数据加载，以及在线线性评估以加速原型设计。
• 基准结果显示多种方法在 CIFAR-10、CIFAR-100 和 ImageNet-100 上的在线线性评估准确率具有竞争力，且对超参数进行了调整以提高相对于原始报告的性能。
• 相比相关库（VISSL、Lightly），solo-learn 支持更多现代的 SSL 方法，面向较小 GPU 预算的研究人员，同时提供额外工具如自动线性评估和 UMAP 可视化。
• 实验还表明在 ImageNet-100 上使用 DALI 进行数据加载时，训练时间更快、内存占用更低。