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

Update SGD to allow per-parameter learning rates and add test coverage

Browse Source
This commit is contained in:
levita_a
2026-07-23 15:00:33 +02:00
parent f6e0654445
commit 3252a0dd8f
5 changed files with 182 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/fancy_ptycho_per_parameter_lrs.py
+1 -1
View File
@@ -34,7 +34,7 @@ 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):
for loss in recon.optimize(50, lr=0.03, batch_size=10):
print(model.report())
model.inspect(min_interval=10)
src/cdtools/models/base.py
+4
View File
@@ -544,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,
@@ -569,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
@@ -595,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
+11 -20
View File
@@ -34,14 +34,15 @@ 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
@@ -90,6 +91,7 @@ class AdamReconstructor(Reconstructor):
str(unused_lr_factors),
stacklevel=3,
)
def print_lrs(self):
"""Prints the current per-parameter learning rates.
@@ -101,13 +103,13 @@ class AdamReconstructor(Reconstructor):
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.
@@ -124,8 +126,6 @@ class AdamReconstructor(Reconstructor):
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,
@@ -139,17 +139,12 @@ class AdamReconstructor(Reconstructor):
param_group['betas'] = betas
param_group['amsgrad'] = amsgrad
param_name = param_group['name']
if isinstance(self.lr_factors, dict):
if param_name not in self.lr_factors:
param_group['lr'] = lr
else:
param_group['lr'] = lr * self.lr_factors[param_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
if verbose:
self.print_lrs()
def optimize(
self,
@@ -165,7 +160,6 @@ class AdamReconstructor(Reconstructor):
thread: bool = True,
calculation_width: int = 10,
shuffle: bool = True,
verbose: bool = False,
):
"""
Runs a round of reconstruction using the Adam optimizer
@@ -195,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).
@@ -216,10 +212,6 @@ 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
@@ -229,7 +221,6 @@ class AdamReconstructor(Reconstructor):
betas=betas,
amsgrad=amsgrad,
lr_factors=lr_factors,
verbose=verbose,
)
# Update the training history
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)