파이토치 GPU 사용하는 방법

파이토치 환경에서 학습을 하다보면 학습 속도가 매우 느릴 때가 있다.

그럴 때 작업 관리자를 켜고 보면 GPU를 사용하지 않고 CPU 사용할 때가 매우 많다.

그래서 이번엔 파이토치를 사용하여 학습을 진행할 때 GPU를 제대로 인식하고 사용하는지에 대한 여부를 확인하는 방법에 대해 기술하려고 한다.

 

(파이토치가 설치되었다는 가정 하에 글을 진행한다.)

당연히 토치는 임포트 돼있어야한다.

 

GPU 사용 가능 여부

print(torch.cuda.is_available())

 

GPU가 제대로 인식되어 있고 토치가 설치되어 있고 gpu를 사용할 수 있다면 True, 그렇지 않으면 False가 출력된다.

 

사용 가능한 GPU 개수

print(torch.cuda.device_count())

 

위 명령어를 입력하면 사용 가능한 gpu의 개수가 출력된다.

 

사용하려는 GPU 지정

device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu')

 

위 명령어는 device라는 변수에 cuda:0을 할당한다는 뜻이다. 즉 첫 번째 gpu(0번)를 지정하여 사용한다고 하는 것이다. 따라서 0번으로 인식되는 첫 번째 gpu를 사용하게 된다. 보통의 경우 gpu가 하나만 장착되어 있기 때문에 0번으로 지정하면 된다. 텐서플로우와 파이토치 모두 사용가능하다.

 

ps. gpu가 gpu0 ~ gpu3까지 4개 존재한다고 가정했을 때, 각 순서에 맞게 cuda:0~cuda:3으로 되는 경우도 있지만 가끔 순서가 바뀔 때가 존재하여 작업관리자를 실행하여 gpu 번호를 직접 확인할 필요가 있다.

 

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

 

위 명령어도 마찬가지로 사용할 gpu를 지정하는 것이다. 마찬가지로 텐서플로우와 파이토치 모두 사용가능하다.