Repository: kymatio/kymatio

Issue:
CUDA<->CPU interface for 3D
Continuation of the discussion from #122. It would be good to have a way to set the `Scattering3D` object to CPU or CUDA modes, just like the 1D and 3D codes with `cpu()` and `cuda()` functions that mutate the object. Note that calling `cuda()` doesn't have to make the object move its filters to GPU (like the 1D and 2D code currently does), but simply gives it the option to. This may be better in 3D due to the memory constraints.

Right now, all this could do is throw and error if the input to `forward()` is not compatible (i.e., sending a CPU tensor to a `Scattering3D` object that is in CUDA mode, or vice versa), but it unifies the interfaces and allows for more complex behavior in the future.


Discussion:

Resolving patch:
diff --git a/kymatio/scattering3d/scattering3d.py b/kymatio/scattering3d/scattering3d.py
--- a/kymatio/scattering3d/scattering3d.py
+++ b/kymatio/scattering3d/scattering3d.py
@@ -40,6 +40,8 @@ def __init__(self, J, shape, L, sigma_0):
         self.L = L
         self.sigma_0 = sigma_0
 
+        self.is_cuda = False
+
         self.build()
 
     def build(self):
@@ -49,6 +51,20 @@ def build(self):
         self.gaussian_filters = gaussian_filter_bank(
                                 self.M, self.N, self.O, self.J + 1, self.sigma_0)
 
+    def cuda(self):
+        """Move to the GPU
+
+        This function prepares the object to accept input Tensors on the GPU.
+        """
+        self.is_cuda = True
+
+    def cpu(self):
+        """Move to the CPU
+
+        This function prepares the object to accept input Tensors on the CPU.
+        """
+        self.is_cuda = False
+
     def _fft_convolve(self, input_array, filter_array):
         """
         Computes the fourier space convolution of the input_array, 
@@ -252,6 +268,14 @@ def _check_input(self, input_array):
                 'The input should be a torch.cuda.FloatTensor, '
                 'a torch.FloatTensor or a torch.DoubleTensor'))
 
+        if self.is_cuda and not input_array.is_cuda:
+            raise(TypeError('This transform is in GPU mode, but the input is '
+                            'on the CPU.'))
+
+        if not self.is_cuda and input_array.is_cuda:
+            raise(TypeError('This transform is in CPU mode, but the input is '
+                            'on the GPU.'))
+
         if (not input_array.is_contiguous()):
             input_array = input_array.contiguous()
 


Test patch:
diff --git a/kymatio/scattering3d/tests/test_scattering3d.py b/kymatio/scattering3d/tests/test_scattering3d.py
--- a/kymatio/scattering3d/tests/test_scattering3d.py
+++ b/kymatio/scattering3d/tests/test_scattering3d.py
@@ -72,8 +72,10 @@ def test_against_standard_computations():
     for device in devices:
         if device == 'cpu':
             x = x.cpu()
+            scattering.cpu()
         else:
             x = x.cuda()
+            scattering.cuda()
         order_0 = compute_integrals(x, integral_powers)
         orders_1_and_2 = scattering(
             x, order_2=True, method='integral',
@@ -177,6 +179,48 @@ def test_scattering_methods():
     Sx = scattering(x, method='local', points=points)
     Sx = scattering(x, method='local', points=points, rotation_covariant=False)
 
+def test_cpu_cuda():
+    if backend.NAME == "skcuda":
+        warnings.warn(("The skcuda backend is not yet implemented for 3D "
+            "scattering, but that's ok (for now)."), RuntimeWarning,
+            stacklevel=2)
+        return
+
+    shape = (32, 32, 32)
+    J = 4
+    L = 3
+    sigma_0 = 1
+    x = torch.randn((1,) + shape)
+
+    S = Scattering3D(J=J, shape=shape, L=L, sigma_0=sigma_0)
+
+    assert not S.is_cuda
+
+    Sx = S(x)
+
+    if 'gpu' in devices:
+        x_gpu = x.cuda()
+
+        with pytest.raises(TypeError) as record:
+            Sx_gpu = S(x_gpu)
+        assert "is in CPU mode" in record.value.args[0]
+
+        S.cuda()
+
+        assert S.is_cuda
+
+        Sx_gpu = S(x_gpu)
+
+        assert Sx_gpu.is_cuda
+
+        with pytest.raises(TypeError) as record:
+            Sx = S(x)
+        assert "is in GPU mode" in record.value.args[0]
+
+    S.cpu()
+
+    assert not S.is_cuda
+
 def test_utils():
     # Simple test
     x = np.arange(8)