diff --git a/src/cdtools/models/base.py b/src/cdtools/models/base.py index df347b6..e237bbe 100644 --- a/src/cdtools/models/base.py +++ b/src/cdtools/models/base.py @@ -40,6 +40,7 @@ import time from scipy import io from contextlib import contextmanager from cdtools.tools.data import nested_dict_to_h5, h5_to_nested_dict, nested_dict_to_numpy, nested_dict_to_torch +from cdtools.tools import multigpu from cdtools.reconstructors import AdamReconstructor, LBFGSReconstructor, SGDReconstructor from cdtools.datasets import CDataset from typing import List, Union, Tuple @@ -65,6 +66,16 @@ class CDIModel(t.nn.Module): self.training_history = '' self.epoch = 0 + # This is a flag related to multi-GPU operation which prevents + # saving/plotting functions from being executed on GPUs outside of + # rank 0. + self.rank = multigpu.get_rank() + + # Keep track of the time each loss history point was taken relative to + # the initialization of this model. + self.INITIAL_TIME = time.time() + self.loss_times = [] + def from_dataset(self, dataset): raise NotImplementedError() @@ -197,7 +208,9 @@ class CDIModel(t.nn.Module): *args Accepts any additional args that model.save_results needs, for this model """ - return nested_dict_to_h5(filename, self.save_results(*args)) + # FOR MULTI-GPU: Only run this method if it's called by the rank 0 GPU + if self.rank == 0: + return nested_dict_to_h5(filename, self.save_results(*args)) @contextmanager @@ -219,11 +232,16 @@ class CDIModel(t.nn.Module): """ try: yield - self.save_to_h5(filename, *args) + + # FOR MULTI-GPU: Only run this method if it's called by the rank 0 GPU + if self.rank == 0: + self.save_to_h5(filename, *args) except: - if exception_filename is None: - exception_filename = filename - self.save_to_h5(exception_filename, *args) + # FOR MULTI-GPU: Only run this method if it's called by the rank 0 GPU + if self.rank == 0: + if exception_filename is None: + exception_filename = filename + self.save_to_h5(exception_filename, *args) raise @contextmanager @@ -245,9 +263,11 @@ class CDIModel(t.nn.Module): try: yield except: - self.save_to_h5(filename, *args) - print('Intermediate results saved under name:') - print(filename, flush=True) + # FOR MULTI-GPU: Only run this method if it's called by the rank 0 GPU + if self.rank == 0: + self.save_to_h5(filename, *args) + print('Intermediate results saved under name:') + print(filename, flush=True) raise @@ -270,6 +290,11 @@ class CDIModel(t.nn.Module): return False def save_checkpoint(self, *args, checkpoint_file=None): + # FOR MULTI-GPU: Dont run this block of code if it isn't + # called by the rank 0 GPU + if self.rank != 0: + return + checkpoint = self.save_results(*args) if (hasattr(self, 'current_optimizer') and self.current_optimizer is not None): @@ -332,7 +357,9 @@ class CDIModel(t.nn.Module): subset: Union[int, List[int]] = None, regularization_factor: Union[float, List[float]] = None, thread=True, - calculation_width=10 + calculation_width=10, + rank=None, + world_size=None ): """ Runs a round of reconstruction using the Adam optimizer from @@ -373,14 +400,29 @@ class CDIModel(t.nn.Module): 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. - + rank : int + Optional, GPU rank assigned during multi-GPU operations. If this + parameter is None, it will be redefined based on the `RANK` + environment variable. If this environment variable doesn't exist, + single-GPU operation will be assumed and a rank of 0 will + automatically be assigned. + world_size : int + Optional, the number of participating GPUs during multi-GPU + operations. If this parameter is None, it will be redefined based on + the `WORLD_SIZE` environment variable. If this environment variable + doesn't exist,single-GPU operation will be assumed and a world_size of + 1 will automatically be assigned. """ + self.rank = rank if rank is not None else multigpu.get_rank() + reconstructor = AdamReconstructor( model=self, dataset=dataset, subset=subset, + rank=self.rank, + world_size=world_size ) - + # Run some reconstructions return reconstructor.optimize( iterations=iterations, @@ -578,6 +620,11 @@ class CDIModel(t.nn.Module): Whether to update existing plots or plot new ones """ + # FOR MULTI-GPU: Dont run this block of code if it isn't + # called by the rank 0 GPU + if self.rank != 0: + return + # We find or create all the figures first_update = False if update and hasattr(self, 'figs') and self.figs: @@ -660,7 +707,11 @@ class CDIModel(t.nn.Module): extention : strategy Default is .eps, the file extension to save with. """ - + # FOR MULTI-GPU: Dont run this block of code if it isn't + # called by the rank 0 GPU + if self.rank != 0: + return + if hasattr(self, 'figs') and self.figs: figs = self.figs else: @@ -687,6 +738,10 @@ class CDIModel(t.nn.Module): logarithmic : bool, default: False Whether to plot the diffraction on a logarithmic scale """ + # FOR MULTI-GPU: Dont run this block of code if it isn't + # called by the rank 0 GPU + if self.rank != 0: + return fig, axes = plt.subplots(1,3,figsize=(12,5.3)) fig.tight_layout(rect=[0.02, 0.09, 0.98, 0.96])