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, + )