问题解决BERT模型使用及一个问题：NotFoundError: Key bert_1/embeddings/LayerNorm/beta not found in checkpoint

Posted 2022-05-19 囚生CY

关于BERT模型的调用，这几天基本上是摸得比较清楚了。

模型源码在https://github.com/google-research/bert，截至本文发布，该项目的Readme.md文件中提供了如下图所示的9个模型的下载链接。前两个是区分大小写的英文模型，第三个是中文模型，4589没有用过具体不太清楚，六七是不区分大小写的英文模型（根据Readme.md中的描述，如果对大小写不是很敏感的话用uncased已经完全足够了，但是我觉得像GEC这种任务应该对大小写还是相当敏感的）。文件夹名称中的"bert_"前缀是我自己添加以区分其他模型的。

下载好所有模型分别解压即为上图所得，每个文件夹中的文件命名与数量都是相同的，如下图所示：

一共五个文件，bert_config.json是配置文件，vocab.txt是对应模型使用的token集合，其他三个ckpt文件即为模型。

调用分两步，第一步先把文本转化为BERT模型输入的形式（一共三个输入参数：input_ids，input_mask，token_type_ids），然后调用模型使用即可得到文本的向量表示，如下所示👇

# -*- coding: UTF-8 -*-
# Author: 囚生
# 调用BERT模型的工具函数

import os
import tensorflow as tf

from bert import modeling,tokenization

def text2input(text,tokenizer,											 # 接收三个参数: 超参数, 文本, 分词器
	maxlen=100,															 # 文本token最大数
	return_tensor=True,													 # 是否返回tensor类型的结果: 否则返回list类型
):																		 # 将文本转化为BERT输入
	tokens = tokenizer.tokenize(text)									 # 分词器分词
	if len(tokens)>maxlen-2: tokens = tokens[:maxlen-2]					 # 注意tokens的数量不能超过maxlen-2, 因为头尾还需要加句首与分句标志
	tokens_bert = ["[CLS]"]										 # 存放token的列表: 置入句首标志
	token_type_ids = [0]												 # 标识token属句类别的列表: 置入句首标志的标识
	for token in tokens:												 # 添加token与对应标识
		tokens_bert.append(token)										 # 添加token
		token_type_ids.append(0)										 # 这个表示一般用0,1,2,...表示是第几句话, 该函数一般只接收一个句子, 因此都是0
	tokens_bert.append("[SEP]")											 # 置入分句标志
	token_type_ids.append(0)											 # 置入分句标志的标识
	input_ids = tokenizer.convert_tokens_to_ids(tokens_bert)			 # 将tokens转化为input_ids
	input_mask = [1]*len(input_ids)										 # 设置蒙布
	while len(input_ids)<maxlen:										 # 对input_ids,input_mask,token_type_ids进行padding
		input_ids.append(0)												 # 零填充
		input_mask.append(0)											 # 零填充
		token_type_ids.append(0)										 # 零填充
	if return_tensor: 													 # 若返回tensor类型
		input_ids = tf.convert_to_tensor([input_ids],dtype=tf.int32,name="input_ids")
		input_mask = tf.convert_to_tensor([input_mask],dtype=tf.int32,name="input_mask")
		token_type_ids = tf.convert_to_tensor([token_type_ids],dtype=tf.int32,name="token_type_ids")
	return input_ids,input_mask,token_type_ids							 # 返回BERT输入的三个参数

