Transformers快速入门代码解析（一）：注意力机制——Attention：Scaled Dot-product Attention

引言

请注意！！！本博客使用了教程Transformers快速入门中的全部代码！！！

只在我个人理解的基础上为代码添加了注释！！！

详细教程请查看Transformers快速入门！！！

万分感谢！！！

自用！！！

Scaled Dot-product Attention代码

# torch.nn：从 PyTorch 中导入神经网络模块，用于定义嵌入层。
from torch import nn 
# transformers.AutoConfig：从 Hugging Face 的 transformers 库中导入 AutoConfig，用于加载预训练模型的配置信息（如词汇大小、隐藏层大小等）。
from transformers import AutoConfig 
# transformers.AutoTokenizer：用于加载与模型匹配的分词器
from transformers import AutoTokenizerimport torch
from math import sqrtimport torch.nn.functional as F# 设置模型检查点
# model_ckpt：指定 Hugging Face 模型的名称或路径，这里是 bert-base-uncased。
# bert-base-uncased 是一个预训练的 BERT 模型，大小适中，分词器会将所有文本转换为小写（uncased）。
# 你也可以选择其他模型，如 bert-large-uncased 或 distilbert-base-uncased。
model_ckpt = "bert-base-uncased"
# 加载分词器
# AutoTokenizer.from_pretrained()：根据 model_ckpt 加载与预训练模型对应的分词器。
# 分词器用于将输入的自然语言文本转换为模型可以处理的 token ID。
tokenizer = AutoTokenizer.from_pretrained(model_ckpt)# 对文本进行分词和编码
# text：输入的文本数据。
text = "time flies like an arrow"
# 将输入文本转换为 token ID，返回一个 PyTorch 张量（return_tensors="pt"）。
# add_special_tokens=False：不添加特殊 token（如 [CLS] 和 [SEP]）。如果为 True，分词器会自动在句首和句末添加这些特殊符号。
inputs = tokenizer(text, return_tensors="pt", add_special_tokens=False)
# 每个数字表示一个 token 的 ID，具体含义可通过 tokenizer.decode 转换回原始单词。
print("==================")
print("inputs的类型，来自transformers：")
print(type(inputs))
print("==================")
print("inputs的数据，来自transformers：")
print(inputs)
print("==================")
print("inputs.input_ids的数据，pytorch类型：")
print(inputs.input_ids)
print("==================")
print("inputs.input_ids的size：")
print(inputs.input_ids.size())
print("==================")# AutoConfig.from_pretrained()：
# 加载与 model_ckpt 对应的模型配置。
# config 包含模型的所有超参数（如隐藏层大小、词汇表大小、层数等）。
config = AutoConfig.from_pretrained(model_ckpt)
# 关键信息：
# config.vocab_size：模型的词汇表大小，这里为 30522（BERT 的词汇表大小）。
# config.hidden_size：每个 token 的嵌入维度，这里为 768（BERT 的默认隐藏层维度）。
# 使用 PyTorch 的 nn.Embedding 定义嵌入层。
# 输入大小：config.vocab_size，即 30522，表示词汇表中所有可能的 token。
# 输出大小：config.hidden_size，即 768，表示每个 token 的嵌入向量维度。
# 嵌入层的作用：
# 将输入的 token ID 转换为高维连续向量表示，这些向量包含了每个 token 的语义信息。
token_emb = nn.Embedding(config.vocab_size, config.hidden_size)
print("token数量以及其特征数量：")
print(token_emb)
print("==================")# inputs.input_ids：输入文本的 token ID，形状为 [1, 5]（1 表示批量大小，5 表示句子中有 5 个 token）。
# inputs_embeds：将 token ID 映射到嵌入向量，嵌入的形状为 [1, 5, 768]。
# =================================================================================================
# 这相当于inputs.input_ids是序列号，要在 [30522 768]的矩阵中找到对应的序列号
# token ID，形状为 [1, 5]表示1（一句话："time flies like an arrow"），5（这句话中有5个单词/token 的ID）
# =================================================================================================
# 1：表示批量大小。
# 5：表示句子中 token 的个数。
# 768：每个 token 的嵌入向量维度。
inputs_embeds = token_emb(inputs.input_ids)
print("给定句子的对应token及其特征：")
print(inputs_embeds)
print("==================")
print("给定句子的对应token及其特征的size：")
print(inputs_embeds.size())
print("==================")# 在自注意力机制中，每个输入向量都会被映射到三个表示：Query（查询），Key（键），Value（值）。
# Q, K, V 通常通过线性层计算，但在这段代码中，直接将 inputs_embeds 用作它们的值。-> [1, 5, 768]
Q = K = V = inputs_embeds
# 在注意力机制中，dim_k 是键向量的维度，用于归一化计算。
# 例如，如果 inputs_embeds 的形状是 [1, 5, 768]（批大小为 1，序列长度为 5，嵌入维度为 768），那么 dim_k = 768。
dim_k = K.size(-1)
# torch.bmm（批量矩阵乘法）:
#### torch.bmm 用于计算批次中的两个三维张量的矩阵乘法。
#### 输入的张量形状必须是 [batch_size, n, m] 和 [batch_size, m, p]，输出形状为 [batch_size, n, p]。
# =================================================================================================
# K.transpose(1,2):
#### K 的形状是 [batch_size, seq_length, embedding_dim] -> [1, 5, 768]。
#### 调用 transpose(1,2) 会交换第二维和第三维，得到形状 [batch_size, embedding_dim, seq_length] -> [1, 768, 5]。
#### 这一步是为了让 K 的嵌入维度与 Q 的序列长度对齐，便于计算矩阵乘法。
# =================================================================================================
# 矩阵乘法:
#### Q 和 K.transpose(1,2) 进行矩阵乘法。
######## Q 的形状是 [batch_size, seq_length, embedding_dim] -> [1, 5, 768]。
######## K.transpose(1,2) 的形状是 [batch_size, embedding_dim, seq_length] -> [1, 768, 5]。
######## 结果 scores 的形状是 [batch_size, seq_length, seq_length] -> [1, 5, 5]。
#### 含义：
######## scores[i, j] 表示输入序列中位置 i 的查询向量 Q[i] 与位置 j 的键向量 K[j] 的点积相似度。
######## 这种相似度用于表示位置 i 关注位置 j 的程度。
# =================================================================================================
# 归一化:
#### 结果被除以 sqrt(dim_k)
#### 这是为了防止嵌入维度较大时点积值过大，导致梯度消失或梯度爆炸的问题。
scores = torch.bmm(Q, K.transpose(1,2)) / sqrt(dim_k)
print("结果（上述text这个句子中，每个单词/token之间的点积相似度）：")
print(scores)
print("==================")
print("结果的size：")
print(scores.size())
print("==================")# 这段代码对前面计算得到的注意力得分矩阵 scores 进行 Softmax 操作，并验证 Softmax 操作后结果的性质（归一化）。以下是逐步的详细解析：
# 对 scores 应用 Softmax
# =================================================================================================
# F.softmax: 这是 PyTorch 中的 Softmax 函数，来自 torch.nn.functional 模块。
# =================================================================================================
# 输入参数:
#### scores 是一个三维张量，形状为 [batch_size, seq_length, seq_length] -> [1, 5, 5]，表示注意力得分矩阵。
#### dim=-1 指定在最后一维（seq_length 维）上应用 Softmax，即对每个查询位置的所有键进行归一化。
# =================================================================================================
# 结果:
#### weights 是一个形状与 scores 相同的张量 [batch_size, seq_length, seq_length] -> [1, 5, 5]，表示注意力权重。
#### 在 dim=-1 维上，每行的值表示一个概率分布，总和为 1。
weights = F.softmax(scores, dim=-1)
print("weights结果（上述text这个句子中，每个单词/token之间的点积相似度归一化之后的结果）")
print(weights)
print("==================")
print("weights结果加和，验证归一化：")
print(weights.sum(dim=-1))
print("==================")# 这段代码是实现 注意力机制（Attention Mechanism） 的最后一步，计算注意力输出（attn_outputs）。以下是逐行的详细解释：
# 背景：自注意力机制的核心步骤
# 在自注意力机制中，核心步骤包括：
# 1. 计算注意力得分（scores）。
# 2. 应用 Softmax 将注意力得分归一化为权重（weights）。
# 3. 将这些权重应用到值向量（V）上，得到注意力输出（attn_outputs）。
# 此代码正是在第 3 步中完成这个操作。
# weights: [batch_size, seq_length, seq_length] -> [1, 5, 5]
# V: [batch_size, seq_length, embedding_dim] -> [1, 5, 768]
# attn_outputs: [batch_size, seq_length, embedding_dim] -> [1, 5, 768]
attn_outputs = torch.bmm(weights, V)
# attn_outputs 是自注意力机制的最终结果。
# 它对输入值向量 V 进行了加权平均，权重来源于注意力权重矩阵 weights。
# 这个输出是序列中每个位置的上下文表示。
print("自注意力机制的最终结果:")
print(attn_outputs)
print("==================")
print("自注意力机制的最终结果的形状:")
print(attn_outputs.shape)
print("==================")

