tensorflow 2.xx multi gpu 사용하기 (gpu, cpu 지정하기, 2.3 이상)

tensorflow에서 여러 개의 GPU에 대해 분산 처리를 지원한다. 만약 내가 Tesla T4 두 개를 가지고 있고 model.fit을 통해 학습을 수행하면, 우선적으로 두 칩의 메모리에 모두 프로세스를 띄우게 되는데, 실제로는 single GPU에서만 잡을 처리하게 된다. 이렇게 되면, 나머지 GPU는 다른 프로세스에서 띄워지지도 않으면서, 사용하지도 않는 한마디로 리소스가 매우 낭비된다. 따라서 모델을 빌드하기 이전에, 먼저 내가 여러 개의 GPU를 사용할 것인지, 하나의 GPU를, 아니면 CPU를 사용할 것인지 결정하고 명시적으로 지정한 뒤에 빌드, 학습, 예측 등을 수행해야 자원을 효율적으로 사용할 수 있다. 그럼 지금부터 이렇게 하는 방법에 대해서 알아보자.

 

목차

 

1. 실행 환경 정보

1) SW 환경

- cuda 에 맞는 드라이버와 텐서플로우 버전은 이 포스팅에서 확인할 수 있다

- !! 참고로 tensorflow 2.3 이하 버전에서는 이와 다른 방식으로 설정해야 하니 다른 글을 참고하자.

- CUDA 12.1
- nvidia driver 530.41.03
- python 3.10
- Tensorflow 2.10

2) HW 환경

- multi-gpu 사용은 기본적으로 gpu 기기가 2개 이상인 환경에서만 제어 가능하다

- Ubuntu 20.04.1
- cpu core 20
- gpu Tesla V100 32GB x2

 

2. 사용가능한 장치 목록 확인

- 내 환경의 경우, CPU 하나(코어 하나가 아님),  GPU 2개를 사용 가능함

import tensorflow as tf
tf.config.get_visible_devices()

[PhysicalDevice(name='/physical_device:CPU:0', device_type='CPU'),
 PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU'),
 PhysicalDevice(name='/physical_device:GPU:1', device_type='GPU')]

여기서 확인이 안되면, 공식 도큐먼트를 참고하여, 각 운영체제 환경에 맞게 tensorflow를 깔아준다. 현재 나는 CUDA가 12.1 인데, 11.8 기준으로 라이브러리를 설치해도 별 문제없이 잘 동작하는 것을 확인했다.

 

 

3. CPU 만 사용하는 방법

1) 방법 1

- 이 방법의 경우, 이 스크립트를 실행할 때에 인식가능한 장치에서 GPU를 완전히 배제하는 방법이며, tensorflow 를 통한 제어가 아닌 python os 를 통해 제어하는 방법이다.

- 모델 빌드 이전 또는 실행할 스크립트의 import문 밑에 아래 명령을 선언해주면, CPU만 사용하게 된다.

import os
import tensorflow as tf

os.environ["CUDA_VISIBLE_DEVICES"] = "-1"

위와 같이 선언한 경우, tensorflow 에서 물리적으로 확인 가능한 디바이스 목록에서 GPU를 완전히 인식하지 못하게 된다

tf.config.get_visible_devices()

PhysicalDevice(name='/physical_device:CPU:0', device_type='CPU')]

2) 방법 2

- 이 방법은 tensorflow 라이브러리를 통해 사용할 장치를 수동으로 제어하는 방식이다. 따라서 tensorflow가 인식하는 물리적 장치의 목록은 바뀌지 않고, with 문 내의 명령은 CPU에 한정하여, with 문 밖의 명령에 대해서는 모든 프로세스에 복사되어 처리된다.

- with 문 내 변수를 CPU 메모리에 올리게 되므로, 모델을 빌드하는 코드 조각이 with문 내에 포함되면 된다.

 

아래 예시는 a, b, c 연산을 모두 CPU에서 처리하는 방법이다.

tf.debugging.set_log_device_placement(True)

