Keras指定GPU训练模式，设置GPU的使用量

缺乏、安全感 2022-10-06 06:47 187阅读 0赞

最近又要用到深度学习模型训练这块的内容了，之前有些操作都忘记了，这里就想着把这些内容记录一下备忘。

我们使用的是阿里云GPU服务器，由于有多个进程同时在训练计算，这里需要限定单个进程里面的GPU使用量，

因为我只有一块显卡，所以这里指定GPU的编号都是 0 ，设定GPU使用量经常用到的方式主要有一下三种：

1、直接指定GPU的使用比例

import os
    import tensorflow as tf
    import keras.backend.tensorflow_backend as KTF
    os.environ["CUDA_VISIBLE_DEVICES"] = "0"
    
    
    
    #设定使用 40% 的GPU资源
    config = tf.ConfigProto()
    config.gpu_options.per_process_gpu_memory_fraction = 0.4
    session = tf.Session(config=config)
    # 设置session
    KTF.set_session(session)

2、设定GPU按照实时计算所需动态增长变化

import tensorflow as tf
    import keras.backend.tensorflow_backend as KTF
    
    config = tf.ConfigProto()  
    config.gpu_options.allow_growth=True   #不全部占满显存, 按需分配
    session = tf.Session(config=config)
    
    # 设置session
    KTF.set_session(session)

3、前面两种方式的综合，限定用量但是可以动态增长变化

import os
    import tensorflow as tf
    import keras.backend.tensorflow_backend as KTF
    
    
    os.environ["CUDA_VISIBLE_DEVICES"] = "0"
    config = tf.ConfigProto()  
    config.gpu_options.allow_growth=True   #不全部占满显存, 按需分配
    session = tf.Session(config=config)
    KTF.set_session(session)

然后，需要查看GPU使用情况的话Nvidia提供了 nvidia-smi 命令，可以很方便地进行查看，如下所示：

![watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L1RvZ2V0aGVyX0Na_size_16_color_FFFFFF_t_70][]

如果想要动态刷新的话，我经常使用的是下面的命令：

watch -n 1 -d nvidia-smi

效果如下所示：

![20210611184907133.gif][]

后续有用的技巧都会持续更新备忘！

[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L1RvZ2V0aGVyX0Na_size_16_color_FFFFFF_t_70]: /images/20221005/712bbacac55b4d39b2931774bc66ea2e.png
[20210611184907133.gif]: /images/20221005/27a49962827b475097d437d5491fe9fe.png