深入浅出 Pytorch — transformer 实现(1)

zidea
492 阅读3分钟

携手创作，共同成长！这是我参与「掘金日新计划 · 8 月更文挑战」的第17天，点击查看活动详情

前言

之前看了一些关于 transformer 的分享文章，还只是停留在理论上。今天我们就来尝试学习到理论知识去联系实际。

BertViz 库

安装 BertViz 库, BertViz 是一个 transformer 中注意力的可视化工具，可以选择基于 transfomers 的模型。

!pip install bertviz

trans_002.png

from transformers import AutoTokenizer

from bertviz.transformers_neuron_view import BertModel
from bertviz.neuron_view import show
  • 从 transformers 导入 AutoTokenizer 也就是拿到词表，AutoTokenizer 具体如何定义词表，这个是根据具体模型而定，不同模型对词表定义不同，也就是将每一个词对应到词表中一个位置，根据需要添加一些特殊符号。
  • 从 transformers_neuron_view 模块导入一个 transformer 模型
  • 从 neuron_view 导入 show 方法可以更直观观察模型
model_ckpt = 'bert-base-uncased'
tokenizer = AutoTokenizer.from_pretrained(model_ckpt)

来简单查看一下 tokenizer 信息，这里是对词表，词表提供从 word 到 id 映射，也提供 id 到 word 的映射，vocab_size 表示词表中一共有 30522 个词，其中包括一些特殊 token，例如 UNK 表示不清楚的 token、PAD 填充 token 等这些特殊的 token

PreTrainedTokenizerFast(name_or_path='bert-base-uncased', vocab_size=30522, model_max_len=512, is_fast=True, padding_side='right', truncation_side='right', special_tokens={'unk_token': '[UNK]', 'sep_token': '[SEP]', 'pad_token': '[PAD]', 'cls_token': '[CLS]', 'mask_token': '[MASK]'})

text = "Deep learning is a subset of machine learning, which is essentially a neural network with three or more layers"

model = BertModel.from_pretrained(model_ckpt)
show(model,'bert',tokenizer,text,display_mode='light',layer=0,head=8)

trans_001.png

准备一个句子用于输入到模型，根据名字 model_ckpt 通过 BerModel 模块提供的 from_pretrained 方法根据名称model_ckpt 可以模型

inputs = tokenizer(text,return_tensors='pt',add_special_tokens=False)
inputs.input_ids

使用 tokenizer 然后可以将输入句子转换为索引 id，通过 add_special_tokens 设置为 False 表示不会给特殊符号id

tensor([[ 2784, 4083, 2003, 1037, 16745, 1997, 3698, 4083, 1010, 2029, 2003, 7687, 1037, 15756, 2897, 2007, 2093, 2030, 2062, 9014]])

config = AutoConfig.from_pretrained(model_ckpt)
config

BertConfig {
"_name_or_path": "bert-base-uncased",
"architectures": [ "BertForMaskedLM" ],
"attention_probs_dropout_prob": 0.1,
"classifier_dropout": null, 
"gradient_checkpointing": false, 
"hidden_act": "gelu", 
"hidden_dropout_prob": 0.1,
"hidden_size": 768, 
"initializer_range": 0.02,
"intermediate_size": 3072, 
"layer_norm_eps": 1e-12, 
"max_position_embeddings": 512,
"model_type": "bert", 
"num_attention_heads": 12, 
"num_hidden_layers": 12,
"pad_token_id": 0, 
"position_embedding_type": "absolute", 
"transformers_version": "4.21.1",
"type_vocab_size": 2,
"use_cache": true, 
"vocab_size": 30522 
}

获取模型配置文件 vocab_size 30522

embedding 层

这里使用 pytorch 提供的 Embedding 模块来创建 token_emb 用于生成 Embedding 层，关于 embedding

import torch
import torch.nn as nn
token_emb = nn.Embedding(config.vocab_size,config.hidden_size)
token_emb

