CUDA_VISIBLE_DEVICES 활용하여 GPU 선택하기

Develop My Life·2023년 6월 9일
0

GPU 선택하기

  • 기존이 이걸 알기 전에는 코드 레벨에서 직접 GPU를 선택할 수 있도록 수정하여 사용했다. 하지만 이는 많은 시간과 여러가지 변수를 고려해야했다.

  • CUDA_VISIBLE_DEVICES를 활용하여 환경변수를 설정함으로써 해당 파일이 실행 될 때 활성화하는 GPU 번호를 선택할 수 있어 별다른 수정없이 특정 GPU에서 학습이 가능하다.

예제

  • 실행 시 사용하는 방법
CUDA_VISIBLE_DEVICES=1 python ./main.py
  • 코드 내에서 추가하는 방법
import os
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"   
os.environ["CUDA_VISIBLE_DEVICES"]="1"

위 코드를 코드 맨 위에 추가한다.
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"를 하는 이유는 nvidia-smi를 실행했을 때 나오는 GPU 인덱스와 동일하게 정렬해주기 위해서다.

0개의 댓글