[Tensorflow] 텐서플로우 GPU 확인, 특정 GPU 사용 방법

텐서플로우 GPU 사용 가능 여부 체크, 원하는 GPU 지정 하기

안녕하세요.

이번 글에서는 tensorflow 2 버전에서 gpu가 연결되었는지를 점검하고,

원하는 gpu를 사용할 수 있도록 하는 방법에 대해서 간단히 설명해보도록 하겠습니다.

 

 

Tensorflow gpu 연결 확인 방법

아래의 코드를 작성해서 손쉽게 gpu 사용 가능 여부를 확인할 수 있습니다.

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

만일, CPU만 연결된 경우는 아래와 같이 메시지가 출력됩니다.

반면, GPU가 1개 이상 연결된 경우는 아래쪽에 GPU:0, GPU:1, ... 식으로

번호가 할당된 디바이스가 추가로 탐색됩니다.

 

 

텐서플로우에서 원하는 gpu 사용 방법

여러개의 gpu 중에서 하나를 선택하여 특정 코드 부분을 실행하기를 원하거나,

gpu가 있는 상황에서도 cpu 사용을 원한다면

아래와 같이 디바이스의 번호를 지정해주시면 됩니다.

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

GPU:n번을 사용하려면 번호를 n으로 지정해주시면 되며,

(위의 예시에서는 GPU:0번이 사용됩니다.)

CPU 강제 사용을 원하신다면 -1로 번호를 선택해주시면 됩니다.

 

 

혹은 with 구문을 통해서도 특정 부분의 코드에서만 디바이스를 지정할 수 있습니다.

# GPU 사용을 원하는 경우
with tf.device('/device:GPU:0'): 
    # 원하는 코드 작성(들여쓰기 필수)

# CPU 사용을 원하는 경우
with tf.device('/cpu:0'): 
    # 원하는 코드 작성(들여쓰기 필수)