AI/MachineLearning

[tensorflow] multi gpu 사용하기

향식이 2022. 1. 13. 14:53

 multi gpu 사용하기

모델링 할 때 gpu를 사용하면 속도가 몇만배이상 빨라진다.

그러나 모델이 복잡해지면 복잡해질수록 single gpu가 아닌 multi gpu를 사용하게 되는데 

이 때는 약간의 코드만 추가해주면 된다. 

 

예시 코드

* tensorflow 기준

 

import tensorflow as tf

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

gpu 0번과 1번을 사용하겠다는 예시이다.

 

나는 실행시킬 때,

$ CUDA_DEVICE_ORDER=PCI_BUS_ID CUDA_VISIBLE_DEVICES=1,2 python3 modeling.py

로 명령하면서 한번 더 설정해준다. 

반응형