Lightning-AI / Lightning-AI/pytorch-lightning

Allow weight reuse in a different lightning module

Open
#21,255 5 comments 0 reactions 0 assignees View on GitHub

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

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 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.