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 <vanpelt@gmail.com>
Co-authored-by: Jirka Borovec <Borda@users.noreply.github.com>
This commit is contained in:
William Falcon
2020-04-03 15:03:00 -04:00
committed by GitHub
co-authored by Chris Van Pelt Jirka Borovec
parent dd5a05926c
commit 3c5530c29d
2 changed files with 11 additions and 5 deletions
+9 -5
View File
@@ -65,10 +65,11 @@ class WandbLogger(LightningLoggerBase):
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 @@ class WandbLogger(LightningLoggerBase):
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 @@ class WandbLogger(LightningLoggerBase):
@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
+2
View File
@@ -50,6 +50,8 @@ def test_wandb_pickle(wandb):
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)