Lightning-AI / Lightning-AI/pytorch-lightning
Can I nest LightningModules inside child modules?
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 31.4k
- Forks
- 3.8k
- Avg merge
- 6d 7h
- Merged PRs (30d)
- 6
Description
### Bug description
Suppose I have a `LightningModule` (parent) that contains a `nn.Module` (child), which in turn contains another `LightningModule` (grandchild). Calling `.log` inside the `LightningModule` (the grandchild) results in the following warning:
> You are trying to `self.log()` but the `self.trainer` reference is not registered on the model yet. This is most likely because the model hasn't been passed to the `Trainer`
The trainer is only set on the direct `children` of the parent `LightningModule`, not all the descendants, since the `trainer.setter` uses `self.children()` rather than `self.modules()`: https://github.com/Lightning-AI/pytorch-lightning/blob/3730e980e388c23f7e9d1f535793e8d614633362/src/lightning/pytorch/core/module.py#L221-L226
### What version are you seeing the problem on?
master
### How to reproduce the bug
```python
# %%
import lightning as L
import torch
from torch import nn
class GrandChild(L.LightningModule):
def dummy_log(self):
self.log("foo", 1)
class Child(nn.Module):
def __init__(self):
super().__init__()
self.module = nn.Linear(1, 1)
self.grandchild = GrandChild()
def forward(self):
self.grandchild.dummy_log()
return 1
class Parent(L.LightningModule):
def __init__(self):
super().__init__()
self.child = Child()
def training_step(self, batch, batch_idx):
return self.child()
def configure_optimizers(self):
optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
return optimizer
def train_dataloader(self):
return torch.utils.data.DataLoader(
torch.utils.data.TensorDataset(torch.randn(10, 1)), batch_size=1
)
# model
parent = Parent()
# train model
trainer = L.Trainer()
trainer.fit(model=parent)
optimizer = parent.configure_optimizers()
loss = parent.training_step(batch=None, batch_idx=0)
```
### Error messages and logs
```
You are trying to `self.log()` but the `self.trainer` reference is not registered on the model yet. This is most likely because the model hasn't been passed to the `Trainer`
```
### Environment
Current environment
* CUDA:
- GPU:
- NVIDIA A100-SXM4-80GB
- available: True
- version: 12.1
* Lightning:
- lightning: 2.2.1
- lightning-utilities: 0.11.2
- pytorch-lightning: 2.2.1
- torch: 2.3.1
- torchmetrics: 1.3.2
- torchvision: 0.18.1
* System:
- OS: Linux
- architecture:
- 64bit
- ELF
- processor: x86_64
- python: 3.11.9
- release: 5.15.0-113-generic
- version: #123-Ubuntu SMP Mon Jun 10 08:16:17 UTC 2024
### More info
_No response_
cc @carmocca @justusschock @awaelchli @borda
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start at src/lightning/pytorch/core/module.py around the trainer.setter referenced in the issue, then run the supplied Parent, Child, and GrandChild reproduction. Check how trainer registration behaves for nested modules and add regression coverage showing that the grandchild can call self.log() without the warning.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 48/100