Add NNCG to optimizers submodule

[pratikrathore8]

No description provided.

[lululxvi]

Format the code using black https://github.com/psf/black

[pratikrathore8]

I've formatted the code with black in the latest commit!

[lululxvi]

If this only supports pytorch, then it should be in https://github.com/lululxvi/deepxde/tree/master/deepxde/optimizers/pytorch

[pratikrathore8]

Moved to the right folder!

[pratikrathore8]

Made the formatting changes. I noticed there are other lines in the code with > 88 chars per line. Would you like for me to fix those?

[lululxvi]

Made the formatting changes. I noticed there are other lines in the code with > 88 chars per line. Would you like for me to fix those?

Yes.

Could you also fix the Codacy issues?

[pratikrathore8]

The formatting has been fixed.

There are still some Codacy issues, but these are mainly due to "too many local variables" and the __init__ function in the optimizer, which is just following PyTorch conventions.

What do we need do next to further integrate NNCG into the codebase?

[lululxvi]

Need to add the NNCG to https://github.com/lululxvi/deepxde/blob/master/deepxde/optimizers/pytorch/optimizers.py

[pratikrathore8]

I've added NNCG to optimizers.py and NNCG_options to config.py

[pratikrathore8]

Fixed!

[pratikrathore8]

Fixed!

[lululxvi]

Have you tried using NNCG in some demo examples?

[pratikrathore8]

A quick update:

I've started thinking about how to make a demo example that uses NNCG. The plan would be to modify some of the existing examples (e.g. Burgers.py) and add on NNCG after Adam and L-BFGS.

I'll have to make some modifications to model.py to full integrate NNCG into the training pipeline. I'm aiming to have an example fully working by the end of next week!

[pratikrathore8]

@lululxvi Sorry for not responding earlier -- I was away from Stanford during the summer. My plan is to take a look at the code next week!

[pratikrathore8]

@lululxvi My apologies for taking so long. I've updated the code -- I moved most of the complexity from model.py into nncg.py.

The main thing to note is that I added an argument perform_backward to train_step for the PyTorch backend -- this argument is True by default and is set to False only in _train_pytorch_nncg. The reason for adding this argument is because NNCG requires Hessian-vector products (hvps), which necessitates a slightly differently implementation from other optimizers, as I have explained below:

We can obtain hvps by either using .backward() with create_graph=True followed by torch.autograd.grad or by simply using torch.autograd.grad twice.

However, I find that the approach using .backward() runs out of memory; the memory of NNCG increases with the number of iterations instead of remaining constant, which is undesirable. I believe this is because using create_graph=True with .backward() leads to the computation graph being retained across iterations instead of being deleted at every iteration.

On the other hand using torch.autograd.grad twice to compute hvps leads to constant memory across iterations.

The memory issues with .backward() might be fixable with a lot of debugging, but I don't think it's worth it -- even PyTorch throws a warning about possible memory leaks when using the .backward() approach to calculate hvps.

Let me know if you'd like for me to make more changes -- I'd be interested in knowing if the tutorial Burgers_NNCG.py in examples works on your end.

The errors in the Codacy analysis are from nncg.py, but these are minor issues related to PyTorch optimizer conventions and too many local variables (which is difficult to avoid in the implementation).

[lululxvi]

Could you fix the other Codacy issues of dict and super-with-arguments?

[pratikrathore8]

@lululxvi I've removed the unnecessary error checking and fixed the Codacy issues of dict and super-with-arguments.

[pratikrathore8]

@lululxvi Made the changes. Should we consider merging _train_pytorch_nncg into _train_sgd by adding an argument to _train_sgd that decides whether or not to perform a backward pass on the loss? Or would you prefer we leave the functions separate, despite the code duplication?

[pratikrathore8]

@lululxvi I added a train_step_nncg function as you suggested. NNCG now runs using _train_sgd.

Thanks for your patience with the refactoring steps so far!

[pratikrathore8]

@lululxvi Done!

[lululxvi]

Add a document for the example similar to this https://deepxde.readthedocs.io/en/latest/demos/pinn_forward/burgers.html

[pratikrathore8]

@lululxvi I added the demo. It's basically the same as the one you linked, except I added the part with NNCG to the end.

Is there a way to check if it's been properly integrated into the docs?

[lululxvi]

Please resolve the conflict.

You can look at the generated doc by clicking "Details" at the right side of "docs/readthedocs.org:deepxde" in in the check list.

[pratikrathore8]

@lululxvi I've fixed the merge conflict!

I took a look at the doc for the demo, and it seems good to me.