[Pytorch] 특정 GPU 사용하기 / 여러개의 multi GPU parallel 하게 사용하기

Neural network를 train을 하다보면 성능을 비교하기 위해 loss function, optimizer 등이 상이한 여러 가지 model을 돌려볼 일이 수도 없이 많다.

 

이 때, 한 장의 GPU만 있다면 어쩔 수 없지만 여러장의 GPU가 존재한다면

A model은 GPU 0번에서, B model은 GPU 1번에서 동시에 두 개의 model을 돌려 빠르게 결과를 보고 싶을 것이다.

 

이를 위하여 각 model을 특정 GPU에서 돌리는 법을 알아보고자 한다.

 

추가적으로, network의 weight param이 크거나 input dataset의 크기가 커서 하나의 GPU memory로는 model을 돌리기에 부족하여 CUDA out of memory error가 뜨는 경우가 많다.

 

이를 해결하기 위하여 여분의 GPU를 함께 사용하여 memory를 parallel하게 사용한다면 out of memory error를 피해 성공적으로 model을 돌릴 수 있다.

 

위 2개 (model을 특정 GPU에서 돌리는 법, 여러 GPU를 parallel하게 사용하는 법)는 거의 비슷한 방법이므로 함께 살펴보도록 하자.

 

0. "nvidia-smi"를 통해 사용할 수 있는 GPU 살펴보기

먼저, ubuntu terminal 혹은 anaconda terminal 등에 "nvidia-smi" 명령어를 입력하여 자신이 model을 돌리고자 하는 환경에서 사용가능한 GPU를 살펴본다.

$nvidia-smi

nvidia-smi 명령어 결과

현재 본인의 경우, 총 8장의 GPU가 available하며 GPU#0, #1은 이미 사용중인 것을 확인할 수 있다.

 

1. "CUDA_VISIBLE_DEVICES"를 통해 cuda가 볼 수 있는 GPU 제한하기

별도의 설정을 하지 않으면 cuda는 GPU 0번을 사용하려 한다. 그러나 현재 GPU 0번에는 memory 자리가 없다.

따라서 사용하지 않고 있는 GPU #2를 사용하여야 한다.

 

이를 위해 "CUDA_VISIBLE_DEVICES"라는 envirable variable를 변경하여

cuda가 볼 수 있는 GPU를 내가 사용하고자 하는 다른 GPU로 제한하여야 한다.

(cuda는 cuda입장에서 visible한 gpu를 사용하는데, 인위적으로 visible devices를 변경함으로서 cuda가 내가 명시한 GPU를 사용하도록 하는 원리이다.)

 

보통 python을 실행하는 경우는 다음 2가지로 나뉘므로 각각 설명하겠다.

1-1) Terminal 사용하여 "~.py" file 돌리는 경우

CUDA_VISIBLE_DEVICES는 위에서 말했듯 envirable variable이라서 "~.py" 와 같은 python file을 terminal을 통해 실행할때는 다음과 같이 명령어 앞에 설정해서 사용할 수 있다.

CUDA_VISIBLE_DEVICES=2 python train.py

즉, 위 명령어는 GPU #2를 사용하여 train.py python file을 실행하라는 의미이다.

 

만약, 여러 GPU를 할당하여 parallel하게 사용하고 싶다면 다음과 같이 comma를 이용하여 2개를 할당하면 된다.

CUDA_VISIBLE_DEVICES=2,3 python train.py

 

단, 뒤에서 말하겠지만 nn.DataParallel()을 작성해주어야 한다.

 

1-2) Jupyter notebook 등의 python script "~.ipynb" file 내에서 돌리는 경우

"~.ipynb" 와 같이 python script 내에서 돌리는 경우에는 다음과 같이 os.environ[ ] code를 활용하여 environment를 설정하여 실행할 수 있다.

os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"  # Arrange GPU devices starting from 0
os.environ["CUDA_VISIBLE_DEVICES"]= "2"  # Set the GPU 2 to use

 

마찬가지로 여러개의 GPU를 할당하고 싶은 경우에는 다음과 같이 comma를 이용한다.

