PyTorch에서 GPU 여러 개 사용하기(1)

PyTorch 공식 튜토리얼 정리 - 라이브러리 살펴보기

DistributedDataParallel

• batch size는 GPU 수보다 커야 한다.
• DistributedDataParallel은 torch.nn.DataParallel보다 상당히 빠르다.
• N개의 GPU가 있는 host에서 DistributedDataParallel를 사용하려면 N개의 프로세스가 있어야 하고, 각 프로세스는 0번에서 N-1번까지의 GPU 각각에서 동작한다. 이것은 CUDA_VISIBLE_DEVICES를 설정하거나 모든 프로세스에 torch.cuda.set_device(i)를 부르면 된다. i는 0에서 N-1까지의 수 이다.
• 각 프로세스에서, 이 모듈을 구성하려면 다음을 참조해야 한다.

>>> torch.distributed.init_process_group(
>>>     backend='nccl', world_size=N, init_method='...'
>>> )
>>> model = DistributedDataParallel(model, device_ids=[i], output_device=i)

• 노드 당 여러 프로세스를 생성하려면 torch.distributed.launch 또는 torch.multiprocessing.spawn을 사용할 수 있다.

Parameters