[Pytorch] cuda 선택 방법

import torch

DEVICE = torch.device('cuda')

- 보통 Pytorch에서 그래픽 카드를 사용하는 방법이다.

- 다중 그래픽 카드가 있는 경우 이와 같이 사용하면 가장 처음 설정된 0번 그래픽 카드를 사용한다.

 

nvidia-smi로 확인한 모습

 

- 혼자서 그래픽 카드를 다 사용하는 경우는 상관없지만 공용 서버에서 다른 사람이 그래픽 카드를 쓰고 있다면 메모리 문제 등으로 내가 실행하고자 하는 모델을 실행시키지 못할 수 있다.

- 몇 주 전, 연구실 선배님이 실험을 돌려서 내 실험을 못 돌리게 되자 찾은 방법. 매우 단순하다.

 

import torch

DEVICE = torch.device('cuda:1')

- 이와 같이 cuda 뒤에 :(id)를 붙여주면 된다. 이 때, id는 그래픽 카드의 아이디를 의미한다.

- 사용 중인 그래픽 카드의 아이디는 `nvidia-smi`와 같은 명령어로 확인할 수 있고, torch에서 명령어를 통해 찾을 수도 있다.