From aa684f27eb447ae7efd56bed0652943c682cca6c Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Mon, 13 Oct 2025 19:20:42 +0200 Subject: [PATCH] Change the pattern for Reconstructor so that the optimizer is defined at object creation, and move the dataloader creation logic to the base optimize() function as it was reused in all subclasses --- src/cdtools/reconstructors/adam.py | 45 ++++++++++------ src/cdtools/reconstructors/base.py | 83 +++++++++++++++++++++++------ src/cdtools/reconstructors/lbfgs.py | 33 ++++++------ src/cdtools/reconstructors/sgd.py | 39 +++++++------- 4 files changed, 133 insertions(+), 67 deletions(-) diff --git a/src/cdtools/reconstructors/adam.py b/src/cdtools/reconstructors/adam.py index 4c54f1f..ac11ee5 100644 --- a/src/cdtools/reconstructors/adam.py +++ b/src/cdtools/reconstructors/adam.py @@ -51,10 +51,12 @@ class AdamReconstructor(Reconstructor): dataset: Ptycho2DDataset, subset: List[int] = None): - super().__init__(model, dataset, subset) - # Define the optimizer for use in this subclass - self.optimizer = t.optim.Adam(self.model.parameters()) + optimizer = t.optim.Adam(model.parameters()) + + super().__init__(model, dataset, optimizer, subset=subset) + + def adjust_optimizer(self, lr: int = 0.005, @@ -79,11 +81,13 @@ class AdamReconstructor(Reconstructor): param_group['betas'] = betas param_group['amsgrad'] = amsgrad + def optimize(self, iterations: int, batch_size: int = 15, lr: float = 0.005, betas: Tuple[float] = (0.9, 0.999), + custom_data_loader = None, schedule: bool = False, amsgrad: bool = False, regularization_factor: Union[float, List[float]] = None, @@ -99,6 +103,12 @@ class AdamReconstructor(Reconstructor): (formerly `CDIModel.AD_optimize`) to run a round of reconstruction once the dataloader and optimizer hyperparameters have been set up. + + The `batch_size` parameter sets the batch size for the default + dataloader. If a custom data loader is desired, it can be passed + in to the `custom_data_loader` argument, which will override the + `batch_size` and `shuffle` parameters + Parameters ---------- @@ -115,6 +125,9 @@ class AdamReconstructor(Reconstructor): schedule : bool Optional, create a learning rate scheduler (torch.optim.lr_scheduler._LRScheduler). + custom_data_loader : t.utils.data.DataLoader + Optional, a custom DataLoader to use. If set, will override + batch_size. amsgrad : bool Optional, whether to use the AMSGrad variant of this algorithm. regularization_factor : float or list(float) @@ -138,18 +151,13 @@ class AdamReconstructor(Reconstructor): f'{regularization_factor}, and schedule = {schedule}.\n' ) - # 1) The subset statement is contained in Reconstructor.__init__ - - # 2) Set up / re-initialize the data laoder - self.setup_dataloader(batch_size=batch_size, shuffle=shuffle) - - # 3) The optimizer is created in self.__init__, but the - # hyperparameters need to be set up with self.adjust_optimizer + # The optimizer is created in self.__init__, but the + # hyperparameters need to be set up with self.adjust_optimizer self.adjust_optimizer(lr=lr, betas=betas, amsgrad=amsgrad) - # 4) Set up the scheduler + # Set up the scheduler if schedule: self.scheduler = \ t.optim.lr_scheduler.ReduceLROnPlateau(self.optimizer, @@ -158,8 +166,13 @@ class AdamReconstructor(Reconstructor): else: self.scheduler = None - # 5) This is analagous to making a call to CDIModel.AD_optimize - return super(AdamReconstructor, self).optimize(iterations, - regularization_factor, - thread, - calculation_width) + # Now, we run the optimize routine defined in the base class + return super(AdamReconstructor, self).optimize( + iterations, + batch_size=batch_size, + custom_data_loader=custom_data_loader, + regularization_factor=regularization_factor, + thread=thread, + calculation_width=calculation_width, + shuffle=shuffle, + ) diff --git a/src/cdtools/reconstructors/base.py b/src/cdtools/reconstructors/base.py index 1937661..bae992b 100644 --- a/src/cdtools/reconstructors/base.py +++ b/src/cdtools/reconstructors/base.py @@ -40,6 +40,8 @@ class Reconstructor: Model for CDI/ptychography reconstruction dataset: CDataset The dataset to reconstruct against + optimizer: torch.optim.Optimizer + The optimizer to use for the reconstruction subset : list(int) or int Optional, a pattern index or list of pattern indices to use @@ -55,31 +57,32 @@ class Reconstructor: def __init__(self, model: CDIModel, dataset: CDataset, + optimizer: t.optim.Optimizer, subset: Union[int, List[int]] = None): + # Store parameters as attributes of Reconstructor - self.subset = subset - - # Initialize attributes that must be defined by the subclasses - self.optimizer = None - self.scheduler = None - self.data_loader = None - - # Store the original model self.model = model + self.optimizer = optimizer - # Store the dataset + # Store the dataset, clipping it to a subset if needed if subset is not None: # if subset is just one pattern, turn into a list for convenience if isinstance(subset, int): subset = [subset] dataset = td.Subset(dataset, subset) + self.dataset = dataset + # Initialize attributes that must be defined by the subclasses + self.scheduler = None + self.data_loader = None + + def setup_dataloader(self, batch_size: int = None, shuffle: bool = True): """ - Sets up / re-initializes the dataloader. + Sets up or re-initializes the dataloader. Parameters ---------- @@ -96,6 +99,7 @@ class Reconstructor: else: self.data_loader = td.Dataloader(self.dataset) + def adjust_optimizer(self, **kwargs): """ Change hyperparameters for the utilized optimizer. @@ -105,7 +109,8 @@ class Reconstructor: """ raise NotImplementedError() - def _run_epoch(self, + + def run_epoch(self, stop_event: threading.Event = None, regularization_factor: Union[float, List[float]] = None, calculation_width: int = 10): @@ -133,6 +138,18 @@ class Reconstructor: diffraction pattern intensity """ + # Setting this as an explicit catch makes me feel more comfortable + # exposing it as a public method. This way a user won't be confused + # if they try to use this directly + if self.data_loader is None: + raise RuntimeError( + 'No data loader was defined. Please run ' + 'Reconstructor.setup_dataloader() before running ' + 'Reconstructor.run_epoch(), or use Reconstructor.optimize(), ' + 'which does it automatically.' + ) + + # Initialize some tracking variables normalization = 0 loss = 0 @@ -212,9 +229,12 @@ class Reconstructor: def optimize(self, iterations: int, + batch_size: int = 1, + custom_data_loader = None, regularization_factor: Union[float, List[float]] = None, thread: bool = True, - calculation_width: int = 10): + calculation_width: int = 10, + shuffle=True): """ Runs a round of reconstruction using the provided optimizer @@ -233,10 +253,25 @@ class Reconstructor: the plots. This behavior can be turned off by setting the keyword argument 'thread' to False. + The `batch_size` parameter sets the batch size for the default + dataloader. If a custom data loader is desired, it can be passed + in to the `custom_data_loader` argument, which will override the + `batch_size` and `shuffle` parameters + + Please see `AdamReconstructor.optimize()` for an example of how to + override this function when designing a subclass + Parameters ---------- iterations : int How many epochs of the algorithm to run. + batch_size : int + Optional, the batch size to use. Default is 1. This is typically + overridden by subclasses with an appropriate default for the + specific optimizer. + custom_data_loader : torch.utils.data.DataLoader + Optional, a custom DataLoader to use. Will override batch_size + if set. regularization_factor : float or list(float) Optional, if the model has a regularizer defined, the set of parameters to pass the regularizer method. @@ -247,6 +282,10 @@ class Reconstructor: Default 10, how many translations to pass through at once for each round of gradient accumulation. This does not affect the result, but may affect the calculation speed. + shuffle : bool + Optional, enable/disable shuffling of the dataset. This option + is intended for diagnostic purposes and should be left as True. + Yields ------ @@ -255,6 +294,11 @@ class Reconstructor: diffraction pattern intensity. """ + if custom_data_loader is None: + self.setup_dataloader(batch_size=batch_size, shuffle=shuffle) + else: + self.data_loader = custom_data_loader + # We store the current optimizer as a model parameter so that # it can be saved and loaded for checkpointing self.current_optimizer = self.optimizer @@ -270,8 +314,10 @@ class Reconstructor: yield float('nan') continue - yield self._run_epoch(regularization_factor=regularization_factor, # noqa - calculation_width=calculation_width) + yield self.run_epoch( + regularization_factor=regularization_factor, # noqa + calculation_width=calculation_width, + ) # But if we do want to thread, it's annoying: else: @@ -282,9 +328,12 @@ class Reconstructor: def target(): try: result_queue.put( - self._run_epoch(stop_event=stop_event, - regularization_factor=regularization_factor, # noqa - calculation_width=calculation_width)) + self.run_epoch( + stop_event=stop_event, + regularization_factor=regularization_factor, # noqa + calculation_width=calculation_width, + ) + ) except Exception as e: # If something bad happens, put the exception into the # result queue diff --git a/src/cdtools/reconstructors/lbfgs.py b/src/cdtools/reconstructors/lbfgs.py index 60c6f97..8907782 100644 --- a/src/cdtools/reconstructors/lbfgs.py +++ b/src/cdtools/reconstructors/lbfgs.py @@ -53,10 +53,16 @@ class LBFGSReconstructor(Reconstructor): dataset: Ptycho2DDataset, subset: List[int] = None): - super().__init__(model, dataset, subset) - # Define the optimizer for use in this subclass - self.optimizer = t.optim.LBFGS(self.model.parameters()) + optimizer = t.optim.LBFGS(model.parameters()) + + super().__init__( + model, + dataset, + optimizer, + subset=subset, + ) + def adjust_optimizer(self, lr: int = 0.005, @@ -121,20 +127,17 @@ class LBFGSReconstructor(Reconstructor): round of gradient accumulation. Does not affect the result, only the calculation speed. """ - # 1) The subset statement is contained in Reconstructor.__init__ - # 2) Set up / re-initialize the data loader. For LBFGS, we load - # all the data at once. - self.setup_dataloader(batch_size=len(self.dataset)) - - # 3) The optimizer is created in self.__init__, but the - # hyperparameters need to be set up with self.adjust_optimizer + # The optimizer is created in self.__init__, but the + # hyperparameters need to be set up with self.adjust_optimizer self.adjust_optimizer(lr=lr, history_size=history_size, line_search_fn=line_search_fn) - # 4) This is analagous to making a call to CDIModel.AD_optimize - return super(LBFGSReconstructor, self).optimize(iterations, - regularization_factor, - thread, - calculation_width) + # Now, we run the optimize routine defined in the base class + return super(LBFGSReconstructor, self).optimize( + iterations, + batch_size=len(self.dataset), + regularization_factor=regularization_factor, + thread=thread, + calculation_width=calculation_width) diff --git a/src/cdtools/reconstructors/sgd.py b/src/cdtools/reconstructors/sgd.py index 5b7bf4e..cb1e26b 100644 --- a/src/cdtools/reconstructors/sgd.py +++ b/src/cdtools/reconstructors/sgd.py @@ -49,11 +49,17 @@ class SGDReconstructor(Reconstructor): dataset: Ptycho2DDataset, subset: List[int] = None): - super().__init__(model, dataset, subset) - # Define the optimizer for use in this subclass - self.optimizer = t.optim.SGD(self.model.parameters()) + optimizer = t.optim.SGD(model.parameters()) + super().__init__( + model, + dataset, + optimizer, + subset=subset, + ) + + def adjust_optimizer(self, lr: int = 0.005, momentum: float = 0, @@ -88,7 +94,7 @@ class SGDReconstructor(Reconstructor): def optimize(self, iterations: int, - batch_size: int = None, + batch_size: int = 15, lr: float = 2e-7, momentum: float = 0, dampening: float = 0, @@ -139,25 +145,20 @@ class SGDReconstructor(Reconstructor): Optional, enable/disable shuffling of the dataset. This option is intended for diagnostic purposes and should be left as True. """ - # 1) The subset statement is contained in Reconstructor.__init__ - # 2) Set up / re-initialize the data laoder - if batch_size is not None: - self.setup_dataloader(batch_size=batch_size, shuffle=shuffle) - else: - # Use default torch dataloader parameters - self.setup_dataloader(batch_size=1, shuffle=False) - - # 3) The optimizer is created in self.__init__, but the - # hyperparameters need to be set up with self.adjust_optimizer + # The optimizer is created in self.__init__, but the + # hyperparameters need to be set up with self.adjust_optimizer self.adjust_optimizer(lr=lr, momentum=momentum, dampening=dampening, weight_decay=weight_decay, nesterov=nesterov) - # 4) This is analagous to making a call to CDIModel.AD_optimize - return super(SGDReconstructor, self).optimize(iterations, - regularization_factor, - thread, - calculation_width) + # Now, we run the optimize routine defined in the base class + return super(SGDReconstructor, self).optimize( + iterations, + batch_size=batch_size, + regularization_factor=regularization_factor, + thread=thread, + calculation_width=calculation_width, + )