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

Created Reconstructors class to replace optimization methods in CDIModel

Browse Source
This commit is contained in:
yoshikisd
2025-08-01 21:37:05 +00:00
parent 25ca3a9317
commit 0785eae552
6 changed files with 794 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/cdtools/__init__.py
+2 -2
View File
@@ -4,9 +4,9 @@ import warnings
warnings.filterwarnings("ignore",
message='To copy construct from a tensor, ')
__all__ = ['tools', 'datasets', 'models']
__all__ = ['tools', 'datasets', 'models', 'reconstructors']
from cdtools import tools
from cdtools import datasets
from cdtools import models
from cdtools import reconstructors
src/cdtools/reconstructors/__init__.py
+16
View File
@@ -0,0 +1,16 @@
"""This module contains optimizers for performing reconstructions
"""
# We define __all__ to be sure that import * only imports what we want
__all__ = [
'Reconstructor',
'Adam',
'LBFGS',
'SGD'
]
from cdtools.reconstructors.base import Reconstructor
from cdtools.reconstructors.adam import Adam
from cdtools.reconstructors.lbfgs import LBFGS
from cdtools.reconstructors.sgd import SGD
src/cdtools/reconstructors/adam.py
+160
View File
@@ -0,0 +1,160 @@
"""This module contains the Adam Reconstructor subclass for performing
optimization ('reconstructions') on ptychographic/CDI models using
the Adam optimizer.
The Reconstructor class is designed to resemble so-called
'Trainer' classes that (in the language of the AI/ML folks) handles
the 'training' of a model given some dataset and optimizer.
"""
import torch as t
from cdtools.datasets.ptycho_2d_dataset import Ptycho2DDataset
from cdtools.models import CDIModel
from typing import Tuple, List, Union
from cdtools.reconstructors import Reconstructor
__all__ = ['Adam']
class Adam(Reconstructor):
"""
The Adam Reconstructor subclass handles the optimization ('reconstruction')
of ptychographic models and datasets using the Adam optimizer.
Parameters
----------
model: CDIModel
Model for CDI/ptychography reconstruction.
dataset: Ptycho2DDataset
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).
Important attributes:
- **model** -- Always points to the core model used.
- **optimizer** -- This class by default uses `torch.optim.Adam` to perform
optimizations.
- **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):
super().__init__(model, dataset, subset)
# Define the optimizer for use in this subclass
self.optimizer = t.optim.Adam(self.model.parameters())
def adjust_optimizer(self,
lr: int = 0.005,
betas: Tuple[float] = (0.9, 0.999),
amsgrad: bool = False):
"""
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.
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.
"""
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),
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
Formerly `CDIModel.Adam_optimize`
This calls the Reconstructor.optimize superclass method
(formerly `CDIModel.AD_optimize`) to run a round of reconstruction
once the dataloader and optimizer hyperparameters have been
set up.
Parameters
----------
iterations : int
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.
betas : tuple
Optional, the beta_1 and beta_2 to use. Default is (0.9, 0.999).
schedule : bool
Optional, create a learning rate scheduler
(torch.optim.lr_scheduler._LRScheduler).
amsgrad : bool
Optional, whether to use the AMSGrad variant of this algorithm.
regularization_factor : float or list(float)
Optional, if the model has a regularizer defined, the set of
parameters to pass the regularizer method.
thread : bool
Default True, whether to run the computation in a separate thread
to allow interaction with plots during computation.
calculation_width : int
Default 10, how many translations to pass through at once for each
round of gradient accumulation. Does not affect the result, only
the calculation speed.
shuffle : bool
Optional, enable/disable shuffling of the dataset. This option
is intended for diagnostic purposes and should be left as True.
"""
# 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'
)
# 1) The subset statement is contained in Reconstructor.__init__
# 2) Set up / re-initialize the data laoder
self.setup_dataloader(batch_size=batch_size, shuffle=shuffle)
# 3) 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)
# 4) Set up the scheduler
if schedule:
self.scheduler = \
t.optim.lr_scheduler.ReduceLROnPlateau(self.optimizer,
factor=0.2,
threshold=1e-9)
else:
self.scheduler = None
# 5) This is analagous to making a call to CDIModel.AD_optimize
return super(Adam, self).optimize(iterations,
regularization_factor,
thread,
calculation_width)
src/cdtools/reconstructors/base.py
+325
View File
@@ -0,0 +1,325 @@
"""This module contains the base Reconstructor class for performing
optimization ('reconstructions') on ptychographic/CDI models.
The Reconstructor class is designed to resemble so-called
'Trainer' classes that (in the language of the AI/ML folks) handles
the 'training' of a model given some dataset and optimizer.
The subclasses of Reconstructor are required to implement
their own data loaders and optimizer adjusters
"""
import torch as t
from torch.utils import data as td
import threading
import queue
import time
from cdtools.datasets import CDataset
from cdtools.models import CDIModel
from typing import List, Union
__all__ = ['Reconstructor']
class Reconstructor:
"""
Reconstructor handles the optimization ('reconstruction') of ptychographic
models given a CDIModel (or subclass) and corresponding CDataset.
This is a base model that defines all functions Reconstructor subclasses
must implement.
Parameters
----------
model: CDIModel
Model for CDI/ptychography reconstruction
dataset: CDataset
The dataset to reconstruct against
subset : list(int) or int
Optional, a pattern index or list of pattern indices to use
Important attributes:
- **model** -- Always points to the core model used.
- **optimizer** -- A `torch.optim.Optimizer` that must be defined when
initializing the Reconstructor subclass.
- **scheduler** -- A `torch.optim.lr_scheduler` that may be 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: CDataset,
subset: Union[int, List[int]] = None):
# Store parameters as attributes of Reconstructor
self.subset = subset
# Initialize attributes that must be defined by the subclasses
self.optimizer = None
self.scheduler = None
self.data_loader = None
# Store the original model
self.model = model
# Store the dataset
if subset is not None:
# if subset is just one pattern, turn into a list for convenience
if isinstance(subset, int):
subset = [subset]
dataset = td.Subset(dataset, subset)
self.dataset = dataset
def setup_dataloader(self,
batch_size: int = None,
shuffle: bool = True):
"""
Sets up / re-initializes the dataloader.
Parameters
----------
batch_size : int
Optional, the size of the minibatches to use
shuffle : bool
Optional, enable/disable shuffling of the dataset. This option
is intended for diagnostic purposes and should be left as True.
"""
if batch_size is not None:
self.data_loader = td.DataLoader(self.dataset,
batch_size=batch_size,
shuffle=shuffle)
else:
self.data_loader = td.Dataloader(self.dataset)
def adjust_optimizer(self, **kwargs):
"""
Change hyperparameters for the utilized optimizer.
For each optimizer, the keyword arguments should be manually defined
as parameters.
"""
raise NotImplementedError()
def _run_epoch(self,
stop_event: threading.Event = None,
regularization_factor: Union[float, List[float]] = None,
calculation_width: int = 10):
"""
Runs one full epoch of the reconstruction. Intended to be called
by Reconstructor.optimize.
Parameters
----------
stop_event : threading.Event
Default None, causes the reconstruction to stop when an exception
occurs in Optimizer.optimize.
regularization_factor : float or list(float)
Optional, if the model has a regularizer defined, the set of
parameters to pass the regularizer method
calculation_width : int
Default 10, how many translations to pass through at once for each
round of gradient accumulation. This does not affect the result,
but may affect the calculation speed.
Returns
------
loss : float
The summed loss over the latest epoch, divided by the total
diffraction pattern intensity
"""
# Initialize some tracking variables
normalization = 0
loss = 0
N = 0
t0 = time.time()
# The data loader is responsible for setting the minibatch
# size, so each set is a minibatch
for inputs, patterns in self.data_loader:
normalization += t.sum(patterns).cpu().numpy()
N += 1
def closure():
self.optimizer.zero_grad()
# We further break up the minibatch into a set of chunks.
# This lets us use larger minibatches than can fit
# on the GPU at once, while still doing batch processing
# for efficiency
input_chunks = [[inp[i:i + calculation_width]
for inp in inputs]
for i in range(0, len(inputs[0]),
calculation_width)]
pattern_chunks = [patterns[i:i + calculation_width]
for i in range(0, len(inputs[0]),
calculation_width)]
total_loss = 0
for inp, pats in zip(input_chunks, pattern_chunks):
# This check allows for graceful exit when threading
if stop_event is not None and stop_event.is_set():
exit()
# Run the simulation
sim_patterns = self.model.forward(*inp)
# Calculate the loss
if hasattr(self, 'mask'):
loss = self.model.loss(pats,
sim_patterns,
mask=self.model.mask)
else:
loss = self.model.loss(pats,
sim_patterns)
# And accumulate the gradients
loss.backward()
# Normalize the accumulating total loss
total_loss += loss.detach() // self.model.world_size
# If we have a regularizer, we can calculate it separately,
# and the gradients will add to the minibatch gradient
if regularization_factor is not None \
and hasattr(self.model, 'regularizer'):
loss = self.model.regularizer(regularization_factor)
loss.backward()
return total_loss
# This takes the step for this minibatch
loss += self.optimizer.step(closure).detach().cpu().numpy()
loss /= normalization
# We step the scheduler after the full epoch
if self.scheduler is not None:
self.scheduler.step(loss)
self.model.loss_history.append(loss)
self.model.epoch = len(self.model.loss_history)
self.model.latest_iteration_time = time.time() - t0
self.model.training_history += self.model.report() + '\n'
return loss
def optimize(self,
iterations: int,
regularization_factor: Union[float, List[float]] = None,
thread: bool = True,
calculation_width: int = 10):
"""
Runs a round of reconstruction using the provided optimizer
Formerly CDIModel.AD_optimize
This is the basic automatic differentiation reconstruction tool
which all the other, algorithm-specific tools, use. It is a
generator which yields the average loss each epoch, ending after
the specified number of iterations.
By default, the computation will be run in a separate thread. This
is done to enable live plotting with matplotlib during a
reconstruction.
If the computation was done in the main thread, this would freeze
the plots. This behavior can be turned off by setting the keyword
argument 'thread' to False.
Parameters
----------
iterations : int
How many epochs of the algorithm to run.
regularization_factor : float or list(float)
Optional, if the model has a regularizer defined, the set of
parameters to pass the regularizer method.
thread : bool
Default True, whether to run the computation in a separate thread
to allow interaction with plots during computation.
calculation_width : int
Default 10, how many translations to pass through at once for each
round of gradient accumulation. This does not affect the result,
but may affect the calculation speed.
Yields
------
loss : float
The summed loss over the latest epoch, divided by the total
diffraction pattern intensity.
"""
# We store the current optimizer as a model parameter so that
# it can be saved and loaded for checkpointing
self.current_optimizer = self.optimizer
# If we don't want to run in a different thread, this is easy
if not thread:
for it in range(iterations):
if self.model.skip_computation():
self.epoch = self.epoch + 1
if len(self.model.loss_history) >= 1:
yield self.model.loss_history[-1]
else:
yield float('nan')
continue
yield self._run_epoch(regularization_factor=regularization_factor, # noqa
calculation_width=calculation_width)
# But if we do want to thread, it's annoying:
else:
# Here we set up the communication with the computation thread
result_queue = queue.Queue()
stop_event = threading.Event()
def target():
try:
result_queue.put(
self._run_epoch(stop_event=stop_event,
regularization_factor=regularization_factor, # noqa
calculation_width=calculation_width))
except Exception as e:
# If something bad happens, put the exception into the
# result queue
result_queue.put(e)
# And this actually starts and monitors the thread
for it in range(iterations):
if self.model.skip_computation():
self.model.epoch = self.model.epoch + 1
if len(self.model.loss_history) >= 1:
yield self.model.loss_history[-1]
else:
yield float('nan')
continue
calc = threading.Thread(target=target,
name='calculator',
daemon=True)
try:
calc.start()
while calc.is_alive():
if hasattr(self.model, 'figs'):
self.model.figs[0].canvas.start_event_loop(0.01)
else:
calc.join()
except KeyboardInterrupt as e:
stop_event.set()
print('\nAsking execution thread to stop cleanly - ' +
'please be patient.')
calc.join()
raise e
res = result_queue.get()
# If something went wrong in the thead, we'll get an exception
if isinstance(res, Exception):
raise res
yield res
# And finally, we unset the current optimizer:
self.current_optimizer = None
src/cdtools/reconstructors/lbfgs.py
+134
View File
@@ -0,0 +1,134 @@
"""This module contains the LBFGS Reconstructor subclass for performing
optimization ('reconstructions') on ptychographic/CDI models using
the LBFGS optimizer.
The Reconstructor class is designed to resemble so-called
'Trainer' classes that (in the language of the AI/ML folks) handles
the 'training' of a model given some dataset and optimizer.
"""
import torch as t
from cdtools.datasets.ptycho_2d_dataset import Ptycho2DDataset
from cdtools.models import CDIModel
from typing import List, Union
from cdtools.reconstructors import Reconstructor
__all__ = ['LBFGS']
class LBFGS(Reconstructor):
"""
The LBFGS Reconstructor subclass handles the optimization
('reconstruction') of ptychographic models and datasets using the LBFGS
optimizer.
Parameters
----------
model: CDIModel
Model for CDI/ptychography reconstruction.
dataset: Ptycho2DDataset
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).
Important attributes:
- **model** -- Always points to the core model used.
- **optimizer** -- This class by default uses `torch.optim.LBFGS` to
perform optimizations.
- **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):
super().__init__(model, dataset, subset)
# Define the optimizer for use in this subclass
self.optimizer = t.optim.LBFGS(self.model.parameters())
def adjust_optimizer(self,
lr: int = 0.005,
history_size: int = 2,
line_search_fn: str = None):
"""
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.
history_size : int
Optional, the length of the history to use.
line_search_fn : str
Optional, either `strong_wolfe` or None
"""
for param_group in self.optimizer.param_groups:
param_group['lr'] = lr
param_group['history_size'] = history_size
param_group['line_search_fn'] = line_search_fn
def optimize(self,
iterations: int,
lr: float = 0.1,
history_size: int = 2,
regularization_factor: Union[float, List[float]] = None,
thread: bool = True,
calculation_width: int = 10,
line_search_fn: str = None):
"""
Runs a round of reconstruction using the LBFGS optimizer
Formerly `CDIModel.LBFGS_optimize`
This algorithm is often less stable that Adam, however in certain
situations or geometries it can be shockingly efficient. Like all
the other optimization routines, it is defined as a generator
function which yields the average loss each epoch.
NOTE: There is no batch size, because it is a usually a bad idea to use
LBFGS on anything but all the data at onece
Parameters
----------
iterations : int
How many epochs of the algorithm to run.
lr : float
Optional, The learning rate (alpha) to use. Default is 0.1.
history_size : int
Optional, the length of the history to use.
regularization_factor : float or list(float)
Optional, if the model has a regularizer defined, the set of
parameters to pass the regularizer method.
thread : bool
Default True, whether to run the computation in a separate thread
to allow interaction with plots during computation.
calculation_width : int
Default 10, how many translations to pass through at once for each
round of gradient accumulation. Does not affect the result, only
the calculation speed.
"""
# 1) The subset statement is contained in Reconstructor.__init__
# 2) Set up / re-initialize the data loader. For LBFGS, we load
# all the data at once.
self.setup_dataloader(batch_size=len(self.dataset))
# 3) The optimizer is created in self.__init__, but the
# hyperparameters need to be set up with self.adjust_optimizer
self.adjust_optimizer(lr=lr,
history_size=history_size,
line_search_fn=line_search_fn)
# 4) This is analagous to making a call to CDIModel.AD_optimize
return super(LBFGS, self).optimize(iterations,
regularization_factor,
thread,
calculation_width)
src/cdtools/reconstructors/sgd.py
+157
View File
@@ -0,0 +1,157 @@
"""This module contains the SGD Reconstructor subclass for performing
optimization ('reconstructions') on ptychographic/CDI models using
stochastic gradient descent.
The Reconstructor class is designed to resemble so-called
'Trainer' classes that (in the language of the AI/ML folks) handles
the 'training' of a model given some dataset and optimizer.
"""
import torch as t
from cdtools.datasets.ptycho_2d_dataset import Ptycho2DDataset
from cdtools.models import CDIModel
from typing import List, Union
from cdtools.reconstructors import Reconstructor
__all__ = ['SGD']
class SGD(Reconstructor):
"""
The Adam Reconstructor subclass handles the optimization ('reconstruction')
of ptychographic models and datasets using the Adam optimizer.
Parameters
----------
model: CDIModel
Model for CDI/ptychography reconstruction.
dataset: Ptycho2DDataset
The dataset to reconstruct against.
subset : list(int) or int
Optional, a pattern index or list of pattern indices to use.
Important attributes:
- **model** -- Always points to the core model used.
- **optimizer** -- This class by default uses `torch.optim.Adam` to perform
optimizations.
- **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):
super().__init__(model, dataset, subset)
# Define the optimizer for use in this subclass
self.optimizer = t.optim.SGD(self.model.parameters())
def adjust_optimizer(self,
lr: int = 0.005,
momentum: float = 0,
dampening: float = 0,
weight_decay: float = 0,
nesterov: bool = False):
"""
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.
momentum : float
Optional, the length of the history to use.
dampening : float
Optional, dampening for the momentum.
weight_decay : float
Optional, weight decay (L2 penalty).
nesterov : bool
Optional, enables Nesterov momentum. Only applicable when momentum
is non-zero.
"""
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 = None,
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):
"""
Runs a round of reconstruction using the Adam optimizer
Formerly `CDIModel.Adam_optimize`
This calls the Reconstructor.optimize superclass method
(formerly `CDIModel.AD_optimize`) to run a round of reconstruction
once the dataloader and optimizer hyperparameters have been
set up.
Parameters
----------
iterations : int
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 to use. The default is 2e-7.
momentum : float
Optional, the length of the history to use.
dampening : float
Optional, dampening for the momentum.
weight_decay : float
Optional, weight decay (L2 penalty).
nesterov : bool
Optional, enables Nesterov momentum. Only applicable when momentum
is non-zero.
regularization_factor : float or list(float)
Optional, if the model has a regularizer defined, the set of
parameters to pass the regularizer method.
thread : bool
Default True, whether to run the computation in a separate thread
to allow interaction with plots during computation.
calculation_width : int
Default 10, how many translations to pass through at once for each
round of gradient accumulation. Does not affect the result, only
the calculation speed.
shuffle : bool
Optional, enable/disable shuffling of the dataset. This option
is intended for diagnostic purposes and should be left as True.
"""
# 1) The subset statement is contained in Reconstructor.__init__
# 2) Set up / re-initialize the data laoder
if batch_size is not None:
self.setup_dataloader(batch_size=batch_size, shuffle=shuffle)
else:
# Use default torch dataloader parameters
self.setup_dataloader(batch_size=1, shuffle=False)
# 3) 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)
# 4) This is analagous to making a call to CDIModel.AD_optimize
return super(SGD, self).optimize(iterations,
regularization_factor,
thread,
calculation_width)