如何在bert中设置特定的GPU？

Question

ResourceExhaustedError （请参阅上面的回溯）：

  

OOM，分配形状为[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，值= Tensor，_device =“ / job：localhost /副本：0 / task：0 / device：GPU：0”]] p

如何设置 gpu 1或其他名称来运行 bert ？

Answer 1

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

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

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