在Tensorflow中限制GPU設備

Question

我在Python中開發了一個使用Tensorflow的應用程序和使用GPU的另一個模型。 我有一台帶有許多GPU（3xNVIDIA GTX1080）的PC，由於所有型號都嘗試使用所有可用的GPU，導致OUT_OF_MEMORY_ERROR，我發現你可以將特定的GPU分配給Python腳本

os.environ['CUDA_VISIBLE_DEVICES'] = '1'

在這里，我附上了我的FCN課程的片段

class FCN:
  def __init__(self):
    os.environ['CUDA_VISIBLE_DEVICES'] = '1'
    self.keep_probability = tf.placeholder(tf.float32, name="keep_probabilty")
    self.image = tf.placeholder(tf.float32, shape=[None, IMAGE_SIZE, IMAGE_SIZE, 3], name="input_image")
    self.annotation = tf.placeholder(tf.int32, shape=[None, IMAGE_SIZE, IMAGE_SIZE, 1], name="annotation")

    self.pred_annotation, logits = inference(self.image, self.keep_probability)
    tf.summary.image("input_image", self.image, max_outputs=2)
    tf.summary.image("ground_truth", tf.cast(self.annotation, tf.uint8), max_outputs=2)
    tf.summary.image("pred_annotation", tf.cast(self.pred_annotation, tf.uint8), max_outputs=2)
    self.loss = tf.reduce_mean((tf.nn.sparse_softmax_cross_entropy_with_logits(logits=logits,
                                                                               labels=tf.squeeze(self.annotation,
                                                                                                 squeeze_dims=[3]),
                                                                               name="entropy")))
    tf.summary.scalar("entropy", self.loss)

...

在同一個文件FCN.py ，我有一個使用該類的主要內容，當Tensorflow打印輸出時，我可以看到只使用了1個GPU，正如我所料。

if __name__ == "__main__":
  fcn = FCN()
  fcn.train_model()

  images_dir = '/home/super/datasets/MeterDataset/full-dataset-gas-images/'
  for img_file in os.listdir(images_dir):
    fcn.segment(os.path.join(images_dir, img_file))

輸出：

2018-01-09 11:31:57.351029: I tensorflow/core/common_runtime/gpu/gpu_device.cc:955] Found device 0 with properties: 
name: GeForce GTX 1080
major: 6 minor: 1 memoryClockRate (GHz) 1.7335
pciBusID 0000:09:00.0
Total memory: 7.92GiB
Free memory: 7.60GiB
2018-01-09 11:31:57.351047: I tensorflow/core/common_runtime/gpu/gpu_device.cc:976] DMA: 0 
2018-01-09 11:31:57.351051: I tensorflow/core/common_runtime/gpu/gpu_device.cc:986] 0:   Y 
2018-01-09 11:31:57.351057: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1045] Creating TensorFlow device (/gpu:0) -> (device: 0, name: GeForce GTX 1080, pci bus id: 0000:09:00.0)

當我嘗試從另一個腳本實例化FCN對象時出現問題。

def main(args):
  start_time = datetime.now()

  font = cv2.FONT_HERSHEY_SIMPLEX

  results_file = "../results.txt"
  if os.path.exists(results_file):
    os.remove(results_file)

  results_file = open(results_file, "a")

  fcn = FCN()

這里對象的創建總是使用所有3個GPU，而不是使用僅分配到__init__()方法。

這里是不受歡迎的輸出：

2018-01-09 11:41:02.537548: I 

tensorflow/core/common_runtime/gpu/gpu_device.cc:976] DMA: 0 1 2 
2018-01-09 11:41:02.537555: I tensorflow/core/common_runtime/gpu/gpu_device.cc:986] 0:   Y Y Y 
2018-01-09 11:41:02.537558: I tensorflow/core/common_runtime/gpu/gpu_device.cc:986] 1:   Y Y Y 
2018-01-09 11:41:02.537561: I tensorflow/core/common_runtime/gpu/gpu_device.cc:986] 2:   Y Y Y 
2018-01-09 11:41:02.537567: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1045] Creating TensorFlow device (/gpu:0) -> (device: 0, name: GeForce GTX 1080, pci bus id: 0000:0b:00.0)
2018-01-09 11:41:02.537571: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1045] Creating TensorFlow device (/gpu:1) -> (device: 1, name: GeForce GTX 1080, pci bus id: 0000:09:00.0)
2018-01-09 11:41:02.537574: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1045] Creating TensorFlow device (/gpu:2) -> (device: 2, name: GeForce GTX 1080, pci bus id: 0000:05:00.0)

Answer 1

這是你可以做的：

• 使用已設置的CUDA_VISIBLE_DEVICES環境變量運行腳本，如下所述 ：

   CUDA_VISIBLE_DEVICES=1 python another_script.py 

• 為Session構造函數提供顯式配置：

   config = tf.ConfigProto(device_count={'GPU': 1}) sess = tf.Session(config=config) 

  ...強制tensorflow只使用一個GPU，無論有多少GPU可用。 您還可以通過visible_device_list設置細粒度的設備列表（有關詳細信息，請參閱config.proto ）。