Fix the annoying warning printing from interactions and get a minimum usable script going for simpleptycho and distributeddataparallel

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()

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()

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()

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)

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,