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

Change the pattern for Reconstructor so that the optimizer is defined at object creation, and move the dataloader creation logic to the base optimize() function as it was reused in all subclasses

Browse Source
This commit is contained in:
levita_a
2025-10-13 19:20:42 +02:00
parent a089dc3b72
commit aa684f27eb
4 changed files with 133 additions and 67 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/cdtools/reconstructors/adam.py
+29 -16
View File
@@ -51,10 +51,12 @@ class AdamReconstructor(Reconstructor):
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())
optimizer = t.optim.Adam(model.parameters())
super().__init__(model, dataset, optimizer, subset=subset)
def adjust_optimizer(self,
lr: int = 0.005,
@@ -79,11 +81,13 @@ class AdamReconstructor(Reconstructor):
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 = None,
schedule: bool = False,
amsgrad: bool = False,
regularization_factor: Union[float, List[float]] = None,
@@ -99,6 +103,12 @@ class AdamReconstructor(Reconstructor):
(formerly `CDIModel.AD_optimize`) to run a round of reconstruction
once the dataloader and optimizer hyperparameters have been
set up.
The `batch_size` parameter sets the batch size for the default
dataloader. If a custom data loader is desired, it can be passed
in to the `custom_data_loader` argument, which will override the
`batch_size` and `shuffle` parameters
Parameters
----------
@@ -115,6 +125,9 @@ class AdamReconstructor(Reconstructor):
schedule : bool
Optional, create a learning rate scheduler
(torch.optim.lr_scheduler._LRScheduler).
custom_data_loader : t.utils.data.DataLoader
Optional, a custom DataLoader to use. If set, will override
batch_size.
amsgrad : bool
Optional, whether to use the AMSGrad variant of this algorithm.
regularization_factor : float or list(float)
@@ -138,18 +151,13 @@ class AdamReconstructor(Reconstructor):
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
# 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
# Set up the scheduler
if schedule:
self.scheduler = \
t.optim.lr_scheduler.ReduceLROnPlateau(self.optimizer,
@@ -158,8 +166,13 @@ class AdamReconstructor(Reconstructor):
else:
self.scheduler = None
# 5) This is analagous to making a call to CDIModel.AD_optimize
return super(AdamReconstructor, self).optimize(iterations,
regularization_factor,
thread,
calculation_width)
# Now, we run the optimize routine defined in the base class
return super(AdamReconstructor, self).optimize(
iterations,
batch_size=batch_size,
custom_data_loader=custom_data_loader,
regularization_factor=regularization_factor,
thread=thread,
calculation_width=calculation_width,
shuffle=shuffle,
)
src/cdtools/reconstructors/base.py
+66 -17
View File
@@ -40,6 +40,8 @@ class Reconstructor:
Model for CDI/ptychography reconstruction
dataset: CDataset
The dataset to reconstruct against
optimizer: torch.optim.Optimizer
The optimizer to use for the reconstruction
subset : list(int) or int
Optional, a pattern index or list of pattern indices to use
@@ -55,31 +57,32 @@ class Reconstructor:
def __init__(self,
model: CDIModel,
dataset: CDataset,
optimizer: t.optim.Optimizer,
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
self.optimizer = optimizer
# Store the dataset
# Store the dataset, clipping it to a subset if needed
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
# Initialize attributes that must be defined by the subclasses
self.scheduler = None
self.data_loader = None
def setup_dataloader(self,
batch_size: int = None,
shuffle: bool = True):
"""
Sets up / re-initializes the dataloader.
Sets up or re-initializes the dataloader.
Parameters
----------
@@ -96,6 +99,7 @@ class Reconstructor:
else:
self.data_loader = td.Dataloader(self.dataset)
def adjust_optimizer(self, **kwargs):
"""
Change hyperparameters for the utilized optimizer.
@@ -105,7 +109,8 @@ class Reconstructor:
"""
raise NotImplementedError()
def _run_epoch(self,
def run_epoch(self,
stop_event: threading.Event = None,
regularization_factor: Union[float, List[float]] = None,
calculation_width: int = 10):
@@ -133,6 +138,18 @@ class Reconstructor:
diffraction pattern intensity
"""
# Setting this as an explicit catch makes me feel more comfortable
# exposing it as a public method. This way a user won't be confused
# if they try to use this directly
if self.data_loader is None:
raise RuntimeError(
'No data loader was defined. Please run '
'Reconstructor.setup_dataloader() before running '
'Reconstructor.run_epoch(), or use Reconstructor.optimize(), '
'which does it automatically.'
)
# Initialize some tracking variables
normalization = 0
loss = 0
@@ -212,9 +229,12 @@ class Reconstructor:
def optimize(self,
iterations: int,
batch_size: int = 1,
custom_data_loader = None,
regularization_factor: Union[float, List[float]] = None,
thread: bool = True,
calculation_width: int = 10):
calculation_width: int = 10,
shuffle=True):
"""
Runs a round of reconstruction using the provided optimizer
@@ -233,10 +253,25 @@ class Reconstructor:
the plots. This behavior can be turned off by setting the keyword
argument 'thread' to False.
The `batch_size` parameter sets the batch size for the default
dataloader. If a custom data loader is desired, it can be passed
in to the `custom_data_loader` argument, which will override the
`batch_size` and `shuffle` parameters
Please see `AdamReconstructor.optimize()` for an example of how to
override this function when designing a subclass
Parameters
----------
iterations : int
How many epochs of the algorithm to run.
batch_size : int
Optional, the batch size to use. Default is 1. This is typically
overridden by subclasses with an appropriate default for the
specific optimizer.
custom_data_loader : torch.utils.data.DataLoader
Optional, a custom DataLoader to use. Will override batch_size
if set.
regularization_factor : float or list(float)
Optional, if the model has a regularizer defined, the set of
parameters to pass the regularizer method.
@@ -247,6 +282,10 @@ class Reconstructor:
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.
shuffle : bool
Optional, enable/disable shuffling of the dataset. This option
is intended for diagnostic purposes and should be left as True.
Yields
------
@@ -255,6 +294,11 @@ class Reconstructor:
diffraction pattern intensity.
"""
if custom_data_loader is None:
self.setup_dataloader(batch_size=batch_size, shuffle=shuffle)
else:
self.data_loader = custom_data_loader
# We store the current optimizer as a model parameter so that
# it can be saved and loaded for checkpointing
self.current_optimizer = self.optimizer
@@ -270,8 +314,10 @@ class Reconstructor:
yield float('nan')
continue
yield self._run_epoch(regularization_factor=regularization_factor, # noqa
calculation_width=calculation_width)
yield self.run_epoch(
regularization_factor=regularization_factor, # noqa
calculation_width=calculation_width,
)
# But if we do want to thread, it's annoying:
else:
@@ -282,9 +328,12 @@ class Reconstructor:
def target():
try:
result_queue.put(
self._run_epoch(stop_event=stop_event,
regularization_factor=regularization_factor, # noqa
calculation_width=calculation_width))
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
src/cdtools/reconstructors/lbfgs.py
+18 -15
View File
@@ -53,10 +53,16 @@ class LBFGSReconstructor(Reconstructor):
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())
optimizer = t.optim.LBFGS(model.parameters())
super().__init__(
model,
dataset,
optimizer,
subset=subset,
)
def adjust_optimizer(self,
lr: int = 0.005,
@@ -121,20 +127,17 @@ class LBFGSReconstructor(Reconstructor):
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
# 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(LBFGSReconstructor, self).optimize(iterations,
regularization_factor,
thread,
calculation_width)
# Now, we run the optimize routine defined in the base class
return super(LBFGSReconstructor, self).optimize(
iterations,
batch_size=len(self.dataset),
regularization_factor=regularization_factor,
thread=thread,
calculation_width=calculation_width)
src/cdtools/reconstructors/sgd.py
+20 -19
View File
@@ -49,11 +49,17 @@ class SGDReconstructor(Reconstructor):
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())
optimizer = t.optim.SGD(model.parameters())
super().__init__(
model,
dataset,
optimizer,
subset=subset,
)
def adjust_optimizer(self,
lr: int = 0.005,
momentum: float = 0,
@@ -88,7 +94,7 @@ class SGDReconstructor(Reconstructor):
def optimize(self,
iterations: int,
batch_size: int = None,
batch_size: int = 15,
lr: float = 2e-7,
momentum: float = 0,
dampening: float = 0,
@@ -139,25 +145,20 @@ class SGDReconstructor(Reconstructor):
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
# 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(SGDReconstructor, self).optimize(iterations,
regularization_factor,
thread,
calculation_width)
# Now, we run the optimize routine defined in the base class
return super(SGDReconstructor, self).optimize(
iterations,
batch_size=batch_size,
regularization_factor=regularization_factor,
thread=thread,
calculation_width=calculation_width,
)