>>>
==================
inputs的类型，来自transformers：
<class 'transformers.tokenization_utils_base.BatchEncoding'>
==================
inputs的数据，来自transformers：
{'input_ids': tensor([[ 2051, 10029,  2066,  2019,  8612]]), 'token_type_ids': tensor([[0, 0, 0, 0, 0]]), 'attention_mask': tensor([[1, 1, 1, 1, 1]])}
==================
inputs.input_ids的数据，pytorch类型：
tensor([[ 2051, 10029,  2066,  2019,  8612]])
==================
inputs.input_ids的size：
torch.Size([1, 5])
==================
token数量以及其特征数量：
Embedding(30522, 768)
==================
给定句子的对应token及其特征：
tensor([[[ 1.2626, -2.1080, -0.6319,  ...,  1.9245, -0.1595,  0.8960],[-2.1495, -1.4998,  1.2840,  ...,  1.1923, -1.4115,  1.1000],[ 2.7451, -0.0587, -0.5572,  ..., -0.4284, -0.2758,  0.3217],[ 0.6120,  0.7005, -0.1534,  ...,  0.4005,  1.0884,  1.3846],[ 0.0194, -0.3392,  0.6146,  ...,  0.4700, -2.5397, -0.1677]]],grad_fn=<EmbeddingBackward0>)
==================
给定句子的对应token及其特征的size：
torch.Size([1, 5, 768])
==================
结果（上述text这个句子中，每个单词/token之间的点积相似度）：
tensor([[[28.7071,  0.4433,  0.7403, -0.4044,  0.2598],[ 0.4433, 27.8755, -0.9110,  0.0327,  1.2176],[ 0.7403, -0.9110, 27.3893,  1.3638,  0.7354],[-0.4044,  0.0327,  1.3638, 27.6121,  0.3601],[ 0.2598,  1.2176,  0.7354,  0.3601, 27.4275]]],grad_fn=<DivBackward0>)
==================
结果的size：
torch.Size([1, 5, 5])
==================
weights结果（上述text这个句子中，每个单词/token之间的点积相似度归一化之后的结果）
tensor([[[1.0000e+00, 5.3110e-13, 7.1472e-13, 2.2751e-13, 4.4206e-13],[1.2200e-12, 1.0000e+00, 3.1489e-13, 8.0912e-13, 2.6461e-12],[2.6697e-12, 5.1205e-13, 1.0000e+00, 4.9802e-12, 2.6567e-12],[6.8006e-13, 1.0529e-12, 3.9854e-12, 1.0000e+00, 1.4607e-12],[1.5895e-12, 4.1419e-12, 2.5573e-12, 1.7570e-12, 1.0000e+00]]],grad_fn=<SoftmaxBackward0>)
==================
weights结果加和，验证归一化：
tensor([[1., 1., 1., 1., 1.]], grad_fn=<SumBackward1>)
==================
自注意力机制的最终结果:
tensor([[[ 1.2626, -2.1080, -0.6319,  ...,  1.9245, -0.1595,  0.8960],[-2.1495, -1.4998,  1.2840,  ...,  1.1923, -1.4115,  1.1000],[ 2.7451, -0.0587, -0.5572,  ..., -0.4284, -0.2758,  0.3217],[ 0.6120,  0.7005, -0.1534,  ...,  0.4005,  1.0884,  1.3846],[ 0.0194, -0.3392,  0.6146,  ...,  0.4700, -2.5397, -0.1677]]],grad_fn=<BmmBackward0>)
==================
自注意力机制的最终结果的形状:
torch.Size([1, 5, 768])
==================

