pytorch mask attention 机制

原创

mob649e8162c013 2024-12-11 07:59:08 ©著作权

文章标签 Mask 数据 python 文章分类 PyTorch 人工智能

©著作权归作者所有：来自51CTO博客作者mob649e8162c013的原创作品，请联系作者获取转载授权，否则将追究法律责任

实现 PyTorch 的 Mask Attention 机制

在深度学习中，自注意力机制是一种重要的技术，尤其在处理序列数据时尤为有效。在处理变长序列时，Mask Attention 机制用于确保模型在注意力计算中忽略某些无效位置。本文将带你逐步实现 PyTorch 中的 Mask Attention 机制。

流程概述

我们将分为以下几个步骤来实现 Mask Attention 机制：

步骤	描述
1	准备数据
2	定义 Attention 函数
3	添加 Mask
4	测试 Attention 函数

首先，让我们稍微了解一下每一个步骤需要做什么。

步骤详解

步骤 1: 准备数据

首先需要创建一个输入序列及其对应的 Mask。

import torch

# 创建一个示例序列 (batch size=2, sequence length=3, feature size=4)
x = torch.tensor([[[1, 0, 0, 1],
                    [0, 1, 0, 0],
                    [0, 0, 1, 0]],

                   [[1, 1, 1, 1],
                    [0, 0, 0, 0],
                    [0, 0, 0, 0]]], dtype=torch.float32)

# 创建 Mask (1表示有效，0表示无效)
mask = torch.tensor([[1, 1, 1],
                     [1, 0, 0]], dtype=torch.float32)

步骤 2: 定义 Attention 函数

接下来，我们定义一个简单的线性 Attention 机制。

import torch.nn.functional as F

def attention(q, k, v, mask=None):
    # 计算注意力分数
    scores = torch.matmul(q, k.transpose(-2, -1)) / (k.size(-1) ** 0.5)
    
    # 如果有 Mask, 应用 Mask
    if mask is not None:
        scores = scores.masked_fill(mask == 0, float('-inf'))
    
    # 计算注意力权重
    attn_weights = F.softmax(scores, dim=-1)
    
    # 应用注意力权重到值向量
    output = torch.matmul(attn_weights, v)
    return output, attn_weights

步骤 3: 添加 Mask

在使用 Mask 时，我们需要将其转换为适合于 Attention 的形式。

# 转换 Mask 形状以匹配分数形状
mask = mask.unsqueeze(1).unsqueeze(2)  # (batch size, 1, 1, sequence length)

步骤 4: 测试 Attention 函数

最后，使用我们之前定义的 Attention 函数进行测试。

# 定义查询、键和值（在这里我们简单地将它们都设置为相同的输入）
query = x
key = x
value = x

# 计算输出
output, attn_weights = attention(query, key, value, mask)

print("Output:", output)
print("Attention Weights:", attn_weights)

旅行图示意

以下是一个旅行图，展示了从准备数据到实现完整 Mask Attention 机制的流程。

journey
    title 实现 PyTorch Mask Attention 机制
    section 准备数据
      创建输入序列: 5: 数据准备
      创建 Mask: 5: 数据准备
    section 定义 Attention 函数
      编写注意力计算: 4: 函数定义
      应用 Mask: 5: 函数定义
    section 测试 Attention 函数
      测试输出: 5: 函数测试

序列图示意

以下是一个序列图，用于表示函数调用的顺序：

sequenceDiagram
    participant User
    participant AttentionFunction
    User->>AttentionFunction: call attention(query, key, value, mask)
    AttentionFunction->>AttentionFunction: compute scores
    AttentionFunction->>AttentionFunction: apply mask
    AttentionFunction->>AttentionFunction: compute softmax
    AttentionFunction->>AttentionFunction: multiply by values
    AttentionFunction-->>User: return output, weights