介绍
把 因果注意力 和 dropout 修改应用 大模型原理：实现带可训练权重的自注意力机制-封装类计算所有上下文向量 中开发的 SelfAttention Python类中。
这个类将成为开发多头注意力的基础，而多头注意力是我们最终实现的注意力类。
实现
导包
import torch
import torch.nn as nn
封装类
class CausalAttention(nn.Module):
def __init__(self, d_in, d_out, context_length, dropout, qkv_bias=False):
"""
:param d_in: 输入嵌入维度
:param d_out: 输出嵌入维度
:param context_length: 上下文长度，用于设置掩码矩阵
:param dropout: dropout忽略概率
:param qkv_bias:
"""
super().__init__()
self.d_out = d_out
self.W_query = nn.Linear(d_in, d_out, bias=qkv_bias)
self.W_key = nn.Linear(d_in, d_out, bias=qkv_bias)
self.W_value = nn.Linear(d_in, d_out, bias=qkv_bias)
# 增加 Dropout 层
self.dropout = nn.Dropout(dropout)
self.register_buffer(
'mask',
torch.triu(torch.ones(context_length, context_length), diagonal=1)
)
def forward(self, x):
b, num_tokens, d_in = x.shape
keys = self.W_key(x)
queries = self.W_query(x)
values = self.W_value(x)
# 将维度1和2转置，将批维度保持在第一个位置（0）
attn_scores = queries @ keys.transpose(1, 2)
# 直接修改传入的矩阵，而不是返回副本，减少内存
attn_scores.masked_fill_(
self.mask.bool()[:num_tokens, :num_tokens], -torch.inf)
attn_weights = torch.softmax(
attn_scores / keys.shape[-1]**0.5, dim=-1
)
# 对注意力权重矩阵进行dropout操作
attn_weights = self.dropout(attn_weights)
# 计算上下文向量
context_vec = attn_weights @ values
return context_vec
关键代码说明：
在 __init__ 方法中增加了 self.register_buffer()，生成固定掩码并注册为缓冲区
作用：用于注册无需梯度（不需要训练）、但需随模型保存/迁移设备的张量，是模型 “非训练参数”的标准管理方式；
好处：
- 只需生成一次，后续重复使用
- 掩码无需梯度，随模型保存
准备数据
# 输入的嵌入向量
inputs = torch.tensor(
[[0.43, 0.15, 0.89], # Your (x^1)
[0.55, 0.87, 0.66], # journey (x^2)
[0.57, 0.85, 0.64], # starts (x^3)
[0.22, 0.58, 0.33], # with (x^4)
[0.77, 0.25, 0.10], # one (x^5)
[0.05, 0.80, 0.55]] # step (x^6)
)
# 输入嵌入维度d_in=3
d_in = inputs.shape[1]
# 输出嵌入维度d_out=2
d_out = 2
模拟批量输入
模拟批量输入，因为真正执行时，是通过数据加载器，批量读取的
# 模拟批量输入，因为真正执行时，是通过数据加载器，批量读取的
batch = torch.stack((inputs, inputs), dim=0)
print("批量数据：\n", batch)
print("批量数据的shape：", batch.shape)
执行结果：
批量数据：
tensor([[[0.4300, 0.1500, 0.8900],
[0.5500, 0.8700, 0.6600],
[0.5700, 0.8500, 0.6400],
[0.2200, 0.5800, 0.3300],
[0.7700, 0.2500, 0.1000],
[0.0500, 0.8000, 0.5500]],
[[0.4300, 0.1500, 0.8900],
[0.5500, 0.8700, 0.6600],
[0.5700, 0.8500, 0.6400],
[0.2200, 0.5800, 0.3300],
[0.7700, 0.2500, 0.1000],
[0.0500, 0.8000, 0.5500]]])
批量数据的shape： torch.Size([2, 6, 3])
计算上下文向量
# 计算上下文向量
print("\n\b计算上下文向量--------")
# 为了复现，设置随机种子
torch.manual_seed(123)
# 词元的数量，用于设置掩码矩阵
context_length = batch.shape[1]
print("context_length:", context_length)
# 传入参数：实例化
ca = CausalAttention(d_in, d_out, context_length, 0.0)
# 计算上下文向量
context_vecs = ca(batch)
print("上下文向量：", context_vecs)
print("上下文向量.shape:", context_vecs.shape)
执行结果：
计算上下文向量--------
context_length: 6
上下文向量： tensor([[[-0.4519, 0.2216],
[-0.5874, 0.0058],
[-0.6300, -0.0632],
[-0.5675, -0.0843],
[-0.5526, -0.0981],
[-0.5299, -0.1081]],
[[-0.4519, 0.2216],
[-0.5874, 0.0058],
[-0.6300, -0.0632],
[-0.5675, -0.0843],
[-0.5526, -0.0981],
[-0.5299, -0.1081]]], grad_fn=<UnsafeViewBackward0>)
上下文向量.shape: torch.Size([2, 6, 2])
最终生成的上下文向量是一个三维张量：torch.Size([2, 6, 2])，其中每个词元现在用二维嵌入来表示。