Tensorflow 加载模型 restore 与 init_from_checkpoint的区别

在把BERT从estimator框架中抽出来的过程中,学习整理了一下加载模型的两种方式。
请注意,两种方式中,调用sess.run(tf.global_variables_initializer())的顺序是不一样的,否则会导致: 加载失败,使用随机初始化参数。

方式一

# 先构建网络结构
build_model()
# 初始化变量
sess.run(tf.global_variables_initializer())
# 最后从checkpoint中加载已训练好的参数
saver = tf.train.Saver()
saver.restore(self.sess, init_checkpoint)

方式二

# 先构建网络结构
build_model()
# 调用init_from_checkpoint方法
tvars = tf.trainable_variables()
(assignment_map,
 initialized_variable_names) = modeling.get_assignment_map_from_checkpoint(
    tvars, init_checkpoint)
tf.train.init_from_checkpoint(init_checkpoint, assignment_map)
# 最后初始化变量
sess.run(tf.global_variables_initializer())

猜你喜欢

转载自blog.csdn.net/miangangzhen/article/details/89531378
今日推荐