OpenSelfSup: Open-MMLab自监督表征学习代码库

本文转载自知乎，已获作者授权转载。

https://zhuanlan.zhihu.com/p/148782886

前言

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

然而，相信做这个领域的研究者都深有感触：

1）自监督任务复杂而多样，不同方法各有各的专用训练代码，难以结合、复用和改进；

2）评价方案不统一，不同的方法难以在公平环境下对比；

3）动辄百万千万的训练数据量，训练效率是个大问题。

针对这些问题，我们（香港中文大学多媒体实验室和南洋理工大学）开源了一套统一的自监督学习代码库：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) 等几个可独立设计的模块，每个模块提供了多个可选方案，也可以自己来设计各个模块。例如下面是OpenSelfSup中SimCLR的backbone, neck和head设计。

model = dict(type='SimCLR',pretrained=None,backbone=dict(type='ResNet',depth=50,in_channels=3,out_indices=[4],  # 0: conv-1, x: stage-xnorm_cfg=dict(type='SyncBN')),neck=dict(type='NonLinearNeckV1',in_channels=2048,hid_channels=2048,out_channels=128,with_avg_pool=True),head=dict(type='ContrastiveHead', temperature=0.1))

标准化的评测方案

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下用正则表达式筛选出对应网络参数然后指定这些值，而不需要改动代码。如下是DeepCluster中指定head (fc layer)的momentum为0的设置。

optimizer = dict(type='SGD', lr=0.1, momentum=0.9, weight_decay=0.00001,    paramwise_options={'\Ahead.': dict(momentum=0.)})

2）使用自己的训练数据。OpenSelfSup将数据处理和数据源解耦，使用自己的训练数据，只需要新增一个自己的data source文件，并在config中指定即可。
3）设计自己的算法。高度模块化的设计使得自己设计算法变得非常简单。你只需要设计自己的模块或者复用现有模块，然后新建一个模型文件将这些模块组织起来即可。如果你的算法包含间隔iteration、间隔epoch的操作，例如DeepCluster需要在每个epoch结束后对训练集做聚类，也只需要新建一个hook来定义这些操作，并在config中设置这个hook就行了，hook的调用是自动进行的。如下是DeepClusterHook中初始聚类和每n个epoch后做聚类的操作。

def before_run(self, runner):if self.initial:self.deepcluster(runner)def after_train_epoch(self, runner):if not self.every_n_epochs(runner, self.interval):return        self.deepcluster(runner)