diff --git a/src/cdtools/models/base.py b/src/cdtools/models/base.py index b455f45..41cdcd0 100644 --- a/src/cdtools/models/base.py +++ b/src/cdtools/models/base.py @@ -40,6 +40,9 @@ import time from scipy import io from contextlib import contextmanager from cdtools.tools.data import nested_dict_to_h5, h5_to_nested_dict, nested_dict_to_numpy, nested_dict_to_torch +from cdtools.datasets import CDataset +from typing import List, Union, Tuple +import os __all__ = ['CDIModel'] @@ -316,202 +319,24 @@ class CDIModel(t.nn.Module): self.current_checkpoint_id += 1 - - - def AD_optimize(self, iterations, data_loader, optimizer,\ - scheduler=None, regularization_factor=None, thread=True, - calculation_width=10): - """Runs a round of reconstruction using the provided optimizer - - This is the basic automatic differentiation reconstruction tool - which all the other, algorithm-specific tools, use. It is a - generator which yields the average loss each epoch, ending after - the specified number of iterations. - - By default, the computation will be run in a separate thread. This - is done to enable live plotting with matplotlib during a reconstruction. - If the computation was done in the main thread, this would freeze - the plots. This behavior can be turned off by setting the keyword - argument 'thread' to False. - - Parameters - ---------- - iterations : int - How many epochs of the algorithm to run - data_loader : torch.utils.data.DataLoader - A data loader loading the CDataset to reconstruct - optimizer : torch.optim.Optimizer - The optimizer to run the reconstruction with - scheduler : torch.optim.lr_scheduler._LRScheduler - Optional, a learning rate scheduler 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 - thread : bool - 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. This does not affect the result, but may affect the calculation speed. - - Yields - ------ - loss : float - The summed loss over the latest epoch, divided by the total diffraction pattern intensity - """ - - def run_epoch(stop_event=None): - """Runs one full epoch of the reconstruction.""" - # First, initialize some tracking variables - normalization = 0 - loss = 0 - N = 0 - t0 = time.time() - - # The data loader is responsible for setting the minibatch - # size, so each set is a minibatch - for inputs, patterns in data_loader: - normalization += t.sum(patterns).cpu().numpy() - N += 1 - def closure(): - optimizer.zero_grad() - - # We further break up the minibatch into a set of chunks. - # This lets us use larger minibatches than can fit - # on the GPU at once, while still doing batch processing - # for efficiency - input_chunks = [[inp[i:i + calculation_width] - for inp in inputs] - for i in range(0, len(inputs[0]), - calculation_width)] - pattern_chunks = [patterns[i:i + calculation_width] - for i in range(0, len(inputs[0]), - calculation_width)] - - total_loss = 0 - for inp, pats in zip(input_chunks, pattern_chunks): - # This check allows for graceful exit when threading - if stop_event is not None and stop_event.is_set(): - exit() - - # Run the simulation - sim_patterns = self.forward(*inp) - - # Calculate the loss - if hasattr(self, 'mask'): - loss = self.loss(pats,sim_patterns, mask=self.mask) - else: - loss = self.loss(pats,sim_patterns) - - # And accumulate the gradients - loss.backward() - total_loss += loss.detach() - - # If we have a regularizer, we can calculate it separately, - # and the gradients will add to the minibatch gradient - if regularization_factor is not None \ - and hasattr(self, 'regularizer'): - loss = self.regularizer(regularization_factor) - loss.backward() - - return total_loss - - # This takes the step for this minibatch - loss += optimizer.step(closure).detach().cpu().numpy() - - - loss /= normalization - - # We step the scheduler after the full epoch - if scheduler is not None: - scheduler.step(loss) - - self.loss_history.append(loss) - self.epoch = len(self.loss_history) - self.latest_iteration_time = time.time() - t0 - self.training_history += self.report() + '\n' - return loss - - # We store the current optimizer as a model parameter so that - # it can be saved and loaded for checkpointing - self.current_optimizer = optimizer - - # If we don't want to run in a different thread, this is easy - if not thread: - for it in range(iterations): - if self.skip_computation(): - self.epoch = self.epoch + 1 - if len(self.loss_history) >= 1: - yield self.loss_history[-1] - else: - yield float('nan') - continue - - yield run_epoch() - - - # But if we do want to thread, it's annoying: - else: - # Here we set up the communication with the computation thread - result_queue = queue.Queue() - stop_event = threading.Event() - def target(): - try: - result_queue.put(run_epoch(stop_event)) - except Exception as e: - # If something bad happens, put the exception into the - # result queue - result_queue.put(e) - - # And this actually starts and monitors the thread - for it in range(iterations): - if self.skip_computation(): - self.epoch = self.epoch + 1 - if len(self.loss_history) >= 1: - yield self.loss_history[-1] - else: - yield float('nan') - continue - - calc = threading.Thread(target=target, name='calculator', daemon=True) - try: - calc.start() - while calc.is_alive(): - if hasattr(self, 'figs'): - self.figs[0].canvas.start_event_loop(0.01) - else: - calc.join() - - except KeyboardInterrupt as e: - stop_event.set() - print('\nAsking execution thread to stop cleanly - please be patient.') - calc.join() - raise e - - res = result_queue.get() - - # If something went wrong in the thead, we'll get an exception - if isinstance(res, Exception): - raise res - - yield res - - # And finally, we unset the current optimizer: - self.current_optimizer = None - def Adam_optimize( self, - iterations, - dataset, - batch_size=15, - lr=0.005, - betas=(0.9, 0.999), - schedule=False, - amsgrad=False, - subset=None, - regularization_factor=None, + iterations: int, + dataset: CDataset, + batch_size: int = 15, + lr: float = 0.005, + betas: Tuple[float] = (0.9, 0.999), + schedule: bool = False, + amsgrad: bool = False, + subset: Union[int, List[int]] = None, + regularization_factor: Union[float, List[float]] = None, thread=True, calculation_width=10 ): - """Runs a round of reconstruction using the Adam optimizer + """ + Runs a round of reconstruction using the Adam optimizer from + cdtools.reconstructors.Adam. This is generally accepted to be the most robust algorithm for use with ptychography. Like all the other optimization routines, @@ -521,125 +346,143 @@ class CDIModel(t.nn.Module): Parameters ---------- iterations : int - How many epochs of the algorithm to run + How many epochs of the algorithm to run. dataset : CDataset - The dataset to reconstruct against + The dataset to reconstruct against. batch_size : int - Optional, the size of the minibatches to use + Optional, the size of the minibatches to use. lr : float - Optional, The learning rate (alpha) to use. Defaultis 0.005. 0.05 is typically the highest possible value with any chance of being stable - betas : tuple + Optional, The learning rate (alpha) to use. Defaultis 0.005. + 0.05 is typically the highest possible value with any chance + of being stable. + betas : tuple(float) Optional, the beta_1 and beta_2 to use. Default is (0.9, 0.999). - schedule : float - Optional, whether to use the ReduceLROnPlateau scheduler + schedule : bool + Optional, whether to use the ReduceLROnPlateau scheduler. + amsgrad : bool + Optional, whether to use the AMSGrad variant of this algorithm. 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 + regularization_factor : float or list(float). + 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 - - """ - - self.training_history += ( - f'Planning {iterations} epochs of Adam, with a learning rate = ' - f'{lr}, batch size = {batch_size}, regularization_factor = ' - f'{regularization_factor}, and schedule = {schedule}.\n' - ) + 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. - - if subset is not None: - # if subset is just one pattern, turn into a list for convenience - if type(subset) == type(1): - subset = [subset] - dataset = torchdata.Subset(dataset, subset) - - # Make a dataloader - data_loader = torchdata.DataLoader(dataset, - batch_size=batch_size, - shuffle=True) - - # Define the optimizer - optimizer = t.optim.Adam( - self.parameters(), - lr = lr, - betas=betas, - amsgrad=amsgrad) - - # Define the scheduler - if schedule: - scheduler = t.optim.lr_scheduler.ReduceLROnPlateau(optimizer, factor=0.2,threshold=1e-9) - else: - scheduler = None - - return self.AD_optimize(iterations, data_loader, optimizer, - scheduler=scheduler, - regularization_factor=regularization_factor, - thread=thread, - calculation_width=calculation_width) + """ + # 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: + if not hasattr(self, 'reconstructor'): + from cdtools.reconstructors import Adam + self.reconstructor = Adam(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) - def LBFGS_optimize(self, iterations, dataset, - lr=0.1,history_size=2, subset=None, - regularization_factor=None, thread=True, - calculation_width=10, line_search_fn=None): - """Runs a round of reconstruction using the L-BFGS optimizer + def LBFGS_optimize(self, + iterations: int, + dataset: CDataset, + lr: float = 0.1, + history_size: int = 2, + subset: Union[int, List[int]] = None, + regularization_factor: Union[float, List[float]] =None, + thread: bool = True, + calculation_width: int = 10, + line_search_fn: str = None): + """ + Runs a round of reconstruction using the L-BFGS optimizer from + cdtools.reconstructors.LBFGS. This algorithm is often less stable that Adam, however in certain situations or geometries it can be shockingly efficient. Like all the other optimization routines, it is defined as a generator function which yields the average loss each epoch. - Note: There is no batch size, because it is a usually a bad idea to use + NOTE: There is no batch size, because it is a usually a bad idea to use LBFGS on anything but all the data at onece Parameters ---------- iterations : int - How many epochs of the algorithm to run + How many epochs of the algorithm to run. dataset : CDataset - The dataset to reconstruct against + The dataset to reconstruct against. lr : float - Optional, the learning rate to use + Optional, the learning rate to use. history_size : int Optional, the length of the history to use. subset : list(int) or int - Optional, a pattern index or list of pattern indices to ues + 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 """ - if subset is not None: - # if just one pattern, turn into a list for convenience - if type(subset) == type(1): - subset = [subset] - dataset = torchdata.Subset(dataset, subset) - - # Make a dataloader. This basically does nothing but load all the - # data at once - data_loader = torchdata.DataLoader(dataset, batch_size=len(dataset)) + # 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: + if not hasattr(self, 'reconstructor'): + from cdtools.reconstructors import LBFGS + self.reconstructor = LBFGS(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) - # Define the optimizer - optimizer = t.optim.LBFGS(self.parameters(), - lr = lr, history_size=history_size, - line_search_fn=line_search_fn) - - return self.AD_optimize(iterations, data_loader, optimizer, - regularization_factor=regularization_factor, - thread=thread, - calculation_width=calculation_width) - - - def SGD_optimize(self, iterations, dataset, batch_size=None, - lr=0.01, momentum=0, dampening=0, weight_decay=0, - nesterov=False, subset=None, regularization_factor=None, - thread=True, calculation_width=10): - """Runs a round of reconstruction using the SGD optimizer + def SGD_optimize(self, + iterations: int, + dataset: CDataset, + batch_size: int = None, + lr: float = 2e-7, + momentum: float = 0, + dampening: float = 0, + weight_decay: float = 0, + nesterov: bool = False, + subset: Union[int, List[int]] = None, + regularization_factor: Union[float, List[float]] = None, + thread: bool = True, + calculation_width: int = 10): + """ + Runs a round of reconstruction using the SGD optimizer from + cdtools.reconstructors.SGD. This algorithm is often less stable that Adam, but it is simpler and is the basic workhorse of gradience descent. @@ -647,51 +490,58 @@ class CDIModel(t.nn.Module): Parameters ---------- iterations : int - How many epochs of the algorithm to run + How many epochs of the algorithm to run. dataset : CDataset - The dataset to reconstruct against + The dataset to reconstruct against. batch_size : int - Optional, the size of the minibatches to use + Optional, the size of the minibatches to use. lr : float - Optional, the learning rate to use + Optional, the learning rate to use. momentum : float Optional, the length of the history to use. + dampening : float + Optional, dampening for the momentum. + weight_decay : float + Optional, weight decay (L2 penalty). + nesterov : bool + Optional, enables Nesterov momentum. Only applicable when momentum + is non-zero. subset : list(int) or int - Optional, a pattern index or list of pattern indices to use + 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 1, how many translations to pass through at once for each round of gradient accumulation + Default 10, how many translations to pass through at once for each + round of gradient accumulation. """ - - if subset is not None: - # if just one pattern, turn into a list for convenience - if type(subset) == type(1): - subset = [subset] - dataset = torchdata.Subset(dataset, subset) - - # Make a dataloader - if batch_size is not None: - data_loader = torchdata.DataLoader(dataset, batch_size=batch_size, - shuffle=True) - else: - data_loader = torchdata.DataLoader(dataset) - - - # Define the optimizer - optimizer = t.optim.SGD(self.parameters(), - lr = lr, momentum=momentum, - dampening=dampening, - weight_decay=weight_decay, - nesterov=nesterov) - - return self.AD_optimize(iterations, data_loader, optimizer, - regularization_factor=regularization_factor, - thread=thread, - calculation_width=calculation_width) + # 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: + if not hasattr(self, 'reconstructor'): + from cdtools.reconstructors import SGD + self.reconstructor = SGD(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) def report(self):