tensorflow设置GPU

一、设置GPU使用

1.一劳永逸的方法，直接在~/.bashrc中设置环境变量
CUDA_VISIBLE_DEVICES

      （1）vim ~/.bashrc
      （2）在.bashrc.最下方添加
export  CUDA_VISIBLE_DEVICES='XX'(这个xx是GPU的序号，比如我这边服务器上有8个GTX1080序号是0~7，可以填入任意多个序号，序号间用逗号隔开)
      （3）source ~/.bashrc
         这样只要每次你用TensorFlow跑程序，都只会占用指定的GPU
2.在终端设置使用的GPU
如果用方法一，虽然方便，但有的时候还是需要指定其他的GPU，这时可以这样，例如

CUDA_VISIBLE_DEVICES=2 python test.py

这样就只会使用序号为2的GPU
3.在程序中指定使用的GPU
import os
os.environ["CUDA_VISIBLE_DEVICES"]=‘6，7’
这里就设置了使用序号为6,7两个的GPU
        
事实上这3种方法都是为了设置环境变量CUDA_VISIBLE_DEVICES，方法一可以分别和方法二，三结合使用，比如虽然方法一设置了CUDA_VISIBLE_DEVICES=2，后面再用方法二设置使用GPU 0，执行的时候只使用GPU 0。

二、设置GPU占用

1、session中设置GPU

(1).定量设置显存

gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.7)
sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options)) 

或者

config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.7
session = tf.Session(config=config)
      这样运行TensorFlow程序时，每个使用的GPU中，占用的显存都不超过总显存的0.7

 

(2).按需设置显存

gpu_options = tf.GPUOptions(allow_growth=True)
sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))   

或者

config = tf.ConfigProto()
config.gpu_options.allow_growth = True
session = tf.Session(config=config)
 这样设置以后，程序就会按需占用GPU显存。

2、tf.estimator中设置GPU

(1).定量设置显存

    session_config = tf.ConfigProto(log_device_placement=True)
    session_config.gpu_options.per_process_gpu_memory_fraction = 0.5
    run_config = tf.estimator.RunConfig().replace(session_config=session_config)

    # Instantiate Estimator
    nn = tf.estimator.Estimator(model_fn=model_fn, config=run_config, params=model_params)

(2).按需设置显存

session_config = tf.ConfigProto(
      log_device_placement=True
      inter_op_parallelism_threads=0,
      intra_op_parallelism_threads=0,
      allow_soft_placement=True)

session_config.gpu_options.allow_growth = True
session_config.gpu_options.allocator_type = 'BFC'

log_device_placement=True
设置为True时，会打印出TensorFlow使用了那种操作

inter_op_parallelism_threads=0
设置线程一个操作内部并行运算的线程数，比如矩阵乘法，如果设置为０，则表示以最优的线程数处理
 
intra_op_parallelism_threads=0
设置多个操作并行运算的线程数，比如 c = a + b，d = e + f . 可以并行运算
 
allow_soft_placement=True
有时候，不同的设备，它的cpu和gpu是不同的，如果将这个选项设置成True，那么当运行设备不满足要求时，会自动分配GPU或者CPU。
 

tf.estimator参考程序:

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function 

import os
import argparse
import sys
import tempfile

# Import urllib
from six.moves import urllib

import numpy as np
import tensorflow as tf

 

FLAGS = None
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
# 开启loggging.
tf.logging.set_verbosity(tf.logging.INFO)
 

# 定义下载数据集.
def maybe_download(train_data, test_data, predict_data):
   """Maybe downloads training data and returns train and test file names."""
    if train_data:
        train_file_name = train_data
    else:
        train_file = tempfile.NamedTemporaryFile(delete=False)
        urllib.request.urlretrieve(
            "http://download.tensorflow.org/data/abalone_train.csv",
            train_file.name)
        train_file_name = train_file.name
        train_file.close()
        print("Training data is downloaded to %s" % train_file_name) 

    if test_data:
        test_file_name = test_data
    else:
        test_file = tempfile.NamedTemporaryFile(delete=False)
        urllib.request.urlretrieve(
            "http://download.tensorflow.org/data/abalone_test.csv", test_file.name)
        test_file_name = test_file.name
        test_file.close()
        print("Test data is downloaded to %s" % test_file_name)


    if predict_data:
        predict_file_name = predict_data
    else:
        predict_file = tempfile.NamedTemporaryFile(delete=False)
        urllib.request.urlretrieve(
            "http://download.tensorflow.org/data/abalone_predict.csv",
            predict_file.name)
        predict_file_name = predict_file.name
        predict_file.close()
        print("Prediction data is downloaded to %s" % predict_file_name) 

    return train_file_name, test_file_name, predict_file_name


