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

load_spawn_weights only in proc rank 0 (#1385)

Browse Source 
Co-authored-by: Alexander Reshytko <areshytko@Alexanders-MacBook-Pro.local>
This commit is contained in:
areshytko
2020-04-06 10:17:16 -04:00
committed by GitHub
co-authored by Alexander Reshytko
parent 4ed3027309
commit 9754c5da55
1 changed files with 11 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pytorch_lightning/trainer/distrib_data_parallel.py
+11 -7
View File
@@ -363,15 +363,19 @@ class TrainerDDPMixin(ABC):
:param model:
:return:
"""
# load weights saved in ddp
path = os.path.join(self.default_save_path, '__temp_weight_ddp_end.ckpt')
loaded_model = original_model.__class__.load_from_checkpoint(path)
# copy loaded weights to old model
original_model.load_state_dict(loaded_model.state_dict())
loaded_model = original_model
# remove ddp weights
os.remove(path)
if self.proc_rank == 0:
# load weights saved in ddp
path = os.path.join(self.default_save_path, '__temp_weight_ddp_end.ckpt')
loaded_model = original_model.__class__.load_from_checkpoint(path)
# copy loaded weights to old model
original_model.load_state_dict(loaded_model.state_dict())
# remove ddp weights
os.remove(path)
return loaded_model