diff --git a/examples/gold_ball_ptycho_spawn.py b/examples/gold_ball_ptycho_spawn.py new file mode 100644 index 0000000..ae13c9e --- /dev/null +++ b/examples/gold_ball_ptycho_spawn.py @@ -0,0 +1,91 @@ +import cdtools +from matplotlib import pyplot as plt +import torch as t + +# If you're noticing that the multi-GPU job is hanging (especially with 100% +# GPU use across all participating devices), you might want to try disabling +# the environment variable NCCL_P2P_DISABLE. +import os +os.environ['NCCL_P2P_DISABLE'] = str(int(True)) + + +# The entire reconstruction script needs to be wrapped in a function +def reconstruct(rank, world_size): + + # In the multigpu setup, we need to explicitly define the rank and + # world_size. The master address and master port should also be + # defined if we don't specify the `init_method` parameter. + cdtools.tools.multigpu.setup(rank=rank, + world_size=world_size, + master_addr='localhost', + master_port='6666') + + filename = 'example_data/AuBalls_700ms_30nmStep_3_6SS_filter.cxi' + dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(filename) + + pad = 10 + dataset.pad(pad) + dataset.inspect() + model = cdtools.models.FancyPtycho.from_dataset( + dataset, + n_modes=3, + probe_support_radius=50, + propagation_distance=2e-6, + units='um', + probe_fourier_crop=pad + ) + model.translation_offsets.data += 0.7 * \ + t.randn_like(model.translation_offsets) + model.weights.requires_grad = False + + # We need to manually define the rank parameter for the model, or else + # all plots will be duplicated by the number of GPUs used. + model.rank = rank + + device = 'cuda' + model.to(device=device) + dataset.get_as(device=device) + + # Rank and world_size also needs to be explicitly defined here + recon = cdtools.reconstructors.AdamReconstructor(model, + dataset, + rank=rank, + world_size=world_size) + + with model.save_on_exception( + 'example_reconstructions/gold_balls_earlyexit.h5', dataset): + + for loss in recon.optimize(20, lr=0.005, batch_size=50): + if rank == 0: + print(model.report()) + if model.epoch % 10 == 0: + model.inspect(dataset) + + for loss in recon.optimize(50, lr=0.002, batch_size=100, + schedule=True): + if rank == 0: + print(model.report()) + + if model.epoch % 10 == 0: + model.inspect(dataset) + + for loss in recon.optimize(100, lr=0.001, batch_size=100, + schedule=True): + if rank == 0: + print(model.report()) + if model.epoch % 10 == 0: + model.inspect(dataset) + + cdtools.tools.multigpu.cleanup() + + model.tidy_probes() + model.save_to_h5('example_reconstructions/gold_balls.h5', dataset) + model.inspect(dataset) + model.compare(dataset) + plt.show() + + +if __name__ == '__main__': + # Specify the number of GPUs we want to use, then spawn the multi-GPU job + ngpus = 2 + t.multiprocessing.spawn(reconstruct, args=(ngpus,), nprocs=ngpus) diff --git a/examples/gold_ball_ptycho_torchrun.py b/examples/gold_ball_ptycho_torchrun.py new file mode 100644 index 0000000..834cb84 --- /dev/null +++ b/examples/gold_ball_ptycho_torchrun.py @@ -0,0 +1,70 @@ +import cdtools +from matplotlib import pyplot as plt +import torch as t + +# At the beginning of the script we need to setup the multi-GPU job +# by initializing the process group and sycnronizing the RNG seed +# across all participating GPUs. +cdtools.tools.multigpu.setup() + +# To avoid redundant print statements, we first grab the GPU "rank" +# (an ID number between 0 and max number of GPUs minus 1). +rank = cdtools.tools.multigpu.get_rank() + +filename = 'example_data/AuBalls_700ms_30nmStep_3_6SS_filter.cxi' +dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(filename) + +pad = 10 +dataset.pad(pad) +dataset.inspect() +model = cdtools.models.FancyPtycho.from_dataset( + dataset, + n_modes=3, + probe_support_radius=50, + propagation_distance=2e-6, + units='um', + probe_fourier_crop=pad +) +model.translation_offsets.data += 0.7 * t.randn_like(model.translation_offsets) +model.weights.requires_grad = False +device = 'cuda' +model.to(device=device) +dataset.get_as(device=device) + +recon = cdtools.reconstructors.AdamReconstructor(model, dataset) + +with model.save_on_exception( + 'example_reconstructions/gold_balls_earlyexit.h5', dataset): + + for loss in recon.optimize(20, lr=0.005, batch_size=50): + # We ensure that only the GPU with rank of 0 runs print statement. + if rank == 0: + print(model.report()) + + # But we don't need to do rank checking for any plotting- or saving- + # related methods; this checking is handled internernally. + if model.epoch % 10 == 0: + model.inspect(dataset) + + for loss in recon.optimize(50, lr=0.002, batch_size=100): + if rank == 0: + print(model.report()) + + if model.epoch % 10 == 0: + model.inspect(dataset) + + for loss in recon.optimize(100, lr=0.001, batch_size=100, schedule=True): + if rank == 0: + print(model.report()) + if model.epoch % 10 == 0: + model.inspect(dataset) + +# After the reconstruction is completed, we need to cleanup things by +# destroying the process group. +cdtools.tools.multigpu.cleanup() + +model.tidy_probes() +model.save_to_h5('example_reconstructions/gold_balls.h5', dataset) +model.inspect(dataset) +model.compare(dataset) +plt.show()