一行命令跑评测,港中文MMLab开源自监督表征学习代码库OpenSelfSup

自监督表征学习发展迅速,但也存在诸多问题。近日,香港中文大学多媒体实验室(MMLab)和南洋理工大学的研究者开源了一套统一的自监督学习代码库 OpenSelfSup。

前言

近几个月来自监督表征学习领域获得了显著突破,特别是随着 Rotation Prediction、DeepCluster、MoCo、SimCLR 等简单有效的方法的诞生,自监督表征学习大有超越有监督表征学习的趋势。

然而,做这个领域的研究者都深有感触:1)自监督任务复杂而多样,不同方法各有各的专用训练代码,难以结合、复用和改进;2)评价方案不统一,不同的方法难以在公平环境下对比;3)动辄百万千万的训练数据量,训练效率是个大问题。

针对这些问题,香港中文大学多媒体实验室(MMLab)和南洋理工大学的研究者最近开源了一套统一的自监督学习代码库:OpenSelfSup。

开源地址:https://github.com/open-mmlab/OpenSelfSup

OpenSelfSup

统一的代码框架和模块化设计

OpenSelfSup 使用 PyTorch 实现,支持基于分类、重建、聚类、memory bank、contrastive learning 的多种自监督学习框架,目前收录了 Relative Location、Rotation Prediction、DeepCluster、OnlineDeepCluster、NPID、MoCo、SimCLR 等一系列表现较好的自监督表征学习方法,后续还将陆续跟进学术界最新算法。

在这个框架下,每个算法被拆解为 backbone、neck、head、memory_bank (optional)、hook (optional) 等多个可独立设计的模块,每个模块均提供多个可选方案,开发者也可以自行设计各个模块。

一行命令跑评测,港中文MMLab开源自监督表征学习代码库OpenSelfSup

SimCLR 中的 backbone、neck 和 head 设计

标准化的评测方案

OpenSelfSup 目前支持 ImageNet/Place205 Linear Classification、ImageNet Semi-Supervised Classification、PASCAL VOC07 Linear SVM、PASCAL VOC / COCO Object Detection 等多个基准的评测方案,可以做到一行命令跑评测,非常方便。

高效率的分布式训练

OpenSelfSup 收录的算法全部实现了多机多卡的分布式训练、提特征和测试。

上手容易

该代码库的环境配置、数据配置均有 from scratch 的脚本或详细指导,简单易行。训练和测试现有算法,只需要一行命令即可搞定。

高度灵活性和可扩展性

1. 改进现有算法。OpenSelfSup 用 config 文件来定义各种参数和模块,方便修改。config 中还支持一些较复杂的调整,比如 data augmentation 的组合、learning rate schedule、独立调整某些网络参数的优化参数等。例如,你希望单独调整 head 中 fully-connected layer 的 momentum 和 learning rate,或者 backbone 中某几层的 weight decay 等等,那么你可以在 config 中 optimizer: paramwise_option 下用正则表达式筛选出对应网络参数,然后指定这些值,而不需要改动代码。

一行命令跑评测,港中文MMLab开源自监督表征学习代码库OpenSelfSup

DeepCluster中指定head(fc layer) 的momentum为0

一行命令跑评测,港中文MMLab开源自监督表征学习代码库OpenSelfSup

2. 使用自己的训练数据。OpenSelfSup 将数据处理和数据源解耦。使用自己的训练数据,只需要新增一个自己的 data source 文件,并在 config 中指定即可。

3. 设计自己的算法。高度模块化的设计使得自行设计算法变得非常简单。你只需要设计自己的模块或者复用现有模块,然后新建一个模型文件将这些模块组织起来即可。如果你的算法包含间隔 iteration、间隔 epoch 的操作(例如 DeepCluster 需要在每个 epoch 结束后对训练集做聚类),也只需要新建一个 hook 来定义这些操作,并在 config 中设置这个 hook 即可。hook 的调用是自动进行的。

一行命令跑评测,港中文MMLab开源自监督表征学习代码库OpenSelfSup

相关推荐