Support arbitrary `Optimizable`s as optimizers #16189

carmocca · 2022-12-23T02:19:59Z

What does this PR do?

This allows us to support works like Gradient Descent: The Ultimate Optimizer (code)

Example script:

import pytorch_lightning as L
import torch
from pytorch_lightning.demos.boring_classes import RandomDataset

# pip install git+https://github.com/kach/gradient-descent-the-ultimate-optimizer
from gradient_descent_the_ultimate_optimizer import gdtuo


# Versions compatible with the requirements of torch optimizers
# Ideally these would be upstreamed so they are not required
class CompatibleAdam(gdtuo.Adam):
    param_groups = None
    defaults = None

    @property
    def state(self):
        return self.cache

    def state_dict(self):
        return {}

    def load_state_dict(self):
        pass


class CompatibleModuleWrapper(gdtuo.ModuleWrapper):
    param_groups = None
    defaults = None

    @property
    def state(self):
        return self.optimizer.state

    def step(self, closure):
        # ignore the closure - optional during manual optimization anyways
        return super().step()

    def state_dict(self):
        return self.optimizer.state_dict()

    def load_state_dict(self):
        return self.optimizer.load_state_dict()


class GDTUOLightningModule(L.LightningModule):
    def __init__(self):
        super().__init__()
        self.layer = torch.nn.Linear(32, 2)
        self.automatic_optimization = False

    def configure_optimizers(self):
        optimizer0 = gdtuo.SGD(1e-5)
        optimizer1 = CompatibleAdam(optimizer=optimizer0)
        mw = CompatibleModuleWrapper(self, optimizer=optimizer1)
        mw.initialize()
        return mw

    def forward(self, x):
        return self.layer(x)

    def training_step(self, batch):
        mw = self.optimizers()
        mw.begin()  # call this before each step, enables gradient tracking on desired params
        mw.zero_grad()

        pred = self(batch)
        loss = torch.nn.functional.mse_loss(pred, torch.ones_like(pred))

        loss.backward(create_graph=True)  # important! use create_graph=True
        mw.step()

        return loss


dataset = RandomDataset(32, 64)
train_dataloader = torch.utils.data.DataLoader(dataset, batch_size=2)
trainer = L.Trainer(max_epochs=10, accelerator="auto")
model = GDTUOLightningModule()
trainer.fit(model, train_dataloader)

Does your PR introduce any breaking changes? If yes, please list them.

None

Before submitting

Was this discussed/approved via a GitHub issue? (not for typos and docs)
Did you read the contributor guideline, Pull Request section?
Did you make sure your PR does only one thing, instead of bundling different changes together?
[n/a] Did you make sure to update the documentation with your changes? (if necessary)
[n/a] Did you write any new necessary tests? (not for typos and docs)
Did you verify new and existing tests pass locally with your changes?
Did you list all the breaking changes introduced by this pull request?
[n/a] Did you update the CHANGELOG? (not for typos, docs, test updates, or minor internal changes/refactors)

PR review

Is this pull request ready for review? (if not, please submit in draft mode)
Check that all items from Before submitting are resolved
Make sure the title is self-explanatory and the description concisely explains the PR
Add labels and milestones (and optionally projects) to the PR so it can be classified

cc @Borda

github-actions · 2022-12-23T02:20:46Z

⚡ Required checks status: All passing 🟢