diff --git a/examples/fancy_ptycho_adjust_lr.py b/examples/fancy_ptycho_adjust_lr.py deleted file mode 100644 index c23c077..0000000 --- a/examples/fancy_ptycho_adjust_lr.py +++ /dev/null @@ -1,54 +0,0 @@ -import cdtools -import torch as t -from matplotlib import pyplot as plt - -filename = 'example_data/lab_ptycho_data.cxi' -dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(filename) - -# FancyPtycho is the workhorse model -model = cdtools.models.FancyPtycho.from_dataset( - dataset, - n_modes=3, # Use 3 incoherently mixing probe modes - oversampling=2, # Simulate the probe on a 2xlarger real-space array - probe_support_radius=120, # Force the probe to 0 outside a radius of 120 pix - propagation_distance=5e-3, # Propagate the initial probe guess by 5 mm - units='mm', # Set the units for the live plots - obj_view_crop=-50, # Expands the field of view in the object plot by 50 pix -) - -if t.cuda.is_available(): - model.to(device='cuda') - dataset.get_as(device='cuda') - -# 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. -# The beta parameters are (0.9, 0.999) by default -# The batch size sets the minibatch size -lr = {'translation_offsets':0.04,'background':0.01} -for loss in recon.optimize(50, lr=lr, batch_size=10, default_lr = 0.03): - print(model.report()) - # Because plotting can be expensive, setting a minimum plotting interval - # (in seconds) can avoid excessive replots. - model.inspect(min_interval=10) - -# It's common to chain several different reconstruction loops. Here, we -# started with an aggressive refinement to find the probe in the previous -# loop, and now we polish the reconstruction with a lower learning rate -# and larger minibatch -for loss in recon.optimize(50, lr=0.005, batch_size=50): - print(model.report()) - model.inspect(min_interval=10) - -# This orthogonalizes the recovered probe modes -model.tidy_probes() - -# Setting replot_all will reopen any windows which were closed earlier -model.inspect(replot_all=True) -model.compare(dataset) -plt.show() diff --git a/examples/fancy_ptycho_per_parameter_lrs.py b/examples/fancy_ptycho_per_parameter_lrs.py new file mode 100644 index 0000000..4c9a0a1 --- /dev/null +++ b/examples/fancy_ptycho_per_parameter_lrs.py @@ -0,0 +1,50 @@ +import cdtools +import torch as t +from matplotlib import pyplot as plt + +filename = 'example_data/lab_ptycho_data.cxi' +dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(filename) + +model = cdtools.models.FancyPtycho.from_dataset( + dataset, + n_modes=3, # Use 3 incoherently mixing probe modes + oversampling=2, # Simulate the probe on a 2xlarger real-space array + probe_support_radius=120, # Force the probe to 0 outside a radius of 120 pix + propagation_distance=5e-3, # Propagate the initial probe guess by 5 mm + units='mm', # Set the units for the live plots + obj_view_crop=-50, # Expands the field of view in the object plot by 50 pix +) + +if t.cuda.is_available(): + model.to(device='cuda') + dataset.get_as(device='cuda') + +# Here, we tune the learning rates of individual parameters. The default +# learning rate factor is 1. Any learning rate factor set here will multiply +# the learning rate for each recon.optimize loop. The dictionary can be passed +# to the reconstructor object at creation time, as done here. It can also be +# updated later with the call to recon.optimize(..., lr_factors=lr_factors). +lr_factors = { + 'translation_offsets' : 1.2, + 'weights' : 0.2, + 'background' : 0.3, +} + +recon = cdtools.reconstructors.AdamReconstructor( + model, dataset, lr_factors=lr_factors) + +# For example, background will get a lr of 0.03 * 0.3 (lr * lr_factor). +for loss in recon.optimize(50, lr=0.03, batch_size=10, lr_factors=lr_factors, verbose=True): + print(model.report()) + model.inspect(min_interval=10) + +# And here background will get a lr of 0.005 * 0.3 (lr * lr_factor). +for loss in recon.optimize(50, lr=0.005, batch_size=50): + print(model.report()) + model.inspect(min_interval=10) + +model.tidy_probes() + +model.inspect(replot_all=True) +model.compare(dataset) +plt.show() diff --git a/src/cdtools/reconstructors/adam.py b/src/cdtools/reconstructors/adam.py index b335166..477e9ef 100644 --- a/src/cdtools/reconstructors/adam.py +++ b/src/cdtools/reconstructors/adam.py @@ -8,6 +8,7 @@ the 'training' of a model given some dataset and optimizer. """ from __future__ import annotations from typing import TYPE_CHECKING +import warnings import torch as t from typing import Tuple, List, Union @@ -46,10 +47,13 @@ class AdamReconstructor(Reconstructor): - **data_loader** -- A torch.utils.data.DataLoader that is defined by calling the `setup_dataloader` method. """ - def __init__(self, - model: CDIModel, - dataset: Ptycho2DDataset, - subset: List[int] = None): + def __init__( + self, + model: CDIModel, + dataset: Ptycho2DDataset, + subset: List[int] = None, + lr_factors: dict = {} + ): # Define the optimizer for use in this subclass param_groups = [] @@ -60,54 +64,109 @@ class AdamReconstructor(Reconstructor): super().__init__(model, dataset, optimizer, subset=subset) + self._set_lr_factors(lr_factors) + + def _set_lr_factors(self, lr_factors): + """Sets the learning rate factors from a provided dictionary - def adjust_optimizer(self, - lr: int | dict = 0.005, - betas: Tuple[float] = (0.9, 0.999), - amsgrad: bool = False, - default_lr : int = 0.005): + This is broken out into it's own function to avoid replicating the + code to emit a warning, and to enable easy future changes as the logic + may need to become more complicated. + + Parameters + ---------- + lr_factors : dict + A dictionary mapping optimizer parameters to adjustment factors for the learning rate. + """ + self.lr_factors = lr_factors + param_group_names = {p['name'] for p in self.optimizer.param_groups} + unused_lr_factors = self.lr_factors.keys() - param_group_names + + if len(unused_lr_factors) != 0: + warnings.warn( + 'The lr_factor dictionary defines some entries ' + + 'which are unused. Check the following entries for typos:' + + str(unused_lr_factors), + stacklevel=3, + ) + + def print_lrs(self): + """Prints the current per-parameter learning rates. + """ + + for param_group in self.optimizer.param_groups: + print( + f"Paramter {param_group['name']} has learning rate " + f"{param_group['lr']}." + ) + + def adjust_optimizer( + self, + lr: int = 0.005, + betas: Tuple[float] = (0.9, 0.999), + amsgrad: bool = False, + lr_factors: dict = None, + verbose: bool = False, + ): """ Change hyperparameters for the utilized optimizer. Parameters ---------- lr : float - Optional, The learning rate (alpha) to use. Default is 0.005. 0.05 + Optional, the learning rate (alpha) to use. Default is 0.005. 0.05 is typically the highest possible value with any chance of being stable. betas : tuple Optional, the beta_1 and beta_2 to use. Default is (0.9, 0.999). amsgrad : bool Optional, whether to use the AMSGrad variant of this algorithm. + lr_factors : dict + Optional, a dictionary mapping optimizer parameters to adjustment factors for the learning rate. + verbose : bool + Default False, whether to print out setup information after adjustment. """ + + # Update the learning rate factors if explicitly given. Otherwise, + # persist the existing dictionary. A common pattern is to set the + # factors once at the start, and then adjust only the learning rate + # afterward. + if lr_factors is not None: + self._set_lr_factors(lr_factors) + for param_group in self.optimizer.param_groups: param_group['betas'] = betas param_group['amsgrad'] = amsgrad param_name = param_group['name'] - if isinstance(lr, dict): - if param_name not in lr: - param_group['lr'] = default_lr + if isinstance(self.lr_factors, dict): + if param_name not in self.lr_factors: + param_group['lr'] = lr else: - param_group['lr'] = lr[param_name] + param_group['lr'] = lr * self.lr_factors[param_name] else: param_group['lr'] = lr - print(f"Found paramter {param_name}, learning rate : {param_group['lr']}") - - def optimize(self, - iterations: int, - batch_size: int = 15, - lr: int | dict = 0.005, - betas: Tuple[float] = (0.9, 0.999), - default_lr : int = 0.005, - custom_data_loader: t.utils.data.DataLoader = None, - schedule: bool = False, - amsgrad: bool = False, - regularization_factor: Union[float, List[float]] = None, - thread: bool = True, - calculation_width: int = 10, - shuffle: bool = True): + if verbose: + self.print_lrs() + + + def optimize( + self, + iterations: int, + batch_size: int = 15, + lr: int = 0.005, + betas: Tuple[float] = (0.9, 0.999), + lr_factors : dict = None, + custom_data_loader: t.utils.data.DataLoader = None, + schedule: bool = False, + amsgrad: bool = False, + regularization_factor: Union[float, List[float]] = None, + thread: bool = True, + calculation_width: int = 10, + shuffle: bool = True, + verbose: bool = False, + ): """ Runs a round of reconstruction using the Adam optimizer @@ -157,21 +216,32 @@ class AdamReconstructor(Reconstructor): shuffle : bool Optional, enable/disable shuffling of the dataset. This option is intended for diagnostic purposes and should be left as True. + lr_factors : dict + Optional, a dictionary mapping optimizer parameters to adjustment factors for the learning rate. + verbose : bool + Default False, whether to print out setup information about the planned run. """ + + # 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, + lr_factors=lr_factors, + verbose=verbose, + ) + # Update the training history self.model.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' ) - - # 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, - default_lr = default_lr) - + self.model.training_history += ( + f'The learning rate factors are {self.lr_factors}, default = 1.\n' + ) + # Set up the scheduler if schedule: self.scheduler = \