diff --git a/odl/contrib/torch/operator.py b/odl/contrib/torch/operator.py index 9064d511b3a..50c6218caef 100644 --- a/odl/contrib/torch/operator.py +++ b/odl/contrib/torch/operator.py @@ -230,6 +230,12 @@ def forward(ctx, operator, input): ctx.op_in_dtype = operator.domain.dtype ctx.op_out_dtype = op_out_dtype + if hasattr(operator, 'gpu_index') and torch.cuda.is_available(): + try: + operator.gpu_index = torch.cuda.current_device() + except AttributeError: + pass + # Evaluate the operator on all inputs in a loop if extra_shape: # Multiple inputs: flatten extra axes, then do one entry at a time @@ -346,6 +352,12 @@ def backward(ctx, grad_output): ''.format(extra_shape + op_out_shape, grad_output_arr.shape) ) + if hasattr(operator, 'gpu_index') and torch.cuda.is_available(): + try: + operator.gpu_index = torch.cuda.current_device() + except AttributeError: + pass + # Evaluate the (derivative) adjoint on all inputs in a loop if extra_shape: # Multiple gradients: flatten extra axes, then do one entry @@ -509,6 +521,13 @@ def __repr__(self): self.__class__.__name__, op_name, op_in_shape, op_out_shape ) + def _replicate_for_data_parallel(self): + replica = super()._replicate_for_data_parallel() + op_replicate = getattr(self.operator, '_replicate', None) + if callable(op_replicate): + replica.operator = op_replicate() + return replica + def copy_if_zero_strides(arr): """Workaround for NumPy issue #9165 with 0 in arr.strides.""" diff --git a/odl/discr/discr_ops.py b/odl/discr/discr_ops.py index 6c5fe065232..1c399afde86 100644 --- a/odl/discr/discr_ops.py +++ b/odl/discr/discr_ops.py @@ -514,18 +514,16 @@ def _resize_discr(discr, newshp, offset, discr_kwargs): for axis, (n_orig, n_new, off, on_bdry) in enumerate(zip( discr.shape, newshp, offset, nodes_on_bdry)): - if not affected[axis]: - new_minpt.append(discr.min_pt[axis]) - new_maxpt.append(discr.max_pt[axis]) - continue - - n_diff = n_new - n_orig - if off is None: - num_r = n_diff // 2 - num_l = n_diff - num_r + if affected[axis]: + n_diff = n_new - n_orig + if off is None: + num_r = n_diff // 2 + num_l = n_diff - num_r + else: + num_r = n_diff - off + num_l = off else: - num_r = n_diff - off - num_l = off + num_l, num_r = 0, 0 try: on_bdry_l, on_bdry_r = on_bdry diff --git a/odl/tomo/backends/astra_cuda.py b/odl/tomo/backends/astra_cuda.py index 96cee858413..94eca67bcac 100644 --- a/odl/tomo/backends/astra_cuda.py +++ b/odl/tomo/backends/astra_cuda.py @@ -46,7 +46,7 @@ class AstraCudaProjectorImpl(object): sino_id = None proj_id = None - def __init__(self, geometry, reco_space, proj_space): + def __init__(self, geometry, reco_space, proj_space, gpu_index=0): """Initialize a new instance. Parameters @@ -58,6 +58,9 @@ def __init__(self, geometry, reco_space, proj_space): projected. proj_space : `DiscreteLp` Projection space, the space of the result. + gpu_index : int, optional + Index of GPU to use. + Default: ``0`` """ assert isinstance(geometry, Geometry) assert isinstance(reco_space, DiscreteLp) @@ -66,6 +69,7 @@ def __init__(self, geometry, reco_space, proj_space): self.geometry = geometry self.reco_space = reco_space self.proj_space = proj_space + self.gpu_index = gpu_index self.create_ids() @@ -159,7 +163,8 @@ def create_ids(self): proj_type = 'cuda' if proj_ndim == 2 else 'cuda3d' self.proj_id = astra_projector( - proj_type, vol_geom, proj_geom, proj_ndim + proj_type, vol_geom, proj_geom, proj_ndim, + gpu_index=self.gpu_index ) self.sino_id = astra_data(proj_geom, @@ -171,7 +176,7 @@ def create_ids(self): # Create algorithm self.algo_id = astra_algorithm( 'forward', proj_ndim, self.vol_id, self.sino_id, - proj_id=self.proj_id, impl='cuda') + proj_id=self.proj_id, impl='cuda', gpu_index=self.gpu_index) def __del__(self): """Delete ASTRA objects.""" @@ -203,7 +208,7 @@ class AstraCudaBackProjectorImpl(object): sino_id = None proj_id = None - def __init__(self, geometry, reco_space, proj_space): + def __init__(self, geometry, reco_space, proj_space, gpu_index=0): """Initialize a new instance. Parameters @@ -214,6 +219,9 @@ def __init__(self, geometry, reco_space, proj_space): Reconstruction space, the space to which the backprojection maps. proj_space : `DiscreteLp` Projection space, the space from which the backprojection maps. + gpu_index : int, optional + Index of GPU to use. + Default: ``0`` """ assert isinstance(geometry, Geometry) assert isinstance(reco_space, DiscreteLp) @@ -222,6 +230,8 @@ def __init__(self, geometry, reco_space, proj_space): self.geometry = geometry self.reco_space = reco_space self.proj_space = proj_space + self.gpu_index = gpu_index + self.create_ids() # Create a mutually exclusive lock so that two callers cant use the @@ -311,7 +321,8 @@ def create_ids(self): proj_type = 'cuda' if proj_ndim == 2 else 'cuda3d' self.proj_id = astra_projector( - proj_type, vol_geom, proj_geom, proj_ndim + proj_type, vol_geom, proj_geom, proj_ndim, + gpu_index=self.gpu_index ) self.vol_id = astra_data(vol_geom, @@ -323,7 +334,7 @@ def create_ids(self): # Create algorithm self.algo_id = astra_algorithm( 'backward', proj_ndim, self.vol_id, self.sino_id, - proj_id=self.proj_id, impl='cuda') + proj_id=self.proj_id, impl='cuda', gpu_index=self.gpu_index) def __del__(self): """Delete ASTRA objects.""" diff --git a/odl/tomo/backends/astra_setup.py b/odl/tomo/backends/astra_setup.py index 658263e59f8..761acb8f2c4 100644 --- a/odl/tomo/backends/astra_setup.py +++ b/odl/tomo/backends/astra_setup.py @@ -615,7 +615,8 @@ def astra_data(astra_geom, datatype, data=None, ndim=2, allow_copy=False): return create(astra_dtype_str, astra_geom) -def astra_projector(astra_proj_type, astra_vol_geom, astra_proj_geom, ndim): +def astra_projector(astra_proj_type, astra_vol_geom, astra_proj_geom, ndim, + gpu_index=0): """Create an ASTRA projector configuration dictionary. Parameters @@ -630,6 +631,9 @@ def astra_projector(astra_proj_type, astra_vol_geom, astra_proj_geom, ndim): ASTRA projection geometry. ndim : {2, 3} Number of dimensions of the projector. + gpu_index : int, optional + Index of GPU to use. + Default: ``0`` Returns ------- @@ -691,7 +695,7 @@ def astra_projector(astra_proj_type, astra_vol_geom, astra_proj_geom, ndim): proj_cfg['type'] = astra_proj_type proj_cfg['VolumeGeometry'] = astra_vol_geom proj_cfg['ProjectionGeometry'] = astra_proj_geom - proj_cfg['options'] = {} + proj_cfg['options'] = {'GPUindex': gpu_index} # Add the approximate 1/r^2 weighting exposed in intermediate versions of # ASTRA @@ -707,7 +711,8 @@ def astra_projector(astra_proj_type, astra_vol_geom, astra_proj_geom, ndim): return astra.projector3d.create(proj_cfg) -def astra_algorithm(direction, ndim, vol_id, sino_id, proj_id, impl): +def astra_algorithm(direction, ndim, vol_id, sino_id, proj_id, impl, + gpu_index=0): """Create an ASTRA algorithm object to run the projector. Parameters @@ -725,6 +730,9 @@ def astra_algorithm(direction, ndim, vol_id, sino_id, proj_id, impl): Handle for the ASTRA projector object. impl : {'cpu', 'cuda'} Implementation of the projector. + gpu_index : int, optional + Index of GPU to use for ``impl=='cuda'``. + Default: ``0`` Returns ------- @@ -757,6 +765,8 @@ def astra_algorithm(direction, ndim, vol_id, sino_id, proj_id, impl): algo_cfg['VolumeDataId'] = vol_id else: algo_cfg['ReconstructionDataId'] = vol_id + if impl == 'cuda': + algo_cfg['option'] = {'GPUindex': gpu_index} # Create ASTRA algorithm object return astra.algorithm.create(algo_cfg) diff --git a/odl/tomo/operators/ray_trafo.py b/odl/tomo/operators/ray_trafo.py index edceb0343b0..b20ed35e2b8 100644 --- a/odl/tomo/operators/ray_trafo.py +++ b/odl/tomo/operators/ray_trafo.py @@ -11,6 +11,7 @@ from __future__ import absolute_import, division, print_function import warnings +from copy import copy import numpy as np @@ -83,6 +84,9 @@ def __init__(self, reco_space, geometry, variant, **kwargs): and on the CPU, since a full volume and a projection dataset are stored. That may be prohibitive in 3D. Default: True + gpu_index : int, optional + Index of GPU to use for ``impl='astra_cuda'``. + Default: ``0`` kwargs Further keyword arguments passed to the projector backend. @@ -211,6 +215,7 @@ def __init__(self, reco_space, geometry, variant, **kwargs): self.__geometry = geometry self.__impl = impl + self.__gpu_index = kwargs.pop('gpu_index', 0) # Generate or check projection space proj_space = kwargs.pop('proj_space', None) @@ -311,6 +316,23 @@ def geometry(self): """Geometry of this operator.""" return self.__geometry + @property + def gpu_index(self): + """Index of GPU for ``'astra_cuda'`` implementation.""" + return self.__gpu_index + + @gpu_index.setter + def gpu_index(self, gpu_index): + """Set index of GPU for ``'astra_cuda'`` implementation.""" + if not isinstance(gpu_index, int): + raise TypeError('`gpu_index` must be an integer') + assert gpu_index >= 0 + if gpu_index != self.__gpu_index: + # Clear cached properties + self._adjoint = None + self._astra_wrapper = None + self.__gpu_index = gpu_index + def _call(self, x, out=None): """Return ``self(x[, out])``.""" if self.domain.is_real: @@ -336,6 +358,16 @@ def _call(self, x, out=None): else: raise RuntimeError('bad domain {!r}'.format(self.domain)) + def _replicate(self): + """Return replica that can be configured independently. + + Useful to create multiple instances on different GPUs.""" + replica = copy(self) + # Clear cached properties + replica._adjoint = None + replica._astra_wrapper = None + return replica + class RayTransform(RayTransformBase): @@ -375,6 +407,9 @@ def __init__(self, domain, geometry, **kwargs): and on the CPU, since a full volume and a projection dataset are stored. That may be prohibitive in 3D. Default: True + gpu_index : int, optional + Index of GPU to use for ``impl='astra_cuda'``. + Default: ``0`` kwargs Further keyword arguments passed to the projector backend. @@ -413,7 +448,7 @@ def _call_real(self, x_real, out_real, **kwargs): if self._astra_wrapper is None: astra_wrapper = AstraCudaProjectorImpl( self.geometry, self.domain.real_space, - self.range.real_space) + self.range.real_space, gpu_index=self.gpu_index) if self.use_cache: self._astra_wrapper = astra_wrapper else: @@ -445,6 +480,7 @@ def adjoint(self): kwargs = self._extra_kwargs.copy() kwargs['domain'] = self.range + kwargs['gpu_index'] = self.gpu_index self._adjoint = RayBackProjection(self.domain, self.geometry, impl=self.impl, use_cache=self.use_cache, @@ -491,6 +527,9 @@ def __init__(self, range, geometry, **kwargs): and on the CPU, since a full volume and a projection dataset are stored. That may be prohibitive in 3D. Default: True + gpu_index : int, optional + Index of GPU to use for ``impl='astra_cuda'``. + Default: ``0`` kwargs Further keyword arguments passed to the projector backend. @@ -527,7 +566,7 @@ def _call_real(self, x_real, out_real, **kwargs): if self._astra_wrapper is None: astra_wrapper = AstraCudaBackProjectorImpl( self.geometry, self.range.real_space, - self.domain.real_space) + self.domain.real_space, gpu_index=self.gpu_index) if self.use_cache: self._astra_wrapper = astra_wrapper else: @@ -559,6 +598,7 @@ def adjoint(self): kwargs = self._extra_kwargs.copy() kwargs['range'] = self.domain + kwargs['gpu_index'] = self.gpu_index self._adjoint = RayTransform(self.range, self.geometry, impl=self.impl, use_cache=self.use_cache,