CVPR | 达摩院开源自监督学习框架CoKe, 单机8卡可训练

团队模型、论文、博文、直播合集，点击此处浏览

一、论文

论文链接: Unsupervised Visual Representation Learning by Online Constrained K-Means

代码链接：https://github.com/idstcv/CoKe

二、背景

        虽然基于instance的自监督学习方法在ImageNet上取得了成功，这些方法往往需要在同一个iteration中优化一个样本的double view，并要求大的batch size/memory bank(e.g., SimCLR, MoCo)或者额外的momentum encoder (e.g., BYOL, MoCo) 来辅助训练, 导致比较大的计算开销.而在有监督学习中，往往比较小的batch size和single view就可以进行有效的训练. 因此，我们设计基于聚类的自监督学习框架CoKe，参考有监督学习，以极简的框架进行高效的训练.同时，为了避免聚类自监督学习中常遇到的collapse问题，我们提出一个新的聚类size下限的限制并开发online算法进行相应的优化. CoKe的示意图如下，利用前一个epoch得到的pseudo label和cluster center，在每个iteration中我们可以利用样本的single view进行训练.

三、从有监督学习到无监督学习

        首先有监督情况下的representation learning可以认为是一个distance metric learning的问题. 给定triplet，优化目标函数为

        其中x是样本的feature，c是不同类别proxy的feature. 在无监督学习中，缺少类别信息y，因此，相应的优化目标可以改写为

        相比有监督学习,新增了一个变量u来表示cluster assignment. 当u固定时，可以发现这个目标就变为了有监督学习. 但是直接优化该目标会导致trivial solution，也就是所有样本collapse到一个类中.不同于现有的balanced assignment约束 (e.g., SwAV), 我们考虑只约束每个cluster的下限，以更好的匹配数据分布. 新的约束可以写为

四、算法

        现在我们来讨论优化提出的目标函数.首先，如果固定c和u，那子问题就是一个标准的有监督学习问题. 考虑到triplet约束和cross entropy的关系 (见 SoftTriple ICCV’19)，该问题可以用标准的cross entropy loss优化.

其中聚类中心c和one-hot标签y来自上一个epoch. 当考虑double-view优化时，当前模型产出的soft label也可以加入优化，具体可见论文.

如果固定x，优化c和u可以写为一个标准的带约束的k-means问题

 K-means问题也可以迭代求解. 如果固定cluster assignment，聚类中心可以有closed-form solution

如果固定c，每个mini-batch的assignment问题为

 通过为每个cluster的constraint提供一个dual variable rho，assignment的更新为

 所有rho也会在每个mini-batch后进行相应的更新.

至此，目标函数的每组变量都可以在SGD的框架下进行优化.

五、结果

        CoKe展示了一个没有memory bank，momentum encoder等复杂结构的自监督学习框架，并可以利用小batch size, i.e., 1024，在单机8卡v100服务器上进行训练. 下面我们展示CoKe在不同任务上的效果.

1、CoKe on ImageNet

        首先，我们在标准的ImageNet linear probing任务上比较single/double/multi-view的CoKe的性能. 可以发现每个iteration只利用每个样本single view的CoKe已经取得了超过72%的Acc. 当配置为double/multi-view时，单机训练的CoKe都可以取得对应SoTA的结果.

 

2、CoKe on Clustering

        CoKe也可以被看作一个single-stage的深度聚类算法，因为我们在CIFAR-10/100上比较了聚类的结果.相比two-stage的SCAN，CoKe也取得了明显的提升. 目前CoKe已开源，期待CoKe在更多不同任务上的表现.

 六、应用

        另外给大家介绍下CV域上的开源免费模型，欢迎大家体验、下载（大部分手机端即可体验）：

ModelScope 魔搭社区https://modelscope.cn/models/damo/cv_resnet50_face-detection_retinaface/summary

ModelScope 魔搭社区https://modelscope.cn/models/damo/cv_resnet101_face-detection_cvpr22papermogface/summary

ModelScope 魔搭社区https://modelscope.cn/models/damo/cv_manual_face-detection_tinymog/summary

ModelScope 魔搭社区https://modelscope.cn/models/damo/cv_manual_face-detection_ulfd/summary

ModelScope 魔搭社区https://modelscope.cn/models/damo/cv_manual_face-detection_mtcnn/summary

ModelScope 魔搭社区https://modelscope.cn/models/damo/cv_resnet_face-recognition_facemask/summary

ModelScope 魔搭社区https://modelscope.cn/models/damo/cv_ir50_face-recognition_arcface/summary

ModelScope 魔搭社区https://modelscope.cn/models/damo/cv_manual_face-liveness_flir/summary

ModelScope 魔搭社区https://modelscope.cn/models/damo/cv_manual_face-liveness_flrgb/summary

ModelScope 魔搭社区https://modelscope.cn/models/damo/cv_manual_facial-landmark-confidence_flcm/summary

ModelScope 魔搭社区https://modelscope.cn/models/damo/cv_vgg19_facial-expression-recognition_fer/summary

ModelScope 魔搭社区https://modelscope.cn/models/damo/cv_resnet34_face-attribute-recognition_fairface/summary