tensorflow 中对 tf.estimator 分配 GPU 方法

本文介绍如何在TensorFlow中指定使用特定的GPU设备及控制GPU资源使用率的方法,包括设置可见GPU、调整GPU内存分配比例及启用按需增长等功能。

摘要生成于 C知道 ,由 DeepSeek-R1 满血版支持, 前往体验 >

1. 指定使用哪一块 GPU: 

先查找有多少块 GPU,并且获得设备编号: 

$ nvidia-smi
 

在代码执行指定使用哪块GPU:

CUDA_VISIBLE_DEVICES=0 ./myapp
 

 

指定使用第0块或者第0,1块GPU:


 
  1. import os
  2. os.environ[ 'CUDA_VISIBLE_DEVICES'] = '0'
  3. os.environ[ 'CUDA_VISIBLE_DEVICES'] = '0,1'

 

2 在tf.estimator() 控制GPU资源使用率:


 
  1. session_config = tf.ConfigProto(log_device_placement= True)
  2. session_config.gpu_options.per_process_gpu_memory_fraction = 0.5
  3. run_config = tf.estimator.RunConfig().replace(session_config=session_config)
  4. # Instantiate Estimator
  5. nn = tf.estimator.Estimator(model_fn=model_fn, config=run_config, params=model_params)

 

3 tf.Config() 还可以这样用:


 
  1. 356 session_config = tf.ConfigProto(
  2. log_device_placement= True
  3. 357 inter_op_parallelism_threads= 0,
  4. 358 intra_op_parallelism_threads= 0,
  5. 359 allow_soft_placement= True)
  6. 360
  7. 361 #lingfeng
  8. 362 session_config.gpu_options.allow_growth = True
  9. 363 session_config.gpu_options.allocator_type = 'BFC'

具体解释


log_device_placement=True

  • 设置为True时,会打印出TensorFlow使用了那种操作

 

inter_op_parallelism_threads=0

  • 设置线程一个操作内部并行运算的线程数,比如矩阵乘法,如果设置为0,则表示以最优的线程数处理

 

intra_op_parallelism_threads=0

  • 设置多个操作并行运算的线程数,比如 c = a + b,d = e + f . 可以并行运算

 

allow_soft_placement=True

  • 有时候,不同的设备,它的cpu和gpu是不同的,如果将这个选项设置成True,那么当运行设备不满足要求时,会自动分配GPU或者CPU。

 

