From 29d14c9537ae023db572ee5e6df71506bb63721c Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Thu, 18 Apr 2024 14:22:14 +0200 Subject: [PATCH] Get a basic version of checkpointing working --- src/cdtools/models/base.py | 102 ++++++++++++++++++++++++++++++++- src/cdtools/tools/data/data.py | 49 ++++++++++++++++ 2 files changed, 148 insertions(+), 3 deletions(-) diff --git a/src/cdtools/models/base.py b/src/cdtools/models/base.py index 8041618..754fe77 100644 --- a/src/cdtools/models/base.py +++ b/src/cdtools/models/base.py @@ -40,7 +40,7 @@ import time from scipy import io from contextlib import contextmanager from .complex_lbfgs import MyLBFGS -from cdtools.tools.data import nested_dict_to_h5 +from cdtools.tools.data import nested_dict_to_h5, h5_to_nested_dict, nested_dict_to_numpy, nested_dict_to_torch __all__ = ['CDIModel'] @@ -173,7 +173,7 @@ class CDIModel(t.nn.Module): results : dict A dictionary containing all the parameters and buffers of the model, i.e. the result of self.state_dict(), converted to numpy. """ - state_dict = {k: v.cpu().numpy() for k, v in self.state_dict().items()} + state_dict = nested_dict_to_numpy(self.state_dict()) return { 'state_dict': state_dict, @@ -195,6 +195,7 @@ class CDIModel(t.nn.Module): Accepts any additional args that model.save_results needs, for this model """ return nested_dict_to_h5(filename, self.save_results(*args)) + @contextmanager def save_on_exit(self, filename, *args, exception_filename=None): @@ -246,7 +247,78 @@ class CDIModel(t.nn.Module): print(filename, flush=True) raise + + def use_checkpoints(self, job_id, checkpoint_file_stem): + self.current_checkpoint_id = 0 + self.job_id = job_id + self.checkpoint_file_stem = checkpoint_file_stem + + def skip_computation(self): + """Returns true if computations should be skipped due to checkpointing + + This is used internally by model.AD_optimize to make the checkpointing + system work, but it is also useful to suppress printing when + computations are being skipped + """ + if (hasattr(self, 'current_checkpoint_id') + and self.current_checkpoint_id != self.job_id): + return True + else: + return False + + def save_checkpoint(self, *args, checkpoint_file=None): + checkpoint = self.save_results(*args) + if (hasattr(self, 'current_optimizer') + and self.current_optimizer is not None): + checkpoint['optimizer_state_dict'] = \ + self.current_optimizer.state_dict() + + if checkpoint_file is None: + checkpoint_file = ( + self.checkpoint_file_stem + + '_' + str(self.current_checkpoint_id) + '.pt' + ) + + t.save(checkpoint, checkpoint_file) + #nested_dict_to_h5(checkpoint_file, checkpoint) + + def load_checkpoint(self, checkpoint_file=None): + if checkpoint_file is None: + checkpoint_file = ( + self.checkpoint_file_stem + + '_' + str(self.current_checkpoint_id) + '.pt' + ) + + checkpoint = t.load(checkpoint_file) + + state_dict = nested_dict_to_torch(checkpoint['state_dict']) + self.load_state_dict(state_dict) + + self.loss_history = list(checkpoint['loss_history']) + self.training_history = checkpoint['training_history'] + self.epoch = checkpoint['epoch'] + + if (hasattr(self, 'current_optimizer') + and self.current_optimizer is not None): + self.current_optimizer.load_state_dict( + checkpoint['optimizer_state_dict']) + + + def checkpoint(self, *args): + if not hasattr(self, 'current_checkpoint_id'): + raise Exception('You must initialize checkpoints with model.use_checkpoints() before calling model.checkpoint()') + + if self.current_checkpoint_id == self.job_id: + self.save_checkpoint(*args) + exit() + elif self.current_checkpoint_id + 1 == self.job_id: + self.load_checkpoint() + + self.current_checkpoint_id += 1 + + + def AD_optimize(self, iterations, data_loader, optimizer,\ scheduler=None, regularization_factor=None, thread=True, calculation_width=10): @@ -358,11 +430,24 @@ class CDIModel(t.nn.Module): self.training_history += self.report() + '\n' return loss + # We store the current optimizer as a model parameter so that + # it can be saved and loaded for checkpointing + self.current_optimizer = optimizer + # If we don't want to run in a different thread, this is easy if not thread: for it in range(iterations): + if self.skip_computation(): + self.epoch = self.epoch + 1 + if len(self.loss_history) >= 1: + yield self.loss_history[-1] + else: + yield float('nan') + continue + yield run_epoch() - + + # But if we do want to thread, it's annoying: else: # Here we set up the communication with the computation thread @@ -378,6 +463,14 @@ class CDIModel(t.nn.Module): # And this actually starts and monitors the thread for it in range(iterations): + if self.skip_computation(): + self.epoch = self.epoch + 1 + if len(self.loss_history) >= 1: + yield self.loss_history[-1] + else: + yield float('nan') + continue + calc = threading.Thread(target=target, name='calculator', daemon=True) try: calc.start() @@ -401,7 +494,10 @@ class CDIModel(t.nn.Module): yield res + # And finally, we unset the current optimizer: + self.current_optimizer = None + def Adam_optimize( self, iterations, diff --git a/src/cdtools/tools/data/data.py b/src/cdtools/tools/data/data.py index aab0a5f..19e7467 100644 --- a/src/cdtools/tools/data/data.py +++ b/src/cdtools/tools/data/data.py @@ -36,6 +36,8 @@ __all__ = ['get_entry_info', 'add_data', 'add_shot_to_shot_info', 'add_ptycho_translations', + 'nested_dict_to_numpy', + 'nested_dict_to_torch', 'nested_dict_to_h5', 'h5_to_nested_dict', ] @@ -864,3 +866,50 @@ def h5_to_nested_dict(h5_file): raise ValueError(f'{value} could not be interpreted sensibly') return d + + +def nested_dict_to_numpy(d): + + new_dict = {} + for key in d.keys(): + value = d[key] + if isinstance(value, numbers.Number): + new_dict[key] = value + # bools are an instance of number, but not np.bool_... + elif isinstance(value, np.bool_): + new_dict[key] = value + elif isinstance(value, np.ndarray): + new_dict[key] = value + elif t.is_tensor(value): + new_dict[key] = value.cpu().numpy() + elif isinstance(value, str): + new_dict[key] = value + elif isinstance(value, Mapping): + new_dict[key] = nested_dict_to_numpy(value) + else: + raise ValueError(f'{value} is not a number, numpy array, torch tensor, or mapping') + + return new_dict + +def nested_dict_to_torch(d): + + new_dict = {} + for key in d.keys(): + value = d[key] + if isinstance(value, numbers.Number): + new_dict[key] = t.as_tensor(value) + # bools are an instance of number, but not np.bool_... + elif isinstance(value, np.bool_): + new_dict[key] = t.as_tensor(value) + elif isinstance(value, np.ndarray): + new_dict[key] = t.as_tensor(value) + elif t.is_tensor(value): + new_dict[key] = value + elif isinstance(value, str): + new_dict[key] = value + elif isinstance(value, Mapping): + new_dict[key] = nested_dict_to_numpy(value) + else: + raise ValueError(f'{value} is not a number, numpy array, torch tensor, or mapping') + + return new_dict