tensorflow 多GPU并行训练 遇到的问题

os.environ['CUDA_VISIBLE_DEVICES'] = "5,6" 
##以下，将graph部署到两个GPU的时候i分别是5和6
with tf.device('/gpu:%d' % i):
            with tf.name_scope('tower_%d' % i) as scope:

在训练的时候发现卡5和卡6占的显存相差巨大。

原因

实际上os.environ那句设置好了使用的卡号之后，在代码里已经自动转为0,1了，所以后面tf.device找不到卡5和卡6；
之所以没有报错是以下语句会在找不到设备的时候自动分配显存。
with tf.Session(config=tf.ConfigProto(allow_soft_placement=True)) as sess:
如果设置成了False，还按照上面的代码训练就会报错了。

解决办法

##以下i分别为0和1.
with tf.device('/gpu:%d' % i):
            with tf.name_scope('tower_%d' % i) as scope: