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

zhangchun3211

于 2020-01-15 13:43:48 发布

阅读量831

点赞数

CC 4.0 BY-SA版权

分类专栏： tensorflow 文章标签： tensorflow 深度学习

本文链接：https://blog.youkuaiyun.com/zhangchun3211/article/details/103987928

tensorflow 专栏收录该内容

1 篇文章

订阅专栏

本文探讨了在使用TensorFlow进行多GPU部署时，显存分配不均的问题及原因。os.environ设置GPU可见设备后，代码内部将其转换为0和1，导致tf.device无法直接识别实际的GPU编号。通过调整配置参数allow_soft_placement，可以实现自动显存分配，避免报错。文章提供了正确的代码示例，帮助读者理解并解决类似问题。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

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占的显存相差巨大。
GPU占用情况

原因

实际上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:

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

原因

解决办法

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