今天想利用Chinese-ELECTRA加载预训练权重来进一步训练,结果出现了下面的错误:
2020-08-11 22:40:26.262591: W tensorflow/core/framework/op_kernel.cc:1502] OP_REQUIRES failed at save_restore_v2_ops.cc:184 : Not found: Key discriminator_predictions/dense/bias/adam_m not found in checkpoint
ERROR:tensorflow:Error recorded from training_loop: Restoring from checkpoint failed. This is most likely due to a Variable name or other graph key that is missing from the checkpoint. Please ensure that you have not altered the graph expected based on the checkpoint. Original error:
Key discriminator_predictions/dense/bias/adam_m not found in checkpoint
[[node save/RestoreV2 (defined at run_pretraining.py:363) ]]
Original stack trace for 'save/RestoreV2':
File "run_pretraining.py", line 404, in <module>
main()
File "run_pretraining.py", line 400, in main
args.model_name, args.data_dir, **hparams))
File "run_pretraining.py", line 363, in train_or_eval
max_steps=config.num_train_steps)
File "/home/test/anaconda3/lib/python3.7/site-packages/tensorflow_estimator/python/estimator/tpu/tpu_estimator.py", line 2871, in train
saving_listeners=saving_listeners)
.....
大概意思是adam不在了,简直把我气哭了眼泪。
解决方法
我下载了exectra-base的预训练模型,放到了models文件夹下面的electra_base文件夹内,然后修改代码configure_pretraining.py:
def model_fn_builder(config: configure_pretraining.PretrainingConfig):
"""Build the model for training."""
def model_fn(features, labels, mode, params):
"""Build the model for training."""
model = PretrainingModel(config, features,
mode == tf.estimator.ModeKeys.TRAIN)
utils.log("Model is built!")
init_checkpoint = config.init_checkpoint
if config is not None:
init_checkpoint = tf.train.latest_checkpoint(config.init_checkpoint)
utils.log("Using checkpoint", init_checkpoint)
tvars = tf.trainable_variables()
scaffold_fn = None
if init_checkpoint:
assignment_map, _ = modeling.get_assignment_map_from_checkpoint(
tvars, init_checkpoint)
# print(assignment_map)
if config.use_tpu:
def tpu_scaffold():
tf.train.init_from_checkpoint(init_checkpoint, assignment_map)
return tf.train.Scaffold()
scaffold_fn = tpu_scaffold
else:
tf.train.init_from_checkpoint(init_checkpoint, assignment_map)
if mode == tf.estimator.ModeKeys.TRAIN:
train_op = optimization.create_optimizer(
model.total_loss, config.learning_rate, config.num_train_steps,
weight_decay_rate=config.weight_decay_rate,
use_tpu=config.use_tpu,
warmup_steps=config.num_warmup_steps,
lr_decay_power=config.lr_decay_power
)
output_spec = tf.estimator.tpu.TPUEstimatorSpec(
mode=mode,
loss=model.total_loss,
train_op=train_op,
training_hooks=[training_utils.ETAHook(
{} if config.use_tpu else dict(loss=model.total_loss),
config.num_train_steps, config.iterations_per_loop,
config.use_tpu)]
)
elif mode == tf.estimator.ModeKeys.EVAL:
output_spec = tf.estimator.tpu.TPUEstimatorSpec(
mode=mode,
loss=model.total_loss,
eval_metrics=model.eval_metrics,
evaluation_hooks=[training_utils.ETAHook(
{} if config.use_tpu else dict(loss=model.total_loss),
config.num_eval_steps, config.iterations_per_loop,
config.use_tpu, is_training=False)])
else:
raise ValueError("Only TRAIN and EVAL modes are supported")
return output_spec
return model_fn
主要添加了加载预训练模型的代码,然后运行命令:
python run_pretraining.py --data-dir data_speech_transformer \
--model-name electra_base_train \
--hparams config.json
注意这里model-name不要填成跟你前面electra_base一样,不然错误会让你调试哭瞎,然后就可以运行了。部分日志为:
333/10000 = 3.3%, SPS: 1.0, ELAP: 5:48, ETA: 2:48:34 - loss: 9.9810
334/10000 = 3.3%, SPS: 1.0, ELAP: 5:49, ETA: 2:48:32 - loss: 9.9372
335/10000 = 3.4%, SPS: 1.0, ELAP: 5:50, ETA: 2:48:30 - loss: 9.9835
336/10000 = 3.4%, SPS: 1.0, ELAP: 5:51, ETA: 2:48:28 - loss: 10.0930
337/10000 = 3.4%, SPS: 1.0, ELAP: 5:52, ETA: 2:48:26 - loss: 10.3831
338/10000 = 3.4%, SPS: 1.0, ELAP: 5:53, ETA: 2:48:23 - loss: 9.8002
339/10000 = 3.4%, SPS: 1.0, ELAP: 5:54, ETA: 2:48:21 - loss: 10.8526
340/10000 = 3.4%, SPS: 1.0, ELAP: 5:55, ETA: 2:48:19 - loss: 9.4941
341/10000 = 3.4%, SPS: 1.0, ELAP: 5:56, ETA: 2:48:17 - loss: 9.4434
342/10000 = 3.4%, SPS: 1.0, ELAP: 5:57, ETA: 2:48:14 - loss: 10.4484
343/10000 = 3.4%, SPS: 1.0, ELAP: 5:58, ETA: 2:48:12 - loss: 10.0355
344/10000 = 3.4%, SPS: 1.0, ELAP: 5:59, ETA: 2:48:10 - loss: 10.0131
参考文献
[1].'adam_m not found in checkpoint ' when further pretraining. https://github.com/google-research/electra/issues/45
[2].Estimator should be able to partially load checkpoints. https://github.com/tensorflow/tensorflow/issues/10155
[3].Chinese-ELECTRA. https://github.com/ymcui/Chinese-ELECTRA