파이썬 코드 실행 시 사용할 GPU를 지정하는 방법(리눅스, 윈도우)

Linux, Window GPU 번호 지정 Python 실행

리눅스, 윈도우 환경에서 사용할 GPU 번호를 지정하여 파이썬 코드 파일을 실행하는

방법에 대해서 간략하게 요약해보도록 하겠습니다.

 

 

리눅스 bash 터미널 환경에서 GPU 할당 방법

리눅스의 터미널 환경에서 사용할 GPU 번호를 지정하려면

python 명령어의 앞 부분에 CUDA_VISIBLE_DEVICES=번호 명령어를 추가해주시면 됩니다.

 

여러 개의 GPU를 지정하려면 콤마(,)를 통하여 원하는 번호들을 나열해주면 되며,

python 실행 구문이 복잡한 경우에도 python의 앞 부분에 CUDA 명령어만 추가하면 됩니다.

# 가장 기본적인 예시(GPU 1개, 여러개 예시)
$ CUDA_VISIBLE_DEVICES=0 python file_name.py
$ CUDA_VISIBLE_DEVICES=0,1,2 python file_name.py

# bash for문, input 인자가 있는 복잡한 예시
$ for var in {1..100}
do
  CUDA_VISIBLE_DEVICES=1 python file_name.py \
    --argument1 $var
done

코드 실행 후, 다른 터미널에서 nvidia-smi를 입력하여 GPU가 잘 사용되고 있는지 여부를

확인해주시면 됩니다.

 

 

윈도우 GPU 할당 방법(python os 라이브러리 이용)

윈도우의 경우, 터미널 종류가 위의 bash 형태와는 달라 os 라이브러리를 이용하는 것을

권장드립니다. 파이썬 파일 내에서 아래의 코드를 윗 부분에 추가해주시면 됩니다.

import os
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"]="0"

# GPU 여러개를 지정하고 싶은 경우는 아래 형태로 사용
os.environ["CUDA_VISIBLE_DEVICES"]="0,1,2"

윈도우에서는 코드 실행 후, 작업 관리자 -> 성능 -> GPU 탭에서 사용 현황을 확인하시면 됩니다.

 

참고로, 리눅스 환경에서도 python 파일 내에 os 라이브러리를 활용하여

사용할 GPU를 지정하는 방법 적용이 가능합니다.