inputs_embeds = token_emb(inputs.input_ids)
print(inputs_embeds)
inputs_embeds.size()

 tensor([[[ 0.1717, -0.4256, -0.6555,  ...,  1.7164, -0.4652,  0.1299],
         [ 0.6570, -1.0636,  1.0804,  ..., -0.3091, -0.0753,  0.3249],
         [ 0.6190, -1.0261,  0.1296,  ...,  0.4145, -0.6006,  0.0333],
         ...,
         [ 1.7625,  0.7076,  0.0977,  ..., -0.9091,  1.4083, -1.4498],
         [ 0.4938,  0.9443,  0.0741,  ..., -0.4177, -0.7573, -1.5135],
         [-1.1800,  0.2995, -0.1345,  ...,  0.6875, -0.4169, -0.1305]]],
       grad_fn=<EmbeddingBackward0>)

当一个句子经过 embedding 层后得到 Tensor 为 batch_size, sequence_length 和 num_features。

torch.Size([1, 20, 768])

自注意力实现

关于自注意力是 Transformer 的核心，

这里我们假设 query 、key 和 value 直接等于输入 inputs_embeds，在自注意力方法是通过 3 个线性变换层 inputs_embeds 转换为为 query 、

from math import sqrt

query = key = value = inputs_embeds
dim_k = key.size(-1)
scores = torch.bmm(query,key.transpose(1,2)) /sqrt(dim_k)
scores.size()

bmm 这个操作是在输入 input(b x n x m) 和 weight (b x m x p) 计算出 (b x n x p)

 import torch.nn.functional as F
 weights = F.softmax(scores,dim=1)

 weights.sum(dim=-1)

att_outputs = torch.bmm(weights,value)
att_outputs.shape
Attention(Q,K,V)=softmax(QKTdk)VAttention(Q,K,V) = softmax (\frac{QK^T}{\sqrt{d_k}}) V
def scaled_dot_product_attention(query, key, value):
  dim_k = key.size(-1)
  scores = torch.bmm(query, key.transpose(1,2)) / sqrt(dim_k)
  weights = F.softmax(scores,dim=-1)
  return torch.bmm(weights,value)
  • 这里除以 dkd_k 是为了防止 QQ 和 KTK^T 内积过大，然后对得到 weights 进行 softmax 得到输入向量间的相关性的得分 score。
  class AttentionHead(nn.Module):
    def __init__(self, embed_dim, head_dim):
      super().__init__()
      self.q = nn.Linear(embed_dim,head_dim)
      self.k = nn.Linear(embed_dim,head_dim)
      self.v = nn.Linear(embed_dim,head_dim)

    def forward(self,hidden_state):
      attn_outputs =  scaled_dot_product_attention(
          self.q(hidden_state),self.k(hidden_state),self.v(hidden_state)
      )
      return attn_outputs

多头自注意力实现

首先解释一下我们为什么需要多头自注意力，所谓多头自注意力，这里通过一个简单比喻来帮助大家理解，例如我们认识一个事物会从不同角度来描述这个事物，可能会从颜色上、从形状上和大小上来描述一个事物，那么多头自注意力就是从不同角度上来计算两个事物之间相关性。

class MultiHeadAttention(nn.Module):
  def __init__(self,config):
    super().__init__()
    embed_dim = config.hidden_size
    num_heads = config.num_attention_heads
    head_dim = embed_dim // num_heads
      
    self.heads = nn.ModuleList(
        [AttentionHead(embed_dim,embed_dim)]
    )

    self.output_linear = nn.Linear(embed_dim,embed_dim)

  def forward(self,hidden_state):
    concatenated_output = torch.cat([h(hidden_state) for h in self.heads],dim=-1)
    concatenated_output = self.output_linear(concatenated_output)
    return concatenated_output
  • embed_dim 是隐藏层的特征维度
  • num_heads 用于指定多头自注意力的头
avatar
文章
阅读
粉丝