【论文学习】Transformer:模型搭建

Transformer由论文Attention Is All You Need提出,该模型舍弃了RNN和CNN,完全基于注意力机制实现了高效、并行化的训练,成为了NLP领域的里程碑式模型。本篇文章介绍Transformer的模型结构并基于PyTorch完成模型的搭建。

本文目录

  • [0 引言](#0 引言)
  • [1 输入部分](#1 输入部分)
    • [1.1 文本嵌入层](#1.1 文本嵌入层)
    • [1.2 位置编码器](#1.2 位置编码器)
  • [2 编码器部分](#2 编码器部分)
    • [2.1 掩码张量](#2.1 掩码张量)
    • [2.2 注意力机制](#2.2 注意力机制)
    • [2.3 多头注意力机制](#2.3 多头注意力机制)
    • [2.4 前馈全连接层](#2.4 前馈全连接层)
    • [2.5 规范化层](#2.5 规范化层)
    • [2.6 子层连接结构(残差连接)](#2.6 子层连接结构(残差连接))
    • [2.7 编码器层](#2.7 编码器层)
    • [2.8 编码器](#2.8 编码器)
  • [3 解码器部分](#3 解码器部分)
    • [3.1 解码器层](#3.1 解码器层)
    • [3.2 解码器](#3.2 解码器)
  • [4 输出部分](#4 输出部分)
  • [5 模型搭建](#5 模型搭建)

0 引言

Transformer基于seq2seq架构实现NLP领域的典型任务,如机器翻译、文本生成等。

在论文中,作者以机器翻译任务为主,详细介绍了Transformer的结构,其总体框架如图1(来源于论文)所示,一共包含四个部分:

  • 输入部分:输入词向量预处理

    • 源文本嵌入层及位置编码器
    • 目标文本嵌入层及位置编码器
  • 编码器部分

    • 多头自注意力机制子层和规范层以及一个残差连接
    • 前馈前连接子层和规范层以及一个残差连接
  • 解码器部分

    • 多头自注意力机制子层和规范层以及一个残差连接
    • 多头注意力机制子层和规范层以及一个残差连接
    • 前馈前连接子层和规范层以及一个残差连接
  • 输出部分:输出向量预处理

    • 线性层

    • softmax层


图1 Transformer模型框架

1 输入部分

1.1 文本嵌入层

文本嵌入层将文本中词汇数字表示转变为高维向量表示,更便于捕捉词汇间的关系。其代码实现如下:

python 复制代码
class Embedding(nn.Module):
    def __init__(self, d_model, vocab):
        '''
        文本嵌入层
        :param d_model: 词嵌入的维度
        :param vocab: 词表的大小
        '''
        super().__init__()
        # 嵌入层
        self.lut = nn.Embedding(vocab, d_model)
        self.d_model = d_model
    def forward(self, x):
        return self.lut(x) * math.sqrt(self.d_model)

1.2 位置编码器

位置编码器在词向量中添加每个词在句子中的位置信息,使模型具有学习词序信息的能力。论文中采用正余弦交替的位置编码,公式如下:
P E t ( i ) = { s i n ( w k t ) , i = 2 k c o s ( w k t ) , i = 2 k + 1 PE_t^{(i)}= \begin{cases} sin(w_kt),i = 2k \\ cos(w_kt),i = 2k+1\ \end{cases} PEt(i)={sin(wkt),i=2kcos(wkt),i=2k+1

w k = 1 1000 0 2 k / d _ m o d e l , k = 0 , 1 , 2 , . . . , d _ m o d e l 2 − 1 w_k={1 \over 10000^{2k/d\_model}},k=0,1,2,...,{d\_model \over 2}-1 wk=100002k/d_model1,k=0,1,2,...,2d_model−1

其中 i i i表示词向量 t t t所处的位置, d _ m o d e l d\_model d_model表示词向量 t t t的维度。

其代码实现如下:

python 复制代码
class PositionEncoding(nn.Module):
    def __init__(self, d_model, dropout=0.2, max_len=5000):
        '''
        位置编码层:正余弦交替编码
        :param d_model: 词向量维度
        :param dropout: 置零比例,正则化,防止模型过拟合
        :param max_len: 每个句子的最大长度
        '''
        super().__init__()
        # dropout正则化层
        self.dropout = nn.Dropout(p=dropout)
        # 初始化一个位置编码矩阵 [max_len, d_model]
        pe = torch.zeros(max_len, d_model)
        # 初始化一个位置矩阵(句子中每个词向量的位置) [max_len, 1]
        position = torch.arange(0, max_len).unsqueeze(1)
        # 构造位置编码中正余弦频率(共用频率) [d_model // 2,]
        w = 1 / torch.pow(10000, torch.arange(0, d_model, 2) / d_model)
        # 构造最终的位置编码:偶数位置使用正弦编码,奇数位置使用余弦编码 [max_len, d_model]
        pe[:, 0::2] = torch.sin(position * w)
        pe[:, 1::2] = torch.cos(position * w)
        # [max_len, d_model] -> [1, max_len, d_model]
        pe = pe.unsqueeze(0)
        # 位置编码注册为模型buffer,其在优化过程中不会更新,在模型保存后重加载时和模型一起加载
        self.register_buffer('pe', pe)
    def forward(self, x):
        '''
        :param x: [bs(句子数量), length(句子长度), d_model] 序列的嵌入(embedding)表示
        :return x_pos:[bs, length, d_model] 加上词序信息的词向量
        '''
        length = x.size(1)  # 句子的长度
        x = x + Variable(self.pe[:, :length, :], requires_grad=False)
        return self.dropout(x)

2 编码器部分

编码器部分结构如图2所示,由N个编码器层组成,每个编码器层由两个子层连接而成:

  • 第一个子层包括一个多头自注意力机制和规范层以及一个残差连接
  • 第二个子层包括一个前馈全连接子层和规划层以及一个残差连接

图2 编码器结构 >

2.1 掩码张量

Transformer中,生成的attention信息中,可能会包含未来的信息,为了避免未来的信息被提前使用,会使用掩码张量对部分信息进行遮掩。其构建代码如下:

python 复制代码
def subsequent_mask(size):
    attn_shape = (1, size, size)

    # 首先构建一个全1矩阵, 再将其形成上三角矩阵(对角线元素也舍弃)
    subsequent_mask = np.triu(np.ones(attn_shape), k=1).astype('uint8')
    # 构造一个下三角矩阵(包含对角线元素)
    return torch.from_numpy(1 - subsequent_mask)

2.2 注意力机制

我们观察事物时,大脑能够快速的把注意力放在事物最具辨识度的部分从而作出判断。基于这种形式,在模型中引入注意力机制,帮助模型聚焦于对当前任务更为关键的信息,降低对其他信息的关注,提高模型对特定任务处理的效率和准确率。

论文中使用的注意力计算方法如图3所示,计算公式如下:
A t t e n t i o n ( Q , K , V ) = s o f t m a x ( Q K T d k ) V Attention(Q,K,V)=softmax({{Q{K^T}} \over {\sqrt {{d_k}} }})V Attention(Q,K,V)=softmax(dk QKT)V

其中 Q ( q u e r y ) , K ( k e y ) , V ( v a l u e ) Q(query),K(key),V(value) Q(query),K(key),V(value)由输入词向量 X X X通过线性变化得到,当 Q = K = V Q=K=V Q=K=V时称为自注意力机制, d k d_k dk为词嵌入维度。

图3 注意力机制结构

注意力计算的实现代码如下:

python 复制代码
def attention(query, key, value, mask=None, dorpout=None):
    '''
    :param query: [bs, head, length, d_k] Q
    :param key: [bs, head, length, d_k] K
    :param value: [bs, head, length, d_k] V
    :return:
    '''
    d_k = query.size(-1)  # d_model/d_k 词嵌入维度
    # 计算Qk^T,得到注意力张量scores  [bs, head, length, length]
    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)
    # 判断是否使用掩码张量
    if mask is not None:
        scores = scores.mask_fill(mask == 0, -1e9)
    # 对scores的最后一维度进行softmax操作 [bs, head, length, length]
    p_attn = F.softmax(scores, dim=-1)
    # 判断是否使用dropout
    if dorpout is not None:
        p_attn = dorpout(p_attn)
    # 得到最终的注意力表达式 [bs, head, length, d_k]
    attn = torch.matmul(p_attn, value)
    return attn, p_attn

2.3 多头注意力机制

多头注意力机制结构如图4所示,其将输入词向量在词嵌入维度上d\_model进行切分,分为 h h h个头,每个头的词嵌入维度为 d _ m o d e l / h d\_model / h d_model/h,分别进行自注意力机制运算后再对结果在通道上进行拼接。

多头注意力机制让每个注意力去优化每个词汇的不同特征部分,从而均衡同一种注意力机制可能产生的偏差,让词义信息具有更丰富的表达。

图4 多头注意力机制结构

多头注意力机制实现代码如下:

python 复制代码
def clones(module, N):
    '''用于生成相同网络层的克隆函数,其参数不共享'''
    return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])

class MultiHeadedAttention(nn.Module):
    def __init__(self, h, d_model, dropout=0.1):
        '''
        多头注意力机制
        :param h: 词向量在词嵌入维度上划分的头数
        :param d_model: 词嵌入维度
        :param dropout: 
        '''
        super().__init__()
        # 每个头的分割词向量维度d_k
        assert d_model % h == 0
        self.d_k = d_model // h
        # 分割头数
        self.h = h
        # 得到线性变化层,分别用于Q,K,V矩阵以及对拼接矩阵的线性操作
        self.linears = clones(nn.Linear(d_model, d_model), 4)
        # 注意力张量
        self.attn = None
        # dropout正则化
        self.dropout = nn.Dropout(p=dropout)
    def forward(self, query, key, value, mask=None):

        # 样本数
        batch_size = query.size(0)

        # 多头处理环节:首先对QKV进行线性操作, 然后按照头数进行划分
        # q,k,v:[bs, length, head, d_k] -> [bs, head, length, d_k]
        query, key, value = \
            [model(x).view(batch_size, -1, self.h, self.d_k).transpose(1, 2)
             for model, x in zip(self.linears, (query, key, value))]
        # 对每个头进行注意力计算
        # x: 注意力张量[bs, head, length, d_k]
        # attn: 注意力张量权重[bs, head, length, length]
        x, self.attn = attention(query, key, value, mask=mask, dorpout=self.dropout)
        # 将多头的注意力张量进行合并 [bs, head, length, d_k] -> [bs, length, head, d_k] -> [bs, length, head*d_k(d_model)]
        x = x.transpose(1, 2).contiguous().view(batch_size, -1, self.h * self.d_k)
        # 多头注意力机制返回结果 [bs, length, d_model]
        return self.linears[-1](x)

2.4 前馈全连接层

前馈全连接层由两个线性操作和ReLU激活函数组成,其公式如下:
F F N ( x ) = W 2 m a x ( 0 , W 1 x + b 1 ) + b 2 FFN(x) = W_2max(0, W_1x+b_1)+b_2 FFN(x)=W2max(0,W1x+b1)+b2

其代码实现如下:

python 复制代码
class PositionwiseFeedForward(nn.Module):
    def __init__(self, d_model, d_hidden, dropout=0.1):
        '''
        前馈全连接层
        :param d_model: 词嵌入维度
        :param d_hidden: 隐藏层维度
        :param dropout:
        '''
        super().__init__()
        self.w1 = nn.Linear(d_model, d_hidden)
        self.w2 = nn.Linear(d_hidden, d_model)
        self.act = F.relu
        self.dropout = nn.Dropout(p=dropout)

    def forward(self, x):
        '''
        :param x: 词向量 (bs, length, d_model)
        :return:
        '''
        y = self.dropout(self.act(self.w1(x)))
        return self.w2(y)

2.5 规范化层

规范化层实现对神经网络数值的标准化,使其特征数值满足均值为0,方差为1的分布,有利于模型的收敛。其代码实现如下:

python 复制代码
class LayerNorm(nn.Module):
    def __init__(self, features, eps=1e-6):
        '''
        规范化层
        :param features: 输入特征的维度
        :param eps:防止分母为0
        '''

        self.a2 = nn.Parameter(torch.ones(features))  # 随着模型训练
        self.b2 = nn.Parameter(torch.zeros(features))  # 随着模型训练
        self.eps = eps

    def forward(self, x):
        '''
        :param x: 词向量 [bs, length, d_model]
        :return:
        '''
        mean = x.mean(-1, keepdim=True)  # 均值
        std = x.std(-1, keepdim=True)  # 方差
        # 标准化 [bs, length, d_model]
        return self.a2 * (x - mean) / (std + self.eps) + self.b2

2.6 子层连接结构(残差连接)

对子层结构进行封装,实现代码如下:

python 复制代码
class SublayerConnection(nn.Module):
    def __init__(self, size, dropout=0.1):
        '''
        :param size: 词嵌入维度
        :param dropout:
        '''
        self.norm = LayerNorm(size)  # 标准化
        self.dropout = nn.Dropout(p=dropout)
    def forward(self, x, sublayer):
        '''
        :param x: 词向量 [bs, length, d_model]
        :sublayer: MultiHeadedAttention / PositionwiseFeedForward
        :return:
        '''
        return x + self.dropout(sublayer(self.norm(x)))

2.7 编码器层

结合以上函数,实现编码器层:

python 复制代码
class EncoderLayer(nn.Module):
    def __init__(self, size, self_attn, feed_forward, dropout):
        '''
        :param size: 词嵌入维度
        :param self_attn: 多头注意力机制
        :param feed_forward: 前馈全连接层
        :param dropout:
        '''
        self.self_attn = self_attn
        self.feed_forward = feed_forward

        # 构造子层连接
        self.sublayer = clones(SublayerConnection(size, dropout), 2)
        self.size = size
    def forward(self, x, mask):
        y = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, mask))  # 第一个子层(多头注意力+标准化+残差连接)
        y = self.sublayer[1](y, self.feed_forward)  # 第二个子层(前馈全连接+标准化+残差连接)
        return y

2.8 编码器

基于编码器层,构造编码器:

python 复制代码
class Encoder(nn.Module):
    def __init__(self, layer, N):
        '''
        :param layer: 编码器层
        :param N: 编码器层数
        '''
        super().__init__()
        self.layers = clones(layer, N)  # 构造N层编码器
        self.norm = LayerNorm(layer.size)  # 构造标准化层
    def forward(self, x, mask):
        for layer in self.layers:
            x = layer(x, mask)
        return self.norm(x)

3 解码器部分

解码器部分结构如图5所示,由N个编码器层组成,每个编码器层由三个子层连接而成:

  • 多头自注意力机制子层和规范层以及一个残差连接
  • 多头注意力机制子层和规范层以及一个残差连接
  • 前馈前连接子层和规范层以及一个残差连接

图5 解码器结构

3.1 解码器层

结合编码器部分中实现的模块,解码器层代码如下:

python 复制代码
class DecoderLayer(nn.Module):
    def __init__(self, size, self_attn, src_attn, feed_forward, dropout=0.1):
        '''
        :param size: 词嵌入维度
        :param self_attn: 多头自注意力机制(Q=K=V)
        :param src_attn: 多头注意力机制(Q!=K=V)
        :param feed_forward: 前馈全连接层
        :param dropout:
        '''
        self.size = size
        self.self_attn = self_attn
        self.src_attn = src_attn
        self.feed_forward = feed_forward
        # 构造子连接层
        self.sublayers = clones(SublayerConnection(size, dropout), 3)
    def forward(self, x, memory, source_mask, target_mask):
        '''
        :param x: 上一层的输出张量 [bs, length, d_model]
        :param memory: 编码器层的输出 [bs, length, d_model]
        :param source_mask: 多头注意力机制的掩码张量
        :param target_mask: 多头自自注意机制的掩码张量
        '''
        m = memory

        # 第一个子层:多头自注意力机制 + 标准化 + 残差连接
        # 目标数据掩码张量:避免对未来信息的使用
        x = self.sublayers[0](x, lambda x:self.self_attn(x, x, x, target_mask))
        # 第二个子层:多头注意力机制 + 标准化 + 残差连接
        # 源数据掩码张量:遮挡对结果无效的信息
        x = self.sublayers[1](x, lambda x:self.src_attn(x, m, m, source_mask))
        # 第三个子层:前馈全连接层 + 标准化 + 残差连接
        x = self.sublayers[2](x, self.feed_forward)

        return x

3.2 解码器

基于解码器层,构造解码器:

python 复制代码
class Decoder(nn.Module):
    def __init__(self, layer, N):
        '''
        :param layer: 解码器层
        :param N: 数量
        '''
        self.layers = clones(layer, N)
        self.norm = LayerNorm(layer.size)
    def forward(self, x, memory, source_mask, target_mask):
        '''
        :param x: 数据的嵌入表示 [bs, length, d_model]
        :param memory: 编码器输出 [bs, length, d_model]
        :param source_mask:
        :param target_mask:
        '''
        for layer in self.layers:
            x = layer(x, memory, source_mask, target_mask)
        return self.norm(x)

4 输出部分

输出层包括线性层和softmax层,其代码如下:

python 复制代码
class Generator(nn.Module):
    def __init__(self, d_model, vocab_size):
        '''
        :param d_model: 词嵌入维度
        :param vocab_size: 词表大小
        '''
        super().__init__()
        # 线性层,将词向量转换为词表预测结果
        self.project = nn.Linear(d_model, vocab_size)

    def forward(self, x):
        '''
        :param x: 解码器的最终输出 [bs, length, d_model]
        :return: 最终预测结果 [bs, length, vocab_size]
        '''
        return F.log_softmax(self.project(x), dim=-1)

5 模型搭建

python 复制代码
class EncoderDecoder(nn.Module):
    def __init__(self, encoder, decoder, source_embed, target_embed, generator):
        '''
        编码器 - 解码器结构
        :param encoder: 编码器
        :param decoder: 解码器
        :param source_embed: 源数据嵌入函数
        :param target_embed: 目标数据嵌入函数
        :param generator: 类别生成器
        '''
        self.encoder = encoder
        self.decoder = decoder
        self.src_embed = source_embed
        self.tgt_embed = target_embed
        self.generator = generator
    def forward(self, source, target, source_mask, target_mask):
        '''
        :param source: 源数据 [bs, length]
        :param target: 目标数据 [bs, length]
        :param source_mask: 源数据掩码张量
        :param target_mask: 目标数据掩码张量
        '''
        return self.decode(self.encode(source, source_mask), source_mask,
                          target, target_mask)
    def encode(self, source, source_mask):
        '''编码器'''
        return self.encoder(self.src_embed(source), source_mask)
    def decode(self, memory, source_mask, target, target_mask):
        return self.decoder(self.tgt_embed(target), memory, source_mask, target_mask)


def make_model(source_vocab, target_vocab, N=6,
               d_model=512, d_hidden=2048, head=8, dropout=0.1):
    '''
    构建模型
    :param source_vocab: 源数据特征(词汇)总数
    :param target_vocab: 目标数据特征(词汇)总数
    :param N: 解码器和编码器个数
    :param d_model: 词嵌入维度
    :param d_hidden: 前馈全连接层隐藏层层数
    :param head: 多头注意力机制头数
    :param dropout:
    '''
    # 深拷贝
    c = copy.deepcopy
    # 实例化多头注意力机制类
    attn = MultiHeadedAttention(head, d_model)
    # 实例化前馈全连接层类
    ff = PositionwiseFeedForward(d_model, dropout)
    # 实例化位置编码类
    position = PositionEncoding(d_model, dropout=dropout)

    # 构造模型
    model = EncoderDecoder(
        encoder=Encoder(EncoderLayer(d_model, c(attn), c(ff), dropout), N),
        decoder=Decoder(DecoderLayer(d_model, c(attn), c(attn), c(ff), dropout), N),
        source_embed=nn.Sequential(Embedding(d_model, source_vocab), c(position)),
        target_embed=nn.Sequential(Embedding(d_model, target_vocab), c(position)),
        generator=Generator(d_model, target_vocab)
    )
    # 初始化参数
    for p in model.parameters():
        if p.dim() > 1:
            nn.init.xavier_uniform(p)
    return model
相关推荐
rommel rain15 小时前
SpecInfer论文阅读
人工智能·语言模型·transformer
Just Jump1 天前
机器翻译基础与模型 之三:基于自注意力的模型
自然语言处理·transformer·机器翻译
cv君2 天前
视频修复技术和实时在线处理
深度学习·音视频·transformer·视频修复
机器学习之心2 天前
POD-Transformer多变量回归预测(Matlab)
matlab·回归·transformer·pod-transformer
regret~2 天前
【论文笔记】LoFLAT: Local Feature Matching using Focused Linear Attention Transformer
论文阅读·深度学习·transformer
迪菲赫尔曼2 天前
即插即用篇 | YOLOv11 引入高效的直方图Transformer模块 | 突破天气障碍:Histoformer引领高效图像修复新路径
人工智能·深度学习·yolo·目标检测·计算机视觉·transformer·注意力机制
Hqst 网络变压器 Andy3 天前
How to connect a 2.5G network transformer to an RJ45 network port and chip
深度学习·5g·transformer
xianghan收藏册3 天前
LLM文档对话 —— pdf解析关键问题
人工智能·深度学习·自然语言处理·chatgpt·transformer
机器白学5 天前
从零开始使用GOT-OCR2.0——多模态通用型OCR(非常具有潜力的开源OCR项目):项目环境安装配置 + 测试使用
ocr·transformer·多模态·视觉语言大模型
_Randy_5 天前
Transformer
人工智能·深度学习·transformer