Alternate lr proposal

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

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

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

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