Lightning-AI / Lightning-AI/pytorch-lightning
Allow weight reuse in a different lightning module
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 31.4k
- Forks
- 3.8k
- Avg merge
- 6d 7h
- Merged PRs (30d)
- 6
Description
@mauvilsa I think we have a different issue than what #21246 fixes, where we want to load the state dict in a different lightning module. E.g. in the following we will have `lr` saved as a hyperparameter in the checkpoint from `TrainingModule`, but the `InferenceModule` does not take it as an argument:
```py
import torch
from lightning.pytorch import LightningModule
class TrainingModule(LightningModule):
def __init__(self, lr: float = 1e-3) -> None:
super().__init__()
self.model = torch.nn.Linear(16, 2)
self.lr = lr
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.model(x)
def configure_optimizers(self):
optimizer = torch.optim.Adam(self.parameters(), lr=self.lr)
return optimizer
def training_step(
self, batch: tuple[torch.Tensor, torch.Tensor], batch_idx: int
) -> torch.Tensor:
x, y = batch
y_hat = self(x)
loss = torch.nn.functional.cross_entropy(y_hat, y)
self.log("train_loss", loss)
return loss
def train_dataloader(self) -> torch.utils.data.DataLoader:
dataset = torch.utils.data.TensorDataset(
torch.rand(32, 16), torch.randint(0, 2, (32,))
)
return torch.utils.data.DataLoader(dataset, batch_size=8)
class InferenceModule(LightningModule):
def __init__(self) -> None:
super().__init__()
self.model = torch.nn.Linear(16, 2)
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.model(x)
def predict_dataloader(self) -> torch.utils.data.DataLoader:
dataset = torch.utils.data.TensorDataset(torch.rand(32, 16))
return torch.utils.data.DataLoader(dataset, batch_size=128)
def predict_step(
self, batch: torch.Tensor, batch_idx: int, dataloader_idx: int = 0
) -> torch.Tensor:
x = batch
y_hat = self(x)
return y_hat
```
_Originally posted by @ziw-liu in https://github.com/Lightning-AI/pytorch-lightning/issues/21116#issuecomment-3336106861_
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 by reproducing the TrainingModule and InferenceModule example from the issue and compare the behavior described in #21246. Trace checkpoint loading and state-dict handling to identify where saved hyperparameters are required, then add coverage showing that InferenceModule can load TrainingModule weights without accepting lr. Done means the different module loads successfully and its parameters are restored.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100