Support for Weights and Biases logger (and possibly others) in Lightning CLI

[ohayonguy]

🚀 Feature

Adding support for weights and biases logger in lightning CLI

Motivation

Currently, instantiating a weights and biases logger from lightning CLI is not supported, and causes an error:

trainer.py: error: Parser key “trainer.logger”: Value “{‘class_path’: ‘pytorch_lightning.loggers.WandbLogger’, ‘init_args’: {‘entity’: ‘gip’, ‘project’: ‘solver’, ‘notes’: ‘notes’}}” does not validate against any of the types in typing.Union[pytorch_lightning.loggers.base.LightningLoggerBase, typing.Iterable[pytorch_lightning.loggers.base.LightningLoggerBase], bool] :: Problem with given class_path “pytorch_lightning.loggers.WandbLogger” :: ‘Configuration check failed :: No action for key “entity” to check its value.’ :: Expected a List but got “{‘class_path’: ‘pytorch_lightning.loggers.WandbLogger’, ‘init_args’: {‘entity’: ‘gip’, ‘project’: ‘solver’, ‘notes’: ‘notes’}}” :: Expected a <class ‘bool’> but got “{‘class_path’: ‘pytorch_lightning.loggers.WandbLogger’, ‘init_args’: {‘entity’: ‘gip’, ‘project’: ‘solver’, ‘notes’: ‘notes’}}”

Pitch

I want to be able to initialize a WandbLogger from the yaml file, and also be able to run wandb sweeps.

[ohayonguy]

@mauvilsa cc

[mauvilsa]

From a quick look at this the problem seems to be only for the cases in which a WandbLogger is given values that should go to its init **kwargs. I will look into this.

[oplatek]

Maybe I am misunderstanding the question,
but for me using instantiating WandBLogger in the config with init_args in the field trainer: logger: class_path works fine.

[ohayonguy]

@oplatek Did you try passing arguments to the wandb logger? For example, passing an 'entity' kwarg.

[oplatek]

@ohayonguy oh... I see, I just used the name, save_dir.

Now, I understand the problem, sorry I misunderstood it.

I experienced this problem with LightningCLI too. I solved it by subclassing the parent and moving the argument from kwargs to regular argument in the inherited __init__ function.

[ohayonguy]

Ahh yup :) that's exactly what we did!
@oplatek

[ohayonguy]

@oplatek We still have a very annoying problem though, when we have nested pl.module classes that takes nn.Module as arguments. These can't be instantiated :(

[oplatek]

I think I solved this too (if I understand you correctly).
I played a bit with LightningCLI on minimal examples to understand it.

See https://github.com/oplatek/pytorch-lightning/blob/d3c5dabd9766cd049f5d21bef7a8dc4845db3814/pl_examples/basic_examples/simple_image_classifier.yml#L3

and
https://github.com/oplatek/pytorch-lightning/blob/d3c5dabd9766cd049f5d21bef7a8dc4845db3814/pl_examples/basic_examples/simple_image_classifier.py#L39

[mauvilsa]

when we have nested pl.module classes that takes nn.Module as arguments. These can't be instantiated :(

This is happening because some init arguments in nn.Module classes don't have type hints. Currently in LightningCLI it is implemented that arguments without a type get the Any type. However looking at this I have noticed that this option is not enabled for some cases. I know why this is and will implement it soon.

[mauvilsa]

From a quick look at this the problem seems to be only for the cases in which a WandbLogger is given values that should go to its init **kwargs. I will look into this.

The issue is that WandbLogger receives **kwargs (wandb.py#L102) which is used for running wandb.init(). jsonargparse is not aware of the arguments that wandb.init() receives. The way jsonargparse works is via class inheritance. If the init of a class has **kwargs then it will look for more arguments in the next class it inherits from. If the following class also receives **kwargs it will continue looking in the next class and so on.

I would propose that we improve the implementation of WandbLogger such that it inherits from a class WandbInitOptions which would include all the arguments that can be accepted. This way the parser can be aware of them and validate against its types. The WandbInitOptions class would not be written manually. It would be a class dynamically generated using some function that inspects the signature of wandb.init(). This way if something changes in wandb it becomes automatically available in lightning.

Please analyze and comment to see if the idea makes sense and to figure out what other issues there could be.

[ohayonguy]

@mauvilsa Your solution seems okay to me, although it's a "workaround". Is there a way to inherit form jsonargparse and change its behavior to our needs? I am not very familiar with it. Changing jsonargparse might be a better solution in the long term, and could give more flexibility (for example, instantiate nn.Modules in subclasses and stuff).

[mauvilsa]

@ohayonguy the nn.Module instantiation is different. This a missing feature that will be implemented in jsonargparse except for maybe a minor modification in LightningCLI.

The proposal for **kwargs is not a workaround. Why do you say it is a workaround? jsonargparse needs to know what arguments it can accept and how to validate them. This is important so that when incorrect values are given it fails early and not after a whole lot of processing has been done. The way to tell jsonargparse what are valid arguments are with the type hints and in the case of classes possibly with class inheritance. This is a change that would be done in pytorch-lightning. From what I understand in the description this will fit your needs.