【CVPR2022】BatchFormer: Learning to Explore Sample Relationships for Robust Representation Learning

【CVPR2022】BatchFormer: Learning to Explore Sample Relationships for Robust Representation Learning

论文:https://arxiv.org/pdf/2203.01522.pdf

代码:https://github.com/zhihou7/BatchFormer

BatchFormer的 V1 版本

这是一个来自 悉尼大学 和 京东 的工作。为了解决当前样本不平衡、零样本、域适应等问题。如下图所示,中间的是鸟(头部类),下面的是同样的鸟类目标,上面的是不同类的鸡、秃鹫目标(尾部类),这些都是外观极为相似的目标。 因为总体外观相似,把知识从样本数量较多的岛类目标,迁移到样本数量较少的秃鹫,是非常有好处的。(吐槽一下,这个图里秃鹫作为数量较少的尾部类可以理解,但很常见的鸡为啥也作为尾部类不是特别理解)

作者提出了一个网络能够从 mini-batch 中学习样本间的关系,叫做 BatchFormer。它可以促进每个 mini-batch 里的样本互相促进学习,比如在样本数量不平衡的长尾识别任务中,利用样本较多的数据促进样本较少类别样本的学习。一个形象的对比图如下所示,Channel attention是在通道上做 attention, VIT 是在空间上做attention,但BatchFormer是在 batch 维度上做 attention。

请添加图片描述

本文提出的特征学习框架如下图所示。特征提取器后面插入 Transformer 。这个Transformer 是沿batch 维度的,把整个 batch 看做一个sequence。同时,在 Transformer 前后都加入了一个分类器。值得注意的是:1)两个分类器是共享的,通过这个共享分类器,能够保持训练和测试的batch不变性。2)Transformer 没有使用位置编码。

请添加图片描述

def BatchFormer(x, y, encoder, is_training):
    # x: input features with the shape [N, C]
    # encoder: TransformerEncoderLayer(C,4,C,0.5)
    if not is_training:
        return x, y
    pre_x = x
    x = encoder(x.unsqueeze(1)).squeeze(1)
    x = torch.cat([pre_x, x], dim=0)
    y = torch.cat([y, y], dim=0)
    return x, y

BatchFormer 的伪代码如上所示。encoder里包括 MSA和 MLP。\(X\in \R^{N\times C}\)代表输入的特征,N 和 C 分别代表序列长度的特征维度(N就是batch里的样本数),因此,encoder的计算过程为:

\[\hat{X}_l=LN(MSA(X_{l-1})+X_{l-1}), \]

\[X_l=LN(MLP(\hat{X}_l)+\hat{X}_l) \]

作者的代码中,直接使用 torch.nn.TransformerEncoderLayer(48, 1, 48, 0.5) 来实现 encoder。四个参数及其含义依次为:

  • d_model – the number of expected features in the input (required).
  • nhead – the number of heads in the multiheadattention models (required).
  • dim_feedforward – the dimension of the feedforward network model (default=2048).
  • dropout – the dropout value (default=0.1).

BatchFormer的输入变为输入的 batch 维度数据,这样自注意力机制就变成不同样本之间的交叉注意力。但是,因为测试数据未知,即样本关系未知,因此BatchFormer前后的特征可能存在差异,所以,不能直接通过移除 BatchFormer 对新的样本进行推断。因此,作者使用了辅助分类器,辅助分类器和最终分类器之间是权重共享的。

BatchFormerV2

作者最近又提出了V2版本的 BatchFormer,将其泛化为一个更通用的模块,来促进一般的目标检测、分割、分类任务。具体来说,将 BatchFormer 插入到两层 Visual Transformer 之间,在每个空间的像素点上面进行Batch Transformer 操作。同时,将每个空间位置上面的 BatchFormer 共享,如下图所示。

从长尾分类的实验结果上看,BatchFormer 可以显著的提升分类性能。当然,作者还做了很多目标检测、图像分割的实验,这里不过多介绍。

posted @ 2022-09-07 01:13  高峰OUC  阅读(257)  评论(0编辑  收藏  举报