mirror of
https://github.com/cdtools-developers/cdtools.git
synced 2026-09-19 09:02:09 +02:00
Changed the names of the reconstructors and updated the old optimizer class documentation
This commit is contained in:
yoshikisd
+84
-81
@@ -319,7 +319,6 @@ class CDIModel(t.nn.Module):
|
||||
|
||||
self.current_checkpoint_id += 1
|
||||
|
||||
|
||||
def Adam_optimize(
|
||||
self,
|
||||
iterations: int,
|
||||
@@ -336,7 +335,7 @@ class CDIModel(t.nn.Module):
|
||||
):
|
||||
"""
|
||||
Runs a round of reconstruction using the Adam optimizer from
|
||||
cdtools.reconstructors.Adam.
|
||||
cdtools.reconstructors.AdamReconstructor.
|
||||
|
||||
This is generally accepted to be the most robust algorithm for use
|
||||
with ptychography. Like all the other optimization routines,
|
||||
@@ -375,45 +374,45 @@ class CDIModel(t.nn.Module):
|
||||
only the calculation speed.
|
||||
|
||||
"""
|
||||
# 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:
|
||||
# We want to have model.Adam_optimize call AND store
|
||||
# cdtools.reconstructors.AdamReconstructor 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)
|
||||
from cdtools.reconstructors import AdamReconstructor
|
||||
self.reconstructor = AdamReconstructor(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)
|
||||
batch_size=batch_size,
|
||||
lr=lr,
|
||||
betas=betas,
|
||||
schedule=schedule,
|
||||
amsgrad=amsgrad,
|
||||
regularization_factor=regularization_factor, # noqa
|
||||
thread=thread,
|
||||
calculation_width=calculation_width)
|
||||
|
||||
|
||||
def LBFGS_optimize(self,
|
||||
iterations: int,
|
||||
def LBFGS_optimize(self,
|
||||
iterations: int,
|
||||
dataset: CDataset,
|
||||
lr: float = 0.1,
|
||||
history_size: int = 2,
|
||||
history_size: int = 2,
|
||||
subset: Union[int, List[int]] = None,
|
||||
regularization_factor: Union[float, List[float]] =None,
|
||||
regularization_factor: Union[float, List[float]] =None,
|
||||
thread: bool = True,
|
||||
calculation_width: int = 10,
|
||||
calculation_width: int = 10,
|
||||
line_search_fn: str = None):
|
||||
"""
|
||||
Runs a round of reconstruction using the L-BFGS optimizer from
|
||||
cdtools.reconstructors.LBFGS.
|
||||
cdtools.reconstructors.LBFGSReconstructor.
|
||||
|
||||
This algorithm is often less stable that Adam, however in certain
|
||||
situations or geometries it can be shockingly efficient. Like all
|
||||
@@ -436,53 +435,55 @@ class CDIModel(t.nn.Module):
|
||||
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.
|
||||
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
|
||||
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.
|
||||
"""
|
||||
# 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:
|
||||
# We want to have model.LBFGS_optimize store
|
||||
# cdtools.reconstructors.LBFGSReconstructor 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)
|
||||
from cdtools.reconstructors import LBFGSReconstructor
|
||||
self.reconstructor = LBFGSReconstructor(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)
|
||||
|
||||
lr=lr,
|
||||
history_size=history_size,
|
||||
regularization_factor=regularization_factor, # noqa
|
||||
thread=thread,
|
||||
calculation_width=calculation_width,
|
||||
line_search_fn=line_search_fn)
|
||||
|
||||
def SGD_optimize(self,
|
||||
iterations: int,
|
||||
dataset: CDataset,
|
||||
iterations: int,
|
||||
dataset: CDataset,
|
||||
batch_size: int = None,
|
||||
lr: float = 2e-7,
|
||||
momentum: float = 0,
|
||||
dampening: float = 0,
|
||||
lr: float = 2e-7,
|
||||
momentum: float = 0,
|
||||
dampening: float = 0,
|
||||
weight_decay: float = 0,
|
||||
nesterov: bool = False,
|
||||
subset: Union[int, List[int]] = None,
|
||||
nesterov: bool = False,
|
||||
subset: Union[int, List[int]] = None,
|
||||
regularization_factor: Union[float, List[float]] = None,
|
||||
thread: bool = True,
|
||||
thread: bool = True,
|
||||
calculation_width: int = 10):
|
||||
"""
|
||||
Runs a round of reconstruction using the SGD optimizer from
|
||||
cdtools.reconstructors.SGD.
|
||||
cdtools.reconstructors.SGDReconstructor.
|
||||
|
||||
This algorithm is often less stable that Adam, but it is simpler
|
||||
and is the basic workhorse of gradience descent.
|
||||
@@ -519,29 +520,31 @@ class CDIModel(t.nn.Module):
|
||||
round of gradient accumulation.
|
||||
|
||||
"""
|
||||
# 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:
|
||||
# We want to have model.SGD_optimize store
|
||||
# cdtools.reconstructors.SGDReconstructor 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)
|
||||
|
||||
from cdtools.reconstructors import SGDReconstructor
|
||||
self.reconstructor = SGDReconstructor(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)
|
||||
batch_size=batch_size,
|
||||
lr=lr,
|
||||
momentum=momentum,
|
||||
dampening=dampening,
|
||||
weight_decay=weight_decay,
|
||||
nesterov=nesterov,
|
||||
regularization_factor=regularization_factor, # noqa
|
||||
thread=thread,
|
||||
calculation_width=calculation_width)
|
||||
|
||||
|
||||
def report(self):
|
||||
|
||||
@@ -1,16 +1,22 @@
|
||||
"""This module contains optimizers for performing reconstructions
|
||||
"""
|
||||
Module `cdtools.tools.reconstructors` contains the `Reconstructor` class and
|
||||
subclasses which run the ptychography reconstructions on a given model and
|
||||
dataset.
|
||||
|
||||
The reconstructors are 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.
|
||||
"""
|
||||
|
||||
# We define __all__ to be sure that import * only imports what we want
|
||||
__all__ = [
|
||||
'Reconstructor',
|
||||
'Adam',
|
||||
'LBFGS',
|
||||
'SGD'
|
||||
'AdamReconstructor',
|
||||
'LBFGSReconstructor',
|
||||
'SGDReconstructor'
|
||||
]
|
||||
|
||||
from cdtools.reconstructors.base import Reconstructor
|
||||
from cdtools.reconstructors.adam import Adam
|
||||
from cdtools.reconstructors.lbfgs import LBFGS
|
||||
from cdtools.reconstructors.sgd import SGD
|
||||
from cdtools.reconstructors.adam import AdamReconstructor
|
||||
from cdtools.reconstructors.lbfgs import LBFGSReconstructor
|
||||
from cdtools.reconstructors.sgd import SGDReconstructor
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
"""This module contains the Adam Reconstructor subclass for performing
|
||||
"""This module contains the AdamReconstructor subclass for performing
|
||||
optimization ('reconstructions') on ptychographic/CDI models using
|
||||
the Adam optimizer.
|
||||
|
||||
@@ -12,10 +12,10 @@ from cdtools.models import CDIModel
|
||||
from typing import Tuple, List, Union
|
||||
from cdtools.reconstructors import Reconstructor
|
||||
|
||||
__all__ = ['Adam']
|
||||
__all__ = ['AdamReconstructor']
|
||||
|
||||
|
||||
class Adam(Reconstructor):
|
||||
class AdamReconstructor(Reconstructor):
|
||||
"""
|
||||
The Adam Reconstructor subclass handles the optimization ('reconstruction')
|
||||
of ptychographic models and datasets using the Adam optimizer.
|
||||
@@ -154,7 +154,7 @@ class Adam(Reconstructor):
|
||||
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)
|
||||
return super(AdamReconstructor, self).optimize(iterations,
|
||||
regularization_factor,
|
||||
thread,
|
||||
calculation_width)
|
||||
|
||||
@@ -178,7 +178,7 @@ class Reconstructor:
|
||||
loss.backward()
|
||||
|
||||
# Normalize the accumulating total loss
|
||||
total_loss += loss.detach()
|
||||
total_loss += loss.detach()
|
||||
|
||||
# If we have a regularizer, we can calculate it separately,
|
||||
# and the gradients will add to the minibatch gradient
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
"""This module contains the LBFGS Reconstructor subclass for performing
|
||||
"""This module contains the LBFGSReconstructor subclass for performing
|
||||
optimization ('reconstructions') on ptychographic/CDI models using
|
||||
the LBFGS optimizer.
|
||||
|
||||
@@ -12,12 +12,12 @@ from cdtools.models import CDIModel
|
||||
from typing import List, Union
|
||||
from cdtools.reconstructors import Reconstructor
|
||||
|
||||
__all__ = ['LBFGS']
|
||||
__all__ = ['LBFGSReconstructor']
|
||||
|
||||
|
||||
class LBFGS(Reconstructor):
|
||||
class LBFGSReconstructor(Reconstructor):
|
||||
"""
|
||||
The LBFGS Reconstructor subclass handles the optimization
|
||||
The LBFGSReconstructor subclass handles the optimization
|
||||
('reconstruction') of ptychographic models and datasets using the LBFGS
|
||||
optimizer.
|
||||
|
||||
@@ -128,7 +128,7 @@ class LBFGS(Reconstructor):
|
||||
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)
|
||||
return super(LBFGSReconstructor, self).optimize(iterations,
|
||||
regularization_factor,
|
||||
thread,
|
||||
calculation_width)
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
"""This module contains the SGD Reconstructor subclass for performing
|
||||
"""This module contains the SGDReconstructor subclass for performing
|
||||
optimization ('reconstructions') on ptychographic/CDI models using
|
||||
stochastic gradient descent.
|
||||
|
||||
@@ -12,13 +12,13 @@ from cdtools.models import CDIModel
|
||||
from typing import List, Union
|
||||
from cdtools.reconstructors import Reconstructor
|
||||
|
||||
__all__ = ['SGD']
|
||||
__all__ = ['SGDReconstructor']
|
||||
|
||||
|
||||
class SGD(Reconstructor):
|
||||
class SGDReconstructor(Reconstructor):
|
||||
"""
|
||||
The Adam Reconstructor subclass handles the optimization ('reconstruction')
|
||||
of ptychographic models and datasets using the Adam optimizer.
|
||||
The SGDReconstructor subclass handles the optimization ('reconstruction')
|
||||
of ptychographic models and datasets using the SGD optimizer.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
@@ -151,7 +151,7 @@ class SGD(Reconstructor):
|
||||
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)
|
||||
return super(SGDReconstructor, self).optimize(iterations,
|
||||
regularization_factor,
|
||||
thread,
|
||||
calculation_width)
|
||||
|
||||
@@ -24,7 +24,7 @@ def test_Adam_gold_balls(gold_ball_cxi, reconstruction_device, show_plot):
|
||||
"""
|
||||
|
||||
print('\nTesting performance on the standard gold balls dataset ' +
|
||||
'with reconstructors.Adam')
|
||||
'with reconstructors.AdamReconstructor')
|
||||
|
||||
dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(gold_ball_cxi)
|
||||
pad = 10
|
||||
@@ -48,11 +48,12 @@ def test_Adam_gold_balls(gold_ball_cxi, reconstruction_device, show_plot):
|
||||
model_recon.to(device=reconstruction_device)
|
||||
dataset.get_as(device=reconstruction_device)
|
||||
|
||||
# ******* Reconstructions with cdtools.reconstructors.Adam.optimize *******
|
||||
print('Running reconstruction using cdtools.reconstructors.Adam.optimize' +
|
||||
# ******* Reconstructions with AdamReconstructor.optimize *******
|
||||
print('Running reconstruction using AdamReconstructor.optimize' +
|
||||
' on provided reconstruction_device,', reconstruction_device)
|
||||
|
||||
recon = cdtools.reconstructors.Adam(model=model_recon, dataset=dataset)
|
||||
recon = cdtools.reconstructors.AdamReconstructor(model=model_recon,
|
||||
dataset=dataset)
|
||||
t.manual_seed(0)
|
||||
|
||||
# Run a reconstruction
|
||||
@@ -85,7 +86,7 @@ def test_Adam_gold_balls(gold_ball_cxi, reconstruction_device, show_plot):
|
||||
model_recon.inspect(dataset)
|
||||
model_recon.compare(dataset)
|
||||
|
||||
# ******* Reconstructions with cdtools.CDIModel.Adam_optimize *******
|
||||
# ******* Reconstructions with CDIModel.Adam_optimize *******
|
||||
print('Running reconstruction using CDIModel.Adam_optimize on provided' +
|
||||
' reconstruction_device,', reconstruction_device)
|
||||
t.manual_seed(0)
|
||||
@@ -150,11 +151,12 @@ def test_LBFGS_RPI(optical_data_ss_cxi,
|
||||
model_recon.to(device=reconstruction_device)
|
||||
dataset.get_as(device=reconstruction_device)
|
||||
|
||||
# ******* Reconstructions with cdtools.reconstructors.LBFGS.optimize ******
|
||||
print('Running reconstruction using cdtools.reconstructors.LBFGS.' +
|
||||
# ******* Reconstructions with LBFGSReconstructor.optimize ******
|
||||
print('Running reconstruction using LBFGSReconstructor.' +
|
||||
'optimize on provided reconstruction_device,', reconstruction_device)
|
||||
|
||||
recon = cdtools.reconstructors.LBFGS(model=model_recon, dataset=dataset)
|
||||
recon = cdtools.reconstructors.LBFGSReconstructor(model=model_recon,
|
||||
dataset=dataset)
|
||||
t.manual_seed(0)
|
||||
|
||||
# Run a reconstruction
|
||||
@@ -178,7 +180,7 @@ def test_LBFGS_RPI(optical_data_ss_cxi,
|
||||
# Check model pointing
|
||||
assert id(model_recon) == id(recon.model)
|
||||
|
||||
# ******* Reconstructions with cdtools.reconstructors.LBFGS.optimize ******
|
||||
# ******* Reconstructions with CDIModel.LBFGS_optimize ******
|
||||
print('Running reconstruction using CDIModel.LBFGS_optimize.' +
|
||||
'optimize on provided reconstruction_device,', reconstruction_device)
|
||||
t.manual_seed(0)
|
||||
@@ -223,7 +225,7 @@ def test_SGD_gold_balls(gold_ball_cxi, reconstruction_device, show_plot):
|
||||
to do some kind of comparative assessment.
|
||||
"""
|
||||
print('\nTesting performance on the standard gold balls dataset ' +
|
||||
'with reconstructors.SGD')
|
||||
'with reconstructors.SGDReconstructor')
|
||||
|
||||
dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(gold_ball_cxi)
|
||||
pad = 10
|
||||
@@ -247,11 +249,12 @@ def test_SGD_gold_balls(gold_ball_cxi, reconstruction_device, show_plot):
|
||||
model_recon.to(device=reconstruction_device)
|
||||
dataset.get_as(device=reconstruction_device)
|
||||
|
||||
# ******* Reconstructions with cdtools.reconstructors.SGD.optimize *******
|
||||
print('Running reconstruction using cdtools.reconstructors.SGD.optimize' +
|
||||
# ******* Reconstructions with SGDReconstructor.optimize *******
|
||||
print('Running reconstruction using SGDReconstructor.optimize' +
|
||||
' on provided reconstruction_device,', reconstruction_device)
|
||||
|
||||
recon = cdtools.reconstructors.SGD(model=model_recon, dataset=dataset)
|
||||
recon = cdtools.reconstructors.SGDReconstructor(model=model_recon,
|
||||
dataset=dataset)
|
||||
t.manual_seed(0)
|
||||
|
||||
# Run a reconstruction
|
||||
|
||||
Reference in New Issue
Block a user