bert 限制GPU使用到特定的GPU设备

mwyxok5s 于 4个月前发布在其他

关注(0)|答案(2)|浏览(69)

要限制run_pretraining.py仅在单个特定的GPU设备上运行，最好的方法是在代码中指定设备。在导入TensorFlow之前，可以插入以下代码：

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"  # 将0替换为所需的GPU设备编号
import tensorflow as tf

将上述代码中的"0"替换为所需的GPU设备编号即可。

bert

来源：https://github.com/google-research/bert/issues/767

2条答案

按热度按时间

w8ntj3qf1#

Having struggled with the same issue. The tpuEstimator uses only the 1st GPU, and I have tried multiples including
os.environ["CUDA_VISIBLE_DEVICES"]="0,1,2,3"
or
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.gpu_options.visible_device_list = '1'
config.graph_options.optimizer_options.global_jit_level = tf.OptimizerOptions.ON_1
and add "session_config=config" in "tf.contrib.tpu.RunConfig"
None of these worked. Desperately in need of help!

赞(0）回复(0）举报 4个月前

jchrr9hc2#

刚刚发现第二个解决方案实际上是有效的，之前我还有一个死掉的内存中的会话，它给我报错了。以下应该可以工作：

config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.gpu_options.visible_device_list = '1,2,3'
config.graph_options.optimizer_options.global_jit_level = tf.OptimizerOptions.ON_1

 # set run_config
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),
        session_config=config
    )

赞(0）回复(0）举报 4个月前

我来回答

bert 限制GPU使用到特定的GPU设备

2条答案

相关问题

热门标签

最新问答