What's the proper way to do logging with nested PL modules?

[siyuanfeng-tri]

Suppose I have pytorch lighning modules A, B. And another pytorch lighning C that holds A, B. I make a trainer out of C, and proceed with the normal workflow. I can run the training and validation loops just fine. But I can't seem to properly log things (e.g. calling self.log()) from A or B. And I think the issue is that A and B are not visible / registered to the trainer, and thus they never get a logger attached. What's the recommended way to resolve this?

Thanks!

[tchaton]

Hey @siyuanfeng-tri.

You would need to override the __setattr__ method of your LightningModule.

    def __setattr__(self, key, value):
        if isinstance(value, (LightningModule, ModuleWrapperBase)):
            self._children.append(key)
        patched_attributes = ["_current_fx_name", "_current_hook_fx_name", "_results"]
        if isinstance(value, Trainer) or key in patched_attributes:
            if hasattr(self, "_children"):
                for child in self._children:
                    setattr(getattr(self, child), key, value)
        super().__setattr__(key, value)

Check this out from Lightning Flash: https://github.com/PyTorchLightning/lightning-flash/blob/master/flash/core/model.py#L91

This is currently not supported yet in Lightning but it is for the Flash Task.

[siyuanfeng-tri]

Awesome, I will give it a try, thanks much