Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We鈥檒l occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix(wandb): allow custom init args #6989

Merged
merged 15 commits into from May 4, 2021
3 changes: 3 additions & 0 deletions CHANGELOG.md
Expand Up @@ -440,6 +440,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
- Fixed `apex` not properly instantiated when running with `ddp` ([#7274](https://github.com/PyTorchLightning/pytorch-lightning/pull/7274))


- Fixed custom init args for `WandbLogger` ([#6989](https://github.com/PyTorchLightning/pytorch-lightning/pull/6989))


## [1.2.7] - 2021-04-06

### Fixed
Expand Down
36 changes: 18 additions & 18 deletions pytorch_lightning/loggers/wandb.py
Expand Up @@ -59,8 +59,7 @@ class WandbLogger(LightningLoggerBase):
log_model: Save checkpoints in wandb dir to upload on W&B servers.
prefix: A string to put at the beginning of metric keys.
experiment: WandB experiment object. Automatically set when creating a run.
\**kwargs: Additional arguments like `entity`, `group`, `tags`, etc. used by
:func:`wandb.init` can be passed as keyword arguments in this logger.
\**kwargs: Arguments passed to :func:`wandb.init` like `entity`, `group`, `tags`, etc.

Raises:
ImportError:
Expand Down Expand Up @@ -93,7 +92,7 @@ def __init__(
save_dir: Optional[str] = None,
offline: Optional[bool] = False,
id: Optional[str] = None,
anonymous: Optional[bool] = False,
anonymous: Optional[bool] = None,
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why change the default?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is actually still the same behavior but is more in line with the W&B documentation so less prone to potential bugs.

version: Optional[str] = None,
project: Optional[str] = None,
log_model: Optional[bool] = False,
Expand Down Expand Up @@ -122,16 +121,25 @@ def __init__(
)

super().__init__()
self._name = name
self._save_dir = save_dir
self._offline = offline
self._id = version or id
self._anonymous = 'allow' if anonymous else None
self._project = project
self._log_model = log_model
self._prefix = prefix
self._experiment = experiment
self._kwargs = kwargs
# set wandb init arguments
anonymous_lut = {True: 'allow', False: None}
self._wandb_init = dict(
name=name,
project=project,
id=version or id,
dir=save_dir,
resume='allow',
anonymous=anonymous_lut.get(anonymous, anonymous)
)
self._wandb_init.update(**kwargs)
# extract parameters
self._save_dir = self._wandb_init.get('dir', None)
self._name = self._wandb_init.get('name', None)
self._id = self._wandb_init.get('id', None)
Borda marked this conversation as resolved.
Show resolved Hide resolved

def __getstate__(self):
state = self.__dict__.copy()
Expand All @@ -158,15 +166,7 @@ def experiment(self) -> Run:
if self._experiment is None:
if self._offline:
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',
**self._kwargs
) if wandb.run is None else wandb.run
self._experiment = wandb.init(**self._wandb_init) if wandb.run is None else wandb.run

# save checkpoints in wandb dir to upload on W&B servers
if self._save_dir is None:
Expand Down
10 changes: 8 additions & 2 deletions tests/loggers/test_wandb.py
Expand Up @@ -37,9 +37,13 @@ def test_wandb_logger_init(wandb, recwarn):

# test wandb.init called when there is no W&B run
wandb.run = None
logger = WandbLogger()
logger = WandbLogger(
name='test_name', save_dir='test_save_dir', version='test_id', project='test_project', resume='never'
)
logger.log_metrics({'acc': 1.0})
wandb.init.assert_called_once()
wandb.init.assert_called_once_with(
name='test_name', dir='test_save_dir', id='test_id', project='test_project', resume='never', anonymous=None
)
wandb.init().log.assert_called_once_with({'acc': 1.0})

# test wandb.init and setting logger experiment externally
Expand All @@ -55,6 +59,8 @@ def test_wandb_logger_init(wandb, recwarn):
wandb.init.reset_mock()
wandb.run = wandb.init()
logger = WandbLogger()
# verify default resume value
assert logger._wandb_init['resume'] == 'allow'
logger.log_metrics({'acc': 1.0}, step=3)
wandb.init.assert_called_once()
wandb.init().log.assert_called_once_with({'acc': 1.0, 'trainer/global_step': 3})
Expand Down