Fancy learning rate

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

src/cdtools/reconstructors/adam.py

@@ -52,41 +52,59 @@ 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.
 
         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['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
 
 
     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,
@@ -116,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).
@@ -155,7 +178,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: