GCN实现人脸聚类

首先介绍一下GCN实现人脸聚类这篇论文,然后记录一下作者开源代码的阅读心得(代码写的很nice!)

Learning to Cluster Faces via Confifidence and Connectivity Estimation CVPR2020
GCN实现人脸聚类

GCN实现人脸聚类

一图胜千言,图2即作者提出的GCN网络来实现人脸聚类的全流程,其实本文的重点不在于GCN模型的构建,而是一种GCN思想在人脸聚类上的应用。GCN实现人脸聚类主要分为四个部分:图的构建,GCN-V实现节点置信度预测,GVN-E实现边连接预测,通过基于树划分的方式获取聚类。

图的构建
我们将一张图看成一个节点,节点的特征由CNN抽取而得,一般为256/512维,我们依据余弦相似度来评估节点之间的相似程度。采用KNN方式先获取所有节点的近邻k个节点并以此构建一个相似图,后续我们GCN处理的图即为此图。

其中,大规模knn计算是非常耗时且耗内存的,作者借用了faiss库来实现k近邻搜索,faiss包含多种距离计算模式,本文采用的内积方式(特征会事先进行归一化,因此内积算出来的即余弦相似度)。关于faiss的介绍,可参考文献https://zhuanlan.zhihu.com/p/210736523

GVN-V
对于密集连接处,节点同属于一个类的概率越大,稀疏处反之。因此,采用节点置信度方式来评估节点是不是属于某特定类是一种可行的方法。作者提出一种节点置信度的计算方式,如下图所示。针对当前节点所有邻接节点,若节点类别一致则累积他们的相似度,若不是同一类则减去节点间相似度,最后除以邻接节点数得出节点平均得分。若得分越高越说明节点所处位置密集连接且邻接节点大多归属于同一类,分越低越说明节点可能处于类边界。GCN-V模型就是来学习预测置信度的,考虑到数据量(11G显卡,58W训练数据)GCN-V为仅单层连接的图卷积网络,输出单个值,采用MSE损失来引导模型训练。GCN训练需要将整个图的特征矩阵和邻接矩阵塞到显卡去跑,因此模型大不了。但是,现在也可通过图分块的方式来扩大训练规模。

GCN实现人脸聚类

GCN-E
待续

聚类
待续

结果
GCN实现人脸聚类
GCN实现人脸聚类

代码笔记
作者开源代码结构清晰,看得出借鉴mmdetection框架的构建思路,采用config方式调节模型构建等,采用mmcv runner实现模型训练。
代码可重点参看一下几个函数
gcn_v.py:GCN-V模型的构建以及计算流程
knn.py:利用faiss搜索k近邻样本
gcn_v_dataset.py:构建图
train_gcn_v.py:利用mmcv runner完成模型训练,顺便看看runner源码,学习hook机制

BaseRunner.py:训练辅助类,定义run,train,val,save_checkpoint函数接口,维护hook列表,定义诸多常用hook注册方法
EpochBasedRunner.py:BaseRunner的继承类,实现逐轮训练常用的步骤
train_gcn_v.py: 训练方法函数,定义batch_processor方法规定模型计算以及loss收集,在_single_train通过配置构建优化器并依据model,batch_processor,optimizer构建runner,runner完成train和val逻辑,runner通过hook完成lr调节,模型保存等操作。

hook:狗子,不不不,是钩子。我们将train过程看成一个流水线操作,before_run,before_epoch,before_iter…,每一个计算节点都可能有相应的操作,例如我们需要在每个epoch开始前看看要不要调节LR,在每个epoch后看要不要存一下模型权重。ok,我们抽象出hook这个概念,hook定义的接口一一对应runner训练所涉及的多个节点(参看hook定义图)。hook依附在runner上,runner在某个计算节点(如after_train_epoch)完成所需操作都是由hook来实现的。

hook接口类定义如下,CheckpointHook继承Hook并重写两个训练后的方法完成一些这个hook在特定位置需要去做的操作。在runner中,例如train方法,我们在某个计算节点调用call_hook(hook下某方法),那么runner下所有的hook都会调用此方法,该干嘛的干嘛,这一步看call_hook的源码就知道了,call_hook会遍历runner下所有注册的hook并挨个调用具体提到的方法。至此,hook机制的大致原理既如此。
GCN实现人脸聚类
GCN实现人脸聚类
GCN实现人脸聚类
GCN实现人脸聚类