tensorflow - 如何在bert中设置特定的gpu？

Question

ResourceExhaustedError（参见上面的回溯）：

分配形状 [768] 的张量和类型 float [[node bert/encoder/layer_0/attention/output/LayerNorm/beta/adam_m/Initializer/zeros（定义在 /home/zyl/souhu/bert/optimization.py： 122) = Const_class=["loc:@bert/encoder/layer_0/attention/output/LayerNorm/beta/adam_m/Assign"], dtype=DT_FLOAT, value=Tensor, _device="/job:localhost/replica:0/任务：0/设备：GPU：0“]]

如何设置gpu 1 或另一个来运行bert？

Answer 1

设置将使用哪些 GPU 的最简单方法是设置CUDA_VISIBLE_DEVICES环境变量。它仍然是GPU:0TensorFlow，不同物理上不同的设备。

如果您在 Python 中使用 BERT（这是一种相当痛苦的方式），您可以使用在块中创建 BERT 图的代码：

with tf.device('/device:GPU:1'):
   model = modeling.BertModel(...)