[ICCV2021] TransReID: Transformer-based Object Re-Identification 行人重识别

[ICCV2021] TransReID:Transformer-based Object Re-Identification

问题
贡献
baseline
- Overlapping Patches
- Position Embedding 和 cls token
- Feature Learning
TransReID
- Side Information Embedding（SIE）
- Jigsaw Patch Module
Experiments
参考

论文：https://arxiv.org/pdf/2102.04378.pdf

Code：heshuting555/TransReID

问题

在ReID任务中，CNN虽然效果很好，但是依然存在缺陷：一是只能在小范围内进行操作，即不能解决长距离依赖问题；另外一点是CNN中的卷积和下采样算子会使信息迅速丢失。
- 在CNN中引入Attention模块来缓解长距离依赖问题，但是大多数注意力模块都作用于网络较深的部分，这并不能从根本上解决CNN感受野受限的问题（同长距离依赖问题）。由于感受野大小受限，Attention不适用于连续的、较大的区域，而且很难从中提取出判别性比较强的特征。
- 细节性的细粒度特征同样也很重要。
基于Transformer的模型可以关注到全局信息，而CNN更多的是关注一些局部信息。全局性的结构化信息对于行人重识别来说也很重要。

图 1 使用Grad-CAM进行注意力图的可视化：(a) 原始图像；(b) 基于CNN的模型；© CNN+Attention模型；(d) 基于Transformer的模型。可以看出基于Transformer的模型能够关注到更多的全局信息

CNN下采样后丢失了一些细节信息（比如下图中的水杯），而Transformer-based的模型不会对图像进行下采样，因而可以保留细节信息。

图2 两种模型对于两个外观极其相似很难区分样本的输出特征图的可视化。与基于CNN的方法相比，基于Transformer的模型能够保留更多的细节信息，如背包、水杯。

Transformer更适合ReID任务的原因：多头注意力可以解决长距离依赖问题，从而使模型可以关注于人体的不同部位；没有下采样操作，所以能保留更多更详细的信息。

贡献

首次将Transformer应用于ReID的方法，并且提出了效果不错的基线模型，效果与CNN相当。经过加tricks之后，效果达到SOTA。
设计了一个区域拼图模块(jigsaw patches module, JPM), 用于将patch embedding进行重排（通过shift and shuffle operation)，以提高模型的鲁棒性和判别能力。
提出辅助信息编码模块(side information embeddings, SIE), 通过可学习的embeddings（其实就是可学习的特征向量）对一些外部信息(比如相机、视角信息)进行编码。

baseline

简要概括以下 ReID 领域的一些工作：

一般 ReID 的 pipeline 是用 CNN 提取特征，使用 cross-entropy (ID loss) 和 triplet loss 来进行训练。当然也有 circle loss(CVPR 2020) 将前面两者很好的结合了起来。

细粒度特征：水平分片、语义分割、姿态估计。

side information：主要就是不同视角下相机拍摄的人变化很大，可以强制让来自不同相机的人的图像学着映射到同一个子空间，从而消除视角的影响。

Transformer in vision：ViT需要大量的数据来pre-train才行；DeiT则利用teacher-student的蒸馏策略来对Vit进行提速。

局部信息能够提升模型的鲁棒性。如果图片划分为不同的块可以应用基于Transformer的模型，但是直接划分可能会破坏图像的全局信息，从而特征破坏长依赖关系。

这篇文章先建立了一个基于 Transformer 的 baseline，再基于此进一步优化。

图 3 基于Transformer的ReID框架

基本架构如上图，输入的图片分 patches，然后经过线性映射后，加入一个[cls] token（图中的*），添加 position embedding 后送入 Transformer 层。最后通过 ID Loss（交叉熵）+ Triplet Loss（三元组损失）进行模型的训练。Transformer 的好处在于不进行下采样操作，每一个Transformer Layer都有全局的感受野，因此可以提取细节信息。

基础模型有3个部分需要注意：重叠采样（Overlapping Patches），位置嵌入（Position Embedding）和 Feature Learning（特征学习）。

Overlapping Patches

