diff --git a/docs/source/examples.rst b/docs/source/examples.rst index 6f356cd..879f86f 100644 --- a/docs/source/examples.rst +++ b/docs/source/examples.rst @@ -110,6 +110,18 @@ If :code:`propagation_distance` is set, it will assume a Fresnel scaling theorem Finally, note the addition of the :code:`panel_plot_mode=True` argument. This is the default mode, and returns the plots in a panel format, good for easily monitoring the progress of a reconstruction. If individual plots are needed for use in presentations, papers, or otherwise, setting :code:`panel_plot_mode=False` will plot each output in it's own window. +Per-Parameter Learning Rates +---------------------------- + +This script shows how the learning rates can be adjusted per parameter, which can sometimes accelerate convergence substantially. + +.. literalinclude:: ../../examples/per_parameter_learning_rates.py + +The major addition is the inclusion of a dictionary, :code:`lr_factors`, which multiplies the main learning rate for each individual parameter. If a specific parameter is not being updated aggressively enough, increase this value from the default of 1. If it is being updated too aggressively and preventing convergence, lower the value. + +This dictionary will persist through all further :code:`recon.optimize()` calls, unles explicitly reset. + + Gold Ball Split --------------- diff --git a/examples/per_parameter_learning_rates.py b/examples/per_parameter_learning_rates.py new file mode 100644 index 0000000..39cf20f --- /dev/null +++ b/examples/per_parameter_learning_rates.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): + 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/models/base.py b/src/cdtools/models/base.py index ebaaf85..943f5ed 100644 --- a/src/cdtools/models/base.py +++ b/src/cdtools/models/base.py @@ -405,6 +405,7 @@ def Adam_optimize( dataset: CDataset, batch_size: int = 15, lr: float = 0.005, + lr_factors: dict = {}, betas: Tuple[float] = (0.9, 0.999), schedule: bool = False, amsgrad: bool = False, @@ -434,6 +435,8 @@ def Adam_optimize( Optional, The learning rate (alpha) to use. Defaultis 0.005. 0.05 is typically the highest possible value with any chance of being stable. + lr_factors : dict + Optional, a dictionary mapping optimizer parameters to adjustment factors for the learning rate. betas : tuple(float) Optional, the beta_1 and beta_2 to use. Default is (0.9, 0.999). schedule : bool @@ -458,6 +461,7 @@ def Adam_optimize( model=self, dataset=dataset, subset=subset, + lr_factors=lr_factors, ) # Run some reconstructions @@ -540,6 +544,7 @@ def SGD_optimize(self, dataset: CDataset, batch_size: int = None, lr: float = 2e-7, + lr_factors : dict = {}, momentum: float = 0, dampening: float = 0, weight_decay: float = 0, @@ -565,6 +570,8 @@ def SGD_optimize(self, Optional, the size of the minibatches to use. lr : float Optional, the learning rate to use. + lr_factors : dict + Optional, a dictionary mapping optimizer parameters to adjustment factors for the learning rate. momentum : float Optional, the length of the history to use. dampening : float @@ -591,6 +598,7 @@ def SGD_optimize(self, model=self, dataset=dataset, subset=subset, + lr_factors=lr_factors, ) # Run some reconstructions diff --git a/src/cdtools/reconstructors/adam.py b/src/cdtools/reconstructors/adam.py index e298bdb..8eaf016 100644 --- a/src/cdtools/reconstructors/adam.py +++ b/src/cdtools/reconstructors/adam.py @@ -8,6 +8,7 @@ """ from __future__ import annotations from typing import TYPE_CHECKING +import warnings import torch as t from typing import Tuple, List, Union @@ -33,67 +34,133 @@ class AdamReconstructor(Reconstructor): The dataset to reconstruct against. subset : list(int) or int Optional, a pattern index or list of pattern indices to use. - schedule : bool - Optional, create a learning rate scheduler - (torch.optim.lr_scheduler._LRScheduler). + lr_factors : dict + Optional, a dictionary mapping optimizer parameters to adjustment factors for the learning rate. Important attributes: - **model** -- Always points to the core model used. - **optimizer** -- This class by default uses `torch.optim.Adam` to perform optimizations. + - **lr_factors** -- A map from optimizer parameters to learning rate + factors. - **scheduler** -- A `torch.optim.lr_scheduler` that is defined during the `optimize` method. - **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 - optimizer = t.optim.Adam(model.parameters()) + param_groups = [] + for name, param in model.named_parameters(): + param_groups.append({'params':[param], 'name':name}) + + optimizer = t.optim.Adam(param_groups) 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 + 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. - def adjust_optimizer(self, - lr: int = 0.005, - betas: Tuple[float] = (0.9, 0.999), - amsgrad: bool = False): + 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, + ): """ 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. """ + + # 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['lr'] = lr 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: 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): + param_name = param_group['name'] + if isinstance(self.lr_factors, dict) and \ + param_name in self.lr_factors: + param_group['lr'] = lr * self.lr_factors[param_name] + else: + param_group['lr'] = lr + + + 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, + ): """ Runs a round of reconstruction using the Adam optimizer @@ -122,6 +189,8 @@ def optimize(self, stable. betas : tuple Optional, the beta_1 and beta_2 to use. Default is (0.9, 0.999). + lr_factors : dict + Optional, a dictionary mapping optimizer parameters to adjustment factors for the learning rate. schedule : bool Optional, create a learning rate scheduler (torch.optim.lr_scheduler._LRScheduler). @@ -144,19 +213,26 @@ def optimize(self, Optional, enable/disable shuffling of the dataset. This option is intended for diagnostic purposes and should be left as True. """ + + # 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, + ) + # 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) - + self.model.training_history += ( + f'The learning rate factors are {self.lr_factors}, default = 1.\n' + ) + # Set up the scheduler if schedule: self.scheduler = \ diff --git a/src/cdtools/reconstructors/sgd.py b/src/cdtools/reconstructors/sgd.py index cb1e26b..c5c495a 100644 --- a/src/cdtools/reconstructors/sgd.py +++ b/src/cdtools/reconstructors/sgd.py @@ -34,23 +34,34 @@ class SGDReconstructor(Reconstructor): The dataset to reconstruct against. subset : list(int) or int Optional, a pattern index or list of pattern indices to use. + lr_factors : dict + Optional, a dictionary mapping optimizer parameters to adjustment factors for the learning rate. Important attributes: - **model** -- Always points to the core model used. - **optimizer** -- This class by default uses `torch.optim.Adam` to perform optimizations. + - **lr_factors** -- A map from optimizer parameters to learning rate + factors. - **scheduler** -- A `torch.optim.lr_scheduler` that is defined during the `optimize` method. - **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 - optimizer = t.optim.SGD(model.parameters()) + param_groups = [] + for name, param in model.named_parameters(): + param_groups.append({'params':[param], 'name':name}) + + optimizer = t.optim.SGD(param_groups) super().__init__( model, @@ -59,13 +70,54 @@ def __init__(self, subset=subset, ) + self._set_lr_factors(lr_factors) + + + def _set_lr_factors(self, lr_factors): + """Sets the learning rate factors from a provided dictionary + + 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, - momentum: float = 0, - dampening: float = 0, - weight_decay: float = 0, - nesterov: bool = False): + def adjust_optimizer( + self, + lr: int = 0.005, + momentum: float = 0, + dampening: float = 0, + weight_decay: float = 0, + nesterov: bool = False, + lr_factors: dict = None, + ): """ Change hyperparameters for the utilized optimizer. @@ -84,26 +136,48 @@ def adjust_optimizer(self, nesterov : bool Optional, enables Nesterov momentum. Only applicable when momentum is non-zero. + lr_factors : dict + Optional, a dictionary mapping optimizer parameters to adjustment factors for the learning rate. + """ + + # 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['lr'] = lr param_group['momentum'] = momentum param_group['dampening'] = dampening param_group['weight_decay'] = weight_decay param_group['nesterov'] = nesterov - def optimize(self, - iterations: int, - batch_size: int = 15, - lr: float = 2e-7, - momentum: float = 0, - dampening: float = 0, - weight_decay: float = 0, - nesterov: bool = False, - regularization_factor: Union[float, List[float]] = None, - thread: bool = True, - calculation_width: int = 10, - shuffle: bool = True): + param_name = param_group['name'] + if isinstance(self.lr_factors, dict) and \ + param_name in self.lr_factors: + param_group['lr'] = lr * self.lr_factors[param_name] + else: + param_group['lr'] = lr + + + def optimize( + self, + iterations: int, + batch_size: int = 15, + lr: float = 2e-7, + momentum: float = 0, + dampening: float = 0, + weight_decay: float = 0, + nesterov: bool = False, + lr_factors : dict = None, + regularization_factor: Union[float, List[float]] = None, + thread: bool = True, + calculation_width: int = 10, + shuffle: bool = True, + ): """ Runs a round of reconstruction using the Adam optimizer @@ -131,6 +205,8 @@ def optimize(self, nesterov : bool Optional, enables Nesterov momentum. Only applicable when momentum is non-zero. + lr_factors : dict + Optional, a dictionary mapping optimizer parameters to adjustment factors for the learning rate. regularization_factor : float or list(float) Optional, if the model has a regularizer defined, the set of parameters to pass the regularizer method. @@ -148,12 +224,27 @@ def optimize(self, # 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) + self.adjust_optimizer( + lr=lr, + momentum=momentum, + dampening=dampening, + weight_decay=weight_decay, + nesterov=nesterov, + lr_factors=lr_factors, + ) + # Update the training history + self.model.training_history += ( + f'Planning {iterations} epochs of SGD, with a learning rate = ' + f'{lr}, batch size = {batch_size}, regularization_factor = ' + f'{regularization_factor}, momentum history length = {momentum},' + f'momemntum dampening = {dampening}, weight_decay = {weight_decay},' + f' and nesterov = {nesterov}.\n' + ) + self.model.training_history += ( + f'The learning rate factors are {self.lr_factors}, default = 1.\n' + ) + # Now, we run the optimize routine defined in the base class return super(SGDReconstructor, self).optimize( iterations, diff --git a/tests/test_reconstructors.py b/tests/test_reconstructors.py index ff34a7a..1ac6731 100644 --- a/tests/test_reconstructors.py +++ b/tests/test_reconstructors.py @@ -20,9 +20,10 @@ def test_Adam_gold_balls(gold_ball_cxi, reconstruction_device, show_plot): 4) Reconstructions performed by `Adam.optimize` and `model.Adam_optimize` calls produce identical results when run over one round of optimization. - 5) The quality of the reconstruction remains below a specified + 5) Checks that the per-parameter learning rates work in both cases + 6) The quality of the reconstruction remains below a specified threshold. - 5) Ensure that the FancyPtycho model works fine and dandy with the + 7) Ensure that the FancyPtycho model works fine and dandy with the Reconstructors. """ @@ -53,12 +54,20 @@ def test_Adam_gold_balls(gold_ball_cxi, reconstruction_device, show_plot): model_recon.to(device=reconstruction_device) dataset.get_as(device=reconstruction_device) + lr_factors = { + 'obj' : 1.1, + 'weights' : 0.5, + } + # ******* Reconstructions with AdamReconstructor.optimize ******* print('Running reconstruction using AdamReconstructor.optimize' + ' on provided reconstruction_device,', reconstruction_device) - recon = cdtools.reconstructors.AdamReconstructor(model=model_recon, - dataset=dataset) + recon = cdtools.reconstructors.AdamReconstructor( + model=model_recon, + dataset=dataset, + lr_factors=lr_factors, + ) t.manual_seed(0) # Run a reconstruction @@ -100,10 +109,13 @@ def test_Adam_gold_balls(gold_ball_cxi, reconstruction_device, show_plot): # We only need to test the first loop to ensure it's identical for i, iterations in enumerate(epoch_tup[:1]): - for loss in model.Adam_optimize(iterations, - dataset, - lr=lr_tup[i], - batch_size=batch_size_tup[i]): + for loss in model.Adam_optimize( + iterations, + dataset, + lr=lr_tup[i], + lr_factors=lr_factors, + batch_size=batch_size_tup[i], + ): print(model.report()) if show_plot: model.inspect(dataset, min_interval=10) @@ -161,8 +173,16 @@ def test_intensity_MSE(gold_ball_cxi, reconstruction_device, show_plot): for loss in recon.optimize(5, lr=.05, batch_size=10): print(model.report()) - # Threshold to be updated after running on a GPU machine - assert model.loss_history[-1] < 1e7 + # Test that Adam optimizer post-creation update of lr_factors works + lr_factors = { + 'background' : 0.3, + 'translation_offsets': 1.2, + } + + for loss in recon.optimize(3, lr=.05, batch_size=10, lr_factors=lr_factors): + print(model.report()) + + assert model.loss_history[-1] < 6.5e6 @pytest.mark.slow @@ -371,3 +391,19 @@ def test_SGD_gold_balls(gold_ball_cxi, reconstruction_device, show_plot): # a threshold of 7.2e-4 for the tested loss. If this value has been # exceeded, the reconstructions have gotten worse. assert model.loss_history[-1] < 0.95 + + print('Testing per-parameter learning rates') + + lr_factors = { + 'background': 0.4, + } + + for loss in model.SGD_optimize(epochs, + dataset, + lr=lr, + lr_factors=lr_factors, + batch_size=batch_size): + print(model.report()) + if show_plot: + model.inspect(dataset, min_interval=10) +