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_kd_v,以及线性层 w_qsw_ksw_vsfc
  • forward 方法首先将输入 qkv 通过线性层映射到多头的维度,然后重塑并转置以便进行并行计算。如果存在掩码,它会被扩展以适应头维度。调用缩放点积注意力模块计算注意力,之后合并多头输出,并应用线性变换和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_1w_2,层归一化 layer_norm,以及dropout层。
  • forward 方法首先通过第一个线性层和ReLU激活函数,然后通过第二个线性层。应用dropout层后,添加残差连接并进行层归一化。

你可能感兴趣的:(机器学习,深度学习,transformer,pytorch)