Skip to content

Unexpected behaviour with SwitchOptimizer #607

Description

@GiovanniCanali

Describe the bug
I would like to switch optimizers during training. However, the switch — handled via the SwitchOptimizer callback — appears to have no effect. Running the same script with or without the callback yields identical results. Even the loss curves on TensorBoard are perfectly aligned, showing no difference in training behavior.

To Reproduce
Below is the code for a simple Poisson problem.

from pina.problem.zoo import Poisson2DSquareProblem as Poisson
from pina.callback import SwitchOptimizer
from pina.domain import CartesianDomain
from pina.optim import TorchOptimizer
from matplotlib import pyplot as plt
from pina.model import FeedForward
from pina.solver import PINN
from pina import Trainer
import numpy as np
import lightning
import torch

# Epochs and seed
epochs = 1000
lightning.seed_everything(37, workers=True)

# Problem
problem = Poisson()
problem.discretise_domain(n=50, mode="grid")

# Model
model = FeedForward(2, 1, 64, 2)

# Solver
optimizer = TorchOptimizer(torch.optim.Adam, lr=1e-3)
solver = PINN(problem=problem, model=model, optimizer=optimizer)

# Callback
new_optim = TorchOptimizer(torch.optim.LBFGS, lr=1.0, line_search_fn='strong_wolfe')
switch_callback = SwitchOptimizer(new_optimizers=new_optim, epoch_switch=800)

# Trainer
trainer = Trainer(
    solver=solver,
    max_epochs=epochs,
    devices=1,
    accelerator="gpu",
    enable_progress_bar=False,
    enable_model_summary=False,
    callbacks=switch_callback, # Comment this line to disable the callback
)

trainer.train()
print("Final optimizer:", solver.optimizer.instance)

Expected behavior
I expect the results and training trajectories to differ when switching optimizers.

Metadata

Metadata

Labels

bugSomething isn't working

Type

Fields

No fields configured for Bug.

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions