diff --git a/examples/gold_ball_ptycho.py b/examples/gold_ball_ptycho.py index 32a05a1..ff0c54f 100644 --- a/examples/gold_ball_ptycho.py +++ b/examples/gold_ball_ptycho.py @@ -50,7 +50,7 @@ recon = cdtools.reconstructors.AdamReconstructor(model, dataset) # This will save out the intermediate results if an exception is thrown # during the reconstruction with model.save_on_exception( - 'example_reconstructions/gold_balls_earlyexit.h5', dataset): + 'example_reconstructions/gold_balls_earlyexit.h5'): for loss in recon.optimize(20, lr=0.005, batch_size=50): print(model.report()) diff --git a/src/cdtools/models/base.py b/src/cdtools/models/base.py index 127d4de..4761b6e 100644 --- a/src/cdtools/models/base.py +++ b/src/cdtools/models/base.py @@ -229,10 +229,10 @@ class CDIModel(t.nn.Module): or produced directly in memory. """ state_dict = nested_dict_to_torch(results_dict['state_dict']) - self.load_state_dict(state_dict) - self.loss_history = list(results_dict['loss_history']) - self.epoch = int(results_dict['epoch']) - self.training_history = str(results_dict['training_history']) + self.load_state_dict(state_dict, strict=False) + self.loss_history = list(results_dict.get('loss_history', [])) + self.epoch = int(results_dict.get('epoch', 0)) + self.training_history = str(results_dict.get('training_history', '')) @classmethod diff --git a/src/cdtools/models/fancy_ptycho.py b/src/cdtools/models/fancy_ptycho.py index c0bfeed..bb9eaa2 100644 --- a/src/cdtools/models/fancy_ptycho.py +++ b/src/cdtools/models/fancy_ptycho.py @@ -1252,7 +1252,6 @@ class FancyPtycho(CDIModel): A fully reconstructed model with all parameters, buffers, and training metadata restored. """ - import numpy as np sd = results_dict['state_dict'] # For optional Parameters (translation_offsets, weights, qe_mask, etc.), @@ -1271,8 +1270,8 @@ class FancyPtycho(CDIModel): obj_basis=sd['obj_basis'], probe_guess=sd['probe'], # normalized; probe_norm restored by _load_results_dict obj_guess=sd['obj'], - surface_normal=sd['surface_normal'], - min_translation=sd['min_translation'], + surface_normal=sd.get('surface_normal', np.array([0., 0., 1.])), + min_translation=sd.get('min_translation', np.array([0., 0.])), background=sd['background'], # sqrt form; restored exactly by _load_results_dict probe_basis=sd.get('probe_basis'), translation_offsets=translation_offsets, # overwritten by _load_results_dict @@ -1281,15 +1280,15 @@ class FancyPtycho(CDIModel): weights=sd.get('weights'), # overwritten by _load_results_dict qe_mask=sd.get('qe_mask'), # overwritten by _load_results_dict saturation=sd.get('saturation'), - translation_scale=float(sd['translation_scale']), - oversampling=int(sd['oversampling']), - fourier_probe=bool(sd['fourier_probe']), + translation_scale=float(sd.get('translation_scale', 1.0)), + oversampling=int(sd.get('oversampling', 1)), + fourier_probe=bool(sd.get('fourier_probe', False)), loss=results_dict.get('loss_function', 'amplitude mse'), - simulate_probe_translation=bool(sd['simulate_probe_translation']), - simulate_finite_pixels=bool(sd['simulate_finite_pixels']), - exponentiate_obj=bool(sd['exponentiate_obj']), - phase_only=bool(sd['phase_only']), - near_field=bool(sd['near_field']), + simulate_probe_translation=bool(sd.get('simulate_probe_translation', False)), + simulate_finite_pixels=bool(sd.get('simulate_finite_pixels', False)), + exponentiate_obj=bool(sd.get('exponentiate_obj', False)), + phase_only=bool(sd.get('phase_only', False)), + near_field=bool(sd.get('near_field', False)), angular_spectrum_propagator=sd.get('angular_spectrum_propagator'), inv_angular_spectrum_propagator=sd.get('inv_angular_spectrum_propagator'), translations=sd.get('original_translations'), diff --git a/tests/models/test_fancy_ptycho.py b/tests/models/test_fancy_ptycho.py index f1bf147..25e5e0c 100644 --- a/tests/models/test_fancy_ptycho.py +++ b/tests/models/test_fancy_ptycho.py @@ -229,3 +229,40 @@ def test_fancy_ptycho_from_results_dict(lab_ptycho_cxi, tmp_path): 'from_results_dict: forward pass output mismatch' assert t.allclose(original_out, loaded_h5_out), \ 'from_results_h5: forward pass output mismatch' + + +def test_fancy_ptycho_from_results_dict_with_missing_keys(lab_ptycho_cxi): + dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(lab_ptycho_cxi) + + t.manual_seed(42) + model = cdtools.models.FancyPtycho.from_dataset(dataset, n_modes=1) + results_dict = model.save_results() + + # Strip top-level training metadata + for key in ('loss_history', 'epoch', 'training_history'): + results_dict.pop(key, None) + + # Strip defaultable state_dict keys + sd_keys_to_strip = ( + 'exponentiate_obj', 'phase_only', 'near_field', 'fourier_probe', + 'simulate_probe_translation', 'simulate_finite_pixels', + 'translation_scale', 'oversampling', 'surface_normal', 'min_translation', + ) + for key in sd_keys_to_strip: + results_dict['state_dict'].pop(key, None) + + loaded = cdtools.models.FancyPtycho.from_results_dict(results_dict) + + assert loaded.loss_history == [] + assert loaded.epoch == 0 + assert loaded.training_history == '' + assert bool(loaded.exponentiate_obj) == False + assert bool(loaded.phase_only) == False + assert bool(loaded.near_field) == False + assert bool(loaded.fourier_probe) == False + assert bool(loaded.simulate_probe_translation) == False + assert bool(loaded.simulate_finite_pixels) == False + assert float(loaded.translation_scale) == 1.0 + assert int(loaded.oversampling) == 1 + assert t.allclose(loaded.surface_normal, t.tensor([0., 0., 1.])) + assert t.allclose(loaded.min_translation, t.tensor([0., 0.]))