Pytorch Bert源码解读

Bert 源码解读

一直在用 Bert 系列做一些自然语言理解任务，包括命名实体识别、关系抽取，文本结构化等信息抽取任务。之前开发任务紧，Bert 拿来就用了，很多代码只知其然而不知其所以然，今日好好读了一下 BERT Pytorch 版本的源码，并逐行做了注释记录，遂开辟 NLP 菜鸟系列，随缘更新一些文章，供基础玩家阅读与学习。

耐心读下来，整个流程可以轻松读懂，源码中一些不常用的函数已经去掉~

BertModel

import torch
import math
import torch.nn as nn
from transformers import BertPreTrainedModel
class BertModel(BertPreTrainedModel):
#继承BertPreTrainedModel类，该类主要包含一些预训练模型的下载地址，权重初始化函数, config配置参数类等
def __init__(self, config):
super().__init__(config) #继承父类__init__
self.config = config #Bert的配置文件
self.embeddings = BertEmbeddings(config) #对输入句子做初步embedding
self.encoder = BertEncoder(config) #Bert的核心构件，由12层BertLayer组成
self.pooler = BertPooler(config) #Bert pool 得到的输出可以用来下游 NSP 任务
self.init_weights()#权重初始化
def forward(self, input_ids, token_type_ids=None, attention_mask=None, position_ids=None, head_mask=None):
'''
:param input_ids: 句子中的每个字查vocab.txt所得到的下标
:param token_type_ids: 有地方称作 segment_ids, 用于Next Sentence Predict任务，Robert后已经取消NSP任务
:param attention_mask: 1代表有用信息，0代表padding这种无用信息
:param position_ids: 代表位置下标
:param head_mask: 处理head_mask，可以用多头剪枝，一般不用/
:return: 大小为[batch_size * seq_length * hidden_size (默认为768)]的tensor
'''
if attention_mask is None:
attention_mask = torch.ones_like(input_ids)#如果未指定attention_mask, 全部设置为1
if token_type_ids is None:
token_type_ids = torch.zeros_like(input_ids)#如果未指定token_type_ids，全部设置为0，即所有句子中“每个字的下标”都为0，不设定NSP任务
extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)#attention_mask 由[batch_size*seq_length] 变成 [batch_size*1*1*seq_length], 后续的 self-attention会用到它
extended_attention_mask = extended_attention_mask.to(dtype=next(self.parameters()).dtype) # fp16 compatibility
extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0 #0代表有用信息, -10000代表padding等无用信息，此操作防止softmax计算溢出
if head_mask is not None: #head 剪枝操作, 略
pass
else: head_mask = [None] * self.config.num_hidden_layers # encoder中的隐藏层数
embedding_output = self.embeddings(input_ids, position_ids=position_ids, token_type_ids=token_type_ids)#大小为[batch_size * seq_length * hidden_size(768)]
##############################跳转到BertEmbeddings()######################################
encoder_outputs = self.encoder(embedding_output, extended_attention_mask, head_mask=head_mask) #Encoder层组成整个transformer
##############################跳转到BertEncoder()#########################################
sequence_output = encoder_outputs[0] #默认为 batch_size * seq_len * 768 的tensor
pooled_output = self.pooler(sequence_output) #可用于下游NSP等任务
outputs = (sequence_output, pooled_output,) + encoder_outputs[1:] #encoder_outputs[1:]代表attention系数
return outputs

BertEmbeddings

