Jinho Tech

tensorflow-gpu, keras, Cuda가 최신 버전으로 잘 설치되어있다면 이 글을 통해 성공적으로 GPU를 사용할 수 있을 것이다.

import tensorflow
import keras
from keras.backend.tensorflow_backend import set_session

# 사용할 수 있는 GPU를 gpus에 저장
gpus = tensorflow.config.experimental.list_physical_devices('GPU')

# 나는 2번 GPU만 사용하도록 설정했다. 
# EX) 0번 GPU가 사용하고 싶다면 gpus[2]를 gpus[0]으로 바꿔주면 된다.
tensorflow.config.experimental.set_visible_devices(gpus[2], 'GPU')

위의 코드를 통해 성공적으로 GPU를 사용할 수 있는지 확인해보자.

코드를 실행한 후에 nvidia-smi 명령어를 통해 GPU가 어떻게 사용되고 있는지 확인해 볼 수 있다.

1, 3번 GPU는 다른 사용자가 사용하는 것이고, 2번 GPU를 성공적으로 사용하고 있는 것을 확인할 수 있다.

그런데 지금 메모리를 쓸데없이 엄청나게 많이 사용하고 있는데(11893MiB/12196iB) 그 이유는 tensor flow에서 기본적으로 사용 가능한 모든 GPU 메모리를 사용하기 때문이다.

필요한 만큼만 메모리를 사용하도록 코드를 추가해주자.

if gpus:
  try:
    tensorflow.config.experimental.set_memory_growth(gpus[2], True)
  except RuntimeError as e:
    print(e)

 나는 2번 GPU만 사용하기 때문에 2번 GPU만 필요한 만큼만 메모리를 사용하도록 코드를 추가해줬다. 추가해주고 다시 코드를 실행하면

메모리 사용량이 11893MiB에서 901MiB로 줄은 것을 확인할 수 있다.