Tensorflow指定GPU训练

以下假设使用第3块GPU训练
1.在代码中添加
第一种方法:

tf.device(‘/gpu:2’)

第二种方法:

import os
os.environ[‘CUDA_VISIBLE_DEVICES’]=’2’

如果要指定1,3块GPU来训练,则:

os.environ[‘CUDA_VISIBLE_DEVICES’]=’0,2’

2.在命令中添加

CUDA_VISIBLE_DEVICES=2 python train.py

猜你喜欢

转载自blog.csdn.net/baidu_27280587/article/details/80430226