def model_fn(features, labels, mode, params):
  """Model function for Estimator."""

  # Connect the first hidden layer to input layer
  # (features["x"]) with relu activation
  first_hidden_layer = tf.layers.dense(features["x"], 10, activation=tf.nn.relu)

  # Connect the second hidden layer to first hidden layer with relu
  second_hidden_layer = tf.layers.dense(
      first_hidden_layer, 10, activation=tf.nn.relu)

  # Connect the output layer to second hidden layer (no activation fn)
  output_layer = tf.layers.dense(second_hidden_layer, 1)

  # Reshape output layer to 1-dim Tensor to return predictions
  predictions = tf.reshape(output_layer, [-1])

 
  # Provide an estimator spec for `ModeKeys.PREDICT`.
  if mode == tf.estimator.ModeKeys.PREDICT:
    return tf.estimator.EstimatorSpec(
        mode=mode,
        predictions={"ages": predictions}) 

  # Calculate loss using mean squared error
  loss = tf.losses.mean_squared_error(labels, predictions)

  # Calculate root mean squared error as additional eval metric
  eval_metric_ops = {
      "rmse": tf.metrics.root_mean_squared_error(
          tf.cast(labels, tf.float64), predictions)
  }

 
  optimizer = tf.train.GradientDescentOptimizer(
      learning_rate=params["learning_rate"])
  train_op = optimizer.minimize(
      loss=loss, global_step=tf.train.get_global_step())

  # Provide an estimator spec for `ModeKeys.EVAL` and `ModeKeys.TRAIN` modes.
  return tf.estimator.EstimatorSpec(
      mode=mode,
      loss=loss,
      train_op=train_op,
      eval_metric_ops=eval_metric_ops)
 

# 创建main()函数，加载train/test/predict数据集.
def main(unused_argv):
    # Load datasets
    abalone_train, abalone_test, abalone_predict = maybe_download(
        FLAGS.train_data, FLAGS.test_data, FLAGS.predict_data)

    # Training examples
    training_set = tf.contrib.learn.datasets.base.load_csv_without_header(
        filename=abalone_train, target_dtype=np.int, features_dtype=np.float64)

    # Test examples
    test_set = tf.contrib.learn.datasets.base.load_csv_without_header(
        filename=abalone_test, target_dtype=np.int, features_dtype=np.float64)

 

    # Set of 7 examples for which to predict abalone ages
    prediction_set = tf.contrib.learn.datasets.base.load_csv_without_header(
        filename=abalone_predict, target_dtype=np.int, features_dtype=np.float64)

    train_input_fn = tf.estimator.inputs.numpy_input_fn(
        x={"x": np.array(training_set.data)},
        y=np.array(training_set.target),
        num_epochs=None,
        shuffle=True)

    LEARNING_RATE = 0.1
    model_params = {"learning_rate": LEARNING_RATE}

    session_config = tf.ConfigProto(log_device_placement=True)
    session_config.gpu_options.per_process_gpu_memory_fraction = 0.5
    run_config = tf.estimator.RunConfig().replace(session_config=session_config)

    # Instantiate Estimator
    nn = tf.estimator.Estimator(model_fn=model_fn, config=run_config, params=model_params)

 
    print("training---")
    nn.train(input_fn=train_input_fn, steps=5000)

    # Score accuracy
    test_input_fn = tf.estimator.inputs.numpy_input_fn(
        x={"x": np.array(test_set.data)},
        y=np.array(test_set.target),
        num_epochs=1,
        shuffle=False)


    ev = nn.evaluate(input_fn=test_input_fn)
    print("Loss: %s" % ev["loss"])
    print("Root Mean Squared Error: %s" % ev["rmse"])


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.register("type", "bool", lambda v: v.lower() == "true")
    parser.add_argument(
        "--train_data", type=str, default="", help="Path to the training data.")
    parser.add_argument(
        "--test_data", type=str, default="", help="Path to the test data.")
    parser.add_argument(
       "--predict_data",
        type=str,
        default="",
        help="Path to the prediction data.")
    FLAGS, unparsed = parser.parse_known_args()
    tf.app.run(main=main, argv=[sys.argv[0]] + unparsed)