TensorflowGPU训练的时候默认选第0块GPU训练,同时占满所有卡的显存。
用tf.device()函数在指定训练时所用GPU
tf.device('/gpu:0')
通过CUDA_VISIBLE_DEVICES来指定(这时只用一块GPU的内存)
import os
os.environ['CUDA_VISIBLE_DEVICES']='0'
# os.environ['CUDA_VISIBLE_DEVICES']='1'
运行python程序前指定:
CUDA_VISIBLE_DEVICES=0 python train.py