파이토치 튜토리얼: Optional: Data Parallelism

Authors: Sung Kim and Jenny Kang

2개이상의 GPU를 DataParallel을 이용하여 사용하는 법

우선 GPU에 모델을 전달하는 방법은 매우 쉬움

device = torch.device("cuda:0")
model.to(device)

그리고, 모든 텐서를 GPU로 복사

mytensor = my_tensor.to(device)

my_tensor.to(device)는 mytensor를 새로 쓰는게 아니라 값을 복사해서 넘김

따라서 이걸 mytensor라는 새로운 변수에 할당해야하는듯

모델의 학습과정을 여러대의 GPU로 이용하는것은 자연스러워 보이지만 파이토치는 1개만 사용하는것을 디폴트로함

따라서 DataParallel 을 이용하여 여러대의 GPU 로 돌릴 수 있음

model = nn.DataParallel(model)

이게 이번 강의의 핵심

Imports and parameters

모듈 임포트 및 변수 정의

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

input_size = 5
output_size = 2

batch_size = 30
data_size = 100

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

Dummy DataSet

임의의 데이터셋을 만듬. getitem만 추가해주면 됨

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]
    
    del __len__(self):
        return self.len

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

Simple Model

이번 강의에 쓸 모델은 단순히 인풋을 넣으면 선형 연산을 하고 아웃풋을 줌 하지만 CNN, RNN, 등 어떤 모델에서도 DataParallel을 쓸 수 있음

print 메소드는 인풋 아웃풋 텐서 사이즈를 알아보기 위해 넣었음. batch rank 0 일때 뭐가 출력되는지 주의해서 볼것

모델 정의:

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

Create Model and DataParallel

이번 강의의 핵심 내용. 먼저, 모델 인스턴스를 만들고 GPU가 여러개인지 체크함

여러개라면 DataParallel로 모델을 덮어쓰고 GPU로 보냄

model = Model(input_size, output_size)
if torch.cuda.device_count() > 1:
    print("Let's use", torch.cuda.device_count(), "GPUs!")
    model = nn.DataParallel(model)

model.to(device)

Run the Model

인풋,아웃풋 텐서들의 사이즈를 볼 수 있음

GPU가 하나라 하지 않았음