First working version

examples/fancy_ptycho_adjust_lr.py

@@ -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()

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()

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 = \