相比于 ViT 将图像分割成 N 个不重叠的小块，本文采用滑动窗口来生成有重叠像素的 patches。假设滑动窗口的步长为 $s$ 个像素，patch 的大小为 $p×p,p=16p\times p, p=16$ ，2个相邻 patch 重叠区域的面积为 $(p−s)×p(p-s)\times p$ 。假设输入图像为 $\times w$ ，最终输出 patch 数量为:
$N=(⌊h−ps⌋+1)×(⌊h−ps⌋+1)N=(\lfloor \frac{h-p}{s} \rfloor + 1) \times (\lfloor \frac{h-p}{s} \rfloor + 1)$
具体实现中，作者通过卷积核为 16 步长 12 的卷积，将重叠采样和特征向量化一起实现。如下：

B = x.shape[0]   # batch size
x = self.patch_embed(x)   # Linear Projection
# self.patch_embed:
# PatchEmbed_overlap(
#   (proj): Conv2d(3, 768, kernel_size=(16, 16), stride=(12, 12))
# )

Overlapping Patches 可以通过改变采样的步长来调整 Patch 的数量。更多的patch通常可以带来更好的性能，同时也会带来更多的计算量。这就需要在性能和计算成本之间进行权衡。

Position Embedding 和 cls token

# 两组可学习的参数
self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))# 加入 cls token
cls_tokens = self.cls_token.expand(B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
x = torch.cat((cls_tokens, x), dim=1)# 引入 Position Embedding
x = x + self.pos_embed    # self.pos_embed.shape == [1, 211, 768]

相对于 CNN 的丰富的空间信息，Transformer 通过位置嵌入来弥补空间信息的缺失。 Position Embedding 通过一组可学习的参数进行空间信息的编码。Position Embedding 和 cls token 都是可学习的模块。

首先将每个 patch 的特征与 cls token（图中的*）叠加，再通过 position embedding 将新的特征进行位置编码（逐元素相加）得到编码后的特征。将其送入后续的 Transformer 层。

cls_token 即0索引对应的 embedding，这个 embedding 用于编码 patch 的整体特征。

Position Embedding 学习到的全局位置信息，图中[0, 1, 2, …, 8]。

这里由于 ReID 任务和 ViT 任务的输入分辨率不同，而位置编码由于位置和输入尺度相关，因此无法直接赋值，这时通过二次线性插值计算获得初始值。

ReID 任务中的输入图像分辨率可能发生变化，CNN 模型中卷积与输入分辨率无关，在最后的分类层之前也可以通过空间金字塔池化 SPP 操作输出固定大小的特征向量，以支持不同分辨率的图像输入。而 Transformer 中 Position Embedding 大小与输入图像尺度相关。为了使用预训练的权重，这里作者使用双线性插值来保值输入大小一致。

v = resize_pos_embed(v, self.pos_embed, self.patch_embed.num_y, self.patch_embed.num_x)

注：为什么需要 class token ？

假设没有 class token，Transformer 层输出了共 N 块 patch（token）的编码信息，为了完成后续的分类任务，我们怎么将这些都是局部 patch 的特征进行利用？
ViT 论文中针对分类问题作者提出有两种方法进行局部信息的整合，1是类似于全局平均池化；2是文中提到的 class token。
class token的好处：
- 不是基于特定图像的特征；
- 随机初始化，参与训练，编码整个数据集的统计信息；
- 聚合所有局部 patch 的信息，并且不基于个别 patch，避免了对 sequence 中某个特定 patch 的偏向性；
- class token 的位置编码固定为 0，避免输出受到位置编码的干扰。
那么全局平均池化是否能够替代 class token 呢？当然是可以的，全局平均池化也能够整个所有的局部 patch 信息。

针对具体任务来说，不同的任务有不同的关注点和处理方法。比如 Rethinking Semantic Segmentation from a Sequence-to-Sequence Perspective with Transformers 中，针对分割任务作者便去掉了 Class token 的设定。

Feature Learning

最后，通过为全局特征构建 ID loss 和 triplet loss 来优化网络，采用的结构使用了 BNNeck 结构，将分类损失和度量损失进行了剥离，之前我们介绍过。正负样本以及 anchor 分别以下标a，p，n标记。度量损失 triplet loss 计算如下：
$LT=log⁡[1+exp⁡(∥fa−fp∥22−∥fa−fn∥22)]\mathcal{L}_T=\log \left[ 1+\exp \left( \lVert f_a-f_p \rVert _{2}^{2}-\lVert f_a-f_n \rVert _{2}^{2} \right) \right]$

# 网络结构
TransReID((patch_embed): PatchEmbed_overlap((proj): Conv2d(3, 768, kernel_size=(16, 16), stride=(12, 12)))(pos_drop): Dropout(p=0.0, inplace=False)(blocks): ModuleList((0): Block((norm1): LayerNorm((768,), eps=1e-06, elementwise_affine=True)(attn): Attention((qkv): Linear(in_features=768, out_features=2304, bias=True)(attn_drop): Dropout(p=0.0, inplace=False)(proj): Linear(in_features=768, out_features=768, bias=True)(proj_drop): Dropout(p=0.0, inplace=False))(drop_path): Identity()(norm2): LayerNorm((768,), eps=1e-06, elementwise_affine=True)(mlp): Mlp((fc1): Linear(in_features=768, out_features=3072, bias=True)(act): GELU()(fc2): Linear(in_features=3072, out_features=768, bias=True)(drop): Dropout(p=0.0, inplace=False)))......)(norm): LayerNorm((768,), eps=1e-06, elementwise_affine=True)(fc): Linear(in_features=768, out_features=1000, bias=True)
)

TransReID

在 ViT 基线的基础上，作者提出了 Side Information Embedding（SIE）和 Jigsaw Patch Module（JPM）两个模块。前者用于引入ReID任务中的不变信息（如相机、视角等）；后者用于进一步挖掘图片的外观信息。

Side Information Embedding（SIE）

在 ReID 中，一个具有挑战性的问题是由于各种摄像机、视角等因素造成的外观偏差。为了解决这个问题，基于 CNN 的框架通常需要修改网络结构或设计特定的损失函数来包含这些非视觉线索(边信息，Side Information)如相机 id 和视角预测等信息。

Transformer 模型非常适合这一类问题，它可以通过将这些 Side Information 编码到 Embedding 中，以此来融合这些 Side Information。类似于 Position Embedding 可以通过一组可学习的参数来编码 Side Information。

具体来说，如果一幅图像的摄像头 ID 为 $C$ ，则其摄像头的 Embedding 可以记为 $S (C)$ 。不同于 Position Embedding 在各 patch 之间的变化，摄像机 Embedding 对于一幅图像的所有 patch 都是相同的。另外，如果物体的视点是可用的无论是通过视点估计算法还是人工标注，都可以将视点标签 $V$ 编码为 $S (V)$ ，然后用于图像的所有 patch。

现在的问题是如何整合两种不同类型的信息。一个可能的方案是直接将2个 Embedding 相加，即 $S (C) + S (V)$ ，但相加的方式可能会使两个 Embedding 项相互抵消。在本文中，作者提出将摄像头 ID $C$ 和视角标签 $V$ 共同编码为 $S (C, V)$ 。换句话说就是对于 $C_N$ 个摄像机 id 和 $V_N$ 个视角标签， $S (C, V)$ 总共有 $Cn×VNC_n \times V_N$ 个不同的值。最后，第 $i$ 个 patch $p_i$ 的输入 Embedding 如下:
$Ei=F(pi)+ρi+λS(C,V)E^i=\mathcal{F}\left( p_i \right) +\rho _i+\lambda S\left( C,V \right)$
上式中， $F\mathcal{F}$ 是学习到的线性映射， $λ\lambda$ 是平衡边信息 $S(C,V)S\left( C,V \right)$ 权重的超参数。由于每个 patch 的位置嵌入 $ρi\rho _i$ 不同，但在不同的图像中可能是相同的，而 $S(C,V)S\left( C,V \right)$ 对于单个图像中每个 patch 都是相同的，但对于不同的图像可能有不同的值。因此 TransReID 能够学习2种不同的映射函数，然后直接添加。整个输入 Embedding 为 $E_0;E_1,E_2,...,E_N]$ ，其中 $E_0$ 是 class token 嵌入。

