cSAXS/cdtools
0
0
Fork 0
mirror of https://github.com/cdtools-developers/cdtools.git synced 2026-09-10 13:32:40 +02:00
Code Issues Packages Projects Releases Wiki Activity

Multi-GPU compatability added for base Reconstructor and AdamReconstructor

Browse Source
This commit is contained in:
yoshikisd
2025-11-08 06:41:48 +00:00
parent 10229e4e28
commit 8d8da151d7
2 changed files with 102 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/cdtools/reconstructors/adam.py
+17 -5
View File
@@ -33,9 +33,18 @@ class AdamReconstructor(Reconstructor):
The dataset to reconstruct against.
subset : list(int) or int
Optional, a pattern index or list of pattern indices to use.
schedule : bool
Optional, create a learning rate scheduler
(torch.optim.lr_scheduler._LRScheduler).
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.
Important attributes:
- **model** -- Always points to the core model used.
@@ -49,12 +58,15 @@ class AdamReconstructor(Reconstructor):
def __init__(self,
model: CDIModel,
dataset: Ptycho2DDataset,
subset: List[int] = None):
subset: List[int] = None,
rank: int = None,
world_size: int = None):
# Define the optimizer for use in this subclass
optimizer = t.optim.Adam(model.parameters())
super().__init__(model, dataset, optimizer, subset=subset)
super().__init__(model, dataset, optimizer,
subset=subset, rank=rank, world_size=world_size)
src/cdtools/reconstructors/base.py
+85 -18
View File
@@ -17,11 +17,13 @@ import threading
import queue
import time
from typing import List, Union
from cdtools.tools import multigpu
from torch.utils.data.distributed import DistributedSampler
if TYPE_CHECKING:
from cdtools.models import CDIModel
from cdtools.datasets import CDataset
__all__ = ['Reconstructor']
@@ -44,6 +46,18 @@ class Reconstructor:
The optimizer to use for the reconstruction
subset : list(int) or int
Optional, a pattern index or list of pattern indices to use
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.
Attributes
----------
@@ -60,8 +74,22 @@ class Reconstructor:
model: CDIModel,
dataset: CDataset,
optimizer: t.optim.Optimizer,
subset: Union[int, List[int]] = None):
subset: Union[int, List[int]] = None,
rank: int = None,
world_size: int = None):
# If we're running multi-GPU jobs, we need to grab some
# information that should be stored as environment variables.
self.rank = rank if rank is not None else multigpu.get_rank()
self.world_size = world_size if world_size is not None \
else multigpu.get_world_size()
self.multi_gpu_used = int(self.world_size) > 1
# Make sure the model and dataset live on the assigned GPUs
if self.multi_gpu_used:
model.to(f'cuda:{self.rank}')
dataset.get_as(f'cuda:{self.rank}')
# Store parameters as attributes of Reconstructor
self.model = model
self.optimizer = optimizer
@@ -79,7 +107,6 @@ class Reconstructor:
self.scheduler = None
self.data_loader = None
def setup_dataloader(self,
batch_size: int = None,
shuffle: bool = True):
@@ -94,14 +121,31 @@ class Reconstructor:
Optional, enable/disable shuffling of the dataset. This option
is intended for diagnostic purposes and should be left as True.
"""
if batch_size is not None:
self.data_loader = td.DataLoader(self.dataset,
batch_size=batch_size,
shuffle=shuffle)
if self.multi_gpu_used:
self.sampler = \
DistributedSampler(self.dataset,
num_replicas=self.world_size,
rank=self.rank,
shuffle=shuffle,
drop_last=False)
# Creating extra threads in children processes may cause problems.
# Leave num_workers at 0.
self.data_loader = \
td.DataLoader(self.dataset,
batch_size=batch_size//self.world_size,
num_workers=0,
drop_last=False,
pin_memory=False,
sampler=self.sampler)
else:
self.data_loader = td.Dataloader(self.dataset)
if batch_size is not None:
self.data_loader = td.DataLoader(self.dataset,
batch_size=batch_size,
shuffle=shuffle)
else:
self.data_loader = td.Dataloader(self.dataset)
def adjust_optimizer(self, **kwargs):
"""
Change hyperparameters for the utilized optimizer.
@@ -111,11 +155,10 @@ class Reconstructor:
"""
raise NotImplementedError()
def run_epoch(self,
stop_event: threading.Event = None,
regularization_factor: Union[float, List[float]] = None,
calculation_width: int = 10):
stop_event: threading.Event = None,
regularization_factor: Union[float, List[float]] = None,
calculation_width: int = 10):
"""
Runs one full epoch of the reconstruction. Intended to be called
by Reconstructor.optimize.
@@ -150,8 +193,18 @@ class Reconstructor:
'Reconstructor.run_epoch(), or use Reconstructor.optimize(), '
'which does it automatically.'
)
# If we're using DistributedSampler (i.e., multi-GPU useage), we need
# to tell it which epoch we're on. Otherwise data shuffling will not
# work properly
if self.multi_gpu_used:
self.data_loader.sampler.set_epoch(self.model.epoch)
# This prevent other GPU rank processes from initializing
# rank 0's GPU (i.e., helps avoid unnecessary memory consumption)
if t.cuda.current_device != self.rank:
t.cuda.set_device(self.rank)
# Initialize some tracking variables
normalization = 0
loss = 0
@@ -201,8 +254,14 @@ class Reconstructor:
# And accumulate the gradients
loss.backward()
# Average and sync gradients + losses for multi-GPU jobs
if self.multi_gpu_used:
multigpu.sync_and_avg_grads(model=self.model,
world_size=self.world_size)
multigpu.sync_loss(loss)
# Normalize the accumulating total loss
total_loss += loss.detach()
total_loss += loss.detach() / self.world_size
# If we have a regularizer, we can calculate it separately,
# and the gradients will add to the minibatch gradient
@@ -212,18 +271,26 @@ class Reconstructor:
loss = self.model.regularizer(regularization_factor)
loss.backward()
# Avg and sync gradients for multi-GPU jobs
if self.multi_gpu_used:
multigpu.sync_and_avg_grads(model=self.model,
world_size=self.world_size)
return total_loss
# This takes the step for this minibatch
loss += self.optimizer.step(closure).detach().cpu().numpy()
loss /= normalization
# We step the scheduler after the full epoch
if self.scheduler is not None:
self.scheduler.step(loss)
if self.multi_gpu_used:
multigpu.sync_lr(rank=self.rank,
optimizer=self.optimizer)
self.model.loss_history.append(loss)
self.model.loss_times.append(time.time() - self.model.INITIAL_TIME)
self.model.epoch = len(self.model.loss_history)
self.model.latest_iteration_time = time.time() - t0
self.model.training_history += self.model.report() + '\n'