[Tensorflow] 텐서플로우 GPU로 실행하도록 설정하는 방법

 

 

tensorflow나 keras를 사용해 ML 코드를 돌릴 때 CPU가 아닌 GPU에서 돌아가도록 설정하는 방법에 대해 소개한다. 방법은 크게 2가지 정도가 있는데 본인이 원하는 상황에 따라 맞게 사용하면 된다.

 

방법 1 : 원하는 부분에만 GPU로 실행하도록 하기

---

# 텐서를 GPU에 할당
with tf.device('/GPU:0'):
  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)
print(c)

위와 같이 with문으로 사용하면 특정 컨텍스트를 묶어 해당 부분만을 특정 GPU로 실행되도록 한다. 이 때, 어떤 GPU를 사용할 것인지 특정해주어야 한다.

 

방법 2 : 전체에 대해 GPU로 실행하도록 하기

---

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)

위의 코드를 처음 부분에 넣어주면 GPU 아래의 모든 부분에 대해 GPU로 실행하도록 한다. 이 때 os.environ["CUDA_VISIBLE_DEVICES"]에 사용하고자 하는 GPU의 넘버를 적어주면 된다.

 

만약 본인이 제대로 설정을 해서 GPU로 돌아가고 있는게 맞는지 확인하려면 nvidia-smi 커맨드를 사용해 자신이 특정한 GPU의 사용량을 보고 확인하면 된다.