论文通过相机信息和视角信息这2个范畴变量来说明 SIE 的使用。然而，SIE 可以扩展为包括范畴变量和数值变量在内的更多种类的信息。在不同实验中，摄像机、视角信息都包含在可用的范围内。

# Initialize SIE Embedding
if camera > 1 and view > 1:self.sie_embed = nn.Parameter(torch.zeros(camera * view, 1, embed_dim))trunc_normal_(self.sie_embed, std=.02)print('camera number is : {} and viewpoint number is : {}'.format(camera, view))print('using SIE_Lambda is : {}'.format(sie_xishu))
elif camera > 1:self.sie_embed = nn.Parameter(torch.zeros(camera, 1, embed_dim))trunc_normal_(self.sie_embed, std=.02)print('camera number is : {}'.format(camera))print('using SIE_Lambda is : {}'.format(sie_xishu))
elif view > 1:self.sie_embed = nn.Parameter(torch.zeros(view, 1, embed_dim))trunc_normal_(self.sie_embed, std=.02)print('viewpoint number is : {}'.format(view))print('using SIE_Lambda is : {}'.format(sie_xishu))# ================
# 引入 Position Embedding, camera 和 view 信息
if self.cam_num > 0 and self.view_num > 0:x = x + self.pos_embed + self.sie_xishu * self.sie_embed[camera_id * self.view_num + view_id]
elif self.cam_num > 0:x = x + self.pos_embed + self.sie_xishu * self.sie_embed[camera_id]
elif self.view_num > 0:x = x + self.pos_embed + self.sie_xishu * self.sie_embed[view_id]
else:x = x + self.pos_embed