封装好的函数：

import torch
import torch.nn.functional as F
from math import sqrt# query_mask, key_mask:
#### 用于屏蔽某些查询或键的位置。如果指定，通常是 [batch_size, seq_length] 的张量。
#### 这些掩码可以屏蔽不需要参与计算的序列位置（例如，填充位置）。
# mask:
#### 更通用的掩码，形状通常为 [batch_size, seq_length, seq_length]，
#### 用于屏蔽具体的查询-键对。
def scaled_dot_product_attention(query, key, value, query_mask=None, key_mask=None, mask=None):dim_k = query.size(-1)scores = torch.bmm(query, key.transpose(1, 2)) / sqrt(dim_k)# 应用掩码（屏蔽不相关的输入）if query_mask is not None and key_mask is not None:# query_mask.unsqueeze(-1) 将查询掩码扩展为 [batch_size, seq_length, 1]。# key_mask.unsqueeze(1) 将键掩码扩展为 [batch_size, 1, seq_length]。# torch.bmm 计算两者的外积，得到形状为 [batch_size, seq_length, seq_length] 的掩码矩阵 mask。# 掩码矩阵中的值为 1 表示对应的查询-键对有效，0 表示无效。mask = torch.bmm(query_mask.unsqueeze(-1), key_mask.unsqueeze(1))if mask is not None:# 如果存在 mask，使用 masked_fill 方法将 mask == 0 的位置填充为 -inf，# 这些位置的得分被屏蔽。# Softmax 在处理 -inf 时会将其归一化为 0，从而屏蔽这些位置。scores = scores.masked_fill(mask == 0, -float("inf"))weights = F.softmax(scores, dim=-1)return torch.bmm(weights, value)