From cb954a60672bdc5f6c167c52ba70c7010df667c3 Mon Sep 17 00:00:00 2001 From: abe Date: Sun, 16 Oct 2022 17:06:44 -0700 Subject: [PATCH] Fix the annoying warning printing from interactions and get a minimum usable script going for simpleptycho and distributeddataparallel --- examples/demo.py | 43 ++++++++++ examples/multi_gpu_simple_ptycho.py | 82 +++++++++++++++++++ examples/simple_ptycho.py | 67 +++++++++++++-- src/cdtools/models/simple_ptycho.py | 13 ++- .../tools/interactions/interactions.py | 18 ++-- 5 files changed, 205 insertions(+), 18 deletions(-) create mode 100644 examples/demo.py create mode 100644 examples/multi_gpu_simple_ptycho.py diff --git a/examples/demo.py b/examples/demo.py new file mode 100644 index 0000000..5d7bf0a --- /dev/null +++ b/examples/demo.py @@ -0,0 +1,43 @@ + + +import torch +import torch.distributed as dist +import torch.nn as nn +import torch.optim as optim + +from torch.nn.parallel import DistributedDataParallel as DDP + +class ToyModel(nn.Module): + def __init__(self): + super(ToyModel, self).__init__() + self.net1 = nn.Linear(10, 10) + self.relu = nn.ReLU() + self.net2 = nn.Linear(10, 5) + + def forward(self, x): + return self.net2(self.relu(self.net1(x))) + + +def demo_basic(): + dist.init_process_group("nccl") + rank = dist.get_rank() + print(f"Start running basic DDP example on rank {rank}.") + + # create model and move it to GPU with id rank + device_id = rank % torch.cuda.device_count() + model = ToyModel().to(device_id) + ddp_model = DDP(model, device_ids=[device_id]) + + loss_fn = nn.MSELoss() + optimizer = optim.SGD(ddp_model.parameters(), lr=0.001) + + for i in range(100): + print(i) + optimizer.zero_grad() + outputs = ddp_model(torch.randn(20, 10)) + labels = torch.randn(20, 5).to(device_id) + loss_fn(outputs, labels).backward() + optimizer.step() + +if __name__ == "__main__": + demo_basic() diff --git a/examples/multi_gpu_simple_ptycho.py b/examples/multi_gpu_simple_ptycho.py new file mode 100644 index 0000000..ad7e646 --- /dev/null +++ b/examples/multi_gpu_simple_ptycho.py @@ -0,0 +1,82 @@ +import cdtools +import torch as t +from cdtools.models.complex_adam import MyAdam +from torch.utils import data as torchdata +import time +import torch.distributed as dist +from torch.nn.parallel import DistributedDataParallel as DDP + + +def demo(): + # First, we load an example dataset from a .cxi file + filename = 'example_data/lab_ptycho_data.cxi' + dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(filename) + + # Next, we create a ptychography model from the dataset + model = cdtools.models.SimplePtycho.from_dataset(dataset) + + + #class MyDataParallel(t.nn.DataParallel): + # def __getattr__(self, name): + # return getattr(self.module, name) + + + print(dist.is_torchelastic_launched()) + dist.init_process_group(backend="nccl", init_method='env://') + rank = dist.get_rank() + print(f"Start running basic DDP example on rank {rank}.") + device_id = rank % t.cuda.device_count() + print('device id', device_id) + model = model.to(device_id) + ddp_model = DDP(model, device_ids=[device_id]) + + + sampler = torchdata.distributed.DistributedSampler(dataset) + # Make a dataloader + #data_loader = torchdata.DataLoader(dataset, batch_size=20, + # shuffle=True) + data_loader = torchdata.DataLoader(dataset, batch_size=20, + shuffle=False, sampler=sampler) + + device = 'cuda:'+str(device_id) + dataset.get_as(device=device) + + # Define the optimizer + optimizer = MyAdam(model.parameters(), lr=0.01) + + normalization=0 + for inputs, patterns in data_loader: + normalization += t.sum(patterns) + + dist.all_reduce(normalization) + + for it in range(100): + sampler.set_epoch(it) + loss = 0 + N = 0 + t0 = time.time() + for inputs, patterns in data_loader: + N += 1 + def closure(): + optimizer.zero_grad() + + sim_patterns = model.forward(*inputs) + + if hasattr(model, 'mask'): + loss = model.loss(patterns,sim_patterns, mask=model.mask) + else: + loss = model.loss(patternss,sim_patterns) + + loss.backward() + return loss.detach() + + loss += optimizer.step(closure).detach()#.cpu().numpy() + #dist.all_reduce(loss) + #loss = dist.all_reduce(loss) + #print(it, 'time', time.time()-t0) + #print(loss / normalization) + + return loss#.cpu().numpy() + +if __name__=='__main__': + demo() diff --git a/examples/simple_ptycho.py b/examples/simple_ptycho.py index e270277..2140f66 100644 --- a/examples/simple_ptycho.py +++ b/examples/simple_ptycho.py @@ -1,22 +1,73 @@ import cdtools from matplotlib import pyplot as plt +import torch as t +from cdtools.models.complex_adam import MyAdam +from torch.utils import data as torchdata +import time # First, we load an example dataset from a .cxi file filename = 'example_data/lab_ptycho_data.cxi' dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(filename) +dataset.translations = t.cat([dataset.translations]*5) +dataset.patterns = t.cat([dataset.patterns]*5) # Next, we create a ptychography model from the dataset model = cdtools.models.SimplePtycho.from_dataset(dataset) -model.to(device='cuda:0') -dataset.get_as(device='cuda:0') +#class MyDataParallel(t.nn.DataParallel): +# def __getattr__(self, name): +# return getattr(self.module, name) -# Now, we run a short reconstruction from the dataset! -for loss in model.Adam_optimize(20, dataset): - print(model.report()) +model = t.nn.DataParallel(model, device_ids=[0,1,2,3]) -print(model.save_results().keys()) + +# Make a dataloader +data_loader = torchdata.DataLoader(dataset, batch_size=20, + shuffle=True) + +device = 'cuda' +model.to(device=device) +#dataset.get_as(device='cuda') + +# Define the optimizer +optimizer = MyAdam(model.parameters(), lr=0.01) + +normalization=0 +for inputs, patterns in data_loader: + normalization += t.sum(patterns).cpu().numpy() + +def run_iteration(stop_event=None): + loss = 0 + N = 0 + t0 = time.time() + for inputs, patterns in data_loader: + inp = (x.to(device) for x in inputs)#inputs.to(device) + pats = patterns.to(device) + N += 1 + def closure(): + optimizer.zero_grad() + + sim_patterns = model.forward(*inp) + + if hasattr(model, 'mask'): + loss = model.module.loss(pats,sim_patterns, mask=model.module.mask) + else: + loss = model.module.loss(pats,sim_patterns) + + loss.backward() + return loss.detach() + + loss += optimizer.step(closure).detach().cpu().numpy() + print('time', time.time()-t0) + return loss / normalization + + +for it in range(20): + print(run_iteration()) + +print(model.module.save_results().keys()) +exit() # Finally, we plot the results -model.inspect(dataset) -model.compare(dataset) +model.module.inspect(dataset) +model.module.compare(dataset) plt.show() diff --git a/src/cdtools/models/simple_ptycho.py b/src/cdtools/models/simple_ptycho.py index 8090f5a..c4d85be 100644 --- a/src/cdtools/models/simple_ptycho.py +++ b/src/cdtools/models/simple_ptycho.py @@ -44,9 +44,16 @@ class SimplePtycho(CDIModel): # object self.register_buffer('probe_norm', t.max(t.abs(probe_guess))) - self.probe = t.nn.Parameter(probe_guess / self.probe_norm) - self.obj = t.nn.Parameter(obj_guess) + self.probe_data = t.nn.Parameter(t.view_as_real(probe_guess / self.probe_norm)) + self.obj_data = t.nn.Parameter(t.view_as_real(obj_guess)) + @property + def probe(self): + return t.view_as_complex(self.probe_data) + + @property + def obj(self): + return t.view_as_complex(self.obj_data) @classmethod def from_dataset(cls, dataset): @@ -99,11 +106,11 @@ class SimplePtycho(CDIModel): def interaction(self, index, translations): + pix_trans = tools.interactions.translations_to_pixel(self.probe_basis, translations, surface_normal=self.surface_normal) pix_trans -= self.min_translation - return tools.interactions.ptycho_2D_round(self.probe_norm * self.probe, self.obj, pix_trans) diff --git a/src/cdtools/tools/interactions/interactions.py b/src/cdtools/tools/interactions/interactions.py index 6043116..0b56d4f 100644 --- a/src/cdtools/tools/interactions/interactions.py +++ b/src/cdtools/tools/interactions/interactions.py @@ -44,13 +44,17 @@ def translations_to_pixel(basis, translations, surface_normal=t.Tensor([0.,0.,1. A Jx2 stack of translations in internal (i,j) pixel-space, or a single translation """ - projection_1 = t.Tensor([[1,0,0], - [0,1,0], - [0,0,0]]).to(device=translations.device,dtype=translations.dtype) - projection_2 = t.inverse(t.Tensor([[1,0,0], - [0,1,0], - -surface_normal/ - surface_normal[2]])).to(device=translations.device,dtype=translations.dtype) + projection_1 = t.as_tensor(np.array([[1,0,0], + [0,1,0], + [0,0,0]]), + device=translations.device, + dtype=translations.dtype) + projection_2 = projection_1.clone() + projection_2[2] = t.as_tensor(-surface_normal / surface_normal[2], + dtype=projection_2.dtype, + device=projection_2.device) + projection_2 = t.inverse(projection_2) + basis_vectors_inv = t.pinverse(basis).to(device=translations.device, dtype=translations.dtype) projection = t.mm(basis_vectors_inv,