Jigsaw Patch Module

论文中将 ViT 的最后一层改为2个独立的分支 Global Branch 和 Jigsaw Branch，以此来学习全局特征和局部特征。

假设输入到最后一层的隐藏特征记为 $Zl−1=[zl−10;zl−11,...,zl−1N]Z_{l-1} = [z^0_{l-1};z^1_{l-1},...,z^N_{l-1}]$ 。全局分支是将 $Z_{l-1}$ 编码成的标准 transformer 输出，其中 $f_g$ 可以看成全局特征。为了学习细粒度特征，作者将 $[zl−11,zl−12,...,zl−1N][z^1_{l-1},z^2_{l-1},...,z^N_{l-1}]$ 分成 $k$ 组，每一个组共用 cls token 嵌入 $z_{l-1}^0$ ，然后将组特征输入到 transformer 层，以学习 $k$ 个局部特征，表示为 $f_l^1,f_l^2,...,f_l^k$ 。

最近的两项研究表明，Token Embedding 主要由其附近的 Token 决定。换而言之，一组邻近的 patches embedding 主要是观察一个有限的连续区域，类似于 CNN 模型中感受野大小受限。

Although transformers can capture the global information in the image very well, a patch token still has a strong correlation with the corresponding patch. ViT-FRCNN [1] shows that the output embeddings of the last layer can be reshaped as a spatial feature map that includes location information. In other words, if we directly divide the original patch embeddings into k parts, each part may only consider a part of the continuous patch embeddings. Therefore, to better capture the long-range dependencies, we rearrange the patch embeddings and then re-group them into different parts, each of which contains several random patch embeddings of an entire image. In this way, the JPM module help to learn robust features with improved discrimination ability and more diversified coverage.

为了解决这个问题，作者提出了Jigsaw Patch Module (JPW) 在 patch 分组前执行 Shuffle 操作由 shift 操作和 patch Shuffle 操作实现，具体如下:

Step 1: shift 操作，将 $[zl−11,zl−12,...,zl−1N][z^1_{l-1},z^2_{l-1},...,z^N_{l-1}]$ 中的前 m 个元素移动至最后，得到 $[zl−1m+1,zl−1m+2,...,zl−1N,zl−11,zl−12,...,zl−1m][z^{m+1}_{l-1},z^{m+2}_{l-1},...,z^N_{l-1},z^1_{l-1},z^2_{l-1},...,z^m_{l-1}]$ .

Step 2: Patch Shuffle 操作，将 shift 操作后的特征分成 k 组。

