TensorFlow 실행시 GPU 지정하기

CUDA 를 이용하여 TensorFlow 를 실행하다보면 GPU 를 지정하여 실행하고 싶은 경우가 생긴다.

이렇때는 어떻게 해야 할까? 처음에는 TensorFlow 에 GPU 를 지정하는 코드를 작성해야 한다고 생각했었다.

 

하지만!!! 아래와 같이 CUDA_VISIBLE_DEVICES 명령어와 python 을 함께 실행하면 된다

CUDA_VISIBLE_DEVICES='0,1,2,3[GPU ID]' python main.py

 

아주 아주 간단하다!