mirror of
https://github.com/cdtools-developers/cdtools.git
synced 2026-09-10 05:22:41 +02:00
Rebased CDIModel to use Reconstructors for reconstructions.
This commit is contained in:
yoshikisd
+161
-311
@@ -40,6 +40,9 @@ import time
|
||||
from scipy import io
|
||||
from contextlib import contextmanager
|
||||
from cdtools.tools.data import nested_dict_to_h5, h5_to_nested_dict, nested_dict_to_numpy, nested_dict_to_torch
|
||||
from cdtools.datasets import CDataset
|
||||
from typing import List, Union, Tuple
|
||||
import os
|
||||
|
||||
__all__ = ['CDIModel']
|
||||
|
||||
@@ -316,202 +319,24 @@ class CDIModel(t.nn.Module):
|
||||
|
||||
self.current_checkpoint_id += 1
|
||||
|
||||
|
||||
|
||||
def AD_optimize(self, iterations, data_loader, optimizer,\
|
||||
scheduler=None, regularization_factor=None, thread=True,
|
||||
calculation_width=10):
|
||||
"""Runs a round of reconstruction using the provided optimizer
|
||||
|
||||
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
|
||||
data_loader : torch.utils.data.DataLoader
|
||||
A data loader loading the CDataset to reconstruct
|
||||
optimizer : torch.optim.Optimizer
|
||||
The optimizer to run the reconstruction with
|
||||
scheduler : torch.optim.lr_scheduler._LRScheduler
|
||||
Optional, a learning rate scheduler 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. 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
|
||||
"""
|
||||
|
||||
def run_epoch(stop_event=None):
|
||||
"""Runs one full epoch of the reconstruction."""
|
||||
# First, 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 data_loader:
|
||||
normalization += t.sum(patterns).cpu().numpy()
|
||||
N += 1
|
||||
def closure():
|
||||
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.forward(*inp)
|
||||
|
||||
# Calculate the loss
|
||||
if hasattr(self, 'mask'):
|
||||
loss = self.loss(pats,sim_patterns, mask=self.mask)
|
||||
else:
|
||||
loss = self.loss(pats,sim_patterns)
|
||||
|
||||
# And accumulate the gradients
|
||||
loss.backward()
|
||||
total_loss += loss.detach()
|
||||
|
||||
# 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, 'regularizer'):
|
||||
loss = self.regularizer(regularization_factor)
|
||||
loss.backward()
|
||||
|
||||
return total_loss
|
||||
|
||||
# This takes the step for this minibatch
|
||||
loss += optimizer.step(closure).detach().cpu().numpy()
|
||||
|
||||
|
||||
loss /= normalization
|
||||
|
||||
# We step the scheduler after the full epoch
|
||||
if scheduler is not None:
|
||||
scheduler.step(loss)
|
||||
|
||||
self.loss_history.append(loss)
|
||||
self.epoch = len(self.loss_history)
|
||||
self.latest_iteration_time = time.time() - t0
|
||||
self.training_history += self.report() + '\n'
|
||||
return loss
|
||||
|
||||
# We store the current optimizer as a model parameter so that
|
||||
# it can be saved and loaded for checkpointing
|
||||
self.current_optimizer = 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.skip_computation():
|
||||
self.epoch = self.epoch + 1
|
||||
if len(self.loss_history) >= 1:
|
||||
yield self.loss_history[-1]
|
||||
else:
|
||||
yield float('nan')
|
||||
continue
|
||||
|
||||
yield run_epoch()
|
||||
|
||||
|
||||
# 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(run_epoch(stop_event))
|
||||
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.skip_computation():
|
||||
self.epoch = self.epoch + 1
|
||||
if len(self.loss_history) >= 1:
|
||||
yield self.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, 'figs'):
|
||||
self.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
|
||||
|
||||
|
||||
def Adam_optimize(
|
||||
self,
|
||||
iterations,
|
||||
dataset,
|
||||
batch_size=15,
|
||||
lr=0.005,
|
||||
betas=(0.9, 0.999),
|
||||
schedule=False,
|
||||
amsgrad=False,
|
||||
subset=None,
|
||||
regularization_factor=None,
|
||||
iterations: int,
|
||||
dataset: CDataset,
|
||||
batch_size: int = 15,
|
||||
lr: float = 0.005,
|
||||
betas: Tuple[float] = (0.9, 0.999),
|
||||
schedule: bool = False,
|
||||
amsgrad: bool = False,
|
||||
subset: Union[int, List[int]] = None,
|
||||
regularization_factor: Union[float, List[float]] = None,
|
||||
thread=True,
|
||||
calculation_width=10
|
||||
):
|
||||
"""Runs a round of reconstruction using the Adam optimizer
|
||||
"""
|
||||
Runs a round of reconstruction using the Adam optimizer from
|
||||
cdtools.reconstructors.Adam.
|
||||
|
||||
This is generally accepted to be the most robust algorithm for use
|
||||
with ptychography. Like all the other optimization routines,
|
||||
@@ -521,125 +346,143 @@ class CDIModel(t.nn.Module):
|
||||
Parameters
|
||||
----------
|
||||
iterations : int
|
||||
How many epochs of the algorithm to run
|
||||
How many epochs of the algorithm to run.
|
||||
dataset : CDataset
|
||||
The dataset to reconstruct against
|
||||
The dataset to reconstruct against.
|
||||
batch_size : int
|
||||
Optional, the size of the minibatches to use
|
||||
Optional, the size of the minibatches to use.
|
||||
lr : float
|
||||
Optional, The learning rate (alpha) to use. Defaultis 0.005. 0.05 is typically the highest possible value with any chance of being stable
|
||||
betas : tuple
|
||||
Optional, The learning rate (alpha) to use. Defaultis 0.005.
|
||||
0.05 is typically the highest possible value with any chance
|
||||
of being stable.
|
||||
betas : tuple(float)
|
||||
Optional, the beta_1 and beta_2 to use. Default is (0.9, 0.999).
|
||||
schedule : float
|
||||
Optional, whether to use the ReduceLROnPlateau scheduler
|
||||
schedule : bool
|
||||
Optional, whether to use the ReduceLROnPlateau scheduler.
|
||||
amsgrad : bool
|
||||
Optional, whether to use the AMSGrad variant of this algorithm.
|
||||
subset : list(int) or int
|
||||
Optional, a pattern index or list of pattern indices 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
|
||||
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
|
||||
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
|
||||
|
||||
"""
|
||||
|
||||
self.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'
|
||||
)
|
||||
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.
|
||||
|
||||
|
||||
if subset is not None:
|
||||
# if subset is just one pattern, turn into a list for convenience
|
||||
if type(subset) == type(1):
|
||||
subset = [subset]
|
||||
dataset = torchdata.Subset(dataset, subset)
|
||||
|
||||
# Make a dataloader
|
||||
data_loader = torchdata.DataLoader(dataset,
|
||||
batch_size=batch_size,
|
||||
shuffle=True)
|
||||
|
||||
# Define the optimizer
|
||||
optimizer = t.optim.Adam(
|
||||
self.parameters(),
|
||||
lr = lr,
|
||||
betas=betas,
|
||||
amsgrad=amsgrad)
|
||||
|
||||
# Define the scheduler
|
||||
if schedule:
|
||||
scheduler = t.optim.lr_scheduler.ReduceLROnPlateau(optimizer, factor=0.2,threshold=1e-9)
|
||||
else:
|
||||
scheduler = None
|
||||
|
||||
return self.AD_optimize(iterations, data_loader, optimizer,
|
||||
scheduler=scheduler,
|
||||
regularization_factor=regularization_factor,
|
||||
thread=thread,
|
||||
calculation_width=calculation_width)
|
||||
"""
|
||||
# We want to have model.Adam_optimize call AND store cdtools.reconstructors.Adam
|
||||
# to perform reconstructions without creating a new reconstructor each time we
|
||||
# update the hyperparameters.
|
||||
#
|
||||
# The only way to do this is to make the Adam reconstructor an attribute
|
||||
# of the model. But since the Adam reconstructor also depends on CDIModel,
|
||||
# this seems to give rise to a circular import error unless
|
||||
# we import cdtools.reconstructors within this method:
|
||||
if not hasattr(self, 'reconstructor'):
|
||||
from cdtools.reconstructors import Adam
|
||||
self.reconstructor = Adam(model=self,
|
||||
dataset=dataset,
|
||||
subset=subset)
|
||||
|
||||
# Run some reconstructions
|
||||
return self.reconstructor.optimize(iterations=iterations,
|
||||
batch_size=batch_size,
|
||||
lr=lr,
|
||||
betas=betas,
|
||||
schedule=schedule,
|
||||
amsgrad=amsgrad,
|
||||
regularization_factor=regularization_factor,
|
||||
thread=thread,
|
||||
calculation_width=calculation_width)
|
||||
|
||||
|
||||
def LBFGS_optimize(self, iterations, dataset,
|
||||
lr=0.1,history_size=2, subset=None,
|
||||
regularization_factor=None, thread=True,
|
||||
calculation_width=10, line_search_fn=None):
|
||||
"""Runs a round of reconstruction using the L-BFGS optimizer
|
||||
def LBFGS_optimize(self,
|
||||
iterations: int,
|
||||
dataset: CDataset,
|
||||
lr: float = 0.1,
|
||||
history_size: int = 2,
|
||||
subset: Union[int, List[int]] = None,
|
||||
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 L-BFGS optimizer from
|
||||
cdtools.reconstructors.LBFGS.
|
||||
|
||||
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
|
||||
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
|
||||
How many epochs of the algorithm to run.
|
||||
dataset : CDataset
|
||||
The dataset to reconstruct against
|
||||
The dataset to reconstruct against.
|
||||
lr : float
|
||||
Optional, the learning rate to use
|
||||
Optional, the learning rate to use.
|
||||
history_size : int
|
||||
Optional, the length of the history to use.
|
||||
subset : list(int) or int
|
||||
Optional, a pattern index or list of pattern indices to ues
|
||||
Optional, a pattern index or list of pattern indices 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
|
||||
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.
|
||||
|
||||
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
|
||||
"""
|
||||
if subset is not None:
|
||||
# if just one pattern, turn into a list for convenience
|
||||
if type(subset) == type(1):
|
||||
subset = [subset]
|
||||
dataset = torchdata.Subset(dataset, subset)
|
||||
|
||||
# Make a dataloader. This basically does nothing but load all the
|
||||
# data at once
|
||||
data_loader = torchdata.DataLoader(dataset, batch_size=len(dataset))
|
||||
# We want to have model.LBFGS_optimize store cdtools.reconstructors.LBFGS
|
||||
# as an attribute to run reconstructions without generating new reconstructors
|
||||
# each time CDIModel.LBFGS_optimize is called.
|
||||
#
|
||||
# Since the LBFGS reconstructor also depends on CDIModel, a circular import error
|
||||
# arises unless we import cdtools.reconstructors within this method:
|
||||
if not hasattr(self, 'reconstructor'):
|
||||
from cdtools.reconstructors import LBFGS
|
||||
self.reconstructor = LBFGS(model=self,
|
||||
dataset=dataset,
|
||||
subset=subset)
|
||||
|
||||
# Run some reconstructions
|
||||
return self.reconstructor.optimize(iterations=iterations,
|
||||
lr=lr,
|
||||
history_size=history_size,
|
||||
regularization_factor=regularization_factor,
|
||||
thread=thread,
|
||||
calculation_width=calculation_width,
|
||||
line_search_fn = line_search_fn)
|
||||
|
||||
|
||||
# Define the optimizer
|
||||
optimizer = t.optim.LBFGS(self.parameters(),
|
||||
lr = lr, history_size=history_size,
|
||||
line_search_fn=line_search_fn)
|
||||
|
||||
return self.AD_optimize(iterations, data_loader, optimizer,
|
||||
regularization_factor=regularization_factor,
|
||||
thread=thread,
|
||||
calculation_width=calculation_width)
|
||||
|
||||
|
||||
def SGD_optimize(self, iterations, dataset, batch_size=None,
|
||||
lr=0.01, momentum=0, dampening=0, weight_decay=0,
|
||||
nesterov=False, subset=None, regularization_factor=None,
|
||||
thread=True, calculation_width=10):
|
||||
"""Runs a round of reconstruction using the SGD optimizer
|
||||
def SGD_optimize(self,
|
||||
iterations: int,
|
||||
dataset: CDataset,
|
||||
batch_size: int = None,
|
||||
lr: float = 2e-7,
|
||||
momentum: float = 0,
|
||||
dampening: float = 0,
|
||||
weight_decay: float = 0,
|
||||
nesterov: bool = False,
|
||||
subset: Union[int, List[int]] = None,
|
||||
regularization_factor: Union[float, List[float]] = None,
|
||||
thread: bool = True,
|
||||
calculation_width: int = 10):
|
||||
"""
|
||||
Runs a round of reconstruction using the SGD optimizer from
|
||||
cdtools.reconstructors.SGD.
|
||||
|
||||
This algorithm is often less stable that Adam, but it is simpler
|
||||
and is the basic workhorse of gradience descent.
|
||||
@@ -647,51 +490,58 @@ class CDIModel(t.nn.Module):
|
||||
Parameters
|
||||
----------
|
||||
iterations : int
|
||||
How many epochs of the algorithm to run
|
||||
How many epochs of the algorithm to run.
|
||||
dataset : CDataset
|
||||
The dataset to reconstruct against
|
||||
The dataset to reconstruct against.
|
||||
batch_size : int
|
||||
Optional, the size of the minibatches to use
|
||||
Optional, the size of the minibatches to use.
|
||||
lr : float
|
||||
Optional, the learning rate to use
|
||||
Optional, the learning rate to use.
|
||||
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.
|
||||
subset : list(int) or int
|
||||
Optional, a pattern index or list of pattern indices to use
|
||||
Optional, a pattern index or list of pattern indices 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
|
||||
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
|
||||
Default True, whether to run the computation in a separate thread
|
||||
to allow interaction with plots during computation.
|
||||
calculation_width : int
|
||||
Default 1, how many translations to pass through at once for each round of gradient accumulation
|
||||
Default 10, how many translations to pass through at once for each
|
||||
round of gradient accumulation.
|
||||
|
||||
"""
|
||||
|
||||
if subset is not None:
|
||||
# if just one pattern, turn into a list for convenience
|
||||
if type(subset) == type(1):
|
||||
subset = [subset]
|
||||
dataset = torchdata.Subset(dataset, subset)
|
||||
|
||||
# Make a dataloader
|
||||
if batch_size is not None:
|
||||
data_loader = torchdata.DataLoader(dataset, batch_size=batch_size,
|
||||
shuffle=True)
|
||||
else:
|
||||
data_loader = torchdata.DataLoader(dataset)
|
||||
|
||||
|
||||
# Define the optimizer
|
||||
optimizer = t.optim.SGD(self.parameters(),
|
||||
lr = lr, momentum=momentum,
|
||||
dampening=dampening,
|
||||
weight_decay=weight_decay,
|
||||
nesterov=nesterov)
|
||||
|
||||
return self.AD_optimize(iterations, data_loader, optimizer,
|
||||
regularization_factor=regularization_factor,
|
||||
thread=thread,
|
||||
calculation_width=calculation_width)
|
||||
# We want to have model.SGD_optimize store cdtools.reconstructors.SGD
|
||||
# as an attribute to run reconstructions without generating new reconstructors
|
||||
# each time CDIModel.SGD_optimize is called.
|
||||
#
|
||||
# Since the SGD reconstructor also depends on CDIModel, a circular import error
|
||||
# arises unless we import cdtools.reconstructors within this method:
|
||||
if not hasattr(self, 'reconstructor'):
|
||||
from cdtools.reconstructors import SGD
|
||||
self.reconstructor = SGD(model=self,
|
||||
dataset=dataset,
|
||||
subset=subset)
|
||||
|
||||
# Run some reconstructions
|
||||
return self.reconstructor.optimize(iterations=iterations,
|
||||
batch_size=batch_size,
|
||||
lr=lr,
|
||||
momentum=momentum,
|
||||
dampening=dampening,
|
||||
weight_decay=weight_decay,
|
||||
nesterov=nesterov,
|
||||
regularization_factor=regularization_factor,
|
||||
thread=thread,
|
||||
calculation_width=calculation_width)
|
||||
|
||||
|
||||
def report(self):
|
||||
|
||||
Reference in New Issue
Block a user