From d98e79940440f8eb7997efce3f66f6276217d8fa Mon Sep 17 00:00:00 2001 From: William Falcon Date: Fri, 7 Jun 2019 15:06:22 -0400 Subject: [PATCH] adding dataparallel --- pytorch_lightning/models/trainer.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/pytorch_lightning/models/trainer.py b/pytorch_lightning/models/trainer.py index 6a83356f..2a9a5404 100644 --- a/pytorch_lightning/models/trainer.py +++ b/pytorch_lightning/models/trainer.py @@ -5,6 +5,7 @@ from pytorch_lightning.root_module.memory import get_gpu_memory_map import traceback from pytorch_lightning.root_module.model_saving import TrainerIO from torch.optim.lr_scheduler import MultiStepLR +from torch.nn import DataParallel import pdb try: @@ -62,6 +63,8 @@ class Trainer(TrainerIO): self.lr_schedulers = [] self.amp_level = amp_level self.check_grad_nans = check_grad_nans + self.data_parallel = True + self.data_parallel_device_ids = [0, 1, 2, 3, 4, 5, 6, 7] # training state self.optimizers = None @@ -242,7 +245,10 @@ class Trainer(TrainerIO): # put on gpu if needed if self.on_gpu: - model = model.cuda() + if self.data_parallel: + model = DataParallel(model, device_ids=self.data_parallel_device_ids) + else: + model = model.cuda() # run tiny validation to make sure program won't crash during val _ = self.validate(model, self.val_dataloader, max_batches=self.nb_sanity_val_steps)