[ Python ] 윈도우에서 GPU 할당하기

Linux에서 GPU 할당하기

리눅스에서 GPU를 정하는 것은 쉽다. 

코드를 돌릴 때, 앞에 CUDA_VISIBLE_DEVICES 를 붙이기만 하면된다.

아래와 같이 숫자는 각각 0부터 하나씩 지정이 되어 있을테니, 숫자를 바꿔서 각각 GPU에서 다른 코드를 돌려볼 수 있다.

$ CUDA_VISIBLE_DEICES=0 python main.py

 

Windows에서 GPU 할당하기

기본적으로 윈도우에서는 리눅스 명령어가 먹히지 않으므로, 저 명령어가 먹히지 않는다. 

따라서 실행시킬 코드 위에 아래 코드를 적어 놓는다.

당연히, 리눅스에서도 아래의 코드를 미리 적어놓으면 앞에 CUDA_VISIBLE_DEVICES를 붙이지 않고, 바로 GPU를 할당할 수 있다.

위와 마찬가지로 숫자를 다르게 써서, 선택할 수 있다.

import os

os.environ["CUDA_VISIBLE_DEVICES"]="0"

아래 사진에서 보면 알 수 있듯, 그냥 두개의 코드를 돌렸을 때는 하나의 GPU에서 코드가 다 돌아가 메모리 사용량이 두배로 뛰지만, 위처럼 할당을 하고 코드를 돌리면 GPU가 서로 병렬적으로 사용되는 것을 알 수 있다.

두 개의 코드를 그냥 돌렸을 때

위와 같이 설정을 해주고 두 개의 코드를 돌렸을 때