[TensorFlow] 텐서플로우 설치 & GPU 사용 체크

1. 텐서플로우(TensorFlow) 설치

pip install tensorflow

 

텐서플로우는 pip 명령어를 통해 설치가 가능하다. 그냥 텐서플로우 설치 시 최신 버전이 자동으로 설치가 되며, 특정 버전을 다운로드 하고 싶다면 아래와 같이 버전을 지정할 수 있다.

pip install tensorflow == 2.10 # 원하는 버전

 

보통 텐서플로우는  GPU 연산을 활용하기 때문에 CUDA 버전에 맞는 버전을 설치해야 하는데, 이는 아래의 페이지에서 확인할 수 있다.

https://www.tensorflow.org/install/source_windows#tested_build_configurations

Windows의 소스에서 빌드,Windows의 소스에서 빌드  |  TensorFlow

 

CUDA를 설치된 텐서플로우 버전에 맞춰 설치할 수도, 그 반대일 수도 있는데, 나는 예전에 CUDA를 설치해둔적이 있어 CUDA 버전에 맞는 텐서플로우를 설치

2. GPU Check

from tensorflow.python.client import device_lib
print(device_lib.list_local_devices()[0])
print(device_lib.list_local_devices()[1])

 

텐서플로우를 설치했다면, GPU를 사용하고 있는지 체크해야한다. device_lib 모듈을 불러온 후, 로컬 다비아스를 체크.

 

device의 인덱스 0은 보통 CPU,1 이후는 GPU로, 오른쪽 사진처럼 GPU에 대한 정보가 뜬다면 성공. GPU에 대한 정보가 조회되지 않는다면 CUDA부터 재설치를 해야한다.

 

3. GPU 최적화

gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
  try:
    # Currently, memory growth needs to be the same across GPUs
    for gpu in gpus:
      tf.config.experimental.set_memory_growth(gpu, True)
    logical_gpus = tf.config.experimental.list_logical_devices('GPU')
    print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
  except RuntimeError as e:
    # Memory growth must be set before GPUs have been initialized
    print(e)

 

위의 코드를 통해 GPU의 메모리 할당을 최적화할 수 있다. 중요한 부분들을 체크해보면

gpus = tf.config.experimental.list_physical_devices('GPU')

 

GPU에 해당하는 디바이스의 목록을 반환받은 후, gpus라는 변수에 저장한 후

 

if gpus:
  try:
    # Currently, memory growth needs to be the same across GPUs
    for gpu in gpus:
      tf.config.experimental.set_memory_growth(gpu, True)

 

텐서플로우가 처음부터 모든 GPU 메모리를 할당하지 않고, 필요에 따라 점진적으로 할당하도록 설정해준다. 이를 통해 메모리 부족 문제를 해결하는데 도움을 줄 수 있다.