您现在的位置是：首页 >技术教程 >Transformer代码学习网站首页技术教程

在init函数里写 需要用到的东西，如softmax、buffer；在forward里写怎么执行的，具体怎么做的

导入库

import numpy as np
import torch
from torch import nn
from einops import rearrange, repeat
from einops.layers.torch import Rearrange

求位置编码的类

思路：
给每个位置加上一个不同的偏置值，以此区别相同的词在不同位置上的不同含义，给输入带上位置信息。得到所有元素的位置编码表（是一系列预生成的东西，不是通过训练得到的）。

如何计算位置编码？----每个位置的x加上当前位置的偏置值
偏置值的计算：
相邻维度（偶数用正弦，奇数用余弦）共享同一频率，但相位不同，使模型能捕捉位置间的相对关系。分母 $10000^{2i/d_{h}id}$ 随着维度索引 i = hid_j//2 增大而指数级增长，导致位置编码的频率逐渐降低。（其中10000是经验值）

class PositionalEncoding(nn.Module):
	# d_hid（隐藏状态的维度）和 n_position（位置的数量）
    def __init__(self, d_hid, n_position=200):
    	super(PositionalEncoding, self).__init__() # 继承父类nn.Module
    	self.register_buffer('pos_table', self._get_sinusoid_encoding_table(n_position, d_hid))
    # 求所有输入元素位置编码表的方法
    def _get_sinusoid_encoding_table(self, n_position, d_hid):
    	def get_position_angle_vec(position):
            return [position / np.power(10000, 2 * (hid_j // 2) / d_hid) for hid_j in range(d_hid)]

        sinusoid_table = np.array([get_position_angle_vec(pos_i) for pos_i in range(n_position)])
        sinusoid_table[:, 0::2] = np.sin(sinusoid_table[:, 0::2])  # dim 2i
        sinusoid_table[:, 1::2] = np.cos(sinusoid_table[:, 1::2])  # dim 2i+1

        return torch.FloatTensor(sinusoid_table).unsqueeze(0)#(1,N,d)
    	def forward(self, x):
    	# 前N个元素（n为序列长度）+自己的偏置值
    	# .clone()创建一个与原张量 self.pos_table 数据相同但存储独立的新张量，避免与原始缓冲区共享内存。.detach()将克隆后的张量从当前计算图中分离。
    		return x + self.pos_table[:, :x.size(1)].clone().detach()
        

前向传播的类（MLP）

Transformer的Encoder中包括两个子类，多头注意力机制和MLP结构，其中MLP结构如下

class FeedForward(nn.Module):
    def __init__(self, dim, hidden_dim, dropout=0.):
	    self.net = nn.Sequential(
	            nn.LayerNorm(dim),
	            nn.Linear(dim, hidden_dim),
	            nn.GELU(),
	            nn.Dropout(dropout),
	            nn.Linear(hidden_dim, dim),
	            nn.Dropout(dropout)
	        )
	def forward(self, x):
        return self.net(x)

多头注意力机制的类

注意力机制：本质是将一个qurey和一个key-value对映射成输出的函数
标准的自注意力机制中，每个头处理输入的不同部分，然后将所有头的输出拼接起来，通过线性变换还原到原始维度。（将每个头得到的z拼起来）比如，如果每个头的维度是64，有8个头，那么拼接后的总维度就是64*8=512。

注意力分数的计算
Attention(Q, K, V) = Softmax(QK^T / sqrt(dim_head)) * V

所以在forward函数中实现：

1. 将输入向量 x 通过 LayerNorm 层进行归一化。
2. 使用线性层 to_qkv 将归一化后的输入向量投影到 Query、Key 和 Value 空间，并使用 chunk 方法将结果切分为三部分，分别表示 Query、Key 和 Value。
3. 使用 rearrange 函数将每个部分重新排列为 (batch_size, num_heads, sequence_length, head_dim) 的形状。
4. 计算注意力分数（做点积）并通过缩放因子进行缩放，得到注意力得分矩阵 dots。使用 Softmax 层计算注意力权重，并应用 Dropout。
5. 将注意力权重与 Value 矩阵相乘得到加权的 Value。

class Attention(nn.Module):
	def __init__(self, dim, heads=8, dim_head=64, dropout=0.):
        super().__init__()
        inner_dim = dim_head * heads
        project_out = not (heads == 1 and dim_head == dim)

        self.heads = heads
        self.scale = dim_head ** -0.5 # 防止数值过大softmax函数梯度消失

        self.norm = nn.LayerNorm(dim) 

        self.attend = nn.Softmax(dim=-1) # 沿着最后一个维度（即序列的每个位置）进行softmax
        self.dropout = nn.Dropout(dropout)

        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias=False)
        # qkv实际是一组参数再乘积，节约网络参数，减小空间，乘积后切分即可

        self.to_out = nn.Sequential(
            nn.Linear(inner_dim, dim),
            nn.Dropout(dropout)
        ) if project_out else nn.Identity()
	def forward(self, x):
		x = self.norm(x)
        qkv = self.to_qkv(x).chunk(3, dim=-1)
        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h=self.heads), qkv)
        dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale # 矩阵乘法QxK转置
        attn = self.attend(dots)
        attn = self.dropout(attn)

        out = torch.matmul(attn, v)
        out = rearrange(out, 'b h n d -> b n (h d)')
        return self.to_out(out)
        

        

Transformer架构

class Transformer(nn.Module):
    def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout=0.):
        super().__init__()
        self.norm = nn.LayerNorm(dim)
        self.layers = nn.ModuleList([])
        # Encoder模块  原文中depth=6
        for _ in range(depth):
            self.layers.append(nn.ModuleList([
                Attention(dim, heads=heads, dim_head=dim_head, dropout=dropout),
                FeedForward(dim, mlp_dim, dropout=dropout)
            ]))
        def forward(self, x): # 残差结构
	        for attn, ff in self.layers:
	            x = attn(x) + x
	            x = ff(x) + x
	        return self.norm(x)