added single node distdataparallel

This commit is contained in:
William Falcon
2019-07-03 15:22:31 -04:00
parent d67e80bf16
commit 251c2e964f
+7 -1
View File
@@ -285,10 +285,16 @@ class Trainer(TrainerIO):
if self.on_gpu:
rank = 0
self.model = model
mp.spawn(dummy, nprocs=len(self.data_parallel_device_ids), args=(rank, ))
mp.spawn(dummy, nprocs=len(self.data_parallel_device_ids), args=(self, ))
else:
self.__run_pretrain_routine(model)
def __getstate__(self):
""" This is called before pickling. """
state = self.__dict__.copy()
del state['namelist']
return {}
def __dp_train(self, gpu_nb, proc_rank):
"""
Entry point into a DP thread