Transfomer重要源码解析：缩放点击注意力，多头自注意力，前馈网络

本文是对Transfomer重要模块的源码解析，完整笔记链接点这里！

缩放点积自注意力 (Scaled Dot-Product Attention)

缩放点积自注意力是一种自注意力机制，它通过查询（Query）、键（Key）和值（Value）的关系来计算注意力权重。该机制的核心在于先计算查询和所有键的点积，然后进行缩放处理，应用softmax函数得到最终的注意力权重，最后用这些权重对值进行加权求和。

源码解析：

import torch
import torch.nn as nn
import torch.nn.functional as F

class ScaledDotProductAttention(nn.Module):
    ''' Scaled Dot-Product Attention '''

    def __init__(self, temperature, attn_dropout=0.1):
        super().__init__()
        self.temperature = temperature  # 温度参数，用于缩放点积
        self.dropout = nn.Dropout(attn_dropout)  # Dropout层

    def forward(self, q, k, v, mask=None):
        attn = torch.matmul(q / self.temperature, k.transpose(2, 3))  # 计算缩放后的点积
        if mask is not None:
            attn = attn.masked_fill(mask == 0, -1e9)  # 掩码操作，将需要忽略的位置设置为一个非常小的值
        attn = self.dropout(F.softmax(attn, dim=-1))  # 应用softmax函数并进行dropout
        output = torch.matmul(attn, v)  # 使用注意力权重对值（v）进行加权求和

        return output, attn

• __init__ 方法中的 temperature 参数用于缩放点积，通常设置为键（Key）维度的平方根。attn_dropout 是在应用softmax函数后进行dropout的比例。
• forward 方法计算缩放点积自注意力。首先，它计算查询（q）和键（k）的点积，并通过除以 temperature 进行缩放。如果提供了 mask，则会使用 masked_fill 将掩码位置的注意力权重设为一个非常小的负数（这里是 -1e9），使得softmax后这些位置的权重接近于0。之后，应用dropout和softmax函数得到最终的注意力权重。最后，使用这些权重对值（v）进行加权求和得到输出。

多头注意力 (Multi-Head Attention)

多头注意力通过将输入分割成多个头，让每个头在不同的子空间表示上计算注意力，然后将这些头的输出合并。这样做可以让模型在多个子空间中捕获丰富的信息。

源码解析：

import torch.nn as nn
import torch.nn.functional as F
from transformer.Modules import ScaledDotProductAttention

class MultiHeadAttention(nn.Module):
    ''' Multi-Head Attention module '''

    def __init__(self, n_head, d_model, d_k, d_v, dropout=0.1):
        super().__init__()
        self.n_head = n_head  # 头的数量
        self.d_k = d_k  # 键/查询的维度
        self.d_v = d_v  # 值的维度
        self.w_qs = nn.Linear(d_model, n_head * d_k, bias=False)  # 查询的线性变换
        self.w_ks = nn.Linear(d_model, n_head * d_k, bias=False)  # 键的线性变换
        self.w_vs = nn.Linear(d_model, n_head * d_v, bias=False)  # 值的线性变换
        self.fc = nn.Linear(n_head * d_v, d_model, bias=False)  # 输出的线性变换
        self.attention = ScaledDotProductAttention(temperature=d_k ** 0.5)  # 缩放点积注意力模块
        self.dropout = nn.Dropout(dropout)  # Dropout层
        self.layer_norm = nn.LayerNorm(d_model, eps=1e-6)  # 层归一化

    def forward(self, q, k, v, mask=None):
        # 保存输入以便后面进行残差连接
        residual = q

        # 线性变换并重塑以准备多头计算
        q = self.w_qs(q).view(sz_b, len_q, n_head, d_k)
        k = self.w_ks(k).view(sz_b, len_k, n_head, d_k)
        v = self.w_vs(v).view(sz_b, len_v, n_head, d_v)

        # 转置以将头维度提前，便于并行计算
        q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)

        # 如果存在掩码，则扩展掩码以适应头维度
        if mask is not None:
            mask = mask.unsqueeze(1)   # 为头维度广播掩码

        # 调用缩放点积注意力模块
        q, attn = self.attention(q, k, v, mask=mask)

        # 转置并重塑以合并多头
        q = q.transpose(1, 2).contiguous().view(sz_b, len_q, -1)
        # 应用线性变换和dropout
        q = self.dropout(self.fc(q))
        # 添加残差连接并进行层归一化
        q += residual
        q = self.layer_norm(q)

        # 返回多头注意力的输出和注意力权重
        return q, attn

• __init__ 方法初始化了多头注意力的参数，包括头的数量 n_head，查询/键/值的维度 d_k 和 d_v，以及线性层 w_qs、w_ks、w_vs 和 fc。
• forward 方法首先将输入 q、k、v 通过线性层映射到多头的维度，然后重塑并转置以便进行并行计算。如果存在掩码，它会被扩展以适应头维度。调用缩放点积注意力模块计算注意力，之后合并多头输出，并应用线性变换和dropout。最后，添加残差连接和层归一化。

前馈网络 (Positionwise FeedForward)

前馈网络（FFN）在自注意力层之后应用，用于进行非线性变换，增加模型的复杂度和表达能力。

源码解析：

import torch.nn as nn
import torch.nn.functional as F

class PositionwiseFeedForward(nn.Module):
    ''' A two-feed-forward-layer module '''

    def __init__(self, d_in, d_hid, dropout=0.1):
        super().__init__()
        self.w_1 = nn.Linear(d_in, d_hid)  # 第一个线性层
        self.w_2 = nn.Linear(d_hid, d_in)  # 第二个线性层
        self.layer_norm = nn.LayerNorm(d_in, eps=1e-6)  # 层归一化
        self.dropout = nn.Dropout(dropout)  # Dropout层

    def forward(self, x):
        # 保存输入以便后面进行残差连接
        residual = x

        # 通过第一个线性层，然后应用ReLU激活函数
        x = self.w_1(x)
        x = F.relu(x)
        # 通过第二个线性层
        x = self.w_2(x)
        # 应用dropout
        x = self.dropout(x)
        # 添加残差连接并进行层归一化
        x += residual
        x = self.layer_norm(x)

        # 返回输出
        return x

• __init__ 方法初始化了两个线性层 w_1 和 w_2，层归一化 layer_norm，以及dropout层。
• forward 方法首先通过第一个线性层和ReLU激活函数，然后通过第二个线性层。应用dropout层后，添加残差连接并进行层归一化。