Configure GPU index for 'astra_cuda', select GPU currently used by PyTorch in OperatorModule

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."""

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

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."""

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)

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,