经过 shuffle 操作，局部特征 $f_l^k$ 可以覆盖不同车身部位或车辆部位的 Patch。JPM 通过共享 transformer 层将其编码为 k 个局部特征 $f_l^1,f_l^2,...,f_l^k$ ，最终利用 k 个局部特征和全局特征 $f_g$ 训练模型。总损失的计算方法如下:
$L=LID(fg)+LT(fg)+1k∑(LID(fli)+LT(fli))\mathcal{L}=\mathcal{L}_{\text{ID}}\left( f_g \right) +\mathcal{L}_{\text{T}}\left( f_g \right) +\frac{1}{k}\sum{\left( \mathcal{L}_{\text{ID}}\left( f_{l}^{i} \right) +\mathcal{L}_{\text{T}}\left( f_{l}^{i} \right) \right)}$
在推理过程中，Concat全局特征和局部特征 $f_g,f_l^1,f_l^2,...,f_l^k]$ 作为最终的特征表示。如果仅使用 $f_g$ 计算成本较低但是性能略有下降。

# JPM模块, patch embedding 重排
def shuffle_unit(features, shift, group, begin=1):batchsize = features.size(0)dim = features.size(-1)# Shift Operation   [: 5:] + [: 1:5]feature_random = torch.cat([features[:, begin - 1 + shift:], features[:, begin:begin - 1 + shift]], dim=1)x = feature_random# Patch Shuffle Operationtry:x = x.view(batchsize, group, -1, dim)except:x = torch.cat([x, x[:, -2:-1, :]], dim=1)x = x.view(batchsize, group, -1, dim)x = torch.transpose(x, 1, 2).contiguous()x = x.view(batchsize, -1, dim)return x

Experiments

作者对于比较主流的公开数据集都进行了实验，数据集列表如下：

下面是一些主流 Backbone 与基于 Transformer 的对比，推理时间以 ResNet50 为单位。从结果可以看出，把图像分为12个 patch 得到的语义嵌入的结果比 ResNet50、101、152以及 ResNeSt50、200 都要好：

对于 JPM 模块的消融实验。‘w/o rearrange’：不适用 shift 和 shuffle 操作的 JPM；‘w/o local’：只是用全局特征。

通过 Grad-CAM 可视化注意力热图。（a）输入图像,（b）Baseline,（c）JPM w/o rearrange,（d）JPM.

SIE模块的消融实验。因为只有 VeRi-766 数据集提供视角信息，所以仅在其上进行实验。

在 Baseline 和 TransReID 上的消融实验。

下表是与 SOTA 模型的对比表，可以看出 TransReID 已经达到 SOTA 水平，同时也超越了很多主流的 ReID 模型，诸如，CBN、OSNet、MGN、HOReID 等。同时对于具有遮挡的 ReID 任务也有比较好的效果，同时 TransReID 在车辆 ReID 也达到了 SOTA 水平。

Trained Models and logs (Size 256)

Datasets	MSMT17	Market	Duke	OCC_Duke	VeRi	VehicleID
Model	mAP \| R1	mAP \| R1	mAP \| R1	mAP \| R1	mAP \| R1	R1 \| R5
Baseline(ViT)	61.8 \| 81.8	87.1 \| 94.6	79.6 \| 89.0	53.8 \| 61.1	79.0 \| 96.6	83.5 \| 96.7
Baseline(ViT)	model \| log	model \| log	model \| log	model \| log	model \| log	model \| test
*TransReID^(ViT)**	67.8 \| 85.3	89.0 \| 95.1	82.2 \| 90.7	59.5 \| 67.4	82.1 \| 97.4	85.2 \| 97.4
*TransReID^(ViT)**	model \| log	model \| log	model \| log	model \| log	model \| log	model \| test
*TransReID^(DeiT)**	66.3 \| 84.0	88.5 \| 95.1	81.9 \| 90.7	57.7 \| 65.2	82.4 \| 97.1	86.0 \| 97.6
*TransReID^(DeiT)**	model \| log	model \| log	model \| log	model \| log	model \| log	model \| test

参考

[1].TransReID: Transformer-based Object Re-Identification.

[2].Transformer 系列| Transformer又搞事情！TransReID首次在ReID中应用，结果喜人(文末获取论文)-技术圈 (proginn.com)

[3].[damo-cv/TransReID: ICCV-2021] TransReID: Transformer-based Object Re-Identification (github.com)