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 函数。