class BertEmbeddings(nn.Module):
def __init__(self, config):
super(BertEmbeddings, self).__init__()
self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=0)#word embeding
self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)#位置embedding, 其中max_position_embeddings等于max_sequence_length
self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)#type_vocab_size=2, 0代表当前句子, 1代表下一句, 用作NSP任务
self.LayerNorm = torch.nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)#层归一化
self.dropout = nn.Dropout(config.hidden_dropout_prob)#dropout
def forward(self, input_ids, token_type_ids, position_ids):
'''
:param input_ids: 句子中的每个字查vocab.txt所得到的下标，大小为[batch_size * seq_length]
:param token_type_ids: 大小为[batch_size * seq_length]
:param position_ids: 大小为[batch_size * seq_length]
:return:
'''
seq_length = input_ids.size(1) #代表句子长度
if position_ids is None: #如果没给出位置下标
position_ids = torch.arange(seq_length, dtype=torch.long, device=input_ids.device)#按照句子长度从0到seq_length给出位置下标
position_ids = position_ids.unsqueeze(0).expand_as(input_ids)#大小为[batch_size * seq_length]
if token_type_ids is None:#如果没给出 token_type_ids, 默认不使用NSP任务
token_type_ids = torch.zeros_like(input_ids)#大小为[batch_size * seq_length]
words_embeddings = self.word_embeddings(input_ids)#大小为[batch_size * seq_length * hidden_size(768)]
position_embeddings = self.position_embeddings(position_ids)#大小为[batch_size * seq_length * hidden_size(768)]
token_type_embeddings = self.token_type_embeddings(token_type_ids)#大小为[batch_size * seq_length * hidden_size(768)]
embeddings = words_embeddings + position_embeddings + token_type_embeddings #embedding之和
embeddings = self.LayerNorm(embeddings)# 层归一化
embeddings = self.dropout(embeddings)# dropout embedding
return embeddings #大小为[batch_size * seq_length * hidden_size(768)]
#回到embedding_output = self.embeddings(input_ids, position_ids=position_ids, token_type_ids=token_type_ids)

BertEncoder—Bert的核心构件

class BertEncoder(nn.Module):
def __init__(self, config):
super(BertEncoder, self).__init__()
self.output_attentions = config.output_attentions #是否输出self-attention的权重系数, 默认是不输出的
self.output_hidden_states = config.output_hidden_states #是否将每一层bert_layer的输出结果添加到all_hidden_states中
self.layer = nn.ModuleList([BertLayer(config) for _ in range(config.num_hidden_layers)]) #num_hidden_layers=12，表示12层的encoder
##############################跳转到BertLayer(),构建BERT########################################
def forward(self, hidden_states, attention_mask, head_mask=None):
'''
一般笔者使用Bert作为embedding的方式, 配置文件中output_hidden_states和output_attentions都不会设置为输出每层的结果
'''
all_hidden_states = () #用于保留每层BertLayer()的输出
all_attentions = ()
#用于保留每层self-attention的系数
for i, layer_module in enumerate(self.layer):
if self.output_hidden_states:
all_hidden_states = all_hidden_states + (hidden_states,)
layer_outputs = layer_module(hidden_states, attention_mask, head_mask[i])
hidden_states = layer_outputs[0]#每一层的输出大小均为[batch_size, seq_len, 768]
if self.output_attentions:#是否输出attention权重
all_attentions = all_attentions + (layer_outputs[1],)
if self.output_hidden_states: #是否将每一层的输出结果都保留下来
all_hidden_states = all_hidden_states + (hidden_states,)
outputs = (hidden_states,)#单独的最后一层输出
if self.output_hidden_states:
outputs = outputs + (all_hidden_states,) #(最后一层的输出, 所有层输出的集合)
if self.output_attentions:
outputs = outputs + (all_attentions,)
#(最后一层的输出, 所有层权重系数的集合)
return outputs
# last-layer hidden state, (all hidden states), (all attentions)
#回到 BertModel 中

Bert Layer—bert encoder中的每一层，核心的self-attention包含于此

class BertLayer(nn.Module):
def __init__(self, config):
super(BertLayer, self).__init__()
self.attention = BertAttention(config) #self-attention, 大小为[batch_size * seq_len * 768]
#################################跳转到BertAttention()##########################################
self.intermediate = BertIntermediate(config)#self-attention 中间层 [batch_size * seq_len * 更大的dim]
#################################跳转到BertIntermediate()#######################################
self.output = BertOutput(config) #输出层
#################################跳转到BertOutput()#############################################
def forward(self, hidden_states, attention_mask, head_mask=None):
attention_outputs = self.attention(hidden_states, attention_mask, head_mask) #self-attention, 第一层的hidden_states为“word_embedding + position_embedding + token_type_embedding”
attention_output = attention_outputs[0]#(attention_outputs[1]代表q, k相乘得到的权重系数, 默认是不返回的)
intermediate_output = self.intermediate(attention_output)#往高维映射
layer_output = self.output(intermediate_output, attention_output)
outputs = (layer_output,) + attention_outputs[1:]
#输出每一层的[batch*seq_len*768]及权重系数(默认不返回)
return outputs
# 回到BertEncoder中

