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

enabling gpu size = 1 to run without data parallel

Browse Source
This commit is contained in:
William Falcon
2019-07-13 22:09:17 -04:00
parent 9ccfc7bd33
commit 7c688fbf2e
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pytorch_lightning/models/trainer.py
+2 -2
View File
@@ -96,7 +96,7 @@ class Trainer(TrainerIO):
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = ','.join([str(x) for x in self.data_parallel_device_ids]) os.environ["CUDA_VISIBLE_DEVICES"] = ','.join([str(x) for x in self.data_parallel_device_ids])
self.data_parallel = self.data_parallel_device_ids is not None and len(self.data_parallel_device_ids) > 0 self.data_parallel = self.data_parallel_device_ids is not None and len(self.data_parallel_device_ids) > 1
# process info # process info
self.proc_rank = 0 self.proc_rank = 0
@@ -291,7 +291,7 @@ class Trainer(TrainerIO):
# when using gpus, first thing we do is spawn a new process between each worker # when using gpus, first thing we do is spawn a new process between each worker
# applies to single gpu, multi-gpu and multi-nodes # applies to single gpu, multi-gpu and multi-nodes
if self.on_gpu and len(self.data_parallel_device_ids) > 1: if self.data_parallel:
self.experiment = self.experiment.get_meta_copy() self.experiment = self.experiment.get_meta_copy()
mp.spawn(self.dp_train, nprocs=len(self.data_parallel_device_ids), args=(model, )) mp.spawn(self.dp_train, nprocs=len(self.data_parallel_device_ids), args=(model, ))
else: else: