Python 원하는 GPU 사용하기

$ CUDA_VISIBLE_DEVICES=[GPU_NUMBER] python [PY_FILE]

ex)
$ CUDA_VISIBLE_DEVICES=0 python main.py

위 명령은 터미널 상에서 GPU 선정 방법이다.

import os

os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"   
os.environ["CUDA_VISIBLE_DEVICES"]="[GPU_NUMBER]"

ex)
import os

os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"   
os.environ["CUDA_VISIBLE_DEVICES"]="0" 
os.environ["CUDA_VISIBLE_DEVICES"]="1"
os.environ["CUDA_VISIBLE_DEVICES"]="0,1"