From 63ebf0d25c0af65b6963abbaccee091e54855ddc Mon Sep 17 00:00:00 2001 From: William Falcon Date: Fri, 3 Apr 2020 15:03:00 -0400 Subject: [PATCH] Wandb bug/wandb multi (#1360) * Allow reinits in sub procs * Dont create an experiment on pickle, name, or project * Comments consistency * Fix test * Apply suggestions from code review Co-authored-by: Chris Van Pelt Co-authored-by: Jirka Borovec --- pytorch_lightning/loggers/wandb.py | 14 +++++++++----- tests/loggers/test_wandb.py | 2 ++ 2 files changed, 11 insertions(+), 5 deletions(-) diff --git a/pytorch_lightning/loggers/wandb.py b/pytorch_lightning/loggers/wandb.py index 39bfcc3cd21e8e..8cabfae52e07a8 100644 --- a/pytorch_lightning/loggers/wandb.py +++ b/pytorch_lightning/loggers/wandb.py @@ -65,10 +65,11 @@ def __init__(self, name: Optional[str] = None, save_dir: Optional[str] = None, def __getstate__(self): state = self.__dict__.copy() + # args needed to reload correct experiment + state['_id'] = self._experiment.id if self._experiment is not None else None + # cannot be pickled state['_experiment'] = None - # args needed to reload correct experiment - state['_id'] = self.experiment.id return state @property @@ -87,7 +88,7 @@ def experiment(self) -> Run: os.environ['WANDB_MODE'] = 'dryrun' self._experiment = wandb.init( name=self._name, dir=self._save_dir, project=self._project, anonymous=self._anonymous, - id=self._id, resume='allow', tags=self._tags, entity=self._entity) + reinit=True, id=self._id, resume='allow', tags=self._tags, entity=self._entity) # save checkpoints in wandb dir to upload on W&B servers if self._log_model: self.save_dir = self._experiment.dir @@ -109,8 +110,11 @@ def log_metrics(self, metrics: Dict[str, float], step: Optional[int] = None) -> @property def name(self) -> str: - return self.experiment.project_name() + # don't create an experiment if we don't have one + name = self._experiment.project_name() if self._experiment else None + return name @property def version(self) -> str: - return self.experiment.id + # don't create an experiment if we don't have one + return self._experiment.id if self._experiment else None diff --git a/tests/loggers/test_wandb.py b/tests/loggers/test_wandb.py index bd430ca21c4175..23cdb6fb3156c2 100644 --- a/tests/loggers/test_wandb.py +++ b/tests/loggers/test_wandb.py @@ -50,6 +50,8 @@ class Experiment: trainer_options = dict(max_epochs=1, logger=logger) trainer = Trainer(**trainer_options) + # Access the experiment to ensure it's created + trainer.logger.experiment pkl_bytes = pickle.dumps(trainer) trainer2 = pickle.loads(pkl_bytes)