def load_model(input_ids,input_mask,token_type_ids,cpath,mpath,
):																		 # 模型载入
	config = modeling.BertConfig.from_json_file(cpath)					 # 载入配置文件
	config_session = tf.ConfigProto()									 # 创建对象配置session运行参数
	config_session.gpu_options.allow_growth = True						 # 动态申请显存
	with tf.Session(config=config_session).as_default() as session:
		model = modeling.BertModel(										 # 载入模型
			config=config,												 # BERT配置信息
			is_training=True,											 # 训练模式
			input_ids=input_ids,										 # 输入参数: 输入token的索引
			input_mask=input_mask,										 # 输入参数: 蒙布
			token_type_ids=token_type_ids,								 # 输入参数: 
			use_one_hot_embeddings=False,								 # 不使用one-hot编码
		)
		saver = tf.train.Saver()										 # 训练保存器
		session.run(tf.global_variables_initializer())					 # 先初始化, 再加载参数，否则会把BERT的参数重新初始化
		saver.restore(session,mpath)									 # 保存模型到ckpt文件
		
	sequence_output = model.get_sequence_output()						 # 获取每个token的输出: shape(batch_size,sequence_length,embedding_size)
	pooled_output = model.get_pooled_output()							 # 获取每个分句的输出: shape(batch_size,embedding_size)
	layers = model.all_encoder_layers									 # 获取所有层的输出: shape(batch_size,sequence_length,embedding_size)
	embedding_output = model.get_embedding_output()
	embedding_table = model.get_embedding_table()

	'''
	with tf.Session() as session:
		session.run(tf.global_variables_initializer())

		sequence_output = session.run(sequence_output)
		pooled_output = session.run(pooled_output)
		embedding_output = session.run(embedding_output)
		embedding_table = session.run(embedding_table)

		print("sequence_output: ".format(sequence_output.shape))		 # (1,32,768)
		print("pooled_output: ".format(pooled_output.shape))			 # (1,768)
		print("embedding_output: ".format(embedding_output.shape))	 # (1,32,768)
		print("embedding_table: ".format(embedding_table.shape))		 # (28996,768)
		
		for layer in layers:
			print(layer.shape)											 # (1,32,768)
	'''
	
	return layers,embedding_output,pooled_output,embedding_table

if __name__ == "__main__":
	# 以下4个路径变量请根据自己的实际情况修改
	root = "otherdata/model/bert_cased_L-12_H-768_A-12"
	vpath = os.path.join(root,"vocab.txt")								 # 词汇表文件
	cpath = os.path.join(root,"bert_config.json")
	mpath = os.path.join(root,"bert_model.ckpt")						 # 这个文件其实不存在, 但是就得这么写, 我也不知道为什么
	
	tokenizer = tokenization.FullTokenizer(vpath)						 # 这个老版本中可能是CharTokenizer类, 目前源码中不存在该类了
	text = "This will , if not already , cause problems as there is very limited space for us ."
	input_ids,input_mask,token_type_ids = text2input(hp,text,tokenizer,32)
	load_model(input_ids,input_mask,token_type_ids,cpath,mpath)
	

注意代码中调用到的bert库，既可以直接从https://github.com/google-research/bert下载后放到同一目录下，因为该项目本身就是bert库的源码，也可以直接pip install bert，其也是从github上下载源码。

上述代码中的text2input是针对单句文字的，多句文字可以参考其他人的博客，因为我自己查下来关于BERT模型调用还是有不少人在详细写过的。

我特别说一下模型输出的问题，通过查看modeling.py中的源码，可以看到BertModel类中定义了5个实例函数：

	sequence_output = model.get_sequence_output()						 # 获取每个token的输出: shape(batch_size,sequence_length,embedding_size)
	pooled_output = model.get_pooled_output()							 # 获取每个分句的输出: shape()
	layers = model.all_encoder_layers									 # 获取所有层的输出: shape()
	embedding_output = model.get_embedding_output()
	embedding_table = model.get_embedding_table()

第三个输出是12层encoders输出的列表，该列表中最后一个元素即为sequence_output，其余不多做解释。

我主要说一下在实际使用中的问题，当我在需要多次调用模型对多个句子进行模型输出时，出现了如下的报错：

