[메모] PyTorch GPU 여러 개 선택해서 쓰기

PyTorch framework 사용해서 머신러닝 알고리즘 돌릴 때 GPU 여러 개를 지정해주는 방법에 대해서 알아본다.

참고한 글: PyTorch 튜토리얼, Stack Overflow: How to use multiple GPUs in pytorch?

이 글을 보면 GPU를 여러 개 가지고 있을 때, (1)가지고 있는 걸 전부 다 쓰거나 (2)갖고 있는 것 중에 특정 GPU를 골라서 쓰는 코드를 알려주고 있다.

1. 전부 다 사용

먼저 설치된 GPU를 전부 다 쓰는 방법이다.

# [1]
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = CreateModel()

model= nn.DataParallel(model)
model.to(device)

# [2]
torch.cuda.device_count()  # Returns the number of GPUs available.

[1]을 실행하고 [2]를 실행하면 내가 가지고 있는 총 GPU 개수가 출력될 것이다.
예를 들어 나의 경우엔 서버 PC에 4개의 GPU가 설치되어 있으므로 4가 출력된다.

2-1. 특정 GPU 사용 (torch 설정)

PyTorch 속성을 이용해서 GPU를 지정하는 방법이다.

device = torch.device("cuda:1,2" if torch.cuda.is_available() else "cpu") ## specify the GPU id's, GPU id's start from 0.

model = CreateModel()

model= nn.DataParallel(model,device_ids = [1, 2])
model.to(device)

GPU의 ID는 항상 0부터 시작함에 유의한다. 위의 코드에서 1, 2를 쓴다는 건 두번째, 세번째를 쓰겠다는 의미가 된다.

그런데 나는 위 코드의 첫번째 줄에서 에러가 났다. RuntimeError: Invalid device string: ‘cuda:1,2’ 에러였다.
구글에 검색해도 이런 사례는 보이지 않았다.

2-2. 특정 GPU 사용 (OS 환경변수 설정)

나는 os 환경변수를 먼저 설정해주는 것으로 해결했다.

import os

os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "1,2"

위의 코드를 제일 먼저 실행해줘야 한다. 이 코드를 실행하고 torch.cuda.device_count()의 반환값을 출력해보면 2가 나온다.
근데 torch.cuda.device_count()를 먼저 출력한 다음에 저 설정을 바꾸고 다시 해보면 2가 아닌 4가 출력된다. 아직 그 이유는 찾지 못했지만 일단은 os 환경변수는 프로그램 상에서 가장 먼저 설정해줘야 하는 것으로 이해했다.

그 다음이 중요한데,

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = CreateModel()

model= nn.DataParallel(model)
model.to(device)

모델을 불러올 때 nn.DataParallel에서 device_ids를 지정해주지 않는다.
이렇게 모델까지 설정한 상태에서 torch.cuda.device_count()를 해보면 역시 2가 출력되어야 한다.

참고

이제 예제를 통해서 원하는 GPU 상에서 코드가 잘 돌아가는지 체크할 건데 그 전에 그걸 확인할 수 있는 방법을 정리해 본다.

$ watch -d -n 0.5 nvidia-smi

터미널에 위 명령어를 치면 GPU 사용 현황을 실시간으로 볼 수 있다.

	OPTION	DESCRIPTION
nvidia-smi	-i, --id=	특정 GPU만 볼 수 있다.
watch	-d, --differences	사용량이 바뀌는 부분을 강조(highlight)한다.
	-n, --interval	지정 시간마다 새로고침 한다. [단위: 초]

예제

예제에 나와있는 대로 하면 데이터 크기가 작아서 사용량이 잘 안 보이길래 숫자만 조금 바꿔서 실행해 보았다.

from torch.utils.data import Dataset, DataLoader

input_size = 1000
output_size = 2

batch_size = 30
data_size = 1000

# Dummy Dataset
class RandomDataset(Dataset):

    def __init__(self, size, length):
        self.len = length
        self.data = torch.randn(length, size)

    def __getitem__(self, index):
        return self.data[index]

    def __len__(self):
        return self.len

rand_loader = DataLoader(dataset=RandomDataset(input_size, data_size), batch_size=batch_size, shuffle=True)

class Model(nn.Module):
    # 우리의 모델

    def __init__(self, input_size, output_size):
        super(Model, self).__init__()
        self.fc = nn.Linear(input_size, output_size)

    def forward(self, input):
        output = self.fc(input)
        print("\tIn Model: input size", input.size(),
              "output size", output.size())

        return output

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = Model(input_size, output_size)

model= nn.DataParallel(model)
model.to(device)

print("Let's use", torch.cuda.device_count(), "GPUs!")
for data in rand_loader:
    input = data.to(device)
    output = model(input)
    print("Outside: input size", input.size(),
          "output_size", output.size())

출력 결과:

    Let's use 2 GPUs!
    	In Model: input size torch.Size([15, 1000]) output size torch.Size([15, 2])
    	In Model: input size torch.Size([15, 1000]) output size torch.Size([15, 2])
    Outside: input size torch.Size([30, 1000]) output_size torch.Size([30, 2])
    	In Model: input size torch.Size([15, 1000]) output size torch.Size([15, 2])
    	In Model: input size torch.Size([15, 1000]) output size torch.Size([15, 2])
    Outside: input size torch.Size([30, 1000]) output_size torch.Size([30, 2])
    	In Model: input size torch.Size([15, 1000]) output size torch.Size([15, 2])
    	In Model: input size torch.Size([15, 1000]) output size torch.Size([15, 2])
    Outside: input size torch.Size([30, 1000]) output_size torch.Size([30, 2])
    	In Model: input size torch.Size([15, 1000]) output size torch.Size([15, 2])
    	In Model: input size torch.Size([15, 1000]) output size torch.Size([15, 2])
    Outside: input size torch.Size([30, 1000]) output_size torch.Size([30, 2])
    	...