From 14d1329655653aad3226a5d5bed53dff8240378a Mon Sep 17 00:00:00 2001 From: William Falcon Date: Mon, 8 Jul 2019 17:51:07 -0400 Subject: [PATCH] auto distribute datasets across nodes --- pytorch_lightning/models/trainer.py | 17 ++++++++++++++--- 1 file changed, 14 insertions(+), 3 deletions(-) diff --git a/pytorch_lightning/models/trainer.py b/pytorch_lightning/models/trainer.py index 13b1300a..fdef6be9 100644 --- a/pytorch_lightning/models/trainer.py +++ b/pytorch_lightning/models/trainer.py @@ -12,6 +12,7 @@ import torch.distributed as dist import os import subprocess from time import sleep +from torch.utils.data.distributed import DistributedSampler try: from apex import amp @@ -87,6 +88,7 @@ class Trainer(TrainerIO): self.amp_level = amp_level self.print_nan_grads = print_nan_grads self.data_parallel_device_ids = None + self.world_size = 1 # gpus come in as a string. # if gpus = -1 then use all available devices @@ -267,7 +269,16 @@ class Trainer(TrainerIO): self.val_dataloader = model.val_dataloader # when distributed data parallel, we need to distribute the dataset to each node - # TODO: implement + if self.nb_gpu_nodes > 1: + self.tng_dataloader = self.__distribute_dataloader(self.tng_dataloader) + self.test_dataloader = self.__distribute_dataloader(self.test_dataloader) + self.val_dataloader = self.__distribute_dataloader(self.val_dataloader) + + def __distribute_dataloader(self, dataloader): + dataset = dataloader.dataset + dataset = DistributedSampler(dataset, num_replicas=self.world_size, rank=self.proc_rank) + dataloader.dataset = dataset + return dataloader # ----------------------------- # MODEL TRAINING @@ -317,11 +328,11 @@ class Trainer(TrainerIO): # determine which process we are and world size self.proc_rank = node_rank * len(self.data_parallel_device_ids) + gpu_nb - world_size = self.nb_gpu_nodes * len(self.data_parallel_device_ids) + self.world_size = self.nb_gpu_nodes * len(self.data_parallel_device_ids) # set up server using proc 0's ip address ip = self.__get_root_node_ip(self.proc_rank, self.nb_gpu_nodes) - dist.init_process_group("nccl", init_method=f'tcp://{ip}:12001', rank=self.proc_rank, world_size=world_size) + dist.init_process_group("nccl", init_method=f'tcp://{ip}:12001', rank=self.proc_rank, world_size=self.world_size) print(f"GPU: {gpu_nb} - Rank: {self.proc_rank}") # copy model to each gpu