NotFoundError (see above for traceback): Key bert_1/embeddings/LayerNorm/beta not found in checkpoint
         [[Node: save_1/RestoreV2 = RestoreV2[dtypes=[DT_FLOAT, DT_FLOAT, DT_FLOAT, DT_FLOAT, DT_FLOAT, ..., DT_FLOAT, DT_FLOAT, DT_FLOAT, DT_FLOAT, DT_FLOAT], _device="/job:localhost/replica:0/task:0/device:CPU:0"](_arg_save_1/Const_0_0, save_1/RestoreV2/tensor_names, save_1/RestoreV2/shape_and_slices)]]
         [[Node: save_1/RestoreV2/_307 = _Recv[client_terminated=false, recv_device="/job:localhost/replica:0/task:0/device:GPU:0", send_device="/job:localhost/replica:0/task:0/device:CPU:0", send_device_incarnation=1, tensor_name="edge_312_save_1/RestoreV2", tensor_type=DT_FLOAT, _device="/job:localhost/replica:0/task:0/device:GPU:0"]()]]

事实上把上述模型调用的代码仅需把最后一行复制一遍（即调用两次模型），就会如上报错，如下所示👇

if __name__ == "__main__":
	# 以下4个路径变量请根据自己的实际情况修改
	root = "otherdata/model/bert_cased_L-12_H-768_A-12"
	vpath = os.path.join(root,"vocab.txt")								 # 词汇表文件
	cpath = os.path.join(root,"bert_config.json")
	mpath = os.path.join(root,"bert_model.ckpt")						 # 这个文件其实不存在, 但是就得这么写, 我也不知道为什么
	
	tokenizer = tokenization.FullTokenizer(vpath)						 # 这个老版本中可能是CharTokenizer类, 目前源码中不存在该类了
	text = "This will , if not already , cause problems as there is very limited space for us ."
	input_ids,input_mask,token_type_ids = text2input(hp,text,tokenizer,32)
	load_model(input_ids,input_mask,token_type_ids,cpath,mpath)
	load_model(input_ids,input_mask,token_type_ids,cpath,mpath)			 # 调用两次

这个问题确实很让人困扰，只能调用一次的模型未免太鸡肋了。

通过检查排错，最后找到了问题所在：

在30行处添加tf.reset_default_graph()，作用是清除默认图的堆栈，并设置全局图为默认图，如下所示

# -*- coding: UTF-8 -*-
# Author: 囚生
# 调用BERT模型的工具函数

import os
import tensorflow as tf

from bert import modeling,tokenization

def text2input(text,tokenizer,											 # 接收三个参数: 超参数, 文本, 分词器
	maxlen=100,															 # 文本token最大数
	return_tensor=True,													 # 是否返回tensor类型的结果: 否则返回list类型
):																		 # 将文本转化为BERT输入
	tokens = tokenizer.tokenize(text)									 # 分词器分词
	if len(tokens)>maxlen-2: tokens = tokens[:maxlen-2]					 # 注意tokens的数量不能超过maxlen-2, 因为头尾还需要加句首与分句标志
	tokens_bert = ["[CLS]"]										 # 存放token的列表: 置入句首标志
	token_type_ids = [0]												 # 标识token属句类别的列表: 置入句首标志的标识
	for token in tokens:												 # 添加token与对应标识
		tokens_bert.append(token)										 # 添加token
		token_type_ids.append(0)										 # 这个表示一般用0,1,2,...表示是第几句话, 该函数一般只接收一个句子, 因此都是0
	tokens_bert.append("[SEP]")											 # 置入分句标志
	token_type_ids.append(0)											 # 置入分句标志的标识
	input_ids = tokenizer.convert_tokens_to_ids(tokens_bert)			 # 将tokens转化为input_ids
	input_mask = [1]*len(input_ids)										 # 设置蒙布
	while len(input_ids)<maxlen:										 # 对input_ids,input_mask,token_type_ids进行padding
		input_ids.append(0)												 # 零填充
		input_mask.append(0)											 # 零填充
		token_type_ids.append(0)										 # 零填充
	if return_tensor: 													 # 若返回tensor类型
		tf.reset_default_graph()
		input_ids = tf.convert_to_tensor([input_ids],dtype=tf.int32,name="input_ids")
		input_mask = tf.convert_to_tensor([input_mask],dtype=tf.int32,name="input_mask")
		token_type_ids = tf.convert_to_tensor([token_type_ids],dtype=tf.int32,name="token_type_ids")
	return input_ids,input_mask,token_type_ids							 # 返回BERT输入的三个参数

