Learning을 여러 개 GPU에서 사용하기 위해서는 모델을 각 GPU에 복사해서 할당 한다.
매 iteration 마다 mini-batch의 크기를 GPU 개수만큼 나누는데 이 작업을 ‘scatter’라 하며, 실제로 scatter 함수를 이용해서 작업을 수행한다.
Scatter 수행 후, 각 GPU에서 forward를 진행하고 각 GPU의 출력 값(ouput)을 특정 한 GPU로 모아주며 이 작업을 ‘gather’라 한다.
한 GPU로 각 GPU에 있는 모델의 출력 값을 모으는(gather) 이유는 loss function을 사용하여 loss 를 얻기 위해서 이고, 얻은 loss를 가지고 각 GPU에서 Backpropagation을 통해 gradient를 구한 후, 모델을 업데이트 하기 위해 gradient를 다시 한 GPU에 모으고, 모델을 업데이트를 한다.
<Reprlicate model each GPU → make Output per GPU → Output Gather → Compute Loss → Loss Scatter → Make gradient on each GPU → Gradient Gather → Model update(=parameter update)>
import os
import torch.nn as nn
os.environ['CUDA_VISIBLE_DEVICES'] = '0,1,2,3'
model = nn.DataParallel(model, output_deivce=0)
위의 과정과 같이 특정 한 GPU에 연산 결과를 모아주는 작업을 수행함으로 메모리 불균형이 발생하게 됨. 특정 GPU로 연산 결과 값을 몰아주는 이유는 Loss를 계산하기 위해서 이므로 Loss를 각 GPU에서 계산을 하게 된다면 메모리 불균형이 발생하지 않는다는 말이 된다.
그러면 어떻게 Loss를 각 GPU에서 연산할 수 있을까? Pytorch-Encoding package가 있다.
Loss function을 병렬 연산 가능하게 하는 방법은 Model을 병렬 연산으로 만드는 방법과 동일하다.
Pytorch에서 Loss fucntion도 하나의 모듈이므로, 이 모듈을 각 GPU에 Replicate 한다.
후, 데이터 label 또한 각 GPU로 scatter 하게 되면 각 GPU에서 loss 계산을 할 수 있는 상태가 된다.
Loss function을 parallel하게 만들어서 하는 코드를 보면, target을 scatter 후, replicated module(loss function)이 각각 계산을 하고, 계산한 output과 Reduce.apply 를 사용하여 각 GPU에 Backward 연산을 수행한다.
class DataParallelCriterion(DataParallel):
def forward():
targets, kwargs = self.scatter(targets, kwargs, self.device_ids)
replicas = self.replicate(self.module, self.device_ids[:len(input)])
targets = tuple(targets_per_gpu[0] for targets_per_gpu in targets)
outputs = _criterion_parallel_apply(replicas, inputs, targets, kwargs)
return Reduce.apply(*outputs) / len(outputs), targets
DataParallelCriterion을 사용할 경우, DataParallel로 모델을 감싸면 안된다.
DataParallel은 디폴트로 하나의 GPU로 출력을 모으기 때문이다. 따라서, Custom DataParallel 클래스인 DataParallelModel을 사용해야 한다.
사용하는 방법은 Pytorch-Encoding package에서 parallel.py 파일만 가져와 import하여 사용하면 된다.
import torch
import torch.nn as nn
from parallel import DataParallelModel, DataParallelCriterion
model = make_model(args)
model = DataParallelModel(model)
model.cuda()
criterion = make_criterion(args)
criterion = DataParallelCriterion(criterion)
for i, (inputs, labels) in enumerate(train_loader):
outputs = model(inputs)
loss = criterion(outputs, labels)
optimizer.zero_grad()
loss.backward()
optimizer.step()
위와 같은 방법(loss calculate parallel)을 사용하면 메모리 사용량의 차이가 상당히 줄어들어 batch 크기를 기존에 비해 늘릴 수 있어 학습 시간을 줄일 수 있다. 다만, GPU util은 낮게 사용한다.
GPU 성능을 높게 사용하려면? ⇒ DistributedDataParallel 사용
main.py 를 실행하면 main 함수가 실행이 되고, 다시 main 함수가 main_worker 함수들을 multi-processing으로 실행하게 된다.
GPU 4개를 하나의 노드로 보고 world_size를 설정한 후, mp.spawn 함수가 4개의 GPU에서 따로 main_worker를 실행하면 되는데, 이 방법은 main_worker에서 dist.init_process_group을 통해 각 GPU마다 분산 학습을 위한 초기화를 실행.
Multi-GPU 학습을 할 경우, backend로 nccl을 사용해야 하며, init_method에서 freeport에 사용 가능한 port를 적으면 되고, DataParallel 대신 DistributedDataParallel을 사용해야 한다.
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel
def main():
args = parser.parse_args()
ngpus_per_node = torch.cuda.device_count()
args.world_size = ngpus_per_node * args.world_size
mp.spawn(main_worker, nprocs=ngpus_per_node,
args=(ngpus_per_node, args))
def main_worker(gpu, ngpus_per_node, args):
global best_acc1
args.gpu = gpu
torch.cuda.set_device(args.gpu)
print("Use GPU: {} for training".format(args.gpu))
args.rank = args.rank * ngpus_per_node + gpu
dist.init_process_group(backend='nccl',
init_method='tcp://127.0.0.1:FREEPORT',
world_size=args.world_size,
rank=args.rank)
model = make_model(args)
model.cuda(args.gpu)
model = DistributedDataParallel(model, device_ids=[args.gpu])
acc = 0
for i in range(args.num_epochs):
model = train(model)
acc = test(model, acc)
DataLoader가 입력을 각 프로세스에 전달하기 위해서 DistributedSampler를 사용해야 한다.
DistibutedSampler는 DistributedDataParallel과 함께 사용해야 한다. Dataset을 DistributedSampler로 감싸주고 DataLoader에 sampler 인자로 넣어주며 되며 DataLoader를 사용하듯이 똑같이 사용하면 된다.
내부의 코드를 보면 각 Sampler는 전체 데이터를 GPU 개수로 나눈 부분 데이터에서만 데이터를 샘플링하며, 부분 데이터를 만들기 위해 전체 데이터 셋 인덱스 리스트를 셔플하고 인덱스 리스트를 쪼개서 각 GPU sampler에 할당한다.
매 학습 epoch 마다 train_sampler.set_epoch(epoch) 명령어를 꼭 호출 해야 한다.
💡 학습 시키는 모델에 따라 또한 optimizer에 따라 multi-GPU **학습하는 방법을 선택**해야 합니다.