Lightning-AI / Lightning-AI/pytorch-lightning

Can I nest LightningModules inside child modules?

Open
#20,053 4 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

lightningmodule question ver: 2.2.x
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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.