全网首发，Swin Transformer+FaceNet实现人脸识别
目录
2.PatchEmbed = Patch Partition + Linear Embedding
（2）Shifted Window based Self-Attention
一、 简介
与其他的深度学习方法在人脸上的应用不同，FaceNet并没有用传统的softmax的方式去进行分类学习，然后抽取其中某一层作为特征，而是直接进行端对端学习一个从图像到欧式空间的编码方法，然后基于这个编码再做人脸识别、人脸验证和人脸聚类等。
FaceNet主要有两个重点：Backbone和Triplet loss。我们也将主要从这两个方面介绍。
代码：oaifaye/facenet-swim-transformer
二、Swin Transformer作为Backbone
FaceNet论文使用了两个Backbone：Zeiler&Fergus架构和Google的Inception v1。随之时代的发展，已经不会有人再使用这两个模型了，基本所有优秀的模型都可以作为FaceNet的Backbone，想大Inception+ResNet足够了，想小MobileNet很优秀。这些模型都很棒，于是我们今天使用Swin Transformer。
目前Transformer应用到图像领域主要有两大挑战：
1.目标尺寸多变。不像NLP任务中token大小基本相同，目标检测中的目标尺寸不一，用单层级的模型很难有好的效果。
2.图片的高分辨率。高分辨率会使得计算复杂度呈现输入图片大小的二次方增长，这显然是不能接受的。
顾名思义，Hierarchical(多层级)解决第一个问题；Shifted Windows(滑窗)解决第二个问题。
1.Swin Transformer整体结构
先看模型整体结构，论文中的结构图跟代码有削微的差别，我以代码为准重画了结构图。因为原版FaceNet的输入是160x160,输出是128x128，所以我将输入改为了160x160，输出层加了全连接层将768映射成128.
图1
整个模型采取层次化的设计，一共包含4个Stage，每个stage都会缩小输入特征图的分辨率，像CNN一样逐层扩大感受野。
2.PatchEmbed = Patch Partition + Linear Embedding
PatchEmbed包含了两个功能：Patch Partition 和 Linear Embedding，Patch Partition负责将图片切成Patch； Linear Embedding负责对Patch做Embedding。结构图如下：
上面的96代表Embedding的维度，我们使用默认值96，代码及中文注释如下：
class PatchEmbed(nn.Module):
r""" 将图片分割成不重叠的小patch,并做Embedding，尺寸下采样尺寸为patch_size的大小。
Args:
img_size (int): 输入图像大小. 默认: 224.我们使用160
patch_size (int): Patch token 的大小，默认为 4*4.
in_chans (int): 输入图像的通道数，默认为 3.
embed_dim (int): 线性 projection 输出的通道数，默认为 96.
norm_layer (nn.Module, optional): 归一化层， 默认为N None.
"""
def __init__(self, img_size=224, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None):
super().__init__()
img_size = to_2tuple(img_size)
patch_size = to_2tuple(patch_size)
patches_resolution = [img_size[0] // patch_size[0], img_size[1] // patch_size[1]]
self.img_size = img_size
self.patch_size = patch_size
self.patches_resolution = patches_resolution
self.num_patches = patches_resolution[0] * patches_resolution[1]
self.in_chans = in_chans
self.embed_dim = embed_dim
self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
if norm_layer is not None:
self.norm = norm_layer(embed_dim)
else:
self.norm = None
def forward(self, x):
B, C, H, W = x.shape
# FIXME look at relaxing size constraints
assert H == self.img_size[0] and W == self.img_size[1],
f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
x = self.proj(x)
x = x.flatten(2)
x = x.transpose(1, 2) # 结构为 [B, num_patches=patche_h*patche_w, C]
if self.norm is not None:
x = self.norm(x)
return x
def flops(self):
Ho, Wo = self.patches_resolution
flops = Ho * Wo * self.embed_dim * self.in_chans * (self.patch_size[0] * self.patch_size[1])
if self.norm is not None:
flops += Ho * Wo * self.embed_dim
return flops3.Swin Transformer Block
重点来了，先看Swin Transformer Block结构图：
图2
图1中可以看到Stage1、2、4都有两个Swin Transformer Block，Stage3有6个Swin Transformer Block。我们以Stage1为例，输入和输出是一样的，都是1,1600,96。这块的重点两个个部分：
（1）Window Partition
根据window_size分窗，这里window_size=5，那们40x40的一个featuremap就能分成8个window，有点类似于YOLYv5中的focus结构，示意图如下：
相关代码：
def window_partition(x, window_size):
"""将输入分割为多个不重叠窗口
Args:
x: (B, H, W, C)
window_size (int): window size
Returns:
windows: (num_windows*B, window_size, window_size, C)
"""
B, H, W, C = x.shape
x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)
windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)
return windows（2）Shifted Window based Self-Attention
图2中的Attention模块是在每个窗口下计算注意力的，大大减少了计算量，但是也带来了一个问题：每个窗口建的联系是薄弱的， 为了更好的促进window间进行信息交互，Swin Transformer引入了shifted window操作。偶数轮次执行Shifted Window，如下图：
左边是没有重叠的Window Attention，而右边则是将窗口进行移位的Shift Window Attention。可以看到移位后的窗口包含了原本相邻窗口的元素。但这也引入了一个新问题，即window的个数翻倍了，由原本四个窗口变成了9个窗口。
为了在保持非重叠窗口的高效计算的同时引入跨窗口连接，我们提出了一种移动窗口分区方法，在实际代码里，我们是通过对特征图移位，并给Attention设置mask来间接实现的。能在保持原有的window个数下，最后的计算结果等价。示意图如下：
具体实现是通过torch.roll来实现的，torch.roll的用法大概是下面这样： 在Attention中通过设置mask，让Shifted Window Attention在与Window Attention相同的窗口个数下，达到等价的计算结果。首先我们对Shift Window后的每个窗口都给上index，并且做一个roll操作（window_size=2, shift_size=1）
我们希望在计算Attention的时候，让具有相同index QK进行计算，而忽略不同index QK计算结果。最后正确的结果如下图所示：
而要想在原始四个窗口下得到正确的结果，我们就必须给Attention的结果加入一个mask（如上图最右边所示）
相关代码:
class WindowAttention(nn.Module):
r""" 基于有相对位置偏差的多头自注意力窗口，支持移位的(shifted)或者不移位的(non-shifted)窗口.
输入:
dim (int): 输入特征的维度.
window_size (tuple[int]): 窗口的大小.
num_heads (int): 注意力头的个数.
qkv_bias (bool, optional): 给 query, key, value 添加可学习的偏置，默认为 True.
qk_scale (float | None, optional): 重写默认的缩放因子 scale.
attn_drop (float, optional): 注意力权重的丢弃率，默认为 0.0.
proj_drop (float, optional): 输出的丢弃率，默认为 0.0.
"""
def __init__(self, dim, window_size, num_heads, qkv_bias=True, qk_scale=None, attn_drop=0., proj_drop=0.):
super().__init__()
self.dim = dim # 输入特征的维度
self.window_size = window_size # 窗口的高 Wh,宽 Ww
self.num_heads = num_heads # 注意力头的个数
head_dim = dim // num_heads # 注意力头的维度
self.scale = qk_scale or head_dim ** -0.5 # 缩放因子 scale
# 定义相对位置偏移的参数表，结构为 [2*Wh-1 * 2*Ww-1, num_heads]
self.relative_position_bias_table = nn.Parameter(
torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads))
# 获取窗口内每个 token 的成对的相对位置索引
coords_h = torch.arange(self.window_size[0]) # 高维度上的坐标 (0, 7)
coords_w = torch.arange(self.window_size[1]) # 宽维度上的坐标 (0, 7)
coords = torch.stack(torch.meshgrid([coords_h, coords_w])) # 坐标，结构为 [2, Wh, Ww]
coords_flatten = torch.flatten(coords, 1) # 重构张量结构为 [2, Wh*Ww]
relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] # 相对坐标，结构为 [2, Wh*Ww, Wh*Ww]
relative_coords = relative_coords.permute(1, 2, 0).contiguous() # 交换维度，结构为 [Wh*Ww, Wh*Ww, 2]
relative_coords[:, :, 0] += self.window_size[0] - 1 # 第1个维度移位
relative_coords[:, :, 1] += self.window_size[1] - 1 # 第1个维度移位
relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1 # 第1个维度的值乘以 2倍的 Ww，再减 1
relative_position_index = relative_coords.sum(-1) # 相对位置索引，结构为 [Wh*Ww, Wh*Ww]
self.register_buffer("relative_position_index", relative_position_index) # 保存数据，不再更新
self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) # 线性层，特征维度变为原来的 3倍
self.attn_drop = nn.Dropout(attn_drop) # 随机丢弃神经元，丢弃率默认为 0.0
self.proj = nn.Linear(dim, dim) # 线性层，特征维度不变
self.proj_drop = nn.Dropout(proj_drop) # 随机丢弃神经元，丢弃率默认为 0.0
trunc_normal_(self.relative_position_bias_table, std=.02) # 截断正态分布，限制标准差为 0.02
self.softmax = nn.Softmax(dim=-1) # 激活函数 softmax
# 定义前向传播
def forward(self, x, mask=None):
"""
输入:
x: 输入特征图，结构为 [num_windows*B, N, C]
mask: (0/-inf) mask, 结构为 [num_windows, Wh*Ww, Wh*Ww] 或者没有 mask
"""
B_, N, C = x.shape # 输入特征图的结构
# 将特征图的通道维度按照注意力头的个数重新划分，并再做交换维度操作
qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
q, k, v = qkv[0], qkv[1], qkv[2] # 方便后续写代码，重新赋值
# q 乘以缩放因子
q = q * self.scale
# @ 代表常规意义上的矩阵相乘
attn = (q @ k.transpose(-2, -1)) # q 和 k 相乘后并交换最后两个维度
# 相对位置偏移，结构为 [Wh*Ww, Wh*Ww, num_heads]
relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view(
self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1)
# 相对位置偏移交换维度，结构为 [num_heads, Wh*Ww, Wh*Ww]
relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()
attn = attn + relative_position_bias.unsqueeze(0) # 带相对位置偏移的注意力图
if mask is not None: # 判断是否有 mask
nW = mask.shape[0] # mask 的宽
# 注意力图与 mask 相加
attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)
attn = attn.view(-1, self.num_heads, N, N) # 恢复注意力图原来的结构
attn = self.softmax(attn) # 激活注意力图 [0, 1] 之间
else:
attn = self.softmax(attn)
attn = self.attn_drop(attn) # 随机设置注意力图中的部分值为 0
# 注意力图与 v 相乘得到新的注意力图
x = (attn @ v).transpose(1, 2).reshape(B_, N, C)
x = self.proj(x) # 通过线性层
x = self.proj_drop(x) # 随机设置新注意力图中的部分值为 0
return x三、Triplet Loss
损失函数公式：
输入是一个三元组，包括锚（Anchor）示例、正（Positive）示例、负（Negative）示例，通过优化锚示例与正示例的距离小于锚示例与负示例的距离，实现样本之间的相似性计算。
a：anchor，锚示例；
p：positive，与a是同一类别的样本；
n：negative，与a是不同类别的样本；
margin是一个大于0的常数。最终的优化目标是拉近a和p的距离，拉远a和n的距离。设定一个margin常量，可以迫使模型努力学习，能让锚点a和负例n的distance值更大，同时让锚点a和正例p的distance值更小。
由于margin的存在，使得triplets loss多了一个参数，margin的大小需要调参。如果margin太大，则模型的损失会很大，而且学习到最后，loss也很难趋近于0，甚至导致网络不收敛，但是可以较有把握的区分较为相似的样本，即a和p更好区分；如果margin太小，loss很容易趋近于0，模型很好训练，但是较难区分a和p。
Triplet Loss代码实现：
def triplet_loss(alpha = 0.2):
def _triplet_loss(y_pred,Batch_size):
anchor, positive, negative = y_pred[:int(Batch_size)], y_pred[int(Batch_size):int(2*Batch_size)], y_pred[int(2*Batch_size):]
pos_dist = torch.sqrt(torch.sum(torch.pow(anchor - positive,2), axis=-1))
neg_dist = torch.sqrt(torch.sum(torch.pow(anchor - negative,2), axis=-1))
keep_all = (neg_dist - pos_dist < alpha).cpu().numpy().flatten()
hard_triplets = np.where(keep_all == 1)
pos_dist = pos_dist[hard_triplets].cuda()
neg_dist = neg_dist[hard_triplets].cuda()
basic_loss = pos_dist - neg_dist + alpha
loss = torch.sum(basic_loss)/torch.max(torch.tensor(1), torch.tensor(len(hard_triplets[0])))
return loss
return _triplet_loss