os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"  # Arrange GPU devices starting from 0
os.environ["CUDA_VISIBLE_DEVICES"]= "2,3"  # Set the GPUs 2 and 3 to use

 

2. GPU가 할당되었는지 확인하기

GPU를 설정한 이후에 다음과 같은 코드로 device에 cuda (GPU)를 설정하고 device와 current_device()를 출력해봄으로 GPU가 잘 할당되었는지 확인한다.

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

print('Device:', device)
print('Current cuda device:', torch.cuda.current_device())
print('Count of using GPUs:', torch.cuda.device_count())

 

torch.cuda.is_available(): cuda가 사용 가능하면 true를 반환함으로서 device에 cuda를 설정하도록 한다.

torch.cuda.current_device(): 현재 cuda가 사용할 GPU를 명시한다.

torch.cuda.device_count(): 현재 cuda가 사용할 GPU 개수를 명시한다.

위 코드의 출력 결과

이 때, cuda.current_device()로 설정된 GPU를 확인하면 위와 같이 GPU #2로 설정을 하였더라도 current_device()의 결과는 0을 반환한다.

 

이는 system environment에서 설정한 GPU를 pytorch cuda는 0으로 mapping해서 보여주기 때문이다.

즉, system에서의 GPU #2가 cuda에게는 #0으로 할당된 것이다. (이 글의 SimonW 글 참고)

 

그러나 실제로 model을 돌려보면 설정한대로 GPU #2에서 model이 돌아가는 것을 확인할 수 있다.

** multi GPU 사용을 위하여 2, 3으로 설정하였더라도 처음에는 2에만 올라오고 이후에 train이 시작될 때 2, 3이 모두 parallel하게 돌아가는 것을 확인할 수 있음.

 

3. to(device)를 통해 GPU로 model 돌리기

이제 사용할 GPU 할당은 끝났고, .to(device) 코드를 통해 model을 GPU에 연결해서 돌려주기만 하면 된다.

net = ResNet50().to(device)

 

만약 multi GPU로 parallel하게 돌리는 경우에는 다음과 같이 nn.DataParallel()을 사용한다.

_net = ResNet50().cuda()
net = nn.DataParallel(_net).to(device)

multi GPU (GPU#2, 3) 이용하는 모습

---

동시에 두개를 소개하다보니 약간 글이 너저분해진 것 같아, 전체 code snippet을 공유하며 글을 마무리 짓고자 한다.

 

1. 특정 GPU 사용하는 code snippet

import os
import torch
import torch.nn as nn

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

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

print('Device:', device)  # 출력결과: cuda 
print('Count of using GPUs:', torch.cuda.device_count())   #출력결과: 1 (GPU #2 한개 사용하므로)
print('Current cuda device:', torch.cuda.current_device())  # 출력결과: 2 (GPU #2 의미)

net = ResNet50().to(device)

...

 

2. multi GPU 사용하는 code snippet

import os
import torch
import torch.nn as nn

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

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

print('Device:', device)  # 출력결과: cuda 
print('Count of using GPUs:', torch.cuda.device_count())   #출력결과: 2 (2, 3 두개 사용하므로)
print('Current cuda device:', torch.cuda.current_device())  # 출력결과: 2 (2, 3 중 앞의 GPU #2 의미)

_net = ResNet50().cuda()
net = nn.DataParallel(_net).to(device)

...

 

+ Multi-GPU 사용 확인용 code snippet (출처: pytorch tutorial)

   * 아래 코드를 복사하여 .py (혹은 .ipynb)를 만들고 실행해보면  multi-GPU 사용이 잘 되는지 확인하실 수 있습니다.

import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader

# 매개변수와 DataLoaders
input_size = 5
output_size = 2

batch_size = 30
data_size = 100



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



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



model = Model(input_size, output_size)
if torch.cuda.device_count() > 1:
  print("Let's use", torch.cuda.device_count(), "GPUs!")
  # dim = 0 [30, xxx] -> [10, ...], [10, ...], [10, ...] on 3 GPUs
  model = nn.DataParallel(model)

model.to(device)



for data in rand_loader:
    input = data.to(device)
    output = model(input)
    print("Outside: input size", input.size(),
          "output_size", output.size())