diff --git a/docs/source/tutorial.rst b/docs/source/tutorial.rst index 8a2fa86..58f3464 100644 --- a/docs/source/tutorial.rst +++ b/docs/source/tutorial.rst @@ -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 diff --git a/examples/near_field_ptycho.py b/examples/near_field_ptycho.py index af0012d..47da91d 100644 --- a/examples/near_field_ptycho.py +++ b/examples/near_field_ptycho.py @@ -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 ) diff --git a/examples/tutorial_simple_ptycho.py b/examples/tutorial_simple_ptycho.py index 67544d7..e5ff284 100644 --- a/examples/tutorial_simple_ptycho.py +++ b/examples/tutorial_simple_ptycho.py @@ -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 = [ diff --git a/src/cdtools/models/bragg_2d_ptycho.py b/src/cdtools/models/bragg_2d_ptycho.py index 881759f..43f6360 100644 --- a/src/cdtools/models/bragg_2d_ptycho.py +++ b/src/cdtools/models/bragg_2d_ptycho.py @@ -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.) diff --git a/src/cdtools/models/fancy_ptycho.py b/src/cdtools/models/fancy_ptycho.py index a72f3df..ea1ddd7 100644 --- a/src/cdtools/models/fancy_ptycho.py +++ b/src/cdtools/models/fancy_ptycho.py @@ -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') diff --git a/src/cdtools/models/multislice_2d_ptycho.py b/src/cdtools/models/multislice_2d_ptycho.py index d6663be..318b0fc 100644 --- a/src/cdtools/models/multislice_2d_ptycho.py +++ b/src/cdtools/models/multislice_2d_ptycho.py @@ -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) diff --git a/src/cdtools/models/multislice_ptycho.py b/src/cdtools/models/multislice_ptycho.py index 3e79b22..03b1b52 100644 --- a/src/cdtools/models/multislice_ptycho.py +++ b/src/cdtools/models/multislice_ptycho.py @@ -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') diff --git a/src/cdtools/models/rpi.py b/src/cdtools/models/rpi.py index 3d8ac6c..4bce45a 100644 --- a/src/cdtools/models/rpi.py +++ b/src/cdtools/models/rpi.py @@ -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) diff --git a/src/cdtools/models/simple_ptycho.py b/src/cdtools/models/simple_ptycho.py index 734ba20..eb99802 100644 --- a/src/cdtools/models/simple_ptycho.py +++ b/src/cdtools/models/simple_ptycho.py @@ -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 = [ diff --git a/src/cdtools/reconstructors/base.py b/src/cdtools/reconstructors/base.py index 4ad7931..ed8afc1 100644 --- a/src/cdtools/reconstructors/base.py +++ b/src/cdtools/reconstructors/base.py @@ -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: diff --git a/src/cdtools/tools/losses/losses.py b/src/cdtools/tools/losses/losses.py index 44e5c5d..e8b8c19 100644 --- a/src/cdtools/tools/losses/losses.py +++ b/src/cdtools/tools/losses/losses.py @@ -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, diff --git a/tests/models/test_fancy_ptycho.py b/tests/models/test_fancy_ptycho.py index 79a5caf..93b3280 100644 --- a/tests/models/test_fancy_ptycho.py +++ b/tests/models/test_fancy_ptycho.py @@ -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 diff --git a/tests/models/test_simple_ptycho.py b/tests/models/test_simple_ptycho.py index 225e463..bcc5e46 100644 --- a/tests/models/test_simple_ptycho.py +++ b/tests/models/test_simple_ptycho.py @@ -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 diff --git a/tests/test_reconstructors.py b/tests/test_reconstructors.py index 3f9245b..ff34a7a 100644 --- a/tests/test_reconstructors.py +++ b/tests/test_reconstructors.py @@ -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 diff --git a/tests/tools/test_losses.py b/tests/tools/test_losses.py index b36b85d..836413c 100644 --- a/tests/tools/test_losses.py +++ b/tests/tools/test_losses.py @@ -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))