ckpt文件restore的两种方法

朱雀 2021-09-28 12:52 805阅读 0赞

这里以BERT恢复模型为例，说明恢复ckpt的两种常见方法

\# 方法一：

init_checkpoint = "chinese_L-12_H-768_A-12/bert_model.ckpt"
    use_tpu = False
    # 获取模型中所有的训练参数。
    tvars = tf.trainable_variables()
    # 加载BERT模型
    
    (assignment_map, initialized_variable_names) = modeling.get_assignment_map_from_checkpoint(tvars,
                                                                                           init_checkpoint)
    
    tf.train.init_from_checkpoint(init_checkpoint, assignment_map)
    
    tf.logging.info("**** Trainable Variables ****")
    # 打印加载模型的参数
    for var in tvars:
        init_string = ""
        if var.name in initialized_variable_names:
            init_string = ", *INIT_FROM_CKPT*"
        tf.logging.info("  name = %s, shape = %s%s", var.name, var.shape,
                        init_string)
    with tf.Session() as sess:
        sess.run(tf.global_variables_initializer())

\# 方法二：

pathname = "chinese\_L-12\_H-768\_A-12/bert\_model.ckpt" \# 模型地址  
bert\_config = modeling.BertConfig.from\_json\_file("chinese\_L-12\_H-768\_A-12/bert\_config.json")\# 配置文件地址。  
configsession = tf.ConfigProto()  
configsession.gpu\_options.allow\_growth = True  
sess = tf.Session(config=configsession)  
input\_ids = tf.placeholder(shape=\[64, 128\], dtype=tf.int32, name="input\_ids")  
input\_mask = tf.placeholder(shape=\[64, 128\], dtype=tf.int32, name="input\_mask")  
segment\_ids = tf.placeholder(shape=\[64, 128\], dtype=tf.int32, name="segment\_ids")

with sess.as\_default():  
    model = modeling.BertModel(  
        config=bert\_config,  
        is\_training=True,  
        input\_ids=input\_ids,  
        input\_mask=input\_mask,  
        token\_type\_ids=segment\_ids,  
        use\_one\_hot\_embeddings=False)  
    saver = tf.train.Saver()  
    sess.run(tf.global\_variables\_initializer())\# 这里尤其注意，先初始化，在加载参数，否者会把bert的参数重新初始化。这里和demo1是有区别的  
    saver.restore(sess, pathname)