def load_model(input_ids,input_mask,token_type_ids,cpath,mpath,
):																		 # 模型载入
	config = modeling.BertConfig.from_json_file(cpath)					 # 载入配置文件
	config_session = tf.ConfigProto()									 # 创建对象配置session运行参数
	config_session.gpu_options.allow_growth = True						 # 动态申请显存
	with tf.Session(config=config_session).as_default() as session:
		model = modeling.BertModel(										 # 载入模型
			config=config,												 # BERT配置信息
			is_training=True,											 # 训练模式
			input_ids=input_ids,										 # 输入参数: 输入token的索引
			input_mask=input_mask,										 # 输入参数: 蒙布
			token_type_ids=token_type_ids,								 # 输入参数: 
			use_one_hot_embeddings=False,								 # 不使用one-hot编码
		)
		saver = tf.train.Saver()										 # 训练保存器
		session.run(tf.global_variables_initializer())					 # 先初始化, 再加载参数，否则会把BERT的参数重新初始化
		saver.restore(session,mpath)									 # 保存模型到ckpt文件
		
	sequence_output = model.get_sequence_output()						 # 获取每个token的输出: shape(batch_size,sequence_length,embedding_size)
	pooled_output = model.get_pooled_output()							 # 获取每个分句的输出: shape(batch_size,embedding_size)
	layers = model.all_encoder_layers									 # 获取所有层的输出: shape(batch_size,sequence_length,embedding_size)
	embedding_output = model.get_embedding_output()
	embedding_table = model.get_embedding_table()

	'''
	with tf.Session() as session:
		session.run(tf.global_variables_initializer())

		sequence_output = session.run(sequence_output)
		pooled_output = session.run(pooled_output)
		embedding_output = session.run(embedding_output)
		embedding_table = session.run(embedding_table)

		print("sequence_output: ".format(sequence_output.shape))		 # (1,32,768)
		print("pooled_output: ".format(pooled_output.shape))			 # (1,768)
		print("embedding_output: ".format(embedding_output.shape))	 # (1,32,768)
		print("embedding_table: ".format(embedding_table.shape))		 # (28996,768)
		
		for layer in layers:
			print(layer.shape)											 # (1,32,768)
	'''
	
	return layers,embedding_output,pooled_output,embedding_table

if __name__ == "__main__":
	# 以下4个路径变量请根据自己的实际情况修改
	root = "otherdata/model/bert_cased_L-12_H-768_A-12"
	vpath = os.path.join(root,"vocab.txt")								 # 词汇表文件
	cpath = os.path.join(root,"bert_config.json")
	mpath = os.path.join(root,"bert_model.ckpt")						 # 这个文件其实不存在, 但是就得这么写, 我也不知道为什么
	
	tokenizer = tokenization.FullTokenizer(vpath)						 # 这个老版本中可能是CharTokenizer类, 目前源码中不存在该类了
	text = "This will , if not already , cause problems as there is very limited space for us ."
	input_ids,input_mask,token_type_ids = text2input(text,tokenizer,32)
	load_model(input_ids,input_mask,token_type_ids,cpath,mpath)
	load_model(input_ids,input_mask,token_type_ids,cpath,mpath)			 # 调用两次
	

总之tensorflow用起来确实挺麻烦，我使用的版本是tensorflow-gpu v1.9.0，因为gpu版本再高一点的我的显卡不支持，而我又不想卸了用cpu版本的了，建议如果使用tensorflow的话要么就用2.0版本以后的，要么就用v1.12.0这个长期更新的版本。总之tensorflow不同版本间的差异确实很大，别人的代码很多时候不改都是跑不了的。

可能pytorch还是好一点吧，可惜现在主流的几个模型transformer，bert源码都是用的tensorflow架构的，也没什么办法了。

分享学习，共同进步！