1、引言
在本文中，我们将探讨近两年来最具影响力的模型架构之一——Transformer模型。自从2017年Vaswani等人发表的论文《注意力是你所需要的全部》以来，Transformer架构在多个领域持续刷新着性能记录，尤其是在自然语言处理（NLP）领域。这种拥有庞大参数量的Transformer模型能够生成长篇且具有说服力的文章，为人工智能的新应用领域开辟了道路。鉴于Transformer架构的热潮在未来几年内似乎不会消退，理解其工作原理并亲自实现它变得至关重要。
尽管Transformer在自然语言处理领域取得了巨大成功，但我们在本教程中不会涉及NLP。原因有三：首先，人工智能的相关文章中有许多优秀的NLP课题研究，这些研究深入探讨Transformer架构在NLP中的应用。其次，GPT等大模型深入研究了自然语言级别的语言生成，读者可以轻松地将Transformer架构应用于其中。最后，也是最关键的，Transformer架构的应用远不止于此。虽然NLP是Transformer架构最初被提出并产生最大影响的领域，但它也推动了其他领域的研究进展，甚至包括计算机视觉。因此，我们将重点讨论Transformer和自注意力机制为何如此强大。在后续的文章中，我们将讨论Transformer在计算机视觉中的应用。
接下来，我们将导入一些标准库，同时，我们还会使用PyTorch Lightning作为辅助框架。如果你对PyTorch Lightning还不太熟悉，请仔细阅读前一篇博文。
# 导入标准库
import os # 用于操作系统功能
import numpy as np # 用于数值计算
import random # 用于生成随机数
import math # 用于数学运算
import json # 用于处理JSON数据
from functools import partial # 用于函数的部分应用
# 导入绘图相关的库
import matplotlib.pyplot as plt # 用于绘图
plt.set_cmap('cividis') # 设置颜色映射
%matplotlib inline # 使matplotlib图形在Jupyter Notebook中显示
from IPython.display import set_matplotlib_formats # 用于设置matplotlib的输出格式
set_matplotlib_formats('svg', 'pdf') # 设置输出格式为SVG和PDF，便于导出
from matplotlib.colors import to_rgb # 用于颜色转换
import matplotlib
matplotlib.rcParams['lines.linewidth'] = 2.0 # 设置线条宽度
import seaborn as sns # 用于数据可视化
sns.reset_orig() # 重置seaborn的默认设置
# 导入进度条库
from tqdm.notebook import tqdm # 用于显示加载进度条
# 导入PyTorch库
import torch # 用于深度学习
import torch.nn as nn # 用于神经网络
import torch.nn.functional as F # 用于神经网络函数
import torch.utils.data as data # 用于数据加载
import torch.optim as optim # 用于优化算法
# 导入Torchvision库
import torchvision # 用于计算机视觉
from torchvision.datasets import CIFAR100 # 用于加载CIFAR100数据集
from torchvision import transforms # 用于数据转换
# 导入PyTorch Lightning库
try:
import pytorch_lightning as pl # 用于简化训练过程
except ModuleNotFoundError: # 如果未安装PyTorch Lightning，则安装它
!pip install --quiet pytorch-lightning>=1.4
import pytorch_lightning as pl
from pytorch_lightning.callbacks import LearningRateMonitor, ModelCheckpoint # 用于监控学习率和保存模型
# 设置数据集和预训练模型的路径
DATASET_PATH = "../data" # 数据集路径
CHECKPOINT_PATH = "../saved_models/tutorial6" # 预训练模型保存路径
# 设置随机种子
pl.seed_everything(42) # 确保结果可复现
# 确保在GPU上的所有操作都是确定性的（如果使用GPU）
torch.backends.cudnn.deterministic = True # 确保CuDNN的确定性
torch.backends.cudnn.benchmark = False # 关闭CuDNN的基准测试
# 根据是否可用选择GPU或CPU
device = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")
print("Device:", device) # 打印设备信息
在准备运行代码之前，请务必确保已经根据实际需要调整了CHECKPOINT_PATH的路径，因为接下来将下载两个关键的预训练模型。如果此路径尚未设置或需要更新，请立即进行相应调整，以确保代码能够正确执行并加载这些重要的模型文件。
import urllib.request # 导入urllib.request库，用于处理URL请求
from urllib.error import HTTPError # 导入HTTPError，用于处理HTTP请求错误
# 定义存放预训练模型的GitHub URL
base_url = "https://raw.githubusercontent.com/phlippe/saved_models/main/tutorial6/"
# 定义需要下载的预训练文件列表
pretrained_files = ["ReverseTask.ckpt", "SetAnomalyTask.ckpt"]
# 如果检查点路径不存在，则创建它
os.makedirs(CHECKPOINT_PATH, exist_ok=True)
# 对于每个文件，检查它是否已经存在。如果不存在，尝试下载它。
for file_name in pretrained_files:
file_path = os.path.join(CHECKPOINT_PATH, file_name) # 获取文件的完整路径
# 如果文件名中包含"/"，则创建相应的文件夹结构
if "/" in file_name:
os.makedirs(file_path.rsplit("/",1)[0], exist_ok=True)
# 如果文件不存在，则尝试下载
if not os.path.isfile(file_path):
file_url = base_url + file_name # 构建文件的URL
print(f"正在下载 {file_url}...")
try:
urllib.request.urlretrieve(file_url, file_path) # 尝试下载文件
except HTTPError as e:
# 如果下载过程中出现HTTP错误，打印错误信息
print("下载过程中出现问题。请尝试从GDrive文件夹下载文件，或联系作者并提供完整的输出信息，包括以下错误：\n", e)
2、Transformer架构
在本文中，我们将亲自动手实现Transformer模型架构。虽然这种架构已经广为人知，Pytorch已经提供了一个名为nn.Transformer的模块，并且有关如何利用它进行下一个词预测的教程也已存在。但为了深入理解其核心细节，我们仍将从头开始实现它。
当然，网络上关于注意力机制和Transformer的教程数不胜数。如果你对这一主题感兴趣，并希望在阅读完本文后获得不同的视角，以下是一些推荐资源：
- 《Transformer：一种用于语言理解的新型神经网络架构》（Jakob Uszkoreit, 2017） - 这是谷歌关于Transformer论文的原始博客文章，主要聚焦于机器翻译的应用。
- 《图解Transformer》（Jay Alammar, 2018） - 一篇非常受欢迎且易于理解的博客文章，通过许多精美的可视化图表直观地解释了Transformer架构。重点在于自然语言处理（NLP）。
- 《注意力机制解析》（Lilian Weng, 2018）- 一篇总结不同领域（包括视觉）中注意力机制的博客文章。
- 《自注意力机制图解》（Raimi Karim, 2019） - 对自注意力机制步骤的清晰可视化。如果你觉得下面的解释太过抽象，强烈推荐阅读。
- 《Transformer家族》（Lilian Weng, 2020） - 一篇非常详尽的博客文章，回顾了除了原始版本之外的更多Transformer变体。
2.1.注意力机制的定义
注意力机制是近年来在神经网络中引起广泛关注的一类新型网络层，特别是在处理序列任务时。尽管在学术文献中“注意力”有多种定义，但我们在这里采用以下定义：注意力机制指的是基于输入查询和元素键动态计算权重的元素加权平均。具体来说，这意味着什么？我们的目标是计算多个元素的特征平均值。不过，我们希望根据元素的实际值来赋予不同的权重，而不是平等地对待每个元素。换句话说，我们希望动态地决定哪些输入更值得“关注”。具体来说，注意力机制通常包括以下四个部分：
- 查询（Query）：查询是一个特征向量，它描述了我们在序列中寻找的内容，即我们可能想要关注的对象。
- 键（Keys）：对于每个输入元素，都有一个键，这也是一个特征向量。这个特征向量大致描述了元素“提供”的内容，或者它何时可能变得重要。键的设计应使我们能够根据查询识别出我们想要关注的元素。
- 值（Values）：对于每个输入元素，我们还有一个值向量。这个特征向量是我们希望进行平均计算的。
- 得分函数（Score function）：为了评估我们想要关注哪些元素，我们需要定义一个得分函数。得分函数以查询和键为输入，并输出查询-键对的得分或注意力权重。它通常通过简单的相似性度量实现，如点积或小型多层感知器（MLP）。
平均值的权重通过所有得分函数输出的softmax函数计算得出。因此，我们为那些与查询最相似的键对应的值向量赋予更高的权重。如果我们用伪代码来描述这个过程，可以这样写：
α
i
=
exp
(
f
a
t
t
n
(
key
i
,
query
)
)
∑
j
exp
(
f
a
t
t
n
(
key
j
,
query
)
)
,
out
=
∑
i
α
i
⋅
value
i
\alpha_i = \frac{\exp\left(f_{attn}\left(\text{key}_i, \text{query}\right)\right)}{\sum_j \exp\left(f_{attn}\left(\text{key}_j, \text{query}\right)\right)}, \hspace{5mm} \text{out} = \sum_i \alpha_i \cdot \text{value}_i
αi=∑jexp(fattn(keyj,query))exp(fattn(keyi,query)),out=i∑αi⋅valuei
在视觉上，我们可以将注意力机制对一系列单词的关注情况展示如下：
每个单词都对应着一个键向量和一个值向量。通过一个得分函数（在这个场景中是点积）将查询向量与所有键向量进行比较，以此确定各自的权重。为了简化，我们这里没有展示softmax过程。最终，所有单词的值向量会根据这些注意力权重进行加权平均。
大多数注意力机制的不同之处在于它们所采用的查询方式、键和值向量的定义，以及所使用的得分函数。在Transformer模型中应用的注意力机制被称为自注意力。在自注意力机制中，序列中的每个元素都充当键、值和查询的角色。对于序列中的每个元素，我们都会通过一个注意力层来评估其查询与其他所有元素键的相似度，并为每个元素生成一个经过加权平均的新值向量。接下来，我们将通过深入探讨Transformer模型中使用的特定注意力机制——缩放点积注意力——来进一步理解这一概念。
2.2.缩放点积注意力
自注意力机制的核心是缩放点积注意力。我们的目标是构建一种注意力机制，它允许序列中的任意元素都能高效地关注到其他任意元素。点积注意力机制的输入包括一组查询 Q Q Q、键 K K K 和值 V V V，其中 Q , K ∈ R T × d k Q, K \in \mathbb{R}^{T \times d_k} Q,K∈RT×dk 和 V ∈ R T × d v V \in \mathbb{R}^{T \times d_v} V∈RT×dv。这里 T T T 代表序列的长度， d k d_k dk 和 d v d_v dv 分别代表查询/键和值的隐藏维度。为了简化说明，我们这里暂时不考虑批量维度。元素 i i i 到 j j j 之间的注意力值是基于查询 Q i Q_i Qi 和键 K j K_j Kj 的相似度，使用点积作为相似度的度量方式。数学上，点积注意力的计算公式如下：
Attention ( Q , K , V ) = softmax ( Q K T d k ) V \text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V Attention(Q,K,V)=softmax(dkQKT)V
这里的矩阵乘法 Q K T QK^T QKT 计算每一对查询和键的点积，生成一个 T × T T \times T T×T 形状的矩阵。每一行代表了特定元素对序列中所有其他元素的注意力得分。我们对这些得分应用softmax函数，并与值向量相乘，从而得到一个加权平均值（权重由注意力值决定）。下面是一个计算图，它从另一个角度展示了这种注意力机制（图示来源：Vaswani et al., 2017）。
我们尚未讨论的一点是 1 d k \frac{1}{\sqrt{d_k}} dk1 这个缩放因子的重要性。这个缩放因子对于在模型初始化后保持注意力值的适当方差非常关键。我们知道，初始化模型时我们希望各层的方差保持一致，因此查询 Q Q Q 和键 K K K 的方差也可能接近 1。但是，如果两个向量的方差为 σ 2 \sigma^2 σ2，那么它们点积的结果方差将会是 σ 4 ⋅ d k \sigma^4 \cdot d_k σ4⋅dk：
q i ∼ N ( 0 , σ 2 ) , k i ∼ N ( 0 , σ 2 ) → Var ( ∑ i = 1 d k q i ⋅ k i ) = σ 4 ⋅ d k q_i \sim \mathcal{N}(0,\sigma^2), k_i \sim \mathcal{N}(0,\sigma^2) \to \text{Var}\left(\sum_{i=1}^{d_k} q_i\cdot k_i\right) = \sigma^4\cdot d_k qi∼N(0,σ2),ki∼N(0,σ2)→Var(i=1∑dkqi⋅ki)=σ4⋅dk
如果不将方差重新缩放至 σ 2 \sigma^2 σ2 附近，那么在对数几率上应用 softmax 函数时，一个随机元素的值将会饱和至 1，而其他所有元素的值则会接近 0。这将导致通过 softmax 的梯度几乎为零，从而使得我们无法适当地学习模型参数。需要注意的是，方差中的额外 σ 2 \sigma^2 σ2 因子，即 σ 4 \sigma^4 σ4 而不是 σ 2 \sigma^2 σ2，通常不会造成问题，因为我们通常将原始方差 σ 2 \sigma^2 σ2 保持在接近 1 的水平。
在上图表中的掩码模块（mask）表示在注意力矩阵中对特定条目进行可选的掩码操作。例如，当我们将不同长度的多个序列堆叠到一个批次中时，就会使用这种方法。为了在 PyTorch 中仍然能够利用并行计算的优势，我们通常会将句子填充到相同的长度，并在计算注意力值时忽略这些填充标记。这通常是通过将相关的注意力对数几率设置为一个极低的值来实现的。
在详细讨论了缩放点积注意力块的细节之后，我们可以编写一个函数，该函数接受查询、键和值的三元组，并计算输出特征：
def scaled_dot_product(q, k, v, mask=None):
d_k = q.size(-1) # 获取查询向量的维度
attn_logits = torch.matmul(q, k.transpose(-2, -1)) # 计算查询和键的点积
attn_logits = attn_logits / math.sqrt(d_k) # 应用缩放因子
if mask is not None: # 如果提供了掩码
attn_logits = attn_logits.masked_fill(mask == 0, -9e15) # 将掩码位置的值设置为一个非常小的数
attention = F.softmax(attn_logits, dim=-1) # 应用 softmax 函数
values = torch.matmul(attention, v) # 计算加权平均值
return values, attention
注意，上述代码支持在序列长度前添加任何额外的维度，这样它也可以用于处理批次数据。为了更好地理解，让我们生成一些随机的查询、键和值向量，并计算注意力输出：
seq_len, d_k = 3, 2
pl.seed_everything(42) # 设置随机种子以确保结果可复现
q = torch.randn(seq_len, d_k) # 生成随机查询向量
k = torch.randn(seq_len, d_k) # 生成随机键向量
v = torch.randn(seq_len, d_k) # 生成随机值向量
values, attention = scaled_dot_product(q, k, v) # 计算输出特征和注意力权重
print("查询向量 Q\n", q)
print("键向量 K\n", k)
print("值向量 V\n", v)
print("输出特征 Values\n", values)
print("注意力权重 Attention\n", attention)
输出示例：
查询向量 Q
tensor([[ 0.3367, 0.1288],
[ 0.2345, 0.2303],
[-1.1229, -0.1863]])
键向量 K
tensor([[ 2.2082, -0.6380],
[ 0.4617, 0.2674],
[ 0.5349, 0.8094]])
值向量 V
tensor([[ 1.1103, -1.6898],
[-0.9890, 0.9580],
[ 1.3221, 0.8172]])
输出特征 Values
tensor([[ 0.5698, -0.1520],
[ 0.5379, -0.0265],
[ 0.2246, 0.5556]])
注意力权重 Attention
tensor([[0.4028, 0.2886, 0.3086],
[0.3538, 0.3069, 0.3393],
[0.1303, 0.4630, 0.4067]])
在继续之前，请确保你能够理解这里的具体计算过程，并通过手工计算来验证。完全理解缩放点积注意力的计算方式是非常重要的。
2.3.多头注意力机制
缩放点积注意力机制使得网络能够对序列进行关注。但是，序列中的元素往往需要同时关注多个不同的方面，单一的加权平均可能并不是最佳选择。因此，我们扩展了注意力机制，引入了多头注意力，即在相同的特征上应用多组不同的查询-键-值三元组。具体来说，给定查询、键和值矩阵，我们将其转换为 h h h 组子查询、子键和子值，然后独立地对它们应用缩放点积注意力。最后，我们将这些子结果拼接起来，并通过一个最终的权重矩阵进行整合。数学上，这一操作可以表达为：
多头注意力 ( Q , K , V ) = 拼接 ( 头 1 , . . . , 头 h ) W O 其中 头 i = 注意力 ( Q W i Q , K W i K , V W i V ) \begin{split} \text{多头注意力}(Q,K,V) & = \text{拼接}(\text{头}_1,...,\text{头}_h)W^{O}\\ \text{其中 } \text{头}_i & = \text{注意力}(QW_i^Q,KW_i^K, VW_i^V) \end{split} 多头注意力(Q,K,V)其中 头i=拼接(头1,...,头h)WO=注意力(QWiQ,KWiK,VWiV)
我们称这为多头注意力层，其具有可学习的参数 W 1... h Q ∈ R D × d k W_{1...h}^{Q}\in\mathbb{R}^{D\times d_k} W1...hQ∈RD×dk， W 1... h K ∈ R D × d k W_{1...h}^{K}\in\mathbb{R}^{D\times d_k} W1...hK∈RD×dk， W 1... h V ∈ R D × d v W_{1...h}^{V}\in\mathbb{R}^{D\times d_v} W1...hV∈RD×dv，和 W O ∈ R h ⋅ d v × d 输出 W^{O}\in\mathbb{R}^{h\cdot d_v\times d_{\text{输出}}} WO∈Rh⋅dv×d输出（其中 D D D 表示输入的维度）。在计算图中，我们可以如下所示进行可视化（图表来源 - Vaswani et al., 2017）。
在神经网络中，我们如何应用多头注意力层，尤其是在输入不是任意查询、键和值向量的情况下？观察上面的计算图，一个简单而有效的实现方法是将当前的特征映射 X ∈ R B × T × d 模型 X\in\mathbb{R}^{B\times T\times d_{\text{模型}}} X∈RB×T×d模型 作为 Q Q Q， K K K 和 V V V（其中 B B B 表示批量大小， T T T 表示序列长度， d 模型 d_{\text{模型}} d模型 表示 X X X 的隐藏维度）。连续的权重矩阵 W Q W^{Q} WQ， W K W^{K} WK，和 W V W^{V} WV 可以将 X X X 转换为代表输入的查询、键和值的特征向量。采用这种方法，我们可以如下实现多头注意力模块。
# 辅助函数，用于支持不同形状的掩码。
# 输出形状支持（批量大小，头数，序列长度，序列长度）
# 如果是2D：在批量大小和头数上广播
# 如果是3D：在头数上广播
# 如果是4D：保持不变
def expand_mask(mask):
assert mask.ndim >= 2, "掩码至少必须是2维的，具有序列长度 x 序列长度"
if mask.ndim == 3:
mask = mask.unsqueeze(1)
while mask.ndim < 4:
mask = mask.unsqueeze(0)
return mask
class MultiheadAttention(nn.Module):
def __init__(self, input_dim, embed_dim, num_heads):
super().__init__()
assert embed_dim % num_heads == 0, "嵌入维度必须是头数的倍数"
self.embed_dim = embed_dim
self.num_heads = num_heads
self.head_dim = embed_dim // num_heads
# 为了提高效率，将所有权重矩阵1...h堆叠在一起
self.qkv_proj = nn.Linear(input_dim, 3*embed_dim)
self.o_proj = nn.Linear(embed_dim, embed_dim)
self._reset_parameters()
def _reset_parameters(self):
# 使用原始Transformer的初始化方法，详见PyTorch文档
nn.init.xavier_uniform_(self.qkv_proj.weight)
self.qkv_proj.bias.data.fill_(0)
nn.init.xavier_uniform_(self.o_proj.weight)
self.o_proj.bias.data.fill_(0)
def forward(self, x, mask=None, return_attention=False):
batch_size, seq_length, _ = x.size()
if mask is not None:
mask = expand_mask(mask)
qkv = self.qkv_proj(x)
# 从线性输出中分离 Q, K, V
qkv = qkv.reshape(batch_size, seq_length, self.num_heads, 3*self.head_dim)
qkv = qkv.permute(0, 2, 1, 3) # [批量, 头, 序列长度, 维度]
q, k, v = qkv.chunk(3, dim=-1)
# 计算值输出
values, attention = scaled_dot_product(q, k, v, mask=mask)
values = values.permute(0, 2, 1, 3) # [批量, 序列长度, 头, 维度]
values = values.reshape(batch_size, seq_length, self.embed_dim)
o = self.o_proj(values)
if return_attention:
return o, attention
else:
return o
多头注意力的一个关键特性是它对于输入是排列不变的。这意味着如果我们交换序列中的两个输入元素，例如 X 1 ↔ X 2 X_1\leftrightarrow X_2 X1↔X2（暂时忽略批量维度），输出除了元素1和2交换外，完全相同。因此，多头注意力实际上是将输入视为一组元素，而不是序列。这种属性使得多头注意力块和Transformer架构非常强大和广泛适用！但如果输入的顺序对于解决任务实际上很重要，比如语言建模，我们就需要在输入特征中编码位置信息，我们将在后续的 位置编码 主题中进一步探讨。
在继续构建Transformer架构之前，我们可以将自注意力操作与我们处理序列数据的其他常见层进行比较：卷积和循环神经网络。下面你可以找到一张由 Vaswani et al. (2017) 提供的表格，关于每层的复杂性、顺序操作的数量和最大路径长度。复杂性是通过执行操作的上界次数来衡量的，而最大路径长度表示前向或后向信号需要穿越的最大步数以到达任何其他位置。这个长度越低，梯度信号就能更好地对长距离依赖进行反向传播。让我们看看下面的表格：
是序列长度， 是表示维度， 是卷积的核大小。与循环网络相比，自注意力层可以并行化所有操作，使其在较短序列长度下执行速度更快。然而，当序列长度超过隐藏维度时，自注意力比RNNs更昂贵。一种减少长序列计算成本的方法是通过限制自注意力只关注输入的邻域，表示为 。尽管如此，最近有很多关于更高效的Transformer架构的工作，仍然允许长距离依赖，如果你感兴趣，你可以在 Tay et al. (2020) 的论文中找到概述。
2.4.Transformer编码器
接下来，我们来探讨如何在变压器（Transformer）架构中应用多头注意力机制。最初，变压器模型是为机器翻译任务设计的，因此它采用了编码器-解码器结构。在这个结构中，编码器负责接收原始语言的句子并生成基于注意力的表示，而解码器则基于编码器的输出生成翻译后的句子，这一过程类似于标准的递归神经网络（RNN）。尽管这种结构对于需要自回归解码的序列到序列任务非常有用，但在这里，我们将重点放在编码器部分。许多自然语言处理（NLP）领域的进步都是基于纯编码器的变压器模型实现的。例如，BERT系列、视觉变压器（Vision Transformer）等。在我们的教程中，我们也将主要关注编码器部分。一旦您理解了编码器的架构，解码器的实现将会非常简单。完整的变压器架构如下所示（图示来源 - [Vaswani et al., 2017](Attention Is All You Need)）：
编码器由 N N N 个相同的块组成，这些块依次应用于输入。以 x x x 作为输入，首先通过上述实现的多头注意力块。然后将输出与原始输入通过残差连接相加，并对和进行连续的层归一化处理。总的来说，它计算的是 LayerNorm ( x + Multihead ( x , x , x ) ) \text{LayerNorm}(x+\text{Multihead}(x,x,x)) LayerNorm(x+Multihead(x,x,x))（其中 x x x 是传递给注意力层的 Q Q Q、 K K K 和 V V V 输入）。在变压器架构中，残差连接至关重要，原因有二：
- 类似于残差网络（ResNets），Transformer设计得非常深。一些模型在编码器中包含超过24个块。因此，残差连接对于确保模型中梯度的顺畅流动至关重要。
- 没有残差连接，原始序列的信息将会丢失。记住，多头注意力层忽略了序列中元素的位置，并且只能基于输入特征来学习它。如果移除残差连接，这意味着在第一个注意力层之后（在初始化之后），这些信息将会丢失。而且，随机初始化的查询和键向量，位置 的输出向量与其原始输入没有关系。所有注意力的输出可能表示相似或相同的信息，模型没有机会区分哪个信息来自哪个输入元素。残差连接的一个替代选项是至少固定一个头来专注于其原始输入，但这非常低效，并没有改善梯度流动的好处。
FFN ( x ) = max ( 0 , x W 1 + b 1 ) W 2 + b 2 x = LayerNorm ( x + FFN ( x ) ) \begin{split} \text{FFN}(x) & = \max(0, xW_1+b_1)W_2 + b_2\\ x & = \text{LayerNorm}(x + \text{FFN}(x)) \end{split} FFN(x)x=max(0,xW1+b1)W2+b2=LayerNorm(x+FFN(x))
这个MLP为模型增加了额外的复杂性，并允许对每个序列元素单独进行转换。你可以想象这使得模型能够“后处理”前一个多头注意力添加的新信息，并为下一个注意力块做好准备。通常，MLP的内部维度是原始输入 x x x 维度的2-8倍。相比狭窄的多层MLP，更宽的层的优势在于更快的并行化执行。
最后，查看了编码器架构的所有部分后，我们可以开始实现它。我们首先通过实现单个编码器块来开始。除了上述层外，我们还将添加dropout层在MLP和MLP以及多头注意力的输出中进行正则化。
class EncoderBlock(nn.Module):
def __init__(self, input_dim, num_heads, dim_feedforward, dropout=0.0):
"""
构造函数：
input_dim - 输入的维度
num_heads - 注意力块中使用的头数
dim_feedforward - MLP中隐藏层的维度
dropout - dropout层中使用的dropout概率
"""
super().__init__()
# 注意力层
self.self_attn = MultiheadAttention(input_dim, input_dim, num_heads)
# 两层MLP
self.linear_net = nn.Sequential(
nn.Linear(input_dim, dim_feedforward),
nn.Dropout(dropout),
nn.ReLU(inplace=True),
nn.Linear(dim_feedforward, input_dim)
)
# 主层之间的层
self.norm1 = nn.LayerNorm(input_dim)
self.norm2 = nn.LayerNorm(input_dim)
self.dropout = nn.Dropout(dropout)
def forward(self, x, mask=None):
# 注意力部分
attn_out = self.self_attn(x, mask=mask)
x = x + self.dropout(attn_out)
x = self.norm1(x)
# MLP部分
linear_out = self.linear_net(x)
x = x + self.dropout(linear_out)
x = self.norm2(x)
return x
基于这个块，我们可以为完整的Transformer编码器实现一个模块。除了一个通过编码器块序列迭代的前向函数外，我们还提供了一个名为get_attention_maps的函数。这个函数的想法是返回编码器中所有多头注意力块的注意力概率。这有助于我们理解，并在某种意义上，解释模型。然而，注意力概率应该谨慎解释，因为它并不一定反映模型的真实解释（关于这一点有一系列论文，包括注意力不是解释和[注意力不是不解释](Attention is not not Explanation))。
class TransformerEncoder(nn.Module):
def __init__(self, num_layers, **block_args):
super().__init__()
self.layers = nn.ModuleList([EncoderBlock(**block_args) for _ in range(num_layers)])
def forward(self, x, mask=None):
for l in self.layers:
x = l(x, mask=mask)
return x
def get_attention_maps(self, x, mask=None):
attention_maps = []
for l in self.layers:
_, attn_map = l.self_attn(x, mask=mask, return_attention=True)
attention_maps.append(attn_map)
x = l(x)
return attention_maps
2.5.位置编码
我们之前讨论过，多头注意力（Multi-Head Attention）模块是排列不变的，它无法区分输入序列中的元素顺序。然而在诸如语言理解等任务中，位置信息对于解释输入单词非常重要。因此，可以通过输入特征添加位置信息。我们可以为每一个可能的位置学习一个嵌入，但这不会泛化到动态输入序列长度。因此，更好的选择是使用网络能够从特征中识别并可能泛化到更长序列的特征模式。Vaswani等人选择的特定模式是不同频率的正弦和余弦函数，如下所示：
P
E
(
p
o
s
,
i
)
=
{
sin
(
p
o
s
1000
0
i
/
d
model
)
如果
i
mod
2
=
0
cos
(
p
o
s
1000
0
(
i
−
1
)
/
d
model
)
否则
PE_{(pos,i)} = \begin{cases} \sin\left(\frac{pos}{10000^{i/d_{\text{model}}}}\right) & \text{如果}\hspace{3mm} i \text{ mod } 2=0\\ \cos\left(\frac{pos}{10000^{(i-1)/d_{\text{model}}}}\right) & \text{否则}\\ \end{cases}
PE(pos,i)=⎩
⎨
⎧sin(10000i/dmodelpos)cos(10000(i−1)/dmodelpos)如果i mod 2=0否则
P
E
(
p
o
s
,
i
)
PE_{(pos,i)}
PE(pos,i) 表示序列中位置
p
o
s
pos
pos 的位置编码，以及隐藏维度
i
i
i。这些值，对于所有隐藏维度进行拼接，被添加到原始输入特征中（在变压器的可视化中，见“位置编码”），构成了位置信息。我们区分偶数（
i
mod
2
=
0
i \text{ mod } 2=0
i mod 2=0）和奇数（
i
mod
2
=
1
i \text{ mod } 2=1
i mod 2=1）隐藏维度，分别应用正弦/余弦。这种编码背后的直觉是，你可以将
P
E
(
p
o
s
+
k
,
:
)
PE_{(pos+k,:)}
PE(pos+k,:) 表示为
P
E
(
p
o
s
,
:
)
PE_{(pos,:)}
PE(pos,:) 的线性函数，这可能允许模型轻松地关注相对位置。不同维度的波长范围从
2
π
2\pi
2π 到
10000
⋅
2
π
10000\cdot 2\pi
10000⋅2π。
位置编码如下实现。代码取自[PyTorch教程](Language Modeling with)关于自然语言处理中的Transformer，并根据我们的需要进行了调整。
class PositionalEncoding(nn.Module):
def __init__(self, d_model, max_len=5000):
"""
输入
d_model - 输入的隐藏维度。
max_len - 预期的序列的最大长度。
"""
super().__init__()
# 创建表示最大长度输入的位置编码的矩阵 [SeqLen, HiddenDim]
pe = torch.zeros(max_len, d_model)
position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
pe[:, 0::2] = torch.sin(position * div_term)
pe[:, 1::2] = torch.cos(position * div_term)
pe = pe.unsqueeze(0)
# register_buffer => 这是一个不是参数的张量，但应该是模块状态的一部分。
# 用于需要与模块在同一设备上的张量。
# persistent=False 告诉PyTorch不要将缓冲区添加到状态字典中（例如，当我们保存模型时）
self.register_buffer('pe', pe, persistent=False)
def forward(self, x):
x = x + self.pe[:, :x.size(1)]
return x
为了理解位置编码，我们可以在下面对其进行可视化。我们将生成一个位置编码在隐藏维度和序列位置的图像。因此，每个像素代表我们对输入特征进行的编码特定位置的变化。让我们在下面进行。
encod_block = PositionalEncoding(d_model=48, max_len=96)
pe = encod_block.pe.squeeze().T.cpu().numpy()
fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(8,3))
pos = ax.imshow(pe, cmap="RdGy", extent=(1,pe.shape[1]+1,pe.shape[0]+1,1))
fig.colorbar(pos, ax=ax)
ax.set_xlabel("序列中的位置")
ax.set_ylabel("隐藏维度")
ax.set_title("隐藏维度上的位置编码")
ax.set_xticks([1]+[i*10 for i in range(1,1+pe.shape[1]//10)])
ax.set_yticks([1]+[i*10 for i in range(1,1+pe.shape[0]//10)])
plt.show()
你可以清楚地看到在隐藏维度中编码位置的正弦和余弦波，波长不同。具体来说，我们可以单独查看每个隐藏维度的正弦/余弦波，以更好地理解模式。下面我们可视化了隐藏维度 、2、3 和 4的位置编码。
sns.set_theme()
fig, ax = plt.subplots(2, 2, figsize=(12,4))
ax = [a for a_list in ax for a in a_list]
for i in range(len(ax)):
ax[i].plot(np.arange(1,17), pe[i,:16], color=f'C{i}', marker="o", markersize=6, markeredgecolor="black")
ax[i].set_title(f"第 {i+1} 个隐藏维度的编码")
ax[i].set_xlabel("序列中的位置", fontsize=10)
ax[i].set_ylabel("位置编码", fontsize=10)
ax[i].set_xticks(np.arange(1,17))
ax[i].tick_params(axis='both', which='major', labelsize=10)
ax[i].tick_params(axis='both', which='minor', labelsize=8)
ax[i].set_ylim(-1.2, 1.2)
fig.subplots_adjust(hspace=0.8)
sns.reset_orig()
plt.show()
正如我们所看到的，隐藏维度2和 2之间的模式仅在起始角度上有所不同。波长是 ，因此在位置6后重复。隐藏维度 2和 3 的波长大约为两倍。
2.6.学习率预热
训练Transformer常用的一种技术是学习率预热。这意味着我们在最初的几次迭代中，逐渐将学习率从0增加到最初指定的学习率。因此，我们缓慢地开始学习，而不是从一开始就采取非常大的步骤。实际上，如果不进行学习率预热，训练深度Transformer可能会使模型发散，并在训练和测试中获得更差的性能。以Liu等人（2019）的以下图表为例，比较了没有预热的Adam（即原始Adam）与预热的Adam：
显然，预热是Transformer架构中的一个重要超参数。为什么它如此重要？目前有两个常见的解释。首先，Adam使用偏差校正因子，但这可能导致在第一次迭代中自适应学习率的方差增加。像[RAdam]这样的改进优化器已被证明可以克服这个问题，不需要预热即可训练Transformer。其次，逐层应用的层归一化在第一次迭代中可能导致梯度非常高，这可以通过使用[Pre-Layer Normalization]（类似于Pre-Activation ResNet）或用其他技术替换层归一化来解决（[Adaptive Normalization]、[Power Normalization])。
然而，许多应用和论文仍然使用原始的Transformer架构和Adam，因为预热是一种简单但有效的方法，可以解决第一次迭代中的梯度问题。我们可以使用许多不同的调度器。例如，原始的Transformer论文使用了一个带有预热的指数衰减调度器。然而，目前最受欢迎的调度器是余弦预热调度器，它将预热与余弦形状的学习率衰减结合起来。我们可以在下面实现它，并可视化学习率因子在周期内的变化。
class CosineWarmupScheduler(optim.lr_scheduler._LRScheduler):
def __init__(self, optimizer, warmup, max_iters):
self.warmup = warmup
self.max_num_iters = max_iters
super().__init__(optimizer)
def get_lr(self):
lr_factor = self.get_lr_factor(epoch=self.last_epoch)
return [base_lr * lr_factor for base_lr in self.base_lrs]
def get_lr_factor(self, epoch):
lr_factor = 0.5 * (1 + np.cos(np.pi * epoch / self.max_num_iters))
if epoch <= self.warmup:
lr_factor *= epoch * 1.0 / self.warmup
return lr_factor
# 需要初始化学习率调度器
p = nn.Parameter(torch.empty(4,4))
optimizer = optim.Adam([p], lr=1e-3)
lr_scheduler = CosineWarmupScheduler(optimizer=optimizer, warmup=100, max_iters=2000)
# 绘图
epochs = list(range(2000))
sns.set()
plt.figure(figsize=(8,3))
plt.plot(epochs, [lr_scheduler.get_lr_factor(e) for e in epochs])
plt.ylabel("学习率因子")
plt.xlabel("迭代次数（按批次）")
plt.title("余弦预热学习率调度器")
plt.show()
sns.reset_orig()
在前100次迭代中，我们将学习率因子从0增加到1，而在所有后续迭代中，我们使用余弦波衰减它。这种调度器的预实现可以在流行的NLP Transformer库[huggingface]中找到。
2.7.PyTorch Lightning 模块
最终，我们可以将 Transformer 架构整合到 PyTorch Lightning 模块中。正如教程5所介绍的，PyTorch Lightning 能够简化我们的训练和测试代码，并且将代码结构化地组织在不同的函数中。我们将实现一个基于 Transformer 编码器的分类器模板，每个序列元素都会有一个预测输出。如果我们需要对整个序列进行分类，通常的做法是在序列中添加一个额外的 [CLS] 标记，代表分类器标记。但在这里，我们专注于每个元素都有输出的任务。
除了 Transformer 架构，我们还增加了一个小型输入网络（将输入维度映射到模型维度）、位置编码以及一个输出网络（将输出编码转换为预测）。我们还增加了一个学习率调度器，它每次迭代更新一次而不是每个周期更新一次。这对于预热和平滑的余弦衰减是必需的。训练、验证和测试步骤目前是空的，将根据我们特定于任务的模型进行填充。
class TransformerPredictor(pl.LightningModule):
def __init__(self, input_dim, model_dim, num_classes, num_heads, num_layers, lr, warmup, max_iters, dropout=0.0, input_dropout=0.0):
"""
参数：
input_dim - 输入的隐藏维度
model_dim - Transformer 内部使用的隐藏维度
num_classes - 每个序列元素要预测的类别数
num_heads - 多头注意力机制中使用的头数
num_layers - 使用的编码器层数
lr - 优化器中的学习率
warmup - 预热步骤数，通常在50到500之间
max_iters - 模型训练的最大迭代次数，这对于余弦预热调度器是必要的
dropout - 模型内部应用的 dropout 比率
input_dropout - 输入特征上应用的 dropout 比率
"""
super().__init__()
self.save_hyperparameters()
self._create_model()
def _create_model(self):
# 将输入维度转换为模型维度
self.input_net = nn.Sequential(
nn.Dropout(self.hparams.input_dropout),
nn.Linear(self.hparams.input_dim, self.hparams.model_dim)
)
# 序列的位置编码
self.positional_encoding = PositionalEncoding(d_model=self.hparams.model_dim)
# Transformer 编码器
self.transformer = TransformerEncoder(num_layers=self.hparams.num_layers,
input_dim=self.hparams.model_dim,
dim_feedforward=2*self.hparams.model_dim,
num_heads=self.hparams.num_heads,
dropout=self.hparams.dropout)
# 每个序列元素的输出分类器
self.output_net = nn.Sequential(
nn.Linear(self.hparams.model_dim, self.hparams.model_dim),
nn.LayerNorm(self.hparams.model_dim),
nn.ReLU(inplace=True),
nn.Dropout(self.hparams.dropout),
nn.Linear(self.hparams.model_dim, self.hparams.num_classes)
)
def forward(self, x, mask=None, add_positional_encoding=True):
"""
参数：
x - 输入特征，形状为 [Batch, SeqLen, input_dim]
mask - 可选的掩码，应用于注意力输出
add_positional_encoding - 如果为 True，则将位置编码添加到输入中
"""
x = self.input_net(x)
if add_positional_encoding:
x = self.positional_encoding(x)
x = self.transformer(x, mask=mask)
x = self.output_net(x)
return x
@torch.no_grad()
def get_attention_maps(self, x, mask=None, add_positional_encoding=True):
"""
提取整个 Transformer 对于单个批次的注意力矩阵的函数。
输入参数与前向传递相同。
"""
x = self.input_net(x)
if add_positional_encoding:
x = self.positional_encoding(x)
attention_maps = self.transformer.get_attention_maps(x, mask=mask)
return attention_maps
def configure_optimizers(self):
optimizer = optim.Adam(self.parameters(), lr=self.hparams.lr)
# 每步应用学习率调度器
lr_scheduler = CosineWarmupScheduler(optimizer,
warmup=self.hparams.warmup,
max_iters=self.hparams.max_iters)
return [optimizer], [{'scheduler': lr_scheduler, 'interval': 'step'}]
def training_step(self, batch, batch_idx):
raise NotImplementedError
def validation_step(self, batch, batch_idx):
raise NotImplementedError
def test_step(self, batch, batch_idx):
raise NotImplementedError
3.实验
在完成了 Transformer 架构的实现之后，我们可以开始进行实验，并将该架构应用于不同的任务。在本文中，我们将重点关注两个任务：并行序列到序列（Sequence-to-Sequence）任务和集合异常检测。这两个任务分别侧重于 Transformer 架构的不同特性，下面我们将逐一进行讨论。
3.1.序列到序列任务
序列到序列任务是一种输入和输出都是序列的任务，这两个序列的长度不一定相同。在这个领域中，一些流行的任务包括机器翻译和文本摘要。通常，我们会使用一个 Transformer 编码器来解释输入序列，并使用一个解码器以自回归的方式生成输出。然而，在这里，我们将回到一个更为简单的示例任务，并仅使用编码器。给定一个包含 N个介于 0和 M之间的数字的序列，任务的目标是反转输入序列。在 Numpy 符号表示中，如果我们的输入是 ，那么输出应该是 。尽管这个任务听起来非常简单，但循环神经网络（RNN）在处理这类任务时可能会遇到问题，因为它们需要处理长期依赖关系。Transformer 被设计用来支持这种长期依赖关系，因此我们预期它将表现得很好。
首先，让我们在下面创建一个数据集类。
class ReverseDataset(data.Dataset):
def __init__(self, num_categories, seq_len, size):
super().__init__()
self.num_categories = num_categories
self.seq_len = seq_len
self.size = size
# 创建一个随机序列的数据集
self.data = torch.randint(self.num_categories, size=(self.size, self.seq_len))
def __len__(self):
# 返回数据集的大小
return self.size
def __getitem__(self, idx):
# 获取单个样本
inp_data = self.data[idx]
labels = torch.flip(inp_data, dims=(0,)) # 反转标签
return inp_data, labels
我们创建了一系列随机数序列，这些数字介于 0 和 num_categories-1 之间。标签简单地是序列维度翻转的张量。我们可以在下面创建相应的数据加载器。
# 创建数据集函数
dataset = partial(ReverseDataset, 10, 16)
# 创建训练数据加载器
train_loader = data.DataLoader(dataset(50000), batch_size=128, shuffle=True, drop_last=True, pin_memory=True)
# 创建验证数据加载器
val_loader = data.DataLoader(dataset(1000), batch_size=128)
# 创建测试数据加载器
test_loader = data.DataLoader(dataset(10000), batch_size=128)
让我们查看数据集的一个随机样本：
inp_data, labels = train_loader.dataset[0]
print("输入数据:", inp_data)
print("标签:", labels)
在训练过程中，我们将输入序列通过 Transformer 编码器传递，并预测每个输入标记的输出。我们使用标准的交叉熵损失来进行预测。每个数字都表示为一个独热向量。需要注意的是，将类别表示为单个标量会极大地降低模型的表达能力，因为在我们的例子中 0 和 1并不比 0 和 9更接近。使用学习到的嵌入向量是一个替代方案，这可以通过 PyTorch 模块 nn.Embedding 提供。然而，在我们的案例中，使用独热向量和额外的线性层（self.input_net 将独热向量映射到一个密集向量，其中权重矩阵的每一行代表特定类别的嵌入）的效果与嵌入层相同。
为了实现训练动态，我们创建一个新的类，继承自 TransformerPredictor 并重写训练、验证和测试步骤函数。
class ReversePredictor(TransformerPredictor):
def _calculate_loss(self, batch, mode="train"):
# 获取数据并将类别转换为独热向量
inp_data, labels = batch
inp_data = F.one_hot(inp_data, num_classes=self.hparams.num_classes).float()
# 进行预测并计算损失和准确率
preds = self.forward(inp_data, add_positional_encoding=True)
loss = F.cross_entropy(preds.view(-1, preds.size(-1)), labels.view(-1))
acc = (preds.argmax(dim=-1) == labels).float().mean()
# 日志记录
self.log(f"{mode}_loss", loss)
self.log(f"{mode}_acc", acc)
return loss, acc
def training_step(self, batch, batch_idx):
loss, _ = self._calculate_loss(batch, mode="train")
return loss
def validation_step(self, batch, batch_idx):
_ = self._calculate_loss(batch, mode="val")
def test_step(self, batch, batch_idx):
_ = self._calculate_loss(batch, mode="test")
最后，我们可以创建一个训练函数，类似于我们在 PyTorch Lightning 教程5中看到的。我们创建一个 pl.Trainer 对象，运行 N个周期，记录在 TensorBoard 中，并根据验证结果保存我们的最佳模型。之后，我们在测试集上测试我们的模型。我们在这里传递给训练器的一个额外参数是 gradient_clip_val。这在优化器步骤之前剪切所有参数的梯度范数，防止模型在获得非常高的梯度时发散，例如在陡峭的损失表面上。对于 Transformer，梯度裁剪可以在前几次迭代中进一步稳定训练，并且之后也有帮助。在纯 PyTorch 中，您可以通过 torch.nn.utils.clip_grad_norm_(...) 应用梯度裁剪。
def train_reverse(**kwargs):
# 创建一个带有生成回调的 PyTorch Lightning 训练器
root_dir = os.path.join(CHECKPOINT_PATH, "ReverseTask")
os.makedirs(root_dir, exist_ok=True)
trainer = pl.Trainer(default_root_dir=root_dir,
callbacks=[ModelCheckpoint(save_weights_only=True, mode="max", monitor="val_acc")],
accelerator="gpu" if str(device).startswith("cuda") else "cpu",
devices=1,
max_epochs=10,
gradient_clip_val=5) # 设置梯度裁剪值
trainer.logger._default_hp_metric = None # 可选的日志记录参数，我们不需要
# 检查是否存在预训练模型。如果存在，加载它并跳过训练
pretrained_filename = os.path.join(CHECKPOINT_PATH, "ReverseTask.ckpt")
if os.path.isfile(pretrained_filename):
print("找到预训练模型，正在加载...")
model = ReversePredictor.load_from_checkpoint(pretrained_filename)
else:
model = ReversePredictor(max_iters=trainer.max_epochs*len(train_loader), **kwargs)
trainer.fit(model, train_loader, val_loader)
# 在验证和测试集上测试最佳模型
val_result = trainer.test(model, val_loader, verbose=False)
test_result = trainer.test(model, test_loader, verbose=False)
result = {"test_acc": test_result[0]["test_acc"], "val_acc": val_result[0]["test_acc"]}
model = model.to(device)
return model, result
最终，我们可以训练模型。在这个设置中，我们将使用单个编码器块和多头注意力中的单个头。之所以选择这样做，是因为任务的简单性，并且在这种情况下，注意力实际上可以被解释为预测的“解释”。
reverse_model, reverse_result = train_reverse(input_dim=train_loader.dataset.num_categories,
model_dim=32,
num_heads=1,
num_classes=train_loader.dataset.num_categories,
num_layers=1,
dropout=0.0,
lr=5e-4,
warmup=50)
我们可以忽略 PyTorch Lightning 关于工作线程数量的警告。由于数据集非常简单，并且 __getitem__ 完成的时间可以忽略不计，我们不需要子进程来为我们提供数据。实际上，更多的工作线程可能会减慢训练速度，因为进程/线程之间的通信开销。首先，让我们打印结果：
print(f"验证准确率: {(100.0 * reverse_result['val_acc']):4.2f}%")
print(f"测试准确率: {(100.0 * reverse_result['test_acc']):4.2f}%")
正如我们所预期的，Transformer 能够正确解决任务。但是，多头注意力块中的注意力对于任意输入看起来如何呢？让我们在下面尝试可视化它。
data_input, labels = next(iter(val_loader))
inp_data = F.one_hot(data_input, num_classes=reverse_model.hparams.num_classes).float()
inp_data = inp_data.to(device)
attention_maps = reverse_model.get_attention_maps(inp_data)
对象 attention_maps 是一个长度为 N 的列表，其中 N是层数。每个元素是一个形状为 [Batch, Heads, SeqLen, SeqLen] 的张量，我们可以在下面验证。
attention_maps[0].shape
接下来，我们将编写一个绘图函数，该函数接受输入序列、注意力图和指示我们想要可视化哪个批次元素的索引作为输入。我们将创建一个图表，其中行表示不同的层，列显示不同的头。记住，每个行的 softmax 已经被单独应用。
def plot_attention_maps(input_data, attn_maps, idx=0):
if input_data is not None:
input_data = input_data[idx].detach().cpu().numpy()
else:
input_data = np.arange(attn_maps[0][idx].shape[-1])
attn_maps = [m[idx].detach().cpu().numpy() for m in attn_maps]
num_heads = attn_maps[0].shape[0]
num_layers = len(attn_maps)
seq_len = input_data.shape[0]
fig_size = 4 if num_heads == 1 else 3
fig, ax = plt.subplots(num_layers, num_heads, figsize=(num_heads*fig_size, num_layers*fig_size))
if num_layers == 1:
ax = [ax]
if num_heads == 1:
ax = [[a] for a in ax]
for row in range(num_layers):
for column in range(num_heads):
ax[row][column].imshow(attn_maps[row][column], origin='lower', vmin=0)
ax[row][column].set_xticks(list(range(seq_len)))
ax[row][column].set_xticklabels(input_data.tolist())
ax[row][column].set_yticks(list(range(seq_len)))
ax[row][column].set_yticklabels(input_data.tolist())
ax[row][column].set_title(f"层 {row+1}, 头 {column+1}")
fig.subplots_adjust(hspace=0.5)
plt.show()
最后，我们可以绘制我们训练的 Transformer 在反转任务上的注意力图：
plot_attention_maps(data_input, attention_maps, idx=0)
模型学会了关注自身翻转索引上的标记。因此，它实际上做了我们预期它做的事情。我们可以看到，它也关注了翻转索引附近的值。这是因为模型不需要完美的、硬性的注意力来解决这个问题，而是可以容忍这种近似的、嘈杂的注意力图。附近的索引是由位置编码的相似性引起的，这也是我们预期的位置编码的效果。
3.2.集合异常检测
集合是许多应用中常用的另一种数据结构，与序列不同，集合中的元素是无序的。循环神经网络（RNN）在应用于集合时，通常需要假设数据中存在某种顺序，但这可能会使模型偏向于数据中并不存在的一种顺序。Vinyals等人（2015）的研究表明，这种假定的顺序可能会显著影响模型的性能。因此，我们应尽量避免在集合上使用RNN。理想情况下，我们的模型应对排列是等变的，即无论我们如何对集合中的元素进行排序，输出结果都应保持不变。
Transformer模型因其多头注意力机制的排列等变性，提供了完美的解决方案。无论输入的顺序如何，它都能输出相同的值（输入和输出的排列是相等的）。本文中，我们将探讨集合上的一个特定任务——集合异常检测。该任务的目标是找出集合中与其他元素不一致的元素。在研究领域中，异常检测的常见应用是在图像集合上进行的，其中N-1张图像属于同一类别或具有相同的高级特征，而另一张图像则属于另一个类别。需要注意的是，这里的“类别”不一定与标准分类问题中的类别相关，而可能涉及多个特征的组合。例如，在人脸数据集中，这可能涉及到戴眼镜的人、男性、有胡须等特征。下面是一个区分不同动物的示例。前四幅图像展示了狐狸，而最后一幅则展示了另一种动物。我们的目标是识别出最后一幅图像是不同的动物，但具体属于哪一类动物并不重要。
在本文中，我们将使用CIFAR100数据集。CIFAR100包含600张图像，每类100张，分辨率为32x32，与CIFAR10类似。更多的类别数量要求模型关注图像中的特定特征，而不是像CIFAR10中的粗略特征，因此使任务更加困难。我们将向模型展示一组9张同一类别的图像和1张另一类别的图像。任务是找出与其他图像不同类别的图像。
直接将原始图像作为Transformer的输入并不是一个好主意，因为它不像卷积神经网络（CNN）那样具有平移不变性，首先需要学习从高维输入中检测图像特征。相反，我们将使用torchvision包中的预训练ResNet34模型来获取图像的高级、低维特征。ResNet模型已经在包含100万张图像的1000个类别和不同分辨率的ImageNet数据集上进行了预训练。然而，在训练和测试期间，图像通常被缩放到224x224的分辨率，因此我们也相应地将CIFAR图像缩放到这个分辨率。以下，我们将加载数据集，并准备数据以供ResNet模型处理。
# ImageNet 统计数据，用于图像预处理时的均值和标准差
DATA_MEANS = np.array([0.485, 0.456, 0.406]) # ImageNet 数据集的均值
DATA_STD = np.array([0.229, 0.224, 0.225]) # ImageNet 数据集的标准差
# 将 numpy 数组转换为 torch 张量，以便在后续步骤中使用
TORCH_DATA_MEANS = torch.from_numpy(DATA_MEANS).view(1,3,1,1) # 转换均值
TORCH_DATA_STD = torch.from_numpy(DATA_STD).view(1,3,1,1) # 转换标准差
# 定义图像转换操作，包括调整大小到 224x224 像素，并使用 ImageNet 的统计数据进行归一化
transform = transforms.Compose([
transforms.Resize((224,224)), # 调整图像大小
transforms.ToTensor(), # 将图像转换为张量
transforms.Normalize(DATA_MEANS, DATA_STD) # 归一化处理
])
# 加载训练数据集，使用 CIFAR100 数据集，并应用上面定义的转换操作
train_set = CIFAR100(root=DATASET_PATH, train=True, transform=transform, download=True)
# 加载测试数据集，同样使用 CIFAR100 数据集，并应用转换操作
test_set = CIFAR100(root=DATASET_PATH, train=False, transform=transform, download=True)
接下来，我们希望在图像上运行预训练的ResNet模型，并在分类层之前提取特征。这些是最高级的特征，应该足够描述图像。CIFAR100与ImageNet有一些相似之处，因此我们不会以任何形式重新训练ResNet模型。然而，如果你想获得最佳性能并且拥有一个非常大的数据集，最好在训练期间将ResNet添加到计算图中并微调其参数。由于我们没有足够大的数据集，并且希望高效地训练我们的模型，我们将事先提取特征。让我们在下面加载并准备模型。
import os # 导入os模块，用于操作环境变量
# 设置环境变量TORCH_HOME为模型检查点路径
os.environ["TORCH_HOME"] = CHECKPOINT_PATH
# 加载预训练的ResNet34模型，使用ImageNet数据集上预训练的权重
pretrained_model = torchvision.models.resnet34(weights='IMAGENET1K_V1')
# 移除模型的分类层，以便只使用模型的特征提取部分
# 在不同的模型中，分类层可能被称为"fc"或"classifier"
# 将它们都设置为一个空的序列模型，这样就不会对特征进行任何改变
pretrained_model.fc = nn.Sequential()
pretrained_model.classifier = nn.Sequential()
# 将模型转移到GPU上进行加速
pretrained_model = pretrained_model.to(device)
# 设置模型为评估模式，这样在模型中不会计算梯度
pretrained_model.eval()
# 遍历模型的所有参数，将requires_grad设置为False，这样在前向传播中不会计算梯度
for p in pretrained_model.parameters():
p.requires_grad = False
现在，我们将编写一个提取特征的函数。这个单元需要访问GPU，因为模型相当深，图像相对较大。GoogleColab上的GPU足够，但运行这个单元可能需要2-3分钟。一旦运行，特征将被导出到磁盘，这样每次运行笔记本时就不必重新计算。然而，这需要>150MB的空闲磁盘空间。因此，如果有足够的空闲磁盘和GPU，建议只在本地计算机上运行此单元（GoogleColab对此也适用）。如果您没有GPU，可以从[GoogleDrive文件夹]下载特征。
@torch.no_grad() # 装饰器，指示PyTorch在该函数内不计算梯度
def extract_features(dataset, save_file):
# 检查保存的特征文件是否存在，如果不存在，则提取特征并保存
if not os.path.isfile(save_file):
data_loader = data.DataLoader(dataset, batch_size=128, shuffle=False, drop_last=False, num_workers=4)
extracted_features = [] # 用于存储提取的特征
for imgs, _ in tqdm(data_loader): # 遍历数据加载器
imgs = imgs.to(device) # 将图像数据移动到设备（GPU或CPU）
feats = pretrained_model(imgs) # 使用预训练模型提取特征
extracted_features.append(feats) # 将提取的特征添加到列表
extracted_features = torch.cat(extracted_features, dim=0) # 将特征按维度0（批次维度）连接起来
extracted_features = extracted_features.detach().cpu() # 将特征从GPU移动到CPU，并脱离计算图
torch.save(extracted_features, save_file) # 将提取的特征保存到文件
else:
extracted_features = torch.load(save_file) # 如果特征文件已存在，则直接加载
return extracted_features # 返回提取的特征
# 定义训练集特征文件的路径
train_feat_file = os.path.join(CHECKPOINT_PATH, "train_set_features.tar")
# 提取训练集的特征并保存
train_set_feats = extract_features(train_set, train_feat_file)
# 定义测试集特征文件的路径
test_feat_file = os.path.join(CHECKPOINT_PATH, "test_set_features.tar")
# 提取测试集的特征并保存
test_feats = extract_features(test_set, test_feat_file)
让我们在下面验证特征的形状。训练集应该有50k个元素，测试集有10k个图像。ResNet34的特征维度是512。如果您尝试其他模型，可能会看到不同的特征维度。
print("Train:", train_set_feats.shape)
print("Test: ", test_feats.shape)
像往常一样，我们希望创建一个验证集，以检测我们何时应该停止训练。在这种情况下，我们将训练集分成90%的训练集和10%的验证集。然而，这里的难度在于我们需要确保验证集对所有100个标签都有相同数量的图像。否则，我们会有类别不平衡，这对于创建图像集是不利的。因此，我们每个类别取10%的图像，并将它们移动到验证集中。下面的代码正是这样做的。
# 将训练集拆分为训练集和验证集
# 从训练集中获取标签
labels = train_set.targets
# 将标签转换为torch.LongTensor类型，并获取类别数
labels = torch.LongTensor(labels)
num_labels = labels.max() + 1
# 对每个类别的标签进行排序，并获取每个类别图像的索引
sorted_indices = torch.argsort(labels).reshape(num_labels, -1) # 形状为 [类别, 每个类别的图像数]
# 计算每个类别用于验证集的图像数量，这里取10%作为验证集
num_val_exmps = sorted_indices.shape[1] // 10
# 根据计算出的验证集图像数量，获取验证集和训练集的图像索引
val_indices = sorted_indices[:, :num_val_exmps].reshape(-1) # 验证集索引
train_indices = sorted_indices[:, num_val_exmps:].reshape(-1) # 训练集索引
# 根据索引获取对应的图像特征和标签，形成训练集和验证集的特征数据
train_feats, train_labels = train_set_feats[train_indices], labels[train_indices]
val_feats, val_labels = train_set_feats[val_indices], labels[val_indices]
现在我们可以为集合异常任务准备一个数据集类。我们定义一个时期是每个图像恰好一次作为“异常”的序列。因此，数据集的长度是其中的图像数量。对于训练集，每次我们使用__getitem__访问一个项目时，我们会随机采样一个与相应索引idx的图像不同的类别。在第二步中，我们从这个采样的类别中采样N-1个图像。最后返回10个图像的集合。__getitem__中的随机性允许我们在每次迭代中看到稍微不同的集合。然而，我们不能对测试集使用相同的策略，因为我们希望测试数据集在每次迭代时都是相同的。因此，我们在__init__方法中采样集合，并在__getitem__中返回它们。下面的代码实现了这种动态。
class SetAnomalyDataset(data.Dataset):
def __init__(self, img_feats, labels, set_size=10, train=True):
"""
初始化集合异常数据集
参数：
img_feats - 形状为 [num_imgs, img_dim] 的张量，表示图像的高级特征。
labels - 形状为 [num_imgs] 的张量，包含图像的类别标签。
set_size - 集合中的元素数量，其中N-1个元素来自同一个类别，1个元素来自另一个类别。
train - 布尔值，如果为True，则每次调用__getitem__时都会生成一个新的集合。
"""
super().__init__()
self.img_feats = img_feats
self.labels = labels
self.set_size = set_size - 1 # 集合大小减1，因为最后一个元素是异常
self.train = train
# 获取每个类别的图像索引
self.num_labels = labels.max() + 1
self.img_idx_by_label = torch.argsort(self.labels).reshape(self.num_labels, -1)
if not train:
self.test_sets = self._create_test_sets() # 为测试集预先创建集合
def _create_test_sets(self):
"""
为测试集预先生成每个图像的集合
"""
test_sets = []
num_imgs = self.img_feats.shape[0]
np.random.seed(42) # 固定随机种子以保证结果可复现
test_sets = [self.sample_img_set(self.labels[idx]) for idx in range(num_imgs)]
test_sets = torch.stack(test_sets, dim=0)
return test_sets
def sample_img_set(self, anomaly_label):
"""
给定异常标签，采样一个新的图像集合，集合中的图像来自与异常标签不同的类别
"""
# 随机选择一个与异常标签不同的类别
set_label = np.random.randint(self.num_labels - 1)
if set_label >= anomaly_label:
set_label += 1
# 从选定的类别中随机采样图像索引
img_indices = np.random.choice(self.img_idx_by_label.shape[1], size=self.set_size, replace=False)
img_indices = self.img_idx_by_label[set_label, img_indices]
return img_indices
def __len__(self):
"""
返回数据集的长度
"""
return self.img_feats.shape[0]
def __getitem__(self, idx):
"""
获取数据集中的一个样本
参数：
idx - 样本的索引
返回：
img_set - 包含集合中所有图像的特征
indices - 图像的原始索引
label - 异常图像的标签（集合中最后一个图像的索引）
"""
anomaly = self.img_feats[idx]
if self.train: # 如果是训练模式，则随机采样图像集合
img_indices = self.sample_img_set(self.labels[idx])
else: # 如果是测试模式，则使用预先生成的集合
img_indices = self.test_sets[idx]
# 将图像特征连接成一个集合，异常图像总是最后一个
img_set = torch.cat([self.img_feats[img_indices], anomaly[None]], dim=0)
indices = torch.cat([img_indices, torch.LongTensor([idx])], dim=0)
label = img_set.shape[0] - 1 # 异常图像的标签
return img_set, indices, label
接下来，我们可以在下面设置我们的数据集和数据加载器。这里，我们将使用集合大小为10，即9张来自一个类别的图像+1个异常图像。如果你想要尝试不同的大小，可以随意更改。
SET_SIZE = 10 # 定义集合中图像的数量，包括异常图像
test_labels = torch.LongTensor(test_set.targets) # 将测试集的标签转换为长整型张量
# 创建训练、验证和测试的数据集对象
train_anom_dataset = SetAnomalyDataset(train_feats, train_labels, set_size=SET_SIZE, train=True)
val_anom_dataset = SetAnomalyDataset(val_feats, val_labels, set_size=SET_SIZE, train=False)
test_anom_dataset = SetAnomalyDataset(test_feats, test_labels, set_size=SET_SIZE, train=False)
# 创建训练、验证和测试的数据加载器
train_anom_loader = data.DataLoader(train_anom_dataset, batch_size=64, shuffle=True, drop_last=True, num_workers=4, pin_memory=True)
val_anom_loader = data.DataLoader(val_anom_dataset, batch_size=64, shuffle=False, drop_last=False, num_workers=4)
test_anom_loader = data.DataLoader(test_anom_dataset, batch_size=64, shuffle=False, drop_last=False, num_workers=4)
为了更好地理解数据集，我们可以在下面绘制一些来自测试数据集的集合。每一行显示一个不同的输入集合，其中前9个来自同一类别。
def visualize_exmp(indices, orig_dataset):
"""
函数：可视化图像集合中的异常示例
参数：
indices - 一个包含图像索引的张量，这些图像将被用于可视化。
orig_dataset - 原始数据集对象，包含图像数据。
"""
images = [orig_dataset[idx][0] for idx in indices.reshape(-1)] # 获取索引对应的图像
images = torch.stack(images, dim=0) # 将图像堆叠成一个批次
images = images * TORCH_DATA_STD + TORCH_DATA_MEANS # 对图像进行标准化处理
# 创建一个图像网格，每行显示 SET_SIZE 张图像
img_grid = torchvision.utils.make_grid(images, nrow=SET_SIZE, normalize=True, pad_value=0.5, padding=16)
img_grid = img_grid.permute(1, 2, 0) # 调整维度顺序，以便于显示
plt.figure(figsize=(12,8)) # 设置图像显示窗口的大小
plt.title("CIFAR100上的异常示例") # 设置窗口标题
plt.imshow(img_grid) # 显示图像网格
plt.axis('off') # 关闭坐标轴
plt.show() # 显示图像
plt.close() # 关闭图像显示窗口
# 获取测试数据加载器中的一个批次的图像索引和标签
_, indices, _ = next(iter(test_anom_loader))
# 调用 visualize_exmp 函数，可视化前4个异常示例
visualize_exmp(indices[:4], test_set)
我们已经可以看到，对于一些集合来说，任务可能比其他的容易。如果异常图像属于不同的但视觉上相似的类别（例如火车与公共汽车，面粉与蠕虫等），可能会特别困难。
在准备了数据之后，我们可以更仔细地看看模型。这里，我们有一个整个集合的分类。为了使预测是排列等变的，我们将为每个图像输出一个logit。在这些logit上，我们应用softmax，并训练异常图像具有最高得分/概率。这与标准分类层有点不同，因为softmax是应用于图像，而不是传统意义上的输出类别。然而，如果我们交换两个图像在它们的位置，我们就会有效地交换它们在输出softmax中的位置。因此，预测与输入是等变的。我们在下面实现了这个想法的Transformer Lightning模块的子类。
class AnomalyPredictor(TransformerPredictor):
"""
类：AnomalyPredictor
继承自TransformerPredictor，用于集合异常检测任务的模型。
"""
def _calculate_loss(self, batch, mode="train"):
"""
方法：计算损失并记录指标
参数：
batch - 包含图像集合、标签的批次数据
mode - 指明当前是训练、验证还是测试阶段
返回：
loss - 计算得到的损失值
acc - 计算得到的准确率
"""
img_sets, _, labels = batch
preds = self.forward(img_sets, add_positional_encoding=False) # 预测时不使用位置编码
preds = preds.squeeze(dim=-1) # 预测结果的形状为 [Batch_size, set_size]
loss = F.cross_entropy(preds, labels) # 计算交叉熵损失
acc = (preds.argmax(dim=-1) == labels).float().mean() # 计算准确率
self.log(f"{mode}_loss", loss) # 记录损失
self.log(f"{mode}_acc", acc, on_step=False, on_epoch=True) # 记录准确率
return loss, acc
def training_step(self, batch, batch_idx):
"""
方法：训练步骤
参数：
batch - 当前批次的数据
batch_idx - 当前批次的索引
返回：
loss - 训练步骤的损失值
"""
loss, _ = self._calculate_loss(batch, mode="train")
return loss
def validation_step(self, batch, batch_idx):
"""
方法：验证步骤
参数：
batch - 当前批次的数据
batch_idx - 当前批次的索引
"""
self._calculate_loss(batch, mode="val") # 在验证步骤中计算损失和准确率，但不返回
def test_step(self, batch, batch_idx):
"""
方法：测试步骤
参数：
batch - 当前批次的数据
batch_idx - 当前批次的索引
"""
self._calculate_loss(batch, mode="test") # 在测试步骤中计算损失和准确率，但不返回
最后，我们在下面编写我们的训练函数。它的结构与反向任务完全相同，因此不需要太多解释。
def train_anomaly(**kwargs):
"""
函数：训练集合异常检测模型
参数：
**kwargs - 其他关键字参数，用于传递给模型的构造函数
返回：
model - 训练完成的模型
result - 训练、验证和测试的结果
"""
# 创建一个带有生成回调的 PyTorch Lightning 训练器
root_dir = os.path.join(CHECKPOINT_PATH, "SetAnomalyTask") # 定义训练日志和模型检查点的存储路径
os.makedirs(root_dir, exist_ok=True) # 如果路径不存在，则创建
trainer = pl.Trainer(default_root_dir=root_dir, # 使用默认的根目录
callbacks=[ModelCheckpoint(save_weights_only=True, mode="max", monitor="val_acc")], # 回调函数，保存验证集准确率最高的模型
accelerator="gpu" if str(device).startswith("cuda") else "cpu", # 使用 GPU 或 CPU
devices=1, # 使用单个设备进行训练
max_epochs=100, # 最大训练周期数
gradient_clip_val=2) # 梯度裁剪值，防止梯度爆炸
trainer.logger._default_hp_metric = None # 不需要默认的超参数日志记录
# 检查是否存在预训练模型。如果存在，加载它并跳过训练
pretrained_filename = os.path.join(CHECKPOINT_PATH, "SetAnomalyTask.ckpt")
if os.path.isfile(pretrained_filename):
print("找到预训练模型，正在加载...")
model = AnomalyPredictor.load_from_checkpoint(pretrained_filename) # 加载预训练模型
else:
model = AnomalyPredictor(max_iters=trainer.max_epochs*len(train_anom_loader), **kwargs) # 创建模型实例
trainer.fit(model, train_anom_loader, val_anom_loader) # 训练模型
model = AnomalyPredictor.load_from_checkpoint(trainer.checkpoint_callback.best_model_path) # 加载最佳模型
# 在验证集和测试集上测试最佳模型
train_result = trainer.test(model, train_anom_loader, verbose=False)
val_result = trainer.test(model, val_anom_loader, verbose=False)
test_result = trainer.test(model, test_anom_loader, verbose=False)
result = {"test_acc": test_result[0]["test_acc"], "val_acc": val_result[0]["test_acc"], "train_acc": train_result[0]["test_acc"]} # 收集结果
model = model.to(device) # 将模型移动到设备（GPU或CPU）
return model, result
现在，我们终于要训练我们的模型了。我们将构建一个包含4层的结构，每层都配备4个注意力头。模型的隐藏层维度被设定为256，并且在整个模型中应用了0.1的dropout率来增强正则化效果。此外，值得注意的是，我们还在输入特征上应用了dropout，这样做可以使得模型对图像中的噪声更加健壮，从而提升其泛化能力。为了平稳地启动模型训练，我们还将采用warmup策略。
通过调用train_anomaly函数，并传入包括输入维度（基于训练异常数据集的图像特征形状）、模型维度、注意力头数、类别数（这里假设为1，因为可能是二分类问题的一个特例）、层数、dropout比率、输入dropout比率、学习率和warmup步数等参数，我们将开始训练模型，并获取训练结果。
训练完成后，我们可以打印出模型在训练集、验证集和测试集上分别达到的准确率，以评估其性能。
print(f"训练准确率: {(100.0*anomaly_result['train_acc']):4.2f}%")
print(f"验证准确率: {(100.0*anomaly_result['val_acc']):4.2f}%")
print(f"测试准确率: {(100.0*anomaly_result['test_acc']):4.2f}%")
这三行输出将直观地展示模型在不同数据集上的表现，帮助我们了解模型的训练效果和泛化能力。
这个模型在验证和测试阶段展现出了高达约94%的准确率，显示出很强的泛化能力。但值得注意的是，由于不同计算机/设备之间的差异，以及不同平台和NumPy版本间种子设置的不一致性，您运行此代码时可能会看到略有不同的分数。尽管如此，我们可以肯定地说，该模型在大多数情况下都能很好地完成任务。
# 从测试数据加载器中获取一批数据，包括输入数据、索引和标签
inp_data, indices, labels = next(iter(test_anom_loader))
# 将输入数据转移到指定的设备（如GPU）上
inp_data = inp_data.to(device)
# 将模型设置为评估模式
anomaly_model.eval()
# 关闭梯度计算，因为我们只进行前向传播
with torch.no_grad():
# 通过模型进行前向传播，注意这里不添加位置编码（如果模型设计中包含位置编码的话）
preds = anomaly_model.forward(inp_data, add_positional_encoding=False)
# 对输出进行softmax操作，以获取概率分布，并去除最后一个维度（假设是多分类问题中的类别维度）
preds = F.softmax(preds.squeeze(dim=-1), dim=-1)
# 对输入数据进行随机置换
permut = np.random.permutation(inp_data.shape[1]) # 生成一个置换索引数组
perm_inp_data = inp_data[:, permut] # 根据置换索引置换输入数据
# 对置换后的输入数据进行前向传播
perm_preds = anomaly_model.forward(perm_inp_data, add_positional_encoding=False)
# 对置换后的输出也进行softmax操作
perm_preds = F.softmax(perm_preds.squeeze(dim=-1), dim=-1)
# 验证模型是否满足置换等变性，即置换输入数据后，输出预测结果也应该相应地置换
# 通过比较置换前后的预测结果是否几乎一致来验证
assert (preds[:, permut] - perm_preds).abs().max() < 1e-5, "预测结果不满足置换等变性"
# 打印第一个样本的原始预测结果和置换后的预测结果，验证它们的一致性
print("原始预测:\n", preds[0, permut].cpu().numpy()) # 注意这里使用permut是为了与perm_preds对应
print("置换后的预测:\n", perm_preds[0].cpu().numpy()) # 无需置换索引，因为perm_preds已经是置换后的结果
在尝试对模型进行解读之前，我们先来验证一下它是否具备置换等变性，即对于输入数据的不同排列组合，模型能否给出相同的预测结果。为此，我们从测试集中随机抽取一批数据，通过模型得到预测概率。
接下来，我们对这批输入数据进行随机置换，再次通过模型得到置换后的预测概率。然后，我们将原始预测结果按照相同的置换顺序进行调整，并与置换后的预测结果进行对比。通过断言检查两者之间的最大绝对差值是否小于一个非常小的数（如1e-5），来验证模型是否满足置换等变性。
打印出原始预测和置换后调整过的预测的第一个样本，您会发现它们几乎是完全一致的，只有因为网络内部操作的微小数值差异而略有不同。
为了更深入地理解模型的工作原理，我们可以绘制模型内部的注意力图。这将帮助我们直观地看到模型在不同图像之间是如何共享和传递信息的，以及每个注意力头可能代表的特定信息或特征。为了简化分析，我们首先需要为上述测试批次数据提取注意力图，并确定模型的离散预测结果。
正如您所见，预测结果几乎完全相同，只有由于网络操作中的微小数值差异而略有不同。为了更深入地理解模型，我们可以绘制模型内部的注意力图。这将帮助我们了解模型在图像之间共享/传递了哪些信息，以及每个注意力头可能代表什么。首先，我们需要为上述测试批次提取注意力图，并为了简化起见，确定离散预测。
# 提取注意力图
attention_maps = anomaly_model.get_attention_maps(inp_data, add_positional_encoding=False)
# 确定离散预测
predictions = preds.argmax(dim=-1)
接下来，我们编写一个绘图函数，该函数将绘制输入集中的图像、模型的预测结果以及Transformer各层中不同注意力头的注意力图。您可以自由地探索不同输入示例的注意力图。
def visualize_prediction(idx):
# 可视化输入示例
visualize_exmp(indices[idx:idx+1], test_set)
# 打印预测结果
print("Prediction:", predictions[idx].item())
# 绘制注意力图
plot_attention_maps(input_data=None, attn_mapsattention=_maps, idx=idx)
# 调用函数可视化第一个预测结果
visualize_prediction(0)
请注意，根据随机种子的不同，您可能会看到略有不同的输入集。在网站上的版本中，我们将9张树木图像与一座火山进行比较。我们发现，例如，第2层的第1个和第3个注意力头以及第3层的第1个注意力头都集中在最后一张图像（火山）上。此外，第4层的所有注意力头都似乎忽略了最后一张图像，并为其分配了非常低的注意力概率。这表明模型确实已经识别出该图像与其他图像不符，因此预测其为异常。第3层的第2至第4个注意力头似乎对所有图像进行了稍微加权的平均，这可能表明模型正在提取所有图像的“平均”信息，以便与图像本身的特征进行比较。
为了找出模型实际出错的地方，我们可以识别出那些模型预测结果不是9的集合，因为在数据集中，我们确保异常总是出现在集合的最后一个位置。
首先，我们使用torch.where来找到预测结果不是9的索引，并将它们转换为CPU上的NumPy数组以便查看。
mistakes = torch.where(predictions != 9)[0].cpu().numpy()
print("Indices with mistake:", mistakes)
# 输出示例: Indices with mistake: [36 49]
由于我们的模型达到了约94%的准确率，所以在64个集合的批次中，我们只有很少的错误。不过，我们还是来可视化其中一个错误，比如最后一个错误：
visualize_prediction(mistakes[-1])
然后，我们还可以打印出该集合中每个图像被预测为异常的概率：
print("Probabilities:")
for i, p in enumerate(preds[mistakes[-1]].cpu().numpy()):
print(f"Image {i}: {100.0*p:4.2f}%")
这段代码将输出每个图像作为异常被预测的概率百分比。通过查看这些概率，我们可以更深入地了解模型为何会在这个特定的集合中出错。例如，如果某个非异常的图像具有非常高的异常概率，那么这可能表明模型在某些特征上被误导了，或者模型在区分异常和非异常图像时遇到了困难。
在这个案例中，模型将一棵棕榈树错误地识别为建筑物，给这张棕榈树图片赋予了约90%的异常概率，而真正的异常（即建筑物）却只得到了8%的概率。这主要是因为建筑物照片的拍摄角度与棕榈树相似，同时第2张棕榈树图片又呈现出与众不同的颜色，使得模型在此处出现了误判。但总的来说，模型的整体表现还是相当不错的。
4.结论
在本文中，我们深入了解了多头注意力层（Multi-Head Attention layer），它使用查询和键之间的缩放点积来找到输入元素之间的相关性和相似性。Transformer架构基于多头注意力层，并在类似于ResNet的块中应用了多个这样的层。Transformer是一种非常重要的近期架构，可以应用于许多任务和数据集。尽管它因在NLP领域的成功而广为人知，但其潜力远不止于此。我们已经看到了它在序列到序列任务和集合异常检测中的应用。如果我们不提供位置编码，它具有的置换等变性质（permutation-equivariant property）使其能够推广到许多场景。因此，了解该架构固然重要，但同样重要的是要了解其可能存在的问题，例如学习率预热（learning rate warm-up）解决的前几次迭代中的梯度问题。如果您对继续研究Transformer架构感兴趣，请查看本文开头列出的博客文章。