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

作者:Xiaohang Zhan
文章首发知乎: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的脚本或者详细的指导,简单易行。训练和测试现有算法,都只需要一行命令搞定。

高度灵活性和可扩展性

改进现有算法。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.)})`

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

设计自己的算法。高度模块化的设计使得自己设计算法变得非常简单。你只需要设计自己的模块或者复用现有模块,然后新建一个模型文件将这些模块组织起来即可。如果你的算法包含间隔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,为这个领域贡献自己的力量。

推荐阅读:

  • 自监督场景去遮挡
  • CVPR2019|条件运动传播:从运动中学习物体性质
  • 人脸聚类那些事儿:利用无标签数据提升人脸识别性能

欢迎关注我 ,我是 Xiaohang~

发表评论

邮箱地址不会被公开。 必填项已用*标注