# 텐서를 CPU에 할당
with tf.device('/CPU: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)

a: (_Arg): /job:localhost/replica:0/task:0/device:CPU:0
b: (_Arg): /job:localhost/replica:0/task:0/device:CPU:0
MatMul: (MatMul): /job:localhost/replica:0/task:0/device:CPU:0
product_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:CPU:0
a: (_Arg): /job:localhost/replica:0/task:0/device:CPU:0
b: (_Arg): /job:localhost/replica:0/task:0/device:CPU:0
_MklMatMul: (_MklMatMul): /job:localhost/replica:0/task:0/device:CPU:0
product_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:CPU:0
Executing op _MklMatMul in device /job:localhost/replica:0/task:0/device:CPU:0
tf.Tensor(
[[22. 28.]
 [49. 64.]], shape=(2, 2), dtype=float32)

위 출력문을 통해 모든 연산이 CPU에서 이루어짐을 알 수 있음

 

반면 아래 코드는 a, b 를 CPU 메모리에 할당하나, C 연산을 위해 두 변수를 GPU에도 복사하여 연산을 수행하게 된다(그런데 CPU에 할당한 로그가 안나오고 있어서 왜인지는 잘 모르겠다..?)

tf.debugging.set_log_device_placement(True)

# 텐서를 CPU에 할당
with tf.device('/CPU: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)

a: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
b: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
MatMul: (MatMul): /job:localhost/replica:0/task:0/device:GPU:0
product_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
Executing op MatMul in device /job:localhost/replica:0/task:0/device:GPU:0
2023-07-19 12:44:21.334753: I tensorflow/core/common_runtime/placer.cc:114] a: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
2023-07-19 12:44:21.334784: I tensorflow/core/common_runtime/placer.cc:114] b: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
2023-07-19 12:44:21.334793: I tensorflow/core/common_runtime/placer.cc:114] MatMul: (MatMul): /job:localhost/replica:0/task:0/device:GPU:0
2023-07-19 12:44:21.334798: I tensorflow/core/common_runtime/placer.cc:114] product_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
tf.Tensor(
[[22. 28.]
 [49. 64.]], shape=(2, 2), dtype=float32)

 

4. 특정 GPU 장치를 지정하는 방법

- CUDA_VISIBLE_DEVICES 환경변수를 조정하면, CPU만 사용하거나, CPU와 특정 GPU를 사용하거나, 모든 자원을 사용하는 등의 조정이 가능하나, CPU를 사용하지 않는 방법은 모르겠다 (없을수도 ? 있을수도? 현재의 나는 모른다)

- 그래서, 위에서 사용한 with 문을 통해 제어하는 방법만 알고 있고 이는 아래와 같이 사용하면 된다.

tf.debugging.set_log_device_placement(True)

try:
  # 유효한 GPU 장치를 명시
  with tf.device('/device:GPU:1'):
    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)
except RuntimeError as e:
  print(e)

input: (_Arg): /job:localhost/replica:0/task:0/device:GPU:1
Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:1
Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:1
_EagerConst: (_EagerConst): /job:localhost/replica:0/task:0/device:GPU:1
output_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:1
a: (_Arg): /job:localhost/replica:0/task:0/device:GPU:1
b: (_Arg): /job:localhost/replica:0/task:0/device:GPU:1
MatMul: (MatMul): /job:localhost/replica:0/task:0/device:GPU:1
product_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:1
Executing op MatMul in device /job:localhost/replica:0/task:0/device:GPU:1
2023-07-19 12:59:35.831261: I tensorflow/core/common_runtime/placer.cc:114] input: (_Arg): /job:localhost/replica:0/task:0/device:GPU:1
2023-07-19 12:59:35.831298: I tensorflow/core/common_runtime/placer.cc:114] _EagerConst: (_EagerConst): /job:localhost/replica:0/task:0/device:GPU:1
2023-07-19 12:59:35.831305: I tensorflow/core/common_runtime/placer.cc:114] output_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:1
2023-07-19 12:59:35.838457: I tensorflow/core/common_runtime/placer.cc:114] a: (_Arg): /job:localhost/replica:0/task:0/device:GPU:1
2023-07-19 12:59:35.838492: I tensorflow/core/common_runtime/placer.cc:114] b: (_Arg): /job:localhost/replica:0/task:0/device:GPU:1
2023-07-19 12:59:35.838500: I tensorflow/core/common_runtime/placer.cc:114] MatMul: (MatMul): /job:localhost/replica:0/task:0/device:GPU:1
2023-07-19 12:59:35.838505: I tensorflow/core/common_runtime/placer.cc:114] product_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:1

 

