cSAXS/cdtools
0
0
Fork 0
mirror of https://github.com/cdtools-developers/cdtools.git synced 2026-09-15 15:42:37 +02:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #91 from cdtools-developers/alternate_lr_proposal

Browse Source 
Alternate lr proposal
This commit is contained in:
sheng_p1
2026-07-24 14:20:03 +02:00
committed by GitHub
parent 3e63ca2bcd 289182f3f3
commit b3c452f16f
6 changed files with 346 additions and 73 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/source/examples.rst
+12
View File
@@ -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
+50
View File
@@ -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
+8
View File
@@ -405,6 +405,7 @@ class CDIModel(t.nn.Module):
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 @@ class CDIModel(t.nn.Module):
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 @@ class CDIModel(t.nn.Module):
model=self,
dataset=dataset,
subset=subset,
lr_factors=lr_factors,
)
# Run some reconstructions
@@ -540,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,
@@ -565,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
@@ -591,6 +598,7 @@ class CDIModel(t.nn.Module):
model=self,
dataset=dataset,
subset=subset,
lr_factors=lr_factors,
)
# Run some reconstructions
src/cdtools/reconstructors/adam.py
+110 -34
View File
@@ -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
@@ -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
def adjust_optimizer(self,
lr: int = 0.005,
betas: Tuple[float] = (0.9, 0.999),
amsgrad: bool = False):
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,
):
"""
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
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 = 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):
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 @@ 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).
@@ -144,19 +213,26 @@ class AdamReconstructor(Reconstructor):
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 = \
src/cdtools/reconstructors/sgd.py
+120 -29
View File
@@ -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,
tests/test_reconstructors.py
+46 -10
View File
@@ -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)