From 8cf217f5d8a93e4cec7bb0eb72713c41baf34905 Mon Sep 17 00:00:00 2001 From: Pengju Sheng Date: Thu, 23 Jul 2026 11:50:01 +0200 Subject: [PATCH 1/3] Add feature adjust learning rate --- examples/fancy_ptycho_adjust_lr.py | 54 ++++++++++++++++++++++++++++++ 1 file changed, 54 insertions(+) create mode 100644 examples/fancy_ptycho_adjust_lr.py diff --git a/examples/fancy_ptycho_adjust_lr.py b/examples/fancy_ptycho_adjust_lr.py new file mode 100644 index 0000000..c23c077 --- /dev/null +++ b/examples/fancy_ptycho_adjust_lr.py @@ -0,0 +1,54 @@ +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() From d0122b1161d92ab6419c74f1d69604d7b09fa953 Mon Sep 17 00:00:00 2001 From: Pengju Sheng Date: Thu, 23 Jul 2026 11:54:14 +0200 Subject: [PATCH 2/3] adjust lr --- src/cdtools/reconstructors/adam.py | 29 ++++++++++++++++++++++------- 1 file changed, 22 insertions(+), 7 deletions(-) diff --git a/src/cdtools/reconstructors/adam.py b/src/cdtools/reconstructors/adam.py index e298bdb..b335166 100644 --- a/src/cdtools/reconstructors/adam.py +++ b/src/cdtools/reconstructors/adam.py @@ -52,16 +52,21 @@ def __init__(self, subset: List[int] = None): # 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) def adjust_optimizer(self, - lr: int = 0.005, + lr: int | dict = 0.005, betas: Tuple[float] = (0.9, 0.999), - amsgrad: bool = False): + amsgrad: bool = False, + default_lr : int = 0.005): """ Change hyperparameters for the utilized optimizer. @@ -77,16 +82,25 @@ def adjust_optimizer(self, Optional, whether to use the AMSGrad variant of this algorithm. """ for param_group in self.optimizer.param_groups: - param_group['lr'] = lr 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 + else: + param_group['lr'] = lr[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: float = 0.005, + 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, @@ -155,7 +169,8 @@ def optimize(self, # hyperparameters need to be set up with self.adjust_optimizer self.adjust_optimizer(lr=lr, betas=betas, - amsgrad=amsgrad) + amsgrad=amsgrad, + default_lr = default_lr) # Set up the scheduler if schedule: From 0aecd41ba7c7d1fc55c8c8aeafa3ddde835fd11e Mon Sep 17 00:00:00 2001 From: Pengju Sheng Date: Thu, 23 Jul 2026 14:29:51 +0200 Subject: [PATCH 3/3] modified docstrings --- src/cdtools/reconstructors/adam.py | 27 ++++++++++++++++++--------- 1 file changed, 18 insertions(+), 9 deletions(-) diff --git a/src/cdtools/reconstructors/adam.py b/src/cdtools/reconstructors/adam.py index b335166..515805c 100644 --- a/src/cdtools/reconstructors/adam.py +++ b/src/cdtools/reconstructors/adam.py @@ -72,14 +72,19 @@ def adjust_optimizer(self, Parameters ---------- - lr : float - 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. + lr : float or dictionary + Optional, The learning rate (alpha) to use. Default is 0.005 for all the + parameters. 0.05 is typically the highest possible value with any chance of being + stable. And you can also assign different learning rates to different parameters + using a dictionary. Example: {'obj':0.005, 'probe':0.01}. 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. + default_lr : float + Optional, The default learning rate for the parameters that have not + been assigned a specific learning rate when you try to assign different + learning rates to different parameters. """ for param_group in self.optimizer.param_groups: param_group['betas'] = betas @@ -93,7 +98,6 @@ def adjust_optimizer(self, else: param_group['lr'] = lr - print(f"Found paramter {param_name}, learning rate : {param_group['lr']}") def optimize(self, iterations: int, @@ -130,12 +134,17 @@ def optimize(self, How many epochs of the algorithm to run. batch_size : int Optional, the size of the minibatches to use. - lr : float - 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. + lr : float or dictionary + Optional, The learning rate (alpha) to use. Default is 0.005 for all the + parameters. 0.05 is typically the highest possible value with any chance of being + stable. And you can also assign different learning rates to different parameters + using a dictionary. Example: {'obj':0.005, 'probe':0.01}. betas : tuple Optional, the beta_1 and beta_2 to use. Default is (0.9, 0.999). + default_lr : float + Optional, The default learning rate for the parameters that have not + been assigned a specific learning rate when you try to assign different + learning rates to different parameters. schedule : bool Optional, create a learning rate scheduler (torch.optim.lr_scheduler._LRScheduler).