深度学习之选择GPU或CPU方法

1 TensorFlow设置方法

1.1 指定GPU

法1 

import os
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "1"

上边表示使用第二块GPU运行程序,如果要使用多块,如第一块和第三块,可使用如下方法指定

os.environ["CUDA_VISIBLE_DEVICES"] = "0,2"

如果想屏蔽输出信息,使用以下设置

os.environ['TF_CPP_MIN_LOG_LEVEL']='3' #close log output

法2

import tensorflow as tf
 
with tf.device('/gpu:1'):
    v1 = tf.constant([1.0, 2.0, 3.0], shape=[3], name='v1')
    v2 = tf.constant([1.0, 2.0, 3.0], shape=[3], name='v2')
    sumV12 = v1 + v2
 
    with tf.Session(config=tf.ConfigProto(log_device_placement=True)) as sess:
        print sess.run(sumV12)

1.2 设置GPU比例

选择GPU后,可以设置每块占用的比例,因为tensorflow默认占用全部

gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.2)
with tf.Session(graph=detection_graph,
    config=tf.ConfigProto(gpu_options=gpu_options)) as sess:

 也可以设置按需分配显存

gpu_options = tf.GPUOptions(allow_growth=True)
with tf.Session(graph=detection_graph,
    config=tf.ConfigProto(gpu_options=gpu_options)) as sess:

1.3 指定CPU

如果想强制使用CPU运行tensorflow,只需要将CUDA_VISIBLE_DEVICES这个参数设为空即可,如下

os.environ["CUDA_VISIBLE_DEVICES"] = " "

如果想要指定CPU核数,使用如下方法 

cpu_num=10#指定使用的CPU个数
config = tf.ConfigProto(device_count={"CPU": cpu_num},
            inter_op_parallelism_threads = cpu_num,
            intra_op_parallelism_threads = cpu_num,
            log_device_placement=True)
# 开始训练
with tf.Session(config=config) as sess:
    #以下编写自己的代码

log_device_placement = True可以把程序运行时跑的设备情况给输出

2 PyTorch 设置方法

PyTorch默认使用从0开始的GPU,如果GPU0正在运行程序,需要指定其他GPU。

有如下两种方法来指定需要使用的GPU。

2.1 CUDA_VISIBLE_DEVICES。

类似tensorflow指定GPU的方式,使用CUDA_VISIBLE_DEVICES。

2.1.1 直接终端中设定:

CUDA_VISIBLE_DEVICES=1 python my_script.py

2.1.2 python代码中设定:

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "2"

见网址:http://www.cnblogs.com/darkknightzh/p/6591923.html

2.2. 使用函数 set_device

import torch
torch.cuda.set_device(id)

该函数见 pytorch-master\torch\cuda\__init__.py。

不过官方建议使用CUDA_VISIBLE_DEVICES,不建议使用 set_device 函数。

相关推荐
©️2020 CSDN 皮肤主题: 游动-白 设计师:白松林 返回首页