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-x
norm_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)

结语

OpenSelfSup是一个长期维护的开源项目，旨在方便学术界和工业届的研究者推动这个领域继续前行。也希望有兴趣的研究者、开发者帮助继续完善OpenSelfSup，为这个领域贡献自己的力量。

END

备注：自监督

自监督/无监督学习交流群

关注最新最前沿的自监督、无监督学习技术，

若已为CV君其他账号好友请直接私信。

我爱计算机视觉

微信号:aicvml

QQ群:805388940

微博知乎:@我爱计算机视觉

投稿:amos@52cv.net

网站:www.52cv.net

在看，让更多人看到