参考代码:


 
  1. from __future__ import absolute_import
  2. from __future__ import division
  3. from __future__ import print_function
  4. import os
  5. import argparse
  6. import sys
  7. import tempfile
  8. # Import urllib
  9. from six.moves import urllib
  10. import numpy as np
  11. import tensorflow as tf
  12. FLAGS = None
  13. os.environ[ 'CUDA_VISIBLE_DEVICES'] = '0'
  14. # 开启loggging.
  15. tf.logging.set_verbosity(tf.logging.INFO)
  16. # 定义下载数据集.
  17. def maybe_download(train_data, test_data, predict_data):
  18. """Maybe downloads training data and returns train and test file names."""
  19. if train_data:
  20. train_file_name = train_data
  21. else:
  22. train_file = tempfile.NamedTemporaryFile(delete= False)
  23. urllib.request.urlretrieve(
  24. "http://download.tensorflow.org/data/abalone_train.csv",
  25. train_file.name)
  26. train_file_name = train_file.name
  27. train_file.close()
  28. print( "Training data is downloaded to %s" % train_file_name)
  29. if test_data:
  30. test_file_name = test_data
  31. else:
  32. test_file = tempfile.NamedTemporaryFile(delete= False)
  33. urllib.request.urlretrieve(
  34. "http://download.tensorflow.org/data/abalone_test.csv", test_file.name)
  35. test_file_name = test_file.name
  36. test_file.close()
  37. print( "Test data is downloaded to %s" % test_file_name)
  38. if predict_data:
  39. predict_file_name = predict_data
  40. else:
  41. predict_file = tempfile.NamedTemporaryFile(delete= False)
  42. urllib.request.urlretrieve(
  43. "http://download.tensorflow.org/data/abalone_predict.csv",
  44. predict_file.name)
  45. predict_file_name = predict_file.name
  46. predict_file.close()
  47. print( "Prediction data is downloaded to %s" % predict_file_name)
  48. return train_file_name, test_file_name, predict_file_name
  49. def model_fn(features, labels, mode, params):
  50. """Model function for Estimator."""
  51. # Connect the first hidden layer to input layer
  52. # (features["x"]) with relu activation
  53. first_hidden_layer = tf.layers.dense(features[ "x"], 10, activation=tf.nn.relu)
  54. # Connect the second hidden layer to first hidden layer with relu
  55. second_hidden_layer = tf.layers.dense(
  56. first_hidden_layer, 10, activation=tf.nn.relu)
  57. # Connect the output layer to second hidden layer (no activation fn)
  58. output_layer = tf.layers.dense(second_hidden_layer, 1)
  59. # Reshape output layer to 1-dim Tensor to return predictions
  60. predictions = tf.reshape(output_layer, [ -1])
  61. # Provide an estimator spec for `ModeKeys.PREDICT`.
  62. if mode == tf.estimator.ModeKeys.PREDICT:
  63. return tf.estimator.EstimatorSpec(
  64. mode=mode,
  65. predictions={ "ages": predictions})
  66. # Calculate loss using mean squared error
  67. loss = tf.losses.mean_squared_error(labels, predictions)
  68. # Calculate root mean squared error as additional eval metric
  69. eval_metric_ops = {
  70. "rmse": tf.metrics.root_mean_squared_error(
  71. tf.cast(labels, tf.float64), predictions)
  72. }
  73. optimizer = tf.train.GradientDescentOptimizer(
  74. learning_rate=params[ "learning_rate"])
  75. train_op = optimizer.minimize(
  76. loss=loss, global_step=tf.train.get_global_step())
  77. # Provide an estimator spec for `ModeKeys.EVAL` and `ModeKeys.TRAIN` modes.
  78. return tf.estimator.EstimatorSpec(
  79. mode=mode,
  80. loss=loss,
  81. train_op=train_op,
  82. eval_metric_ops=eval_metric_ops)
  83. # 创建main()函数,加载train/test/predict数据集.
  84. def main(unused_argv):
  85. # Load datasets
  86. abalone_train, abalone_test, abalone_predict = maybe_download(
  87. FLAGS.train_data, FLAGS.test_data, FLAGS.predict_data)
  88. # Training examples
  89. training_set = tf.contrib.learn.datasets.base.load_csv_without_header(
  90. filename=abalone_train, target_dtype=np.int, features_dtype=np.float64)
  91. # Test examples
  92. test_set = tf.contrib.learn.datasets.base.load_csv_without_header(
  93. filename=abalone_test, target_dtype=np.int, features_dtype=np.float64)
  94. # Set of 7 examples for which to predict abalone ages
  95. prediction_set = tf.contrib.learn.datasets.base.load_csv_without_header(
  96. filename=abalone_predict, target_dtype=np.int, features_dtype=np.float64)
  97. train_input_fn = tf.estimator.inputs.numpy_input_fn(
  98. x={ "x": np.array(training_set.data)},
  99. y=np.array(training_set.target),
  100. num_epochs= None,
  101. shuffle= True)
  102. LEARNING_RATE = 0.1
  103. model_params = { "learning_rate": LEARNING_RATE}
  104. session_config = tf.ConfigProto(log_device_placement= True)
  105. session_config.gpu_options.per_process_gpu_memory_fraction = 0.5
  106. run_config = tf.estimator.RunConfig().replace(session_config=session_config)
  107. # Instantiate Estimator
  108. nn = tf.estimator.Estimator(model_fn=model_fn, config=run_config, params=model_params)
  109. print( "training---")
  110. nn.train(input_fn=train_input_fn, steps= 5000)
  111. # Score accuracy
  112. test_input_fn = tf.estimator.inputs.numpy_input_fn(
  113. x={ "x": np.array(test_set.data)},
  114. y=np.array(test_set.target),
  115. num_epochs= 1,
  116. shuffle= False)
  117. ev = nn.evaluate(input_fn=test_input_fn)
  118. print( "Loss: %s" % ev[ "loss"])
  119. print( "Root Mean Squared Error: %s" % ev[ "rmse"])
  120. if __name__ == "__main__":
  121. parser = argparse.ArgumentParser()
  122. parser.register( "type", "bool", lambda v: v.lower() == "true")
  123. parser.add_argument(
  124. "--train_data", type=str, default= "", help= "Path to the training data.")
  125. parser.add_argument(
  126. "--test_data", type=str, default= "", help= "Path to the test data.")
  127. parser.add_argument(
  128. "--predict_data",
  129. type=str,
  130. default= "",
  131. help= "Path to the prediction data.")
  132. FLAGS, unparsed = parser.parse_known_args()
  133. tf.app.run(main=main, argv=[sys.argv[ 0]] + unparsed)

 

参考: 使用tf.estimator中创建Estimators


评论 2
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值