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

Merge pull request #70 from cdtools-developers/newpoisson

Browse Source 
Switch to a more sensible normalization strategy for loss metrics
This commit is contained in:
Abe Levitan
2026-04-17 15:12:03 +02:00
committed by GitHub
parent 3f028e2fbf fcd0b31b2c
commit 65a75b5c7e
15 changed files with 481 additions and 75 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/source/tutorial.rst
+10 -6
View File
@@ -196,6 +196,7 @@ Once again, we start with the basic skeleton
.. code-block:: python
from functools import partial
import torch as t
from cdtools.models import CDIModel
from cdtools import tools
@@ -251,7 +252,11 @@ There is no requirement for what the arguments to the initialization function of
self.probe = t.nn.Parameter(probe_guess / self.probe_norm)
self.obj = t.nn.Parameter(obj_guess)
# We register a loss function and an appropriate normalization
self.loss = partial(tools.losses.amplitude_mse, use_sum=True)
self.loss_normalizer = tools.losses.AmplitudeMSENormalizer()
The first thing to notice about this model is that all the fixed, geometric information is stored with the :code:`module.register_buffer()` function. This is what makes it possible to move all the relevant tensors between devices using a single call to :code:`module.to()`, for example. It stores thetensor as an object attribute, but it also registers it so that pytorch is aware that this attribute helps to encode the state of the model.
The supporting information we need is the wavelength of the illumination, the basis of the probe array in real space, and an offset to define the zero point of the translation.
@@ -268,6 +273,8 @@ The Adam optimizer is designed so that the learning rate sets the maximum stepsi
This is important to remember when adding additional error models. Rescaling all the parameters to have a typical amplitude near 1 is the best way to get well-behaved reconstructions.
The final two lines assign a loss function and its associated normalizer. Here we use :code:`amplitude_mse` and its paired :code:`AmplitudeMSENormalizer`. The normalization function :code:`amplitude_mse` computes the the mean squared error between the square roots of the simulated and measured intensities. In this case, we call it with the :code:`use_sum=True` flag, which will actually calculate a sum-square error, which will be normalized afterward to a mean-squared error by the normalizer, :code:`AmplitudeMSENormalizer`. This pattern is used to ensure that the losses from minibatches with different sizes are properly weighted.
Initialization from Dataset
+++++++++++++++++++++++++++
@@ -347,7 +354,7 @@ Here, we take input in the form of an index and a translation. Note that this in
We start by mapping the translation, given in real space, into pixel coordinates. Then, we use an "off-the-shelf" interaction model - :code:`ptycho_2d_round`, which models a standard 2D ptychography interaction, but rounds the translations to the nearest whole pixel (does not attempt subpixel translations).
The next three definitions amount to just choosing an off-the-shelf function to simulate each step in the chain.
The next two definitions amount to just choosing an off-the-shelf function to simulate each step in the chain.
.. code-block:: python
@@ -357,11 +364,8 @@ The next three definitions amount to just choosing an off-the-shelf function to
def measurement(self, wavefields):
return tools.measurements.intensity(wavefields)
def loss(self, sim_data, real_data):
return tools.losses.amplitude_mse(real_data, sim_data)
The forward propagator maps the exit wave to the wave at the surface of the detector, here using a far-field propagator. The measurement maps that exit wave to a measured pixel value, and the loss defines a loss function to attempt to minimize. The loss function we've chosen - the amplitude mean squared error - is the most reliable one, and can also easily be overridden by an end user.
The forward propagator maps the exit wave to the wave at the surface of the detector, here using a far-field propagator. The measurement maps that wavefield to a measured pixel value. The loss function was already assigned in :code:`__init__` as described above.
Plotting
examples/near_field_ptycho.py
+2 -1
View File
@@ -26,7 +26,8 @@ model = cdtools.models.FancyPtycho.from_dataset(
near_field=True,
propagation_distance=3.65e-3, # 3.65 downstream from focus
units='um', # Set the units for the live plots
obj_view_crop=-35,
obj_view_crop=-35, # Expand the view for the live plots
loss="poisson_nll", # Best option for photon-counting detectors
panel_plot_mode=True, # Set to False to get individual figures
)
examples/tutorial_simple_ptycho.py
+5 -3
View File
@@ -1,3 +1,4 @@
from functools import partial
import torch as t
from cdtools.models import CDIModel
from cdtools import tools
@@ -41,6 +42,10 @@ class SimplePtycho(CDIModel):
self.probe = t.nn.Parameter(probe_guess / self.probe_norm)
self.obj = t.nn.Parameter(obj_guess)
# We register a loss function and an appropriate normalization
self.loss = partial(tools.losses.amplitude_mse, use_sum=True)
self.loss_normalizer = tools.losses.AmplitudeMSENormalizer()
@classmethod
def from_dataset(cls, dataset):
@@ -102,9 +107,6 @@ class SimplePtycho(CDIModel):
def measurement(self, wavefields):
return tools.measurements.intensity(wavefields)
def loss(self, real_data, sim_data):
return tools.losses.amplitude_mse(real_data, sim_data)
# This lists all the plots to display on a call to model.inspect()
plot_list = [
src/cdtools/models/bragg_2d_ptycho.py
+20 -5
View File
@@ -1,3 +1,4 @@
from functools import partial
import torch as t
from cdtools.models import CDIModel
from cdtools.datasets import Ptycho2DDataset
@@ -76,6 +77,7 @@ class Bragg2DPtycho(CDIModel):
propagate_probe=True,
correct_tilt=True,
lens=False,
loss='amplitude mse',
units='um',
dtype=t.float32,
obj_view_crop=0,
@@ -237,7 +239,22 @@ class Bragg2DPtycho(CDIModel):
# TODO: probably doesn't support non-float-32 dtypes
self.register_buffer('universal_propagator',
universal_propagator)
# Here we set the appropriate loss function
if (loss.lower().strip() == 'amplitude mse'
or loss.lower().strip() == 'amplitude_mse'):
self.loss = partial(tools.losses.amplitude_mse, use_sum=True)
self.loss_normalizer = tools.losses.AmplitudeMSENormalizer()
elif (loss.lower().strip() == 'poisson nll'
or loss.lower().strip() == 'poisson_nll'):
self.loss = tools.losses.poisson_nll
self.loss_normalizer = tools.losses.SimplePoissonNLLNormalizer()
elif (loss.lower().strip() == 'intensity mse'
or loss.lower().strip() == 'intensity_mse'):
self.loss = partial(tools.losses.intensity_mse, use_sum=True)
self.loss_normalizer = tools.losses.IntensityMSENormalizer()
else:
raise KeyError('Specified loss function not supported')
@classmethod
@@ -257,6 +274,7 @@ class Bragg2DPtycho(CDIModel):
propagate_probe=True,
correct_tilt=True,
lens=False,
loss='amplitude mse',
obj_padding=200,
obj_view_crop=None,
units='um',
@@ -450,6 +468,7 @@ class Bragg2DPtycho(CDIModel):
propagate_probe=propagate_probe,
correct_tilt=correct_tilt,
lens=lens,
loss=loss,
obj_view_crop=obj_view_crop,
units=units,
panel_plot_mode=panel_plot_mode,
@@ -534,10 +553,6 @@ class Bragg2DPtycho(CDIModel):
)
def loss(self, sim_data, real_data, mask=None):
return tools.losses.amplitude_mse(real_data, sim_data, mask=mask)
def sim_to_dataset(self, args_list):
# In the future, potentially add more control
# over what metadata is saved (names, etc.)
src/cdtools/models/fancy_ptycho.py
+8 -1
View File
@@ -1,3 +1,4 @@
from functools import partial
import torch as t
from cdtools.models import CDIModel
from cdtools.datasets import Ptycho2DDataset
@@ -219,10 +220,16 @@ class FancyPtycho(CDIModel):
# Here we set the appropriate loss function
if (loss.lower().strip() == 'amplitude mse'
or loss.lower().strip() == 'amplitude_mse'):
self.loss = tools.losses.amplitude_mse
self.loss = partial(tools.losses.amplitude_mse, use_sum=True)
self.loss_normalizer = tools.losses.AmplitudeMSENormalizer()
elif (loss.lower().strip() == 'poisson nll'
or loss.lower().strip() == 'poisson_nll'):
self.loss = tools.losses.poisson_nll
self.loss_normalizer = tools.losses.SimplePoissonNLLNormalizer()
elif (loss.lower().strip() == 'intensity mse'
or loss.lower().strip() == 'intensity_mse'):
self.loss = partial(tools.losses.intensity_mse, use_sum=True)
self.loss_normalizer = tools.losses.IntensityMSENormalizer()
else:
raise KeyError('Specified loss function not supported')
src/cdtools/models/multislice_2d_ptycho.py
+22 -7
View File
@@ -1,3 +1,4 @@
from functools import partial
import torch as t
from cdtools.models import CDIModel
from cdtools.datasets import Ptycho2DDataset
@@ -46,6 +47,7 @@ class Multislice2DPtycho(CDIModel):
fourier_probe=False,
prevent_aliasing=True,
phase_only=False,
loss='amplitude mse',
units='um',
panel_plot_mode=False,
plot_level=1,
@@ -155,9 +157,25 @@ class Multislice2DPtycho(CDIModel):
self.as_prop = tools.propagators.generate_angular_spectrum_propagator(shape, spacing, self.wavelength, self.dz, self.bandlimit)
# Here we set the appropriate loss function
if (loss.lower().strip() == 'amplitude mse'
or loss.lower().strip() == 'amplitude_mse'):
self.loss = partial(tools.losses.amplitude_mse, use_sum=True)
self.loss_normalizer = tools.losses.AmplitudeMSENormalizer()
elif (loss.lower().strip() == 'poisson nll'
or loss.lower().strip() == 'poisson_nll'):
self.loss = tools.losses.poisson_nll
self.loss_normalizer = tools.losses.SimplePoissonNLLNormalizer()
elif (loss.lower().strip() == 'intensity mse'
or loss.lower().strip() == 'intensity_mse'):
self.loss = partial(tools.losses.intensity_mse, use_sum=True)
self.loss_normalizer = tools.losses.IntensityMSENormalizer()
else:
raise KeyError('Specified loss function not supported')
@classmethod
def from_dataset(cls, dataset, dz, nz, probe_convergence_semiangle, padding=0, n_modes=1, dm_rank=None, translation_scale=1, saturation=None, propagation_distance=None, scattering_mode=None, oversampling=1, auto_center=True, bandlimit=None, replicate_slice=False, subpixel=True, exponentiate_obj=True, units='um', fourier_probe=False, phase_only=False, prevent_aliasing=True, probe_support_radius=None, panel_plot_mode=False, plot_level=1):
def from_dataset(cls, dataset, dz, nz, probe_convergence_semiangle, padding=0, n_modes=1, dm_rank=None, translation_scale=1, saturation=None, propagation_distance=None, scattering_mode=None, oversampling=1, auto_center=True, bandlimit=None, replicate_slice=False, subpixel=True, exponentiate_obj=True, units='um', fourier_probe=False, phase_only=False, prevent_aliasing=True, probe_support_radius=None, panel_plot_mode=False, plot_level=1, loss='amplitude_mse'):
wavelength = dataset.wavelength
det_basis = dataset.detector_geometry['basis']
@@ -301,7 +319,9 @@ class Multislice2DPtycho(CDIModel):
phase_only=phase_only,
prevent_aliasing=prevent_aliasing,
panel_plot_mode=panel_plot_mode,
plot_level=plot_level)
plot_level=plot_level,
loss=loss,
)
def interaction(self, index, translations):
@@ -415,11 +435,6 @@ class Multislice2DPtycho(CDIModel):
oversampling=self.oversampling)
def loss(self, sim_data, real_data, mask=None):
return tools.losses.amplitude_mse(real_data, sim_data, mask=mask)
#return tools.losses.poisson_nll(real_data, sim_data, mask=mask,eps=0.5)
def to(self, *args, **kwargs):
super(Multislice2DPtycho, self).to(*args, **kwargs)
self.wavelength = self.wavelength.to(*args,**kwargs)
src/cdtools/models/multislice_ptycho.py
+9 -2
View File
@@ -1,3 +1,4 @@
from functools import partial
import torch as t
from cdtools.models import CDIModel
from cdtools.datasets import Ptycho2DDataset
@@ -164,14 +165,20 @@ class MultislicePtycho(CDIModel):
self.register_buffer('simulate_finite_pixels',
t.as_tensor(simulate_finite_pixels, dtype=bool))
# Here we set the appropriate loss function
if (loss.lower().strip() == 'amplitude mse'
or loss.lower().strip() == 'amplitude_mse'):
self.loss = tools.losses.amplitude_mse
self.loss = partial(tools.losses.amplitude_mse, use_sum=True)
self.loss_normalizer = tools.losses.AmplitudeMSENormalizer()
elif (loss.lower().strip() == 'poisson nll'
or loss.lower().strip() == 'poisson_nll'):
self.loss = tools.losses.poisson_nll
self.loss_normalizer = tools.losses.SimplePoissonNLLNormalizer()
elif (loss.lower().strip() == 'intensity mse'
or loss.lower().strip() == 'intensity_mse'):
self.loss = partial(tools.losses.intensity_mse, use_sum=True)
self.loss_normalizer = tools.losses.IntensityMSENormalizer()
else:
raise KeyError('Specified loss function not supported')
src/cdtools/models/rpi.py
+24 -1
View File
@@ -1,3 +1,4 @@
from functools import partial
import torch as t
from cdtools.models import CDIModel
from cdtools import tools
@@ -54,6 +55,7 @@ class RPI(CDIModel):
exponentiate_obj=False,
phase_only=False,
propagation_distance=0,
loss='amplitude mse',
units='um',
dtype=t.float32,
panel_plot_mode=True,
@@ -146,6 +148,22 @@ class RPI(CDIModel):
self.register_buffer('prop_dir',
t.as_tensor([0, 0, 1], dtype=dtype))
# Here we set the appropriate loss function
if (loss.lower().strip() == 'amplitude mse'
or loss.lower().strip() == 'amplitude_mse'):
self.loss = partial(tools.losses.amplitude_mse, use_sum=True)
self.loss_normalizer = tools.losses.AmplitudeMSENormalizer()
elif (loss.lower().strip() == 'poisson nll'
or loss.lower().strip() == 'poisson_nll'):
self.loss = tools.losses.poisson_nll
self.loss_normalizer = tools.losses.SimplePoissonNLLNormalizer()
elif (loss.lower().strip() == 'intensity mse'
or loss.lower().strip() == 'intensity_mse'):
self.loss = partial(tools.losses.intensity_mse, use_sum=True)
self.loss_normalizer = tools.losses.IntensityMSENormalizer()
else:
raise KeyError('Specified loss function not supported')
@classmethod
def from_dataset(
@@ -164,6 +182,7 @@ class RPI(CDIModel):
exponentiate_obj=False,
phase_only=False,
probe_threshold=0,
loss='amplitude mse',
dtype=t.float32,
panel_plot_mode=True,
plot_level=1,
@@ -259,7 +278,9 @@ class RPI(CDIModel):
phase_only=phase_only,
weight_matrix=weight_matrix,
panel_plot_mode=panel_plot_mode,
plot_level=plot_level)
plot_level=plot_level,
loss=loss,
)
# I don't love this pattern, where I do the "real" obj initialization
# after creating the rpi object. But, I chose this so that I could
@@ -291,6 +312,7 @@ class RPI(CDIModel):
dtype=t.float32,
panel_plot_mode=True,
plot_level=1,
loss='amplitude_mse',
):
complex_dtype = (t.ones([1], dtype=dtype) +
@@ -336,6 +358,7 @@ class RPI(CDIModel):
phase_only=phase_only,
panel_plot_mode=panel_plot_mode,
plot_level=plot_level,
loss=loss,
)
rpi_object.init_obj(initialization)
src/cdtools/models/simple_ptycho.py
+6 -3
View File
@@ -1,3 +1,4 @@
from functools import partial
import torch as t
from cdtools.models import CDIModel
from cdtools import tools
@@ -41,6 +42,10 @@ class SimplePtycho(CDIModel):
self.probe = t.nn.Parameter(probe_guess / self.probe_norm)
self.obj = t.nn.Parameter(obj_guess)
# We register a loss function and an appropriate normalization
self.loss = partial(tools.losses.amplitude_mse, use_sum=True)
self.loss_normalizer = tools.losses.AmplitudeMSENormalizer()
@classmethod
def from_dataset(cls, dataset):
@@ -99,12 +104,10 @@ class SimplePtycho(CDIModel):
def forward_propagator(self, wavefields):
return tools.propagators.far_field(wavefields)
def measurement(self, wavefields):
return tools.measurements.intensity(wavefields)
def loss(self, real_data, sim_data):
return tools.losses.amplitude_mse(real_data, sim_data)
# This lists all the plots to display on a call to model.inspect()
plot_list = [
src/cdtools/reconstructors/base.py
+16 -3
View File
@@ -162,7 +162,15 @@ class Reconstructor:
# 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()
if hasattr(self.model, 'loss_normalizer') and \
self.model.loss_normalizer is not None:
if hasattr(self.model, 'mask'):
mask = self.model.mask
else:
mask = None
self.model.loss_normalizer.accumulate(patterns, mask=mask)
N += 1
def closure():
@@ -216,9 +224,14 @@ class Reconstructor:
return total_loss
# This takes the step for this minibatch
loss += self.optimizer.step(closure).detach().cpu().numpy()
loss += self.optimizer.step(closure).detach()
loss /= normalization
if hasattr(self.model, 'loss_normalizer') and \
self.model.loss_normalizer is not None:
loss = self.model.loss_normalizer.normalize_loss(loss)
# Make sure to return a scalar value which is fully numpy
loss = loss.cpu().numpy()[()]
# We step the scheduler after the full epoch
if self.scheduler is not None:
src/cdtools/tools/losses/losses.py
+263 -25
View File
@@ -8,10 +8,17 @@ maximum likelihood metric for a system with Poisson statistics.
import torch as t
__all__ = ['amplitude_mse', 'intensity_mse', 'poisson_nll']
__all__ = [
'amplitude_mse',
'AmplitudeMSENormalizer',
'intensity_mse',
'IntensityMSENormalizer',
'poisson_nll',
'SimplePoissonNLLNormalizer',
]
def amplitude_mse(intensities, sim_intensities, mask=None):
def amplitude_mse(intensities, sim_intensities, mask=None, use_sum=False):
""" Returns the mean squared error of a simulated dataset's amplitudes
Calculates the mean squared error between a given set of
@@ -20,17 +27,20 @@ def amplitude_mse(intensities, sim_intensities, mask=None):
This function calculates the mean squared error between their
associated amplitudes. Because this is not well defined for negative
numbers, make sure that all the intensities are >0 before using this
loss. Note that this is actually a sum-squared error, because this
formulation makes it vastly simpler to compare error calculations
between reconstructions with different minibatch size. I hope to
find a better way to do this that is more honest with this
cost function, though.
loss.
It can accept intensity and simulated intensity tensors of any shape
as long as their shapes match, and the provided mask array can be
broadcast correctly along them.
This is empirically the most useful loss function for most cases
This is empirically the most useful loss function for most cases where
a photon counting detector cannot be used.
Note that, when used with the AmplitudeMSENormalizer, this function
should be called with use_sum=True, in order to return the sum-squared
error rather than the mean-squared error. This allows for the
AmplitudeMSENormalizer to properly weight the loss arising from minibatches
which may not have equal length.
Parameters
----------
@@ -40,6 +50,8 @@ def amplitude_mse(intensities, sim_intensities, mask=None):
A tensor of simulated detector intensities
mask : torch.Tensor
A mask with ones for pixels to include and zeros for pixels to exclude
use_sum : bool
Default is False. If set to True, actually performs the sum squared error
Returns
-------
@@ -52,16 +64,65 @@ def amplitude_mse(intensities, sim_intensities, mask=None):
# with all the errors working off of the same inputs
if mask is None:
return t.sum((t.sqrt(sim_intensities) -
t.sqrt(intensities))**2)
if use_sum:
return t.sum((t.sqrt(sim_intensities) -
t.sqrt(intensities))**2)
else:
return t.mean((t.sqrt(sim_intensities) -
t.sqrt(intensities))**2)
else:
masked_intensities = intensities.masked_select(mask)
return t.sum((t.sqrt(sim_intensities.masked_select(mask)) -
t.sqrt(masked_intensities))**2)
if use_sum:
return t.sum((t.sqrt(sim_intensities.masked_select(mask)) -
t.sqrt(masked_intensities))**2)
else:
return t.mean((t.sqrt(sim_intensities.masked_select(mask)) -
t.sqrt(masked_intensities))**2)
class AmplitudeMSENormalizer(object):
""" Normalizer for the amplitude MSE loss, used with recon.optimize
This is a normalizer designed for use with the recon.optimize function. The
normalization is done separately from the loss, in order to make it simple to
use different normalization strategies for different loss metrics and to make it
easier to work with different minibatch sizes.
This normalizer accumulates the total number of pixels across all patterns
during the first epoch, then divides the summed loss by this count to
convert from sum-squared error to mean-squared error.
The normalizer is stateful: it completes its accumulation phase on the
first epoch and then applies the same normalization factor for all
subsequent epochs.
Methods
-------
accumulate(patterns, mask=None)
Accumulate the normalization factor (called once per minibatch).
normalize_loss(loss)
Apply the accumulated normalization (called once per epoch).
"""
def intensity_mse(intensities, sim_intensities, mask=None):
def __init__(self):
self.first_pass_complete = False
self.num_pix = 0
def accumulate(self, patterns, mask=None):
if not self.first_pass_complete:
if mask is None:
self.num_pix += patterns.numel()
else:
self.num_pix += patterns.masked_select(mask).numel()
def normalize_loss(self, loss):
if not self.first_pass_complete:
self.first_pass_complete = True
return loss / self.num_pix
def intensity_mse(intensities, sim_intensities, mask=None, use_sum=False):
""" Returns the mean squared error of a simulated dataset's intensities
Calculates the summed mean squared error between a given set of
@@ -72,6 +133,15 @@ def intensity_mse(intensities, sim_intensities, mask=None):
It can accept intensity and simulated intensity tensors of any shape
as long as their shapes match, and the provided mask array can be
broadcast correctly along them.
This is rarely a good loss function for ptychography, but can occasionally
be useful.
Note that, when used with the IntensityMSENormalizer, this function
should be called with use_sum=True, in order to return the sum-squared
error rather than the mean-squared error. This allows for the
IntensityMSENormalizer to properly weight the loss arising from minibatches
which may not have equal length.
Parameters
----------
@@ -81,6 +151,8 @@ def intensity_mse(intensities, sim_intensities, mask=None):
A tensor of simulated detector intensities
mask : torch.Tensor
A mask with ones for pixels to include and zeros for pixels to exclude
use_sum : bool
Default is False. If set to True, actually performs the sum squared error
Returns
-------
@@ -89,15 +161,86 @@ def intensity_mse(intensities, sim_intensities, mask=None):
"""
if mask is None:
return t.sum((sim_intensities - intensities)**2) \
/ intensities.view(-1).shape[0]
if use_sum:
return t.sum((sim_intensities - intensities)**2)
else:
return t.mean((sim_intensities - intensities)**2)
else:
masked_intensities = intensities.masked_select(mask)
return t.sum((sim_intensities.masked_select(mask) -
masked_intensities)**2) \
/ masked_intensities.shape[0]
if use_sum:
return t.sum((sim_intensities.masked_select(mask) -
intensities.masked_select(mask))**2)
else:
return t.mean((sim_intensities.masked_select(mask) -
intensities.masked_select(mask))**2)
class IntensityMSENormalizer(object):
""" Normalizer for the intensity MSE loss, used with recon.optimize
This is a normalizer designed for use with the recon.optimize function. The
normalization is done separately from the loss, in order to make it simple to
use different normalization strategies for different loss metrics and to make it
easier to work with different minibatch sizes.
This normalizer accumulates the total number of pixels across all patterns
during the first epoch, then divides the summed loss by this count to
convert from sum-squared error to mean-squared error.
The normalizer is stateful: it completes its accumulation phase on the
first epoch and then applies the same normalization factor for all
subsequent epochs.
Methods
-------
accumulate(patterns, mask=None)
Accumulate the normalization factor (called once per minibatch).
normalize_loss(loss)
Apply the accumulated normalization (called once per epoch).
"""
def __init__(self):
self.first_pass_complete = False
self.num_pix = 0
def accumulate(self, patterns, mask=None):
"""Accumulate pixel counts from a batch of patterns.
Parameters
----------
patterns : torch.Tensor
A tensor of measured detector patterns
mask : torch.Tensor, optional
A mask with ones for pixels to include and zeros for pixels to
exclude. If provided, only masked pixels are counted.
"""
if not self.first_pass_complete:
if mask is None:
self.num_pix += patterns.numel()
else:
self.num_pix += patterns.masked_select(mask).numel()
def normalize_loss(self, loss):
"""Convert summed loss to mean loss by dividing by pixel count.
Parameters
----------
loss : torch.Tensor
The accumulated summed loss across minibatches in an epoch
Returns
-------
normalized_loss : torch.Tensor
The loss divided by the total number of pixels
"""
if not self.first_pass_complete:
self.first_pass_complete = True
return loss / self.num_pix
def poisson_nll(
intensities,
@@ -124,6 +267,9 @@ def poisson_nll(
The default value of eps is 1e-6 - a nonzero value here helps avoid
divergence of the log function near zero.
This is generally the best loss metric to use for ptychography when
a photon counting detector is used.
Parameters
----------
@@ -135,6 +281,8 @@ def poisson_nll(
A mask with ones for pixels to include and zeros for pixels to exclude
eps : float
Optional, a small number to add to the simulated intensities
subtract_min : bool
Default is False, whether to subtract a min to produce a nonnegative output
Returns
-------
@@ -144,8 +292,7 @@ def poisson_nll(
"""
if mask is None:
nll = t.sum(sim_intensities+eps -
t.xlogy(intensities,sim_intensities+eps)) \
/ intensities.view(-1).shape[0]
t.xlogy(intensities,sim_intensities+eps))
if subtract_min:
nll -= t.sum(intensities - t.xlogy(intensities,intensities))
@@ -155,17 +302,108 @@ def poisson_nll(
masked_sims = sim_intensities.masked_select(mask)
nll = t.sum(masked_sims + eps - \
t.xlogy(masked_intensities, masked_sims+eps)) \
/ masked_intensities.shape[0]
t.xlogy(masked_intensities, masked_sims+eps))
if subtract_min:
nll -= t.nansum(masked_intensities - \
t.xlogy(masked_intensities, masked_intensities)) \
/ masked_intensities.shape[0]
t.xlogy(masked_intensities, masked_intensities))
return nll
class SimplePoissonNLLNormalizer(object):
""" Normalizer for the intensity MSE loss, used with recon.optimize
This is a normalizer designed for use with the recon.optimize function. The
normalization is done separately from the loss, in order to make it simple to
use different normalization strategies for different loss metrics and to make it
easier to work with different minibatch sizes.
This normalizer converts raw Poisson negative log likelihood values into
a statistic that is more interpretable for comparing reconstructions. It
performs two operations:
1. **Offset subtraction**: Subtracts the NLL calculated when comparing
measured patterns to themselves (i.e., poisson_nll(data, data)). This
represents the best-case scenario and makes the loss non-negative.
2. **Normalization scaling**: Divides by 0.5 times the count of non-zero
pixels in the measured patterns. This is because, roughly, each non-zero
pixel is expected to contribute 0.5 to the Poisson NLL, if Poisson noise
were the only relevant source of noise in the data.
The normalizer is stateful: it completes its accumulation phase on the
first epoch by processing all patterns in the data, then applies the
same normalization factors for all subsequent epochs.
Methods
-------
accumulate(patterns, mask=None)
Accumulate the normalization factor (called once per minibatch).
normalize_loss(loss)
Apply the accumulated normalization (called once per epoch).
"""
def __init__(self):
self.first_pass_complete = False
self.sum_nonzero = 0
self.offset = 0
def accumulate(self, patterns, mask=None):
"""Accumulate statistics needed for normalization from a batch.
During the first epoch, this method counts non-zero pixels and
computes the Poisson NLL comparing patterns to themselves, which
defines the offset baseline for the loss.
Parameters
----------
patterns : torch.Tensor
A tensor of measured detector patterns
mask : torch.Tensor, optional
A mask with ones for pixels to include and zeros for pixels to
exclude. If provided, only masked pixels are counted.
"""
if not self.first_pass_complete:
self.offset += poisson_nll(patterns, patterns, mask=mask)
if mask is None:
self.sum_nonzero += t.sum(patterns >= 1)
else:
masked_pats = patterns.masked_select(mask)
self.sum_nonzero += t.sum(masked_pats >= 1)
def normalize_loss(self, loss):
"""Normalize the Poisson NLL for interpretability across datasets.
Parameters
----------
loss : torch.Tensor
The accumulated Poisson NLL across minibatches in an epoch
Returns
-------
normalized_loss : torch.Tensor
The offset-corrected and scaled loss value
"""
if not self.first_pass_complete:
self.normalization = 0.5 * self.sum_nonzero
self.first_pass_complete = True
return (loss - self.offset) / self.normalization
#
# Note: I have two other ideas for how to normalize the Poisson NLL
#
# Idea 2: Use the mean pattern to estimate the expected error
# Idea 3: Use the simulated intensities to estimate it, but use detach
# so it doesn't hit the backward pass
#
def poisson_plus_fixed_nll(
intensities,
tests/models/test_fancy_ptycho.py
+5 -4
View File
@@ -97,8 +97,8 @@ def test_lab_ptycho(lab_ptycho_cxi, reconstruction_device, show_plot):
print('Running reconstruction on provided reconstruction_device,',
reconstruction_device)
#model.to(device=reconstruction_device)
#dataset.get_as(device=reconstruction_device)
model.to(device=reconstruction_device)
dataset.get_as(device=reconstruction_device)
for loss in model.Adam_optimize(50, dataset, lr=0.02, batch_size=10):
print(model.report())
@@ -124,7 +124,7 @@ def test_lab_ptycho(lab_ptycho_cxi, reconstruction_device, show_plot):
plt.close('all')
# If this fails, the reconstruction has gotten worse
assert model.loss_history[-1] < 0.0013
assert model.loss_history[-1] < 0.38
@pytest.mark.slow
@@ -139,6 +139,7 @@ def test_near_field_ptycho(near_field_ptycho_cxi, reconstruction_device, show_pl
near_field=True,
propagation_distance=3.65e-3, # 3.65 downstream from focus
panel_plot_mode=False, # test without panel plot mode
loss='poisson_nll',
)
print('Running reconstruction on provided reconstruction_device,',
@@ -165,4 +166,4 @@ def test_near_field_ptycho(near_field_ptycho_cxi, reconstruction_device, show_pl
plt.close('all')
# If this fails, the reconstruction has gotten worse
assert model.loss_history[-1] < 0.005
assert model.loss_history[-1] < 18
tests/models/test_simple_ptycho.py
+1 -1
View File
@@ -30,4 +30,4 @@ def test_simple_ptycho(lab_ptycho_cxi, reconstruction_device, show_plot):
plt.close('all')
# If this fails, the reconstruction got worse
assert model.loss_history[-1] < 0.013
assert model.loss_history[-1] < 6.5
tests/test_reconstructors.py
+43 -3
View File
@@ -1,3 +1,4 @@
from numbers import Number
import pytest
import time
import cdtools
@@ -39,6 +40,7 @@ def test_Adam_gold_balls(gold_ball_cxi, reconstruction_device, show_plot):
units='um',
probe_fourier_crop=pad,
panel_plot_mode=False, # At least one check without panel plot mode
loss='amplitude_mse',
)
model.translation_offsets.data += 0.7 * \
@@ -114,6 +116,11 @@ def test_Adam_gold_balls(gold_ball_cxi, reconstruction_device, show_plot):
time.sleep(3)
plt.close('all')
# Check that the losses returned in loss_history are not torch tensors
assert isinstance(model.loss_history[-1], Number) and \
not isinstance(model.loss_history[-1], t.Tensor)
# Ensure equivalency between the model reconstructions during the first
# pass, where they should be identical
assert np.allclose(model_recon.loss_history[:epoch_tup[0]], model.loss_history[:epoch_tup[0]])
@@ -122,7 +129,40 @@ def test_Adam_gold_balls(gold_ball_cxi, reconstruction_device, show_plot):
# comes from running a reconstruction when it was working well and
# choosing a rough value. If it triggers this assertion error, something
# changed to make the final quality worse!
assert model_recon.loss_history[-1] < 0.0001
assert model_recon.loss_history[-1] < 0.13
@pytest.mark.slow
def test_intensity_MSE(gold_ball_cxi, reconstruction_device, show_plot):
"""
This test checks that the intensity_mse loss function works end-to-end
with the AdamReconstructor, using the Au particle dataset.
"""
print('\nTesting performance on the standard gold balls dataset ' +
'with intensity_mse loss')
dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(gold_ball_cxi)
model = cdtools.models.FancyPtycho.from_dataset(
dataset,
n_modes=1,
propagation_distance=-3e-6,
units='nm',
loss='intensity_mse',
)
model.to(device=reconstruction_device)
dataset.get_as(device=reconstruction_device)
recon = cdtools.reconstructors.AdamReconstructor(model=model,
dataset=dataset)
t.manual_seed(0)
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
@pytest.mark.slow
@@ -217,7 +257,7 @@ def test_LBFGS_RPI(optical_data_ss_cxi,
# The final loss when testing this was 2.28607e-3. Based on this, we set
# a threshold of 2.3e-3 for the tested loss. If this value has been
# exceeded, the reconstructions have gotten worse.
assert model_recon.loss_history[-1] < 0.0023
assert model_recon.loss_history[-1] < 0.14
@pytest.mark.slow
@@ -330,4 +370,4 @@ def test_SGD_gold_balls(gold_ball_cxi, reconstruction_device, show_plot):
# The final loss when testing this was 7.12188e-4. Based on this, we set
# 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.00072
assert model.loss_history[-1] < 0.95
tests/tools/test_losses.py
+47 -10
View File
@@ -1,5 +1,6 @@
import numpy as np
import torch as t
from scipy.special import xlogy
from cdtools.tools import losses
@@ -20,15 +21,34 @@ def test_amplitude_mse():
# First, test without a mask
np_result = np.sum((np.sqrt(data) - np.sqrt(sim))**2)
# np_result /= data.size
torch_result = losses.amplitude_mse(t.from_numpy(data), t.from_numpy(sim))
torch_result = losses.amplitude_mse(t.from_numpy(data), t.from_numpy(sim),
use_sum=True)
assert np.isclose(np_result, np.take(torch_result.numpy(), 0))
# Then, test with a mask
np_result = np.sum(mask * (np.sqrt(data) - np.sqrt(sim))**2)
# np_result /= np.count_nonzero(mask * np.ones_like(data))
torch_result = losses.amplitude_mse(t.from_numpy(data), t.from_numpy(sim), mask=t.from_numpy(mask))
torch_result = losses.amplitude_mse(t.from_numpy(data), t.from_numpy(sim),
mask=t.from_numpy(mask), use_sum=True)
assert np.isclose(np_result, np.take(torch_result.numpy(), 0))
# Now, test the version with use_sum=False, the default
# First, test without a mask
np_result = np.mean((np.sqrt(data) - np.sqrt(sim))**2)
# np_result /= data.size
torch_result = losses.amplitude_mse(t.from_numpy(data), t.from_numpy(sim))
assert np.isclose(np_result, np.take(torch_result.numpy(), 0))
# Then, test with a mask. Note that with a mask, the masked pixels
# should not contribute to the denominator for the mean.
np_result = np.sum(mask * (np.sqrt(data) - np.sqrt(sim))**2)
np_result /= np.count_nonzero(mask * np.ones_like(data))
torch_result = losses.amplitude_mse(t.from_numpy(data), t.from_numpy(sim),
mask=t.from_numpy(mask), use_sum=False)
assert np.isclose(np_result, np.take(torch_result.numpy(), 0))
def test_intensity_mse():
# Make some fake data
@@ -38,6 +58,20 @@ def test_intensity_mse():
# and define a simple mask that needs to be broadcast
mask = (np.random.rand(100, 100) > 0.1).astype(bool)
# First, test without a mask
np_result = np.sum((data - sim)**2)
torch_result = losses.intensity_mse(t.from_numpy(data), t.from_numpy(sim),
use_sum=True)
assert np.isclose(np_result, np.take(torch_result.numpy(), 0))
# Then, test with a mask
np_result = np.sum(mask * (data - sim)**2)
torch_result = losses.intensity_mse(t.from_numpy(data), t.from_numpy(sim),
mask=t.from_numpy(mask), use_sum=True)
assert np.isclose(np_result, np.take(torch_result.numpy(), 0))
# Now, test the version with use_sum=False, the default
# First, test without a mask
np_result = np.sum((data - sim)**2)
np_result /= data.size
@@ -47,27 +81,30 @@ def test_intensity_mse():
# Then, test with a mask
np_result = np.sum(mask * (data - sim)**2)
np_result /= np.count_nonzero(mask * np.ones_like(data))
torch_result = losses.intensity_mse(t.from_numpy(data), t.from_numpy(sim), mask=t.from_numpy(mask))
torch_result = losses.intensity_mse(t.from_numpy(data), t.from_numpy(sim),
mask=t.from_numpy(mask), use_sum=False)
assert np.isclose(np_result, np.take(torch_result.numpy(), 0))
def test_poisson_nll():
# Make some fake data
data = np.random.rand(10, 100, 100)
# And add some noise to it
# Make some fake data spread over a realistic photon-count range,
# with ~5% of pixels set to zero
data = 10 * np.random.rand(10, 100, 100)
data[np.random.rand(10, 100, 100) < 0.05] = 0
# Add some noise, but set ~5% of sim pixels to exactly match data
sim = data + 0.1 * np.random.rand(10, 100, 100)
exact_match = np.random.rand(10, 100, 100) < 0.05
sim[exact_match] = data[exact_match]
# and define a simple mask that needs to be broadcast
mask = (np.random.rand(100, 100) > 0.1).astype(bool)
# First, test without a mask
np_result = np.sum(sim - data * np.log(sim))
np_result /= data.size
np_result = np.sum(sim - xlogy(data, sim))
torch_result = losses.poisson_nll(t.from_numpy(data), t.from_numpy(sim), eps=0)
assert np.isclose(np_result, np.take(torch_result.numpy(), 0))
# Then, test with a mask
np_result = np.sum(mask * (sim - data * np.log(sim)))
np_result /= np.count_nonzero(mask * np.ones_like(data))
np_result = np.sum(mask * (sim - xlogy(data, sim)))
torch_result = losses.poisson_nll(t.from_numpy(data), t.from_numpy(sim),
mask=t.from_numpy(mask), eps=0)
assert np.isclose(np_result, np.take(torch_result.numpy(), 0))