From 82ce845385170afbad9845e73be1d487e17b7373 Mon Sep 17 00:00:00 2001 From: yoshikisd Date: Mon, 4 Aug 2025 16:51:24 +0000 Subject: [PATCH] Changed the names of the reconstructors and updated the old optimizer class documentation --- src/cdtools/models/base.py | 165 +++++++++++++------------ src/cdtools/reconstructors/__init__.py | 20 +-- src/cdtools/reconstructors/adam.py | 14 +-- src/cdtools/reconstructors/base.py | 2 +- src/cdtools/reconstructors/lbfgs.py | 16 +-- src/cdtools/reconstructors/sgd.py | 18 +-- tests/test_reconstructors.py | 29 +++-- 7 files changed, 138 insertions(+), 126 deletions(-) diff --git a/src/cdtools/models/base.py b/src/cdtools/models/base.py index 41cdcd0..635cdfb 100644 --- a/src/cdtools/models/base.py +++ b/src/cdtools/models/base.py @@ -319,7 +319,6 @@ class CDIModel(t.nn.Module): self.current_checkpoint_id += 1 - def Adam_optimize( self, iterations: int, @@ -336,7 +335,7 @@ class CDIModel(t.nn.Module): ): """ Runs a round of reconstruction using the Adam optimizer from - cdtools.reconstructors.Adam. + cdtools.reconstructors.AdamReconstructor. This is generally accepted to be the most robust algorithm for use with ptychography. Like all the other optimization routines, @@ -375,45 +374,45 @@ class CDIModel(t.nn.Module): only the calculation speed. """ - # We want to have model.Adam_optimize call AND store cdtools.reconstructors.Adam - # to perform reconstructions without creating a new reconstructor each time we - # update the hyperparameters. - # - # The only way to do this is to make the Adam reconstructor an attribute - # of the model. But since the Adam reconstructor also depends on CDIModel, - # this seems to give rise to a circular import error unless - # we import cdtools.reconstructors within this method: + # We want to have model.Adam_optimize call AND store + # cdtools.reconstructors.AdamReconstructor to perform reconstructions + # without creating a new reconstructor each time we update the + # hyperparameters. + # + # The only way to do this is to make the Adam reconstructor an + # attribute of the model. But since the Adam reconstructor also + # depends on CDIModel, this seems to give rise to a circular import + # error unless we import cdtools.reconstructors within this method: if not hasattr(self, 'reconstructor'): - from cdtools.reconstructors import Adam - self.reconstructor = Adam(model=self, - dataset=dataset, - subset=subset) + from cdtools.reconstructors import AdamReconstructor + self.reconstructor = AdamReconstructor(model=self, + dataset=dataset, + subset=subset) # Run some reconstructions return self.reconstructor.optimize(iterations=iterations, - batch_size=batch_size, - lr=lr, - betas=betas, - schedule=schedule, - amsgrad=amsgrad, - regularization_factor=regularization_factor, - thread=thread, - calculation_width=calculation_width) + batch_size=batch_size, + lr=lr, + betas=betas, + schedule=schedule, + amsgrad=amsgrad, + regularization_factor=regularization_factor, # noqa + thread=thread, + calculation_width=calculation_width) - - def LBFGS_optimize(self, - iterations: int, + def LBFGS_optimize(self, + iterations: int, dataset: CDataset, lr: float = 0.1, - history_size: int = 2, + history_size: int = 2, subset: Union[int, List[int]] = None, - regularization_factor: Union[float, List[float]] =None, + regularization_factor: Union[float, List[float]] =None, thread: bool = True, - calculation_width: int = 10, + calculation_width: int = 10, line_search_fn: str = None): """ Runs a round of reconstruction using the L-BFGS optimizer from - cdtools.reconstructors.LBFGS. + cdtools.reconstructors.LBFGSReconstructor. This algorithm is often less stable that Adam, however in certain situations or geometries it can be shockingly efficient. Like all @@ -436,53 +435,55 @@ class CDIModel(t.nn.Module): subset : list(int) or int Optional, a pattern index or list of pattern indices to use. regularization_factor : float or list(float) - Optional, if the model has a regularizer defined, the set of parameters - to pass the regularizer method. + Optional, if the model has a regularizer defined, the set of + parameters to pass the regularizer method. thread : bool - Default True, whether to run the computation in a separate thread to allow - interaction with plots during computation. + Default True, whether to run the computation in a separate thread + to allow interaction with plots during computation. calculation_width : int - Default 10, how many translations to pass through at once for each round of - gradient accumulation. Does not affect the result, only the calculation speed + Default 10, how many translations to pass through at once for each + round of gradient accumulation. Does not affect the result, only + the calculation speed. """ - # We want to have model.LBFGS_optimize store cdtools.reconstructors.LBFGS - # as an attribute to run reconstructions without generating new reconstructors - # each time CDIModel.LBFGS_optimize is called. - # - # Since the LBFGS reconstructor also depends on CDIModel, a circular import error - # arises unless we import cdtools.reconstructors within this method: + # We want to have model.LBFGS_optimize store + # cdtools.reconstructors.LBFGSReconstructor as an attribute to run + # reconstructions without generating new reconstructors each time + # CDIModel.LBFGS_optimize is called. + # + # Since the LBFGS reconstructor also depends on CDIModel, a circular + # import error arises unless we import cdtools.reconstructors within + # this method: if not hasattr(self, 'reconstructor'): - from cdtools.reconstructors import LBFGS - self.reconstructor = LBFGS(model=self, - dataset=dataset, - subset=subset) + from cdtools.reconstructors import LBFGSReconstructor + self.reconstructor = LBFGSReconstructor(model=self, + dataset=dataset, + subset=subset) # Run some reconstructions return self.reconstructor.optimize(iterations=iterations, - lr=lr, - history_size=history_size, - regularization_factor=regularization_factor, - thread=thread, - calculation_width=calculation_width, - line_search_fn = line_search_fn) - + lr=lr, + history_size=history_size, + regularization_factor=regularization_factor, # noqa + thread=thread, + calculation_width=calculation_width, + line_search_fn=line_search_fn) def SGD_optimize(self, - iterations: int, - dataset: CDataset, + iterations: int, + dataset: CDataset, batch_size: int = None, - lr: float = 2e-7, - momentum: float = 0, - dampening: float = 0, + lr: float = 2e-7, + momentum: float = 0, + dampening: float = 0, weight_decay: float = 0, - nesterov: bool = False, - subset: Union[int, List[int]] = None, + nesterov: bool = False, + subset: Union[int, List[int]] = None, regularization_factor: Union[float, List[float]] = None, - thread: bool = True, + thread: bool = True, calculation_width: int = 10): """ Runs a round of reconstruction using the SGD optimizer from - cdtools.reconstructors.SGD. + cdtools.reconstructors.SGDReconstructor. This algorithm is often less stable that Adam, but it is simpler and is the basic workhorse of gradience descent. @@ -519,29 +520,31 @@ class CDIModel(t.nn.Module): round of gradient accumulation. """ - # We want to have model.SGD_optimize store cdtools.reconstructors.SGD - # as an attribute to run reconstructions without generating new reconstructors - # each time CDIModel.SGD_optimize is called. - # - # Since the SGD reconstructor also depends on CDIModel, a circular import error - # arises unless we import cdtools.reconstructors within this method: + # We want to have model.SGD_optimize store + # cdtools.reconstructors.SGDReconstructor as an attribute to run + # reconstructions without generating new reconstructors each time + # CDIModel.SGD_optimize is called. + # + # Since the SGD reconstructor also depends on CDIModel, a circular + # import error arises unless we import cdtools.reconstructors within + # this method: if not hasattr(self, 'reconstructor'): - from cdtools.reconstructors import SGD - self.reconstructor = SGD(model=self, - dataset=dataset, - subset=subset) - + from cdtools.reconstructors import SGDReconstructor + self.reconstructor = SGDReconstructor(model=self, + dataset=dataset, + subset=subset) + # Run some reconstructions return self.reconstructor.optimize(iterations=iterations, - batch_size=batch_size, - lr=lr, - momentum=momentum, - dampening=dampening, - weight_decay=weight_decay, - nesterov=nesterov, - regularization_factor=regularization_factor, - thread=thread, - calculation_width=calculation_width) + batch_size=batch_size, + lr=lr, + momentum=momentum, + dampening=dampening, + weight_decay=weight_decay, + nesterov=nesterov, + regularization_factor=regularization_factor, # noqa + thread=thread, + calculation_width=calculation_width) def report(self): diff --git a/src/cdtools/reconstructors/__init__.py b/src/cdtools/reconstructors/__init__.py index 84b96ab..84a1a81 100644 --- a/src/cdtools/reconstructors/__init__.py +++ b/src/cdtools/reconstructors/__init__.py @@ -1,16 +1,22 @@ -"""This module contains optimizers for performing reconstructions +""" +Module `cdtools.tools.reconstructors` contains the `Reconstructor` class and +subclasses which run the ptychography reconstructions on a given model and +dataset. +The reconstructors are designed to resemble so-called 'Trainer' classes that +(in the language of the AI/ML folks) handles the 'training' of a model given +some dataset and optimizer. """ # We define __all__ to be sure that import * only imports what we want __all__ = [ 'Reconstructor', - 'Adam', - 'LBFGS', - 'SGD' + 'AdamReconstructor', + 'LBFGSReconstructor', + 'SGDReconstructor' ] from cdtools.reconstructors.base import Reconstructor -from cdtools.reconstructors.adam import Adam -from cdtools.reconstructors.lbfgs import LBFGS -from cdtools.reconstructors.sgd import SGD +from cdtools.reconstructors.adam import AdamReconstructor +from cdtools.reconstructors.lbfgs import LBFGSReconstructor +from cdtools.reconstructors.sgd import SGDReconstructor diff --git a/src/cdtools/reconstructors/adam.py b/src/cdtools/reconstructors/adam.py index 5a489a0..6eecc9c 100644 --- a/src/cdtools/reconstructors/adam.py +++ b/src/cdtools/reconstructors/adam.py @@ -1,4 +1,4 @@ -"""This module contains the Adam Reconstructor subclass for performing +"""This module contains the AdamReconstructor subclass for performing optimization ('reconstructions') on ptychographic/CDI models using the Adam optimizer. @@ -12,10 +12,10 @@ from cdtools.models import CDIModel from typing import Tuple, List, Union from cdtools.reconstructors import Reconstructor -__all__ = ['Adam'] +__all__ = ['AdamReconstructor'] -class Adam(Reconstructor): +class AdamReconstructor(Reconstructor): """ The Adam Reconstructor subclass handles the optimization ('reconstruction') of ptychographic models and datasets using the Adam optimizer. @@ -154,7 +154,7 @@ class Adam(Reconstructor): self.scheduler = None # 5) This is analagous to making a call to CDIModel.AD_optimize - return super(Adam, self).optimize(iterations, - regularization_factor, - thread, - calculation_width) + return super(AdamReconstructor, self).optimize(iterations, + regularization_factor, + thread, + calculation_width) diff --git a/src/cdtools/reconstructors/base.py b/src/cdtools/reconstructors/base.py index d5d34a6..fd2cbb2 100644 --- a/src/cdtools/reconstructors/base.py +++ b/src/cdtools/reconstructors/base.py @@ -178,7 +178,7 @@ class Reconstructor: loss.backward() # Normalize the accumulating total loss - total_loss += loss.detach() + total_loss += loss.detach() # If we have a regularizer, we can calculate it separately, # and the gradients will add to the minibatch gradient diff --git a/src/cdtools/reconstructors/lbfgs.py b/src/cdtools/reconstructors/lbfgs.py index 0b51dfd..8bbd7ae 100644 --- a/src/cdtools/reconstructors/lbfgs.py +++ b/src/cdtools/reconstructors/lbfgs.py @@ -1,4 +1,4 @@ -"""This module contains the LBFGS Reconstructor subclass for performing +"""This module contains the LBFGSReconstructor subclass for performing optimization ('reconstructions') on ptychographic/CDI models using the LBFGS optimizer. @@ -12,12 +12,12 @@ from cdtools.models import CDIModel from typing import List, Union from cdtools.reconstructors import Reconstructor -__all__ = ['LBFGS'] +__all__ = ['LBFGSReconstructor'] -class LBFGS(Reconstructor): +class LBFGSReconstructor(Reconstructor): """ - The LBFGS Reconstructor subclass handles the optimization + The LBFGSReconstructor subclass handles the optimization ('reconstruction') of ptychographic models and datasets using the LBFGS optimizer. @@ -128,7 +128,7 @@ class LBFGS(Reconstructor): line_search_fn=line_search_fn) # 4) This is analagous to making a call to CDIModel.AD_optimize - return super(LBFGS, self).optimize(iterations, - regularization_factor, - thread, - calculation_width) + return super(LBFGSReconstructor, self).optimize(iterations, + regularization_factor, + thread, + calculation_width) diff --git a/src/cdtools/reconstructors/sgd.py b/src/cdtools/reconstructors/sgd.py index f2dd7b0..a8f2ebc 100644 --- a/src/cdtools/reconstructors/sgd.py +++ b/src/cdtools/reconstructors/sgd.py @@ -1,4 +1,4 @@ -"""This module contains the SGD Reconstructor subclass for performing +"""This module contains the SGDReconstructor subclass for performing optimization ('reconstructions') on ptychographic/CDI models using stochastic gradient descent. @@ -12,13 +12,13 @@ from cdtools.models import CDIModel from typing import List, Union from cdtools.reconstructors import Reconstructor -__all__ = ['SGD'] +__all__ = ['SGDReconstructor'] -class SGD(Reconstructor): +class SGDReconstructor(Reconstructor): """ - The Adam Reconstructor subclass handles the optimization ('reconstruction') - of ptychographic models and datasets using the Adam optimizer. + The SGDReconstructor subclass handles the optimization ('reconstruction') + of ptychographic models and datasets using the SGD optimizer. Parameters ---------- @@ -151,7 +151,7 @@ class SGD(Reconstructor): nesterov=nesterov) # 4) This is analagous to making a call to CDIModel.AD_optimize - return super(SGD, self).optimize(iterations, - regularization_factor, - thread, - calculation_width) + return super(SGDReconstructor, self).optimize(iterations, + regularization_factor, + thread, + calculation_width) diff --git a/tests/test_reconstructors.py b/tests/test_reconstructors.py index 0f46239..3995b6f 100644 --- a/tests/test_reconstructors.py +++ b/tests/test_reconstructors.py @@ -24,7 +24,7 @@ def test_Adam_gold_balls(gold_ball_cxi, reconstruction_device, show_plot): """ print('\nTesting performance on the standard gold balls dataset ' + - 'with reconstructors.Adam') + 'with reconstructors.AdamReconstructor') dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(gold_ball_cxi) pad = 10 @@ -48,11 +48,12 @@ def test_Adam_gold_balls(gold_ball_cxi, reconstruction_device, show_plot): model_recon.to(device=reconstruction_device) dataset.get_as(device=reconstruction_device) - # ******* Reconstructions with cdtools.reconstructors.Adam.optimize ******* - print('Running reconstruction using cdtools.reconstructors.Adam.optimize' + + # ******* Reconstructions with AdamReconstructor.optimize ******* + print('Running reconstruction using AdamReconstructor.optimize' + ' on provided reconstruction_device,', reconstruction_device) - recon = cdtools.reconstructors.Adam(model=model_recon, dataset=dataset) + recon = cdtools.reconstructors.AdamReconstructor(model=model_recon, + dataset=dataset) t.manual_seed(0) # Run a reconstruction @@ -85,7 +86,7 @@ def test_Adam_gold_balls(gold_ball_cxi, reconstruction_device, show_plot): model_recon.inspect(dataset) model_recon.compare(dataset) - # ******* Reconstructions with cdtools.CDIModel.Adam_optimize ******* + # ******* Reconstructions with CDIModel.Adam_optimize ******* print('Running reconstruction using CDIModel.Adam_optimize on provided' + ' reconstruction_device,', reconstruction_device) t.manual_seed(0) @@ -150,11 +151,12 @@ def test_LBFGS_RPI(optical_data_ss_cxi, model_recon.to(device=reconstruction_device) dataset.get_as(device=reconstruction_device) - # ******* Reconstructions with cdtools.reconstructors.LBFGS.optimize ****** - print('Running reconstruction using cdtools.reconstructors.LBFGS.' + + # ******* Reconstructions with LBFGSReconstructor.optimize ****** + print('Running reconstruction using LBFGSReconstructor.' + 'optimize on provided reconstruction_device,', reconstruction_device) - recon = cdtools.reconstructors.LBFGS(model=model_recon, dataset=dataset) + recon = cdtools.reconstructors.LBFGSReconstructor(model=model_recon, + dataset=dataset) t.manual_seed(0) # Run a reconstruction @@ -178,7 +180,7 @@ def test_LBFGS_RPI(optical_data_ss_cxi, # Check model pointing assert id(model_recon) == id(recon.model) - # ******* Reconstructions with cdtools.reconstructors.LBFGS.optimize ****** + # ******* Reconstructions with CDIModel.LBFGS_optimize ****** print('Running reconstruction using CDIModel.LBFGS_optimize.' + 'optimize on provided reconstruction_device,', reconstruction_device) t.manual_seed(0) @@ -223,7 +225,7 @@ def test_SGD_gold_balls(gold_ball_cxi, reconstruction_device, show_plot): to do some kind of comparative assessment. """ print('\nTesting performance on the standard gold balls dataset ' + - 'with reconstructors.SGD') + 'with reconstructors.SGDReconstructor') dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(gold_ball_cxi) pad = 10 @@ -247,11 +249,12 @@ def test_SGD_gold_balls(gold_ball_cxi, reconstruction_device, show_plot): model_recon.to(device=reconstruction_device) dataset.get_as(device=reconstruction_device) - # ******* Reconstructions with cdtools.reconstructors.SGD.optimize ******* - print('Running reconstruction using cdtools.reconstructors.SGD.optimize' + + # ******* Reconstructions with SGDReconstructor.optimize ******* + print('Running reconstruction using SGDReconstructor.optimize' + ' on provided reconstruction_device,', reconstruction_device) - recon = cdtools.reconstructors.SGD(model=model_recon, dataset=dataset) + recon = cdtools.reconstructors.SGDReconstructor(model=model_recon, + dataset=dataset) t.manual_seed(0) # Run a reconstruction