[Pytorch] GPU 여부 확인, 사용할 GPU 번호 지정, 모델 및 텐서에 GPU 할당 방법

파이토치 GPU 체크, 지정 및 할당 방법 정리

파이토치에서 GPU 사용 가능 여부 및 연결된 GPU 정보를 확인해보고

특정 GPU만 사용하도록 설정하는 방법과 모델 및 텐서에 원하는 GPU를

할당할 수 있는 방법에 대하여 정리해보도록 하겠습니다.

 

 

GPU 사용 가능 여부 확인

아래의 명령어의 결과가 True로 반환되면 GPU가 연결되어 사용 가능하다는 의미입니다.

import torch

# GPU 사용 가능 -> True, GPU 사용 불가 -> False
print(torch.cuda.is_available())

 

다음과 같이 코드를 작성하면 GPU 여부에 따라서 사용할 디바이스 종류 자동 지정이 가능합니다.

# GPU 사용 가능 -> 가장 빠른 번호 GPU, GPU 사용 불가 -> CPU 자동 지정 예시
device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu')

 

 

연결된 GPU 정보 및 사용 가능 GPU 개수 확인

아래의 코드로 각 번호에 연결된 그래픽 카드의 이름과 총 GPU의 개수를 확인하실 수 있습니다.

# GPU 이름 체크(cuda:0에 연결된 그래픽 카드 기준)
print(torch.cuda.get_device_name(device = 0)) # 'NVIDIA TITAN X (Pascal)'

# 사용 가능 GPU 개수 체크
print(torch.cuda.device_count()) # 3

반응형

특정 번호의 GPU 사용 지정

만일, 특정 번호(들)의 GPU만을 파이토치에서 사용하고 싶은 경우는 os 모듈의 기능을 활용하여

다음 코드처럼 지정해주시면 됩니다.

import os

# 0, 1, 2번의 GPU를 모두 사용하고 싶은 경우 예시(각각 cuda:0, cuda:1, cuda:2에 지정)
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2"

# 2번 GPU만 사용하고 싶은 경우 예시(cuda:0에 지정)
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "2"

여러개의 GPU를 연결한 경우, 그래픽 카드가 연결된 순서대로 각각

cuda:0, cuda:1, ... 디바이스로 인식되도록 지정된다는 점을 참고해주세요.

 

 

모델 및 텐서에 원하는 GPU 할당

to('cuda:0'), 혹은 위에서 지정한 to(device) 등의 형태로 각 모델과 텐서에 할당될

GPU의 번호 혹은 CPU 여부를 지정할 수 있습니다.

model = my_Net() # 모델 예시
tensor1 = torch.ones(10) # 텐서 예시

model.to('cuda:0') # cuda:0에 지정된 GPU에 모델 할당
tensor1.to('cpu') # CPU에 텐서 할당

 

만일, "RuntimeError: Expected all tensors to be on the same device, 

but found at least two devices, cpu and cuda:0!"와 같은 오류가 발생한다면,

모델과 텐서가 다른 디바이스에 할당 중인 것이 원인이므로

아래 예시처럼 모델과 텐서를 같은 디바이스에 할당시켜주면 해결됩니다.

model = my_Net()
tensor1 = torch.ones(10).to('cpu')
tensor2 = torch.ones(10).to('cuda:0')
model.to('cuda:0')

model(tensor1) # 오류 발생(모델 : cuda:0, 텐서 : cpu)
model(tensor2) # 정상 작동(모델 : cuda:0, 텐서 : cuda:0)