将Bert抽取语言表征模型修改为GPU运行

最新推荐文章于 2025-06-29 01:20:51 发布

原创

最新推荐文章于 2025-06-29 01:20:51 发布 · 4.9k 阅读

20 ·

CC 4.0 BY-SA版权

文章标签：

#深度学习 #Bert #NLP #GPU运行

本文档介绍了如何将Bert语言表征模型修改以在GPU上运行，详细阐述了在`run_pretraining.py`中涉及的四个关键修改位置，包括main函数和model_fn函数内的代码替换，旨在优化模型以适应GPU环境。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

tf.estimator.EstimatorSpec优化器类：
1.predict:只需要传入参数mode和predictions
2.eval(评估模式):需要传入参数mode,loss,eval_metric_ops
3.train:需要传入参数mode,loss,train_op

(一)、Bert抽取特征语言模型修改(GPU运行配置)

修改位置1：run_pretraining.py main()函数中：

源码1:

run_config = tf.contrib.tpu.RunConfig(
	cluster=tpu_cluster_resolver,
	master=FLAGS.master,
	model_dir=FLAGS.output_dir,
	save_checkpoints_steps=FLAGS.save_checkpoints_steps,
	tpu_config=tf.contrib.tpu.TPUConfig(
  	  iterations_per_loop=FLAGS.iterations_per_loop,
	  num_shards=FLAGS.num_tpu_cores,
	  per_host_input_for_training=is_per_host))

替换代码1:

# construct estimator
tf.logging.info("load estimator ...")
config = tf.ConfigProto(
  allow_soft_placement=True,log_device_placement=True,gpu_options={"allow_growth":True})
  #config.gpu_options.allow_growth = True 
  #config.gpu_options.per_process_gpu_memory_fraction = 0.99

run_config = tf.