当前位置 主页 > 网站技术 > 代码类 >

    tensorflow指定GPU与动态分配GPU memory设置

    栏目:代码类 时间:2020-02-03 15:10

    在tensorflow中,默认指定占用所有的GPU,如需指定占用的GPU,可以在命令行中:

    export CUDA_VISIBLE_DEVICES=1

    这样便是只占用1号GPU,通过命令

    nvidia-smi

    可以查看各个GPU的使用情况。

    另外,也可以在python程序中指定GPU,并且动态分配memory,代码如下

    import os
    import sys
    os.environ['CUDA_VISIBLE_DEVICES'] = sys.argv[1]
    import tensorflow as tf
    from keras.backend.tensorflow_backend import set_session
    config = tf.ConfigProto()
    config.gpu_options.allow_growth = True
    set_session(tf.Session(config=config))

    这样,占用的GPU就是sys.argv[1]了,并且会按需分配memory。

    以上这篇tensorflow指定GPU与动态分配GPU memory设置就是小编分享给大家的全部内容了,希望能给大家一个参考,也希望大家多多支持IIS7站长之家。