Using the optimizer LBFGS

[sergwio]

Hello, I'm trying to use an optimizer LBFGS:
optimizer = torch.optim.LBFGS(problem.parameters())

I got an error:

File c:\Program Files (x86)\Microsoft Visual Studio\Shared\Python39_64\lib\site-packages\neuromancer\trainer.py:256, in Trainer.train(self)
254 output[self.train_metric].backward()
255 torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.clip)
--> 256 self.optimizer.step()
257 losses.append(output[self.train_metric])
258 self.callback.end_batch(self, output)

File c:\Program Files (x86)\Microsoft Visual Studio\Shared\Python39_64\lib\site-packages\torch\optim\optimizer.py:485, in Optimizer.profile_hook_step..wrapper(*args, **kwargs)
480 else:
481 raise RuntimeError(
482 f"{func} must return None or a tuple of (new_args, new_kwargs), but got {result}."
483 )
--> 485 out = func(*args, **kwargs)
486 self._optimizer_step_code()
488 # call optimizer step post hooks

File c:\Program Files (x86)\Microsoft Visual Studio\Shared\Python39_64\lib\site-packages\torch\utils_contextlib.py:116, in context_decorator..decorate_context(*args, **kwargs)
113 @functools.wraps(func)
114 def decorate_context(*args, **kwargs):
115 with ctx_factory():
--> 116 return func(*args, **kwargs)

TypeError: step() missing 1 required positional argument: 'closure'

[drgona]

Dear @sergwio. Thanks for using Neuromancer!

LBFGS in PyTorch requires a closure that recomputes loss and gradients each step(). It’s also memory-heavy and generally prefers full-batch or very large batches.

Minimal pattern (drop-in idea for your training loop to try):

optimizer = torch.optim.LBFGS(model.parameters(),
                              lr=1.0,           # often OK to start at 1.0
                              max_iter=20,      # per .step() internal iterations
                              history_size=50,  # reduce if OOM
                              line_search_fn='strong_wolfe')  # optional
loss_meter = []
for epoch in range(num_epochs):
    for batch in train_loader:                 # LBFGS works best with big/full batches
        x, y = batch
        def closure():
            optimizer.zero_grad(set_to_none=True)
            yhat = model(x)
            loss = loss_fn(yhat, y) + reg_terms(model)  # include NM penalties/constraints here
            loss.backward()
            return loss
        loss = optimizer.step(closure)         # returns last closure loss
        loss_meter.append(loss.item())

You would need to create a similar closure for the Neuromancer trainer, or use a custom trainer to support this.
Let us know if this helps.

[sergwio-dev]

Dear @drgona, thank you for your answer.
I implemented the training as follows.

optimizer2 = torch.optim.LBFGS(problem.parameters(),
                              lr=1e-0,           # often OK to start at 1.0
                              #max_iter=3000,      # per .step() internal iterations
                              #history_size=3500,  # reduce if OOM
                              line_search_fn='strong_wolfe',
                              #tolerance_grad=1e-100,
                              #tolerance_change=1e-300
                              )  

num_epochs = 50
loss_meter = []
for epoch in range(num_epochs):
    problem.train()
    for batch in train_loader:               
        x = batch
        def closure():
            optimizer2.zero_grad(set_to_none=True)
            loss = problem(x).get('train_loss')
            loss.backward()
            return loss
        torch.nn.utils.clip_grad_norm_(problem.parameters(), 100.0)
        loss = optimizer2.step(closure)         
        loss_meter.append(loss.item())
    print(f"Epoch {epoch}/{num_epochs}, Loss: {loss.item():.15f}")

I hope this is correct.
I have tried different parameters, LBFGS is faster, but AdamW is more accurate for my simple case (Helmholtz equation).
The full text of the program is here

test.ipynb

[drgona]

Great it worked out! I am going to close this issue.
Since this is about developing new feature I am going to transfer this to a discussion. If you are interested feel free to submit PR with your example on how to use LBFGS in Neuromancer.

[Parv621]

I faced a similar problem working on one of the DeepONet examples. A small change in trainer.py (using an if condition for an instance of LBFGS) can enable running L-BFGS, but as of now, I am not sure it's a robust addition to the library.

[Parv621]

Change to be added in trainer.py (replace the inner train loop)

is_lbfgs = isinstance(self.optimizer, torch.optim.LBFGS)

for t_batch in self.train_data:
    t_batch["epoch"] = i
    t_batch = move_batch_to_device(t_batch, self.device)
    
    ## Changes for LBFGS
    if is_lbfgs:
        def closure():
            self.optimizer.zero_grad()
            out = self.model(t_batch)
            loss = out[self.train_metric]
            loss.backward()
            return loss
        loss = self.optimizer.step(closure)         # loss is a Tensor
        output = {"train_loss": loss}               # stash for logging/mean
    else:
        self.optimizer.zero_grad()
        output = self.model(t_batch)
        loss = output[self.train_metric]
        loss.backward()
        torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.clip)
        self.optimizer.step()

    losses.append(loss)
    self.callback.end_batch(self, output)

I am not sure how to deal with batch training here - for LBFGS, full-batch training is preferred.