[Tensorflow] 텐서플로우 GPU이용하여 실행하도록 하는 방법

tensorflow 를 이용하여 머신러닝을 실행할때, cpu나 gpu를 사용하게 된다. 
gpu를 시용하면 cpu만 사용하는 것에 비해 속도가 약 10배 이상 차이가 난다. 그렇기에 모델을 학습시킬때 gpu를 사용하는 것이 훨씬 효율적이며, cpu에 부하도 훨씬 적을 것입니다.

아래에서 gpu로 실행하도록 설정하는 두가지 방법에 대해서 알아보겠습니다.

 

아래 방법 사용 전 gpu device가 잡히는지 확인해야 합니다.

from tensorflow.python.client import device_lib
print(device_lib.list_local_devices())

위 코드에서 gpu가 잡히는게 확인 되었다면 아래 방법을 실행하면 됩니다.

방법 1: 전체 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의 넘버를 적어주면 된다.

 

방법 2 : 원하는 부분에만 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)

어떤 GPU를 사용할 것인지 지정해주어야 한다.