testadmin
/
ParakeetRebeccaRosario
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

fix bugs with multiprocess training.

This commit is contained in:
chenfeiyu 2021-02-18 19:09:54 +08:00
parent 3f60b6e0a3
commit f423323bae
2 changed files with 6 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
parakeet/training/experiment.py
View File

@ -195,17 +195,16 @@ class ExperimentBase(object):
self.save() self.save()
exit(-1) exit(-1)
@mp_tools.rank_zero_only
def setup_output_dir(self): def setup_output_dir(self):
"""Create a directory used for output. """Create a directory used for output.
""" """
# output dir # output dir
output_dir = Path(self.args.output).expanduser() output_dir = Path(self.args.output).expanduser()
output_dir.mkdir(parents=True, exist_ok=True) if dist.get_rank() == 0:
output_dir.mkdir(parents=True, exist_ok=True)
self.output_dir = output_dir self.output_dir = output_dir
@mp_tools.rank_zero_only
def setup_checkpointer(self): def setup_checkpointer(self):
"""Create a directory used to save checkpoints into. """Create a directory used to save checkpoints into.
@ -213,7 +212,8 @@ class ExperimentBase(object):
""" """
# checkpoint dir # checkpoint dir
checkpoint_dir = self.output_dir / "checkpoints" checkpoint_dir = self.output_dir / "checkpoints"
checkpoint_dir.mkdir(exist_ok=True) if dist.get_rank() == 0:
checkpoint_dir.mkdir(exist_ok=True)
self.checkpoint_dir = checkpoint_dir self.checkpoint_dir = checkpoint_dir

3
parakeet/utils/mp_tools.py
View File

@ -20,11 +20,10 @@ __all__ = ["rank_zero_only"]
def rank_zero_only(func): def rank_zero_only(func):
local_rank = dist.get_rank()
@wraps(func) @wraps(func)
def wrapper(*args, **kwargs): def wrapper(*args, **kwargs):
if local_rank != 0: if dist.get_rank() != 0:
return return
result = func(*args, **kwargs) result = func(*args, **kwargs)
return result return result