[Tensorflow] Tensorflow-gpu 확인 및 사용하기

1. 로컬 디바이스 목록 확인

from tensorflow.python.client import device_lib

print(device_lib.list_local_devices())

 

본 코드는 Python에서 Tensorflow 라이브러리를 사용하여 현재 시스템에서 사용 가능한 로컬 디바이스 목록을 확인하는 데 사용된다. 

 

 

tensorflow에서 gpu를 사용할 수 있는 환경을 잘 구축했다면 위 사진과 같이 CPU 및 GPU 정보 모두 뜰 것이다. 

 

 

2. GPU 및 CPU 사용 설정

import os

# GPU 사용시 "0"
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

# CPU 사용시 "-1"
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"

여러개의 GPU가 있다면 GPU 번호를 할당해주면 되고 CPU만을 사용하길 원한다면 -1 로 할당해주면 된다.