Separate optimization methods from CDIModel into separate class Reconstructors

[yoshikisd]

PR Summary

This PR separates ptychography optimization/reconstruction methods from CDIModel to a new class called Reconstructor.

The motivation of this change was to make the CDTools more compatible with other PyTorch-related packages which expect libraries/scripts to follow specific code patterns. Specifically, problems arise if optimization is a method of torch.nn.module/cdtools.models.CDIModel, which was the case as of v0.3.0.pypi1. Additionally, Reconstructor's implementation allows us to re-use an optimizer rather than create a brand new one each time a new reconstruction loop is called (as was done with calls to model.{Adam, LBFGS, SGD}_optimize)

This PR serves to split PR #17, a development PR that ended up creating both Reconstructors as well as multi-GPU capability. This PR was created by checking out all Reconstructors-related stuff from PR #17 and omits any multi-GPU related code.

@allevitan If you can add your reviews from #17, that would be appreciated.

Technical details

Implementation:

cdtools.reconstructors is largely based on methods from cdtools.models.CDIModel. For the base class cdtools.reconstructors.Reconstructor, the method(s):

• _run_epoch and optimize are based on cdtools.models.CDIModel.AD_Optimize.
• setup_dataloader is based on similar blocks of data-loading code present in cdtools.models.CDIModel.{Adam, LBFGS, SGD}_optimize
• adjust_optimizer is not implemented and must be defined within the subclasses to adjust the optimizer hyperparameters.

For the subclasses cdtools.reconstructors.{Adam, LBFGS, SGD}, the method(s):

• __init__ initializes an optimizer attribute as a torch.optim.{Adam, LBFGS, SGD} using the model parameters.
• adjust_optimizer changes the relevant optimizer hyperparameters
• optimize are based on `cdtools.models.CDIModel.{Adam, LBFGS, SGD}

The main change to the code behavior is that the hyperparameters are updated between different reconstruction loops without creating a new torch.optim.{Adam, LBFGS, SGD}. In some cases, this can result in an improvement in reconstruction quality as discussed in (#17 (comment)).

Changes in cdtools.models.CDIModel and backwards compatibility

The following changes have been made in cdtools.models.CDIModel methods:

• AD_Optimize has been removed.
• {Adam, LBFGS, SGD}_Optimize now initialize a reconstructor attribute as a cdtools.reconstructors.{Adam, LBFGS, SGD} and make a call to their optimize method.

The latter change allows scripts written using v0.3.0.pypi1 to be backwards compatible with this update.

However, this comes at the cost of creating a circular dependency between cdtools.models.CDIModel and cdtools.reconstructors. From a reconstruction standpoint it doesn’t seem to cause any issues. However, it’s not clear if this may cause issues in the long term in terms of maintainability. Thus, we should consider deprecating {Adam, LBFGS, SGD}_Optimize in favor of the following code pattern:

  import cdtools

  filename = ...
  dataset = ...
  model = cdtools.models.FancyPtycho.from_dataset(...)

  device = 'cuda'
  model.to(device=device)
  dataset.get_as(device=device)

+ recon = cdtools.optimizer.Adam(model, dataset)

- for loss in model.Adam_optimize(50, dataset, lr=0.02, batch_size=10):
+ for loss in recon.optimize(50, lr=0.02, batch_size=10):
    print(model.report())
    if model.epoch % 10 == 0:
        model.inspect(dataset)

PyTest

Three tests have been created in test_reconstructors.py which function similar to the tests for the models. These three tests run reconstruction loops using both the traditional method (model.XXX_Optimize) and new method (recon.optimize), but use different optimizers (Adam, LBFGS, SGD), datasets (gold ball dataset for Adam and SGD, siemens star optical data for LBFGS), and models (FancyPtycho for Adam and SGD, RPI for LBFGS). These three tests check the following:

• Calls to Reconstructor.adjust_optimizer are updating the hyperparameters
• Ensure that the model attribute of the Reconstructor points to the original model it was provided as an initialization parameter
• Ensure that reconstructions performed by the old and new methods produce exactly identical results given the same random seed
• Ensure that the final loss of the reconstructions remain below a specified threshold

Additionally, the first test test_Adam_gold_balls checks if data loading methods specific to single-GPU operations are working as intended.

[yoshikisd]

I've addressed some of the reviews requested in PR #17.

Comments #17 (comment) and #17 (comment) were addressed in 82ce845.

[allevitan]

Note that I am aware this breaks the checkpointing system, but we will let it break as it is still undocumented and I'd like to update it in any case. see #48 for more info.

[allevitan]

This is awesome, I really like the changes and I want to pull this in ASAP. There were a few specific issues which I wanted to change, the biggest being:

• Handling circular imports in a better way
• Returning to the old behavior when the model._optimize functions are called. I think it will be simpler and more maintainable going forward to do it this way.
• Changing a few details in the instantiation of the classes that I think simplify the structure a bit
• updating the docs and examples
• Making the tests work again

I submitted these changes as a PR to your forked branch. If you have a moment to review those and pull them in to your branch, then I will approve the PR here and we are good to go.

[yoshikisd]

Thanks for making the changes Abe! I'm happy that we've gotten rid of the circular import, and I like the changes you've made. I've made a couple comments regarding hinting/documentation in the code; if you can make those changes I'll merge those changes into this branch/PR.

[allevitan]

I've rechecked the tests, everything is passing. Ready to merge.