BertAttention—self-attention 机制的实现

class BertAttention(nn.Module):
def __init__(self, config):
super(BertAttention, self).__init__()
self.self = BertSelfAttention(config)
#################################跳转到BertSelfAttention()#############################################
self.output = BertSelfOutput(config)
#################################跳转到BertSelfOutput()#############################################
def forward(self, input_tensor, attention_mask, head_mask=None):
self_outputs = self.self(input_tensor, attention_mask, head_mask) #得到self-attention后的input_tensor, 大小为 ([batch_size * seq_len * 768], 0), 0代表不返回attention的系数
attention_output = self.output(self_outputs[0], input_tensor)
#再做一次FC映射及layerNorm及Dropout常规操作
outputs = (attention_output,) + self_outputs[1:]
#大小为 ([batch_size * seq_len * 768], 0), 0代表不返回attention的系数
return outputs
# 回到BertLayer中
class BertSelfAttention(nn.Module):
def __init__(self, config):
super(BertSelfAttention, self).__init__()
self.output_attentions = config.output_attentions #是否输出 self_attentions 系数
self.num_attention_heads = config.num_attention_heads #multi-head 头数: 12
self.attention_head_size = int(config.hidden_size / config.num_attention_heads) #单个head的维度: 768 / 12 = 64
self.all_head_size = self.num_attention_heads * self.attention_head_size #all_head_size 与 hidden_size 大小相等, all_head_size = 12 * 64 = 768
self.query = nn.Linear(config.hidden_size, self.all_head_size) # query矩阵, hidden_size * all_head_size (768*768) 大小的矩阵
self.key = nn.Linear(config.hidden_size, self.all_head_size)
# key矩阵, hidden_size * all_head_size (768*768) 大小的矩阵
self.value = nn.Linear(config.hidden_size, self.all_head_size) # value矩阵, hidden_size * all_head_size (768*768) 大小的矩阵
self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
def transpose_for_scores(self, x):
#输入x: [batch_size * seq_len * 768]
new_x_shape = x.size()[:-1] #大小为[batch_size * seq_len]
new_x_shape = new_x_shape + (self.num_attention_heads, self.attention_head_size) #大小为[batch_size, seq_len, 12, 64]
x = x.view(*new_x_shape)
return x.permute(0, 2, 1, 3) #输出大小为[batch_size, 12, seq_len, 64], 12代表头数，64代表每个头的维度
def forward(self, hidden_states, attention_mask, head_mask=None):
'''
768 代表 hidden_size
:param hidden_states:
输入大小为[batch_size * seq_len * 768]
:param attention_mask: 输入大小为[batch_size * 1 * 1 * 768]
:param head_mask: 默认不进行剪枝操作
:return:
'''
mixed_query_layer = self.query(hidden_states) #输出为 batch_size * seq_len * 768
mixed_key_layer = self.key(hidden_states)
#输出为 batch_size * seq_len * 768
mixed_value_layer = self.value(hidden_states) #输出为 batch_size * seq_len * 768
query_layer = self.transpose_for_scores(mixed_query_layer) #转为multi-head-attention, 大小为batch_size * 12 * seq_len * 64
key_layer = self.transpose_for_scores(mixed_key_layer)
#转为multi-head-attention, 大小为batch_size * 12 * seq_len * 64
value_layer = self.transpose_for_scores(mixed_value_layer) #转为multi-head-attention, 大小为batch_size * 12 * seq_len * 64
attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2)) #输出为 batch_size * 12 * seq_len * seq_len
#self-attention 机制是 q, k 矩阵乘归一化得出系数后, 乘到 value 上
attention_scores = attention_scores / math.sqrt(self.attention_head_size) #归一化
attention_scores = attention_scores + attention_mask #[batch_size * 12 * seq_len * seq_len] 与 [batch_size*1*1*seq_len] 之和
attention_probs = nn.Softmax(dim=-1)(attention_scores) #这样可以mask掉无用信息
attention_probs = self.dropout(attention_probs) #大小为 [batch_size * 12 * seq_len * seq_len]
if head_mask is not None:
attention_probs = attention_probs * head_mask
context_layer = torch.matmul(attention_probs, value_layer)#系数与原值乘积就代表着该注意句子中的哪一部分，大小为 [batch_size * 12 * seq_len * 64]
context_layer = context_layer.permute(0, 2, 1, 3).contiguous()#大小为 [batch_size * seq_len * 12 * 64], 后续view操作需要tensor连续, 需调用contiguous()
new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,) #大小为 [batch_size * seq_len] + (768, )
context_layer = context_layer.view(*new_context_layer_shape) #大小为 [batch_size * seq_len * 768]
outputs = (context_layer, attention_probs) if self.output_attentions else (context_layer,)#选择是否输出attention系数
return outputs
# 回到BertAttention中
class BertSelfOutput(nn.Module):
def __init__(self, config):
super(BertSelfOutput, self).__init__()
self.dense = nn.Linear(config.hidden_size, config.hidden_size) #768 * 768 全连接
# self.LayerNorm = BertLayerNorm(config.hidden_size, eps=config.layer_norm_eps)
self.dropout = nn.Dropout(config.hidden_dropout_prob)
def forward(self, hidden_states, input_tensor):
hidden_states = self.dense(hidden_states)
hidden_states = self.dropout(hidden_states)
# hidden_states = self.LayerNorm(hidden_states + input_tensor)
return hidden_states
# 回到BertAttention中

