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

    tensorflow使用指定gpu的方法

    栏目:代码类 时间:2020-02-04 12:05

    TensorFlow是一个基于数据流编程(dataflow programming)的符号数学系统,被广泛应用于各类机器学习(machine learning)算法的编程实现,其前身是谷歌的神经网络算法库DistBelief [1]  。
    Tensorflow拥有多层级结构,可部署于各类服务器、PC终端和网页并支持GPU和TPU高性能数值计算,被广泛应用于谷歌内部的产品开发和各领域的科学研究 。

    TensorFlow由谷歌人工智能团队谷歌大脑(Google Brain)开发和维护,拥有包括TensorFlow Hub、TensorFlow Lite、TensorFlow Research Cloud在内的多个项目以及各类应用程序接口(Application Programming Interface, API) 。自2015年11月9日起,TensorFlow依据阿帕奇授权协议(Apache 2.0 open source license)开放源代码 。

    持续监控GPU使用情况命令:

    $ watch -n 10 nvidia-smi

    一、指定使用某个显卡

    如果机器中有多块GPU,tensorflow会默认吃掉所有能用的显存, 如果实验室多人公用一台服务器,希望指定使用特定某块GPU。
    可以在文件开头加入如下代码:

    import os
    os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
    os.environ["CUDA_VISIBLE_DEVICES"] = "1"  # 使用第二块GPU(从0开始)

    也可以制定使用某几块GPU

    import os
    os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
    os.environ["CUDA_VISIBLE_DEVICES"] = "0, 2" # 使用第一, 三块GPU

    禁用GPU

    import os
    os.environ["CUDA_VISIBLE_DEVICES"] = "-1"

    支持的设备

    在一套标准系统中通常有多台计算设备。TensorFlow 支持 CPU 和 GPU 这两种设备。它们均用 strings 表示。例如:

    "/cpu:0":机器的 CPU。
    "/device:GPU:0":机器的 GPU(如果有一个)。
    "/device:GPU:1":机器的第二个 GPU(以此类推)。

    如果 TensorFlow 指令中兼有 CPU 和 GPU 实现,当该指令分配到设备时,GPU 设备有优先权。例如,如果 matmul 同时存在 CPU 和 GPU 核函数,在同时有 cpu:0 和 gpu:0 设备的系统中,gpu:0 会被选来运行 matmul。

    记录设备分配方式

    要找出您的指令和张量被分配到哪个设备,请创建会话并将 log_device_placement 配置选项设为 True。

    #Creates a graph.
    a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[2, 3], name='a')
    b = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[3, 2], name='b')
    c = tf.matmul(a, b)
    #Creates a session with log_device_placement set to True.
    sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))
    #Runs the op.
    print(sess.run(c))

    应该会看到以下输出内容:

    Device mapping:
    /job:localhost/replica:0/task:0/device:GPU:0 -> device: 0, name: Tesla K40c, pci bus
    id: 0000:05:00.0
    b: /job:localhost/replica:0/task:0/device:GPU:0
    a: /job:localhost/replica:0/task:0/device:GPU:0
    MatMul: /job:localhost/replica:0/task:0/device:GPU:0
    [[ 22. 28.]
     [ 49. 64.]]

    手动分配设备

    如果您希望特定指令在您选择的设备(而非系统自动为您选择的设备)上运行,您可以使用 with tf.device 创建设备上下文,这个上下文中的所有指令都将被分配在同一个设备上运行。

    # Creates a graph.
    with tf.device('/cpu:0'):
     a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[2, 3], name='a')
     b = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[3, 2], name='b')
    c = tf.matmul(a, b)
    # Creates a session with log_device_placement set to True.
    sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))
    # Runs the op.
    print(sess.run(c))