Update the docs to include the reconstructors class and discuss their use in the examples

This commit is contained in:
2025-10-16 14:43:04 +02:00
parent 3c8de5dc19
commit 8328b852d3
5 changed files with 29 additions and 19 deletions
+8 -2
View File
@@ -31,9 +31,9 @@ When reading this script, note the basic workflow. After the data is loaded, a m
Next, the model is moved to the GPU using the :code:`model.to` function. Any device understood by :code:`torch.Tensor.to` can be specified here. The next line is a bit more subtle - the dataset is told to move patterns to the GPU before passing them to the model using the :code:`dataset.get_as` function. This function does not move the stored patterns to the GPU. If there is sufficient GPU memory, the patterns can also be pre-moved to the GPU using :code:`dataset.to`, but the speedup is empirically quite small.
Once the device is selected, a reconstruction is run using :code:`model.Adam_optimize`. This is a generator function which will yield at every epoch, to allow some monitoring code to be run.
Once the device is selected, a reconstruction is run using :code:`model.Adam_optimize`. This is a generator function which will yield at the end of every epoch, to allow some monitoring code to be run.
Finally, the results can be studied using :code:`model.inspect(dataet)`, which creates or updates a set of plots showing the current state of the model parameters. :code:`model.compare(dataset)` is also called, which shows how the simulated diffraction patterns compare to the measured diffraction patterns in the dataset.
Finally, the results can be studied using :code:`model.inspect(dataset)`, which creates or updates a set of plots showing the current state of the model parameters. :code:`model.compare(dataset)` is also called, which shows how the simulated diffraction patterns compare to the measured diffraction patterns in the dataset.
Fancy Ptycho
@@ -63,6 +63,12 @@ By default, FancyPtycho will also optimize over the following model parameters,
These corrections can be turned off (on) by calling :code:`model.<parameter>.requires_grad = False #(True)`.
Note as well two other changes that are made in this script, when compared to `simple_ptycho.py`. First, a `Reconstructor` object is explicitly created, in this case an `AdamReconstructor`. This object stores a model, dataset, and pytorch optimizer. It is then used to orchestrate the later reconstruction using a call to `Reconstructor.optimize()`.
We use this pattern, instead of the simpler call to `model.Adam_optimize()`, because having the reconstructor store the optimizer as well as the model and dataset allows the moment estimates to persist between multiple rounds of optimization. This leads to the second change: In this script, we run two optimization loops. The first loop aggressively refines the probe, with a low minibatch size and a high learning rate. The second loop has a smaller learning rate and a larger batch size, which allow for a more precise final estimation of the object.
In this case, we used one reconstructor, but it is possible to create additional reconstructors to zero out all the persistant information in the optimizer, if desired, or even to instantiate multiple reconstructors on the same model with different optimization algorithms (e.g. `model.LBFGS_optimize()`).
Gold Ball Ptycho
----------------
+1
View File
@@ -9,6 +9,7 @@
general
datasets
models
reconstructors
tools/index
indices_tables
+5
View File
@@ -0,0 +1,5 @@
Reconstructors
==============
.. automodule:: cdtools.reconstructors
:members:
+5 -9
View File
@@ -19,15 +19,11 @@ device = 'cuda'
model.to(device=device)
dataset.get_as(device=device)
# Now, to do the reconstruction we will use the more flexible pattern of
# creating an explicit reconstructor. This is what is used behind the
# scenes by the convenience functions model.<algorithm>_optimize. For a long
# reconstruction script with multiple steps, it is better to create the
# reconstructor explicitly.
#
# The reconstructor will store the model and dataset and create an appropriate
# optimizer. This allows the optimizer to persist, along with e.g. estimates
# of the moments of individual parameters between loops
# For this script, we use a slightly different pattern where we explicitly
# create a `Reconstructor` class to orchestrate the reconstruction. The
# reconstructor will store the model and dataset and create an appropriate
# optimizer. This allows the optimizer to persist between loops, along with
# e.g. estimates of the moments of individual parameters
recon = cdtools.reconstructors.AdamReconstructor(model, dataset)
# The learning rate parameter sets the alpha for Adam.
+10 -8
View File
@@ -45,14 +45,16 @@ class Reconstructor:
subset : list(int) or int
Optional, a pattern index or list of pattern indices to use
Important attributes:
- **model** -- Always points to the core model used.
- **optimizer** -- A `torch.optim.Optimizer` that must be defined when
initializing the Reconstructor subclass.
- **scheduler** -- A `torch.optim.lr_scheduler` that may be defined during
the `optimize` method.
- **data_loader** -- A torch.utils.data.DataLoader that is defined by
calling the `setup_dataloader` method.
Attributes
----------
model : CDIModel
Points to the core model used.
optimizer : torch.optim.Optimizer
Must be defined when initializing the Reconstructor subclass.
scheduler : torch.optim.lr_scheduler, optional
May be defined during the ``optimize`` method.
data_loader : torch.utils.data.DataLoader
Defined by calling the ``setup_dataloader`` method.
"""
def __init__(self,
model: CDIModel,