BertIntermediate—对self-attention的输出结果做更高维度的映射

class BertIntermediate(nn.Module):
'''
映射self-attention后的输入[batch_size, seq_len, 768]到更大的一个维度, 维度大小为config.intermediate_size=3072
'''
def __init__(self, config):
super(BertIntermediate, self).__init__()
self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
self.intermediate_act_fn = nn.GELU()
def forward(self, hidden_states):
hidden_states = self.dense(hidden_states)
hidden_states = self.intermediate_act_fn(hidden_states)
return hidden_states
# 回到BertLayer中

BertOutput 层—bert encoder 中每层layer的输出

class BertOutput(nn.Module):
'''
BertOutput的输出代表每一层 encoder 的输出
'''
def __init__(self, config):
super(BertOutput, self).__init__()
self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
# self.LayerNorm = BertLayerNorm(config.hidden_size, eps=config.layer_norm_eps)
self.dropout = nn.Dropout(config.hidden_dropout_prob)
def forward(self, hidden_states, input_tensor):
hidden_states = self.dense(hidden_states)
hidden_states = self.dropout(hidden_states)
# hidden_states = self.LayerNorm(hidden_states + input_tensor)
return hidden_states
# 回到BertLayer中

BertPooler—pooler层的输入是encoder最后一层的输出

class BertPooler(nn.Module):
#这一层只是简单地取出了句子的第一个token
def __init__(self, config):
super(BertPooler, self).__init__()
self.dense = nn.Linear(config.hidden_size, config.hidden_size)
self.activation = nn.Tanh()
def forward(self, hidden_states):
#hidden_states 大小为 [batch_size, seq_length, hidden_size]]
#取出每一句的第一个单词，即[CLS]对应的向量，做全连接和tanh激活。得到的输出可以用来下游NSP任务
first_token_tensor = hidden_states[:, 0]
pooled_output = self.dense(first_token_tensor)
pooled_output = self.activation(pooled_output)
return pooled_output

这样整个bert流程便理过来了，主要是bertLayer构建encoder部分比较绕，因为bertlayer中又嵌套了self-attention机制。

最后

以上就是清脆翅膀为你收集整理的Pytorch Bert源码解读的全部内容，希望文章能够帮你解决Pytorch Bert源码解读所遇到的程序开发问题。

如果觉得靠谱客网站的内容还不错，欢迎将靠谱客网站推荐给程序员好友。