From 3252a0dd8f795676063b5f75686f36a3d576ccc7 Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Thu, 23 Jul 2026 15:00:33 +0200 Subject: [PATCH] Update SGD to allow per-parameter learning rates and add test coverage --- examples/fancy_ptycho_per_parameter_lrs.py | 2 +- src/cdtools/models/base.py | 4 + src/cdtools/reconstructors/adam.py | 31 ++--- src/cdtools/reconstructors/sgd.py | 149 +++++++++++++++++---- tests/test_reconstructors.py | 56 ++++++-- 5 files changed, 182 insertions(+), 60 deletions(-) diff --git a/examples/fancy_ptycho_per_parameter_lrs.py b/examples/fancy_ptycho_per_parameter_lrs.py index 4c9a0a1..39cf20f 100644 --- a/examples/fancy_ptycho_per_parameter_lrs.py +++ b/examples/fancy_ptycho_per_parameter_lrs.py @@ -34,7 +34,7 @@ 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): +for loss in recon.optimize(50, lr=0.03, batch_size=10): print(model.report()) model.inspect(min_interval=10) diff --git a/src/cdtools/models/base.py b/src/cdtools/models/base.py index 1beb560..943f5ed 100644 --- a/src/cdtools/models/base.py +++ b/src/cdtools/models/base.py @@ -544,6 +544,7 @@ class CDIModel(t.nn.Module): dataset: CDataset, batch_size: int = None, lr: float = 2e-7, + lr_factors : dict = {}, momentum: float = 0, dampening: float = 0, weight_decay: float = 0, @@ -569,6 +570,8 @@ class CDIModel(t.nn.Module): 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 @@ -595,6 +598,7 @@ class CDIModel(t.nn.Module): 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 477e9ef..8eaf016 100644 --- a/src/cdtools/reconstructors/adam.py +++ b/src/cdtools/reconstructors/adam.py @@ -34,14 +34,15 @@ 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 @@ -90,6 +91,7 @@ class AdamReconstructor(Reconstructor): str(unused_lr_factors), stacklevel=3, ) + def print_lrs(self): """Prints the current per-parameter learning rates. @@ -101,13 +103,13 @@ class AdamReconstructor(Reconstructor): 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. @@ -124,8 +126,6 @@ class AdamReconstructor(Reconstructor): 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, @@ -139,17 +139,12 @@ class AdamReconstructor(Reconstructor): param_group['betas'] = betas param_group['amsgrad'] = amsgrad param_name = param_group['name'] - if isinstance(self.lr_factors, dict): - if param_name not in self.lr_factors: - param_group['lr'] = lr - else: - param_group['lr'] = lr * self.lr_factors[param_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 - if verbose: - self.print_lrs() - def optimize( self, @@ -165,7 +160,6 @@ class AdamReconstructor(Reconstructor): thread: bool = True, calculation_width: int = 10, shuffle: bool = True, - verbose: bool = False, ): """ Runs a round of reconstruction using the Adam optimizer @@ -195,6 +189,8 @@ class AdamReconstructor(Reconstructor): 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). @@ -216,10 +212,6 @@ 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 @@ -229,7 +221,6 @@ class AdamReconstructor(Reconstructor): betas=betas, amsgrad=amsgrad, lr_factors=lr_factors, - verbose=verbose, ) # Update the training history 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 @@ class SGDReconstructor(Reconstructor): 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 - def adjust_optimizer(self, - lr: int = 0.005, - momentum: float = 0, - dampening: float = 0, - weight_decay: float = 0, - nesterov: bool = False): + 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, + lr_factors: dict = None, + ): """ Change hyperparameters for the utilized optimizer. @@ -84,26 +136,48 @@ class SGDReconstructor(Reconstructor): 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 @@ class SGDReconstructor(Reconstructor): 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 @@ class SGDReconstructor(Reconstructor): # 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) +