mirror of
https://github.com/cdtools-developers/cdtools.git
synced 2026-09-09 21:12:42 +02:00
85 lines
2.7 KiB
Python
85 lines
2.7 KiB
Python
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, num_workers=0)
|
|
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)
|
|
optimizer = t.optim.Adam(model.parameters(), lr=0.01)
|
|
|
|
normalization=0
|
|
for inputs, patterns in dataset:#data_loader:
|
|
normalization += t.sum(patterns)
|
|
|
|
dist.all_reduce(normalization)
|
|
|
|
for it in range(10):
|
|
print(it)
|
|
sampler.set_epoch(it)
|
|
loss = 0
|
|
N = 0
|
|
t0 = time.time()
|
|
for inputs, patterns in data_loader:
|
|
#print(inputs)
|
|
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)#.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()
|