멀티 GPU 시스템에서 하나의 GPU만 사용하기

방법 1

tf.debugging.set_log_device_placement(True)

try:
  # 유효하지 않은 GPU 장치를 명시
  with tf.device('/device:GPU:2'):
    a = tf.constant([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
    b = tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
    c = tf.matmul(a, b)
except RuntimeError as e:
  print(e)

ref: https://www.tensorflow.org/guide/gpu?hl=ko#%EB%A9%80%ED%8B%B0_gpu_%EC%8B%9C%EC%8A%A4%ED%85%9C%EC%97%90%EC%84%9C_%ED%95%98%EB%82%98%EC%9D%98_gpu%EB%A7%8C_%EC%82%AC%EC%9A%A9%ED%95%98%EA%B8%B0

GPU 사용하기  |  TensorFlow Core

 

 

방법 2

 

아래의 코드를 맨 처음 부분에 넣어주면 됨.

Ex> os.environ["CUDA_VISIBLE_DEVICES"]="0" ,    GPU 0를 사용하겠음 선언.

import os

os.environ["CUDA_VISIBLE_DEVICES"]="0"
gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
    try:
        tf.config.experimental.set_memory_growth(gpus[0], True)
    except RuntimeError as e:
        print(e)