deepxde.optimizers.pytorch

deepxde.optimizers.pytorch.nncg module

class deepxde.optimizers.pytorch.nncg.NNCG(params, lr=1.0, rank=10, mu=0.0001, update_freq=20, chunk_size=1, cg_tol=1e-16, cg_max_iters=1000, line_search_fn=None, verbose=False)[source]

Bases: Optimizer

Implementation of NysNewtonCG, a damped Newton-CG method

that uses Nyström preconditioning.

Rathore et al. Challenges in Training PINNs: A Loss Landscape Perspective. Preprint, 2024. <https://arxiv.org/abs/2402.01868>

Warning

This optimizer doesn’t support per-parameter options and parameter groups (there can be only one).

NOTE: This optimizer is currently a beta version.

Our implementation is inspired by the PyTorch implementation of L-BFGS <https://pytorch.org/docs/stable/_modules/torch/optim/lbfgs.html#LBFGS>.

The parameters rank and mu will probably need to be tuned for your specific problem. If the optimizer is running very slowly, you can try one of the following: - Increase the rank (this should increase the accuracy of the Nyström approximation in PCG) - Reduce cg_tol (this will allow PCG to terminate with a less accurate solution) - Reduce cg_max_iters (this will allow PCG to terminate after fewer iterations)

Parameters:
  • params (iterable) – iterable of parameters to optimize or dicts defining parameter groups

  • lr (float, optional) – learning rate (default: 1.0)

  • rank (int, optional) – rank of the Nyström approximation (default: 10)

  • mu (float, optional) – damping parameter (default: 1e-4)

  • update_freq (int, optional) – frequency of updating the preconditioner

  • chunk_size (int, optional) – number of Hessian-vector products to be computed in parallel (default: 1)

  • cg_tol (float, optional) – tolerance for PCG (default: 1e-16)

  • cg_max_iters (int, optional) – maximum number of PCG iterations (default: 1000)

  • line_search_fn (str, optional) – either ‘armijo’ or None (default: None)

  • verbose (bool, optional) – verbosity (default: False)

step(closure)[source]

Perform a single optimization step.

Parameters:

closure (callable) – A closure that reevaluates the model and returns the loss w.r.t. the parameters.