[PyTorch] Multi-GPU

Multi-GPU

Model parallel

• 다중 GPU에서 학습 분산 방법
  • 모델 나누기 (AlexNet)
  • 데이터 나누기

(출처) http://www.idris.fr/eng/ia/model-parallelism-pytorch-eng.html

Data parallel

• 데이터를 나눠서 GPU에 각각 할당한 후, 결과의 평균을 구함

forward

• 데이터를 나눠서 GPU에 각각 할당
• 모델들을 각각의 GPU에 복사
• 연산 실행 -> 결과를 한 GPU에 모아서 loss 계산

backward

• 각각의 GPU에 나눔
• weight의 새로운 그래디언트 값 구함
• 모아서 하나의 GPU가 평균을 냄

-> GPU 사용 불균형 문제 발생, Batch 사이즈 감소, GIL

DistributedDataParallel

• 모으지 않고 각각 개별적으로 연산한 후 평균을 구함

# DataParallel
parallel_model = torch.nn.DataParallel(model)

# DistributedDataParallel
# sampler 사용 == syntax를 어떻게 사용하는가
# pin_memory : 메모리에 데이터를 바로 올릴 수 있도록 절차를 간소화하는 데이터 저장 방법 (메모리 -> GPU를 빠르게)
# num_workers : GPU의 개수
​
train_sampler = torch.utils.data.distributed.DistributedSampler(train_data)
shuffle = False
pin_memory = True
​
trainloader = torch.utils.data.DataLoader(
					train_data, batch_size=20, shuffle = shuffle, 
                    pin_memory=pin_memory, num_workers=3, sampler=train_sampler)

main

def main():
  n_gpus = torch.cuda.device_count()
  torch.multiprocessing.spawn(main_worker, nprocs=n_gpus, args=(n_gpus, ))
​
 def main_worker(gpu, n_gpus):
  image_size = 224
  batch_size = 512
  num_worker = 8
  epochs = ...

  batch_size = int(batch_size / n_gpus)
  num_worker = int(num_worker / n_gpus)
​
  torch.distributed.init_process_group(
    backend='nccl', 'init_method='tcp://127.0.0.1:2568’, world_size=n_gpus, rank=gpu)
    model = MODEL
    torch.cuda.set_device(gpu)
    model = model.cuda(gpu)
    model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[gpu])


# Python의 멀티프로세싱 코드
from multiprocessing import Pool
def f(x):
    return x*x
if __name__ == '__main__':
    with Pool(5) as p:
      print(p.map(f,[1,2,3]))