만약 유효하지 않은 장치를 명시하게 되면, 0번 장치를 사용하도록 되어 있다.

tf.debugging.set_log_device_placement(True)

try:
  # 유효하지 않은 GPU 장치를 명시
  with tf.device('/device:GPU:2'):
    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)
except RuntimeError as e:
  print(e)

input: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
_EagerConst: (_EagerConst): /job:localhost/replica:0/task:0/device:GPU:0
output_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
a: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
b: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
MatMul: (MatMul): /job:localhost/replica:0/task:0/device:GPU:0
product_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
Executing op MatMul in device /job:localhost/replica:0/task:0/device:GPU:0
2023-07-19 13:01:28.714505: I tensorflow/core/common_runtime/placer.cc:114] input: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
2023-07-19 13:01:28.714540: I tensorflow/core/common_runtime/placer.cc:114] _EagerConst: (_EagerConst): /job:localhost/replica:0/task:0/device:GPU:0
2023-07-19 13:01:28.714546: I tensorflow/core/common_runtime/placer.cc:114] output_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0
2023-07-19 13:01:28.726749: I tensorflow/core/common_runtime/placer.cc:114] a: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
2023-07-19 13:01:28.726783: I tensorflow/core/common_runtime/placer.cc:114] b: (_Arg): /job:localhost/replica:0/task:0/device:GPU:0
2023-07-19 13:01:28.726791: I tensorflow/core/common_runtime/placer.cc:114] MatMul: (MatMul): /job:localhost/replica:0/task:0/device:GPU:0
2023-07-19 13:01:28.726797: I tensorflow/core/common_runtime/placer.cc:114] product_RetVal: (_Retval): /job:localhost/replica:0/task:0/device:GPU:0

 

5. 모든 GPU를 사용하는 방법

- 모델을 빌드하는 코드를 with문 안쪽에 작성하면, 해당 모델에 대한 동작이 multi GPU에 대해 분산 처리되어 작동한다.

mirrored_strategy = tf.distribute.MirroredStrategy()

with mirrored_strategy.scope():
  model = tf.keras.Sequential([tf.keras.layers.Dense(1, input_shape=(1,))])

model.compile(loss='mse', optimizer='sgd')

 

만약 일부 GPU에 대해서만 동작시키고 싶으면 다음과 같이 scope 내에 파라미터를 주면 된다.

mirrored_strategy = tf.distribute.MirroredStrategy(devices=["/gpu:0", "/gpu:1"])

with mirrored_strategy.scope():
  model = tf.keras.Sequential([tf.keras.layers.Dense(1, input_shape=(1,))])

model.compile(loss='mse', optimizer='sgd')

 

6. 마무리

이상 끝~~ 지금까지 tensorflow 연산을 cpu, single-gpu, 그리고 multi-gpu  환경에서 동작시키는 방법을 알아보았다. 필요할 때마다 찾아서 복붙하면서 썼는데, 앞으로는 이 자료를 참고해서 써야겠군.

 

7. 참고

• 공식 도큐먼트: TensorFlow를 사용한 분산 훈련 https://www.tensorflow.org/guide/distributed_training?hl=ko 
• 공식 도큐먼트: pip 로 Tensorflow 설치 https://www.tensorflow.org/install/pip?hl=ko#linux
• 공식 도큐먼트: GPU 사용하기 https://www.tensorflow.org/guide/gpu?hl=ko