wassname/pytorch-lightning
1
0
Fork 0
mirror of https://github.com/wassname/pytorch-lightning.git synced 2026-09-09 11:32:07 +08:00
Code Issues Packages Projects Releases Wiki Activity

updated dist sampler

Browse Source
This commit is contained in:
William Falcon
2019-07-08 19:26:51 -04:00
parent d596ff2039
commit 96314cbf46
1 changed files with 5 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/new_project_templates/lightning_module_template.py
+5 -13
View File
@@ -159,19 +159,11 @@ class LightningTemplateModel(LightningModule):
dataset = MNIST(root=self.hparams.data_root, train=train, transform=transform, download=True)
# when using multi-node we need to add the datasampler
print('-'*100)
print(self.trainer.world_size, self.trainer.proc_rank)
print('-'*100)
train_sampler = DistributedSampler(dataset, num_replicas=self.trainer.world_size, rank=self.trainer.proc_rank)
# try:
# if self.hparams.nb_gpu_nodes > 1:
# print('-'*100)
# print(self.trainer.world_size, self.trainer.proc_rank)
# print('-'*100)
# train_sampler = DistributedSampler(dataset, num_replicas=self.trainer.world_size, rank=self.trainer.proc_rank)
# except Exception as e:
# print('no sampler')
# train_sampler = None
try:
if self.hparams.nb_gpu_nodes > 1:
train_sampler = DistributedSampler(dataset, num_replicas=self.trainer.world_size, rank=self.trainer.proc_rank)
except Exception as e:
train_sampler = None
should_shuffle = train_sampler is None
loader = DataLoader(