Repository: kymatio/kymatio

Issue:
scattering3d needs a skcuda backend
because scattering3d doesn't have that backend.

Ideally we would make scattering just not import 3D if that backend is chosen. Or we could make it be not functional and raise an error when it is executed.

 It shouldn't cause a global importerror though


Discussion:
fyi `KYMATIO_BACKEND_2D=skcuda` works and even `KYMATIO_BACKEND=skcuda KYMATIO_BACKEND_3D=torch` works
I would vote for (2), that is, make it non-functional when `skcuda` backend is chosen and raise an error when instantiating (or calling `forward()`).
I'm also voting for (2).

Otherwise, if we choose (1), then skcuda users will suddenly get errors of the like of
```
AttributeError: module 'kymatio' has no attribute 'scattering3d'
```
which is not informative enough for them to change their backend.
If you want I can add skcuda backend for 3D. I did it on pyscatharm, I have a speedup of 15%. It only concerns the Cdgmm function. It'a also possible to remove skcuda dependency if you copy the few necessary files as I did in pyscatharm : https://github.com/louity/pyscatharm/tree/master/scatharm/skcuda_utils 
@louity that would probably be the best solution, assuming it doesn't take too much time. The 15% reduction is similar to what we see in 1D and can always be improved upon later.
@eickenberg  should I make a fork from eickenberg/kymatio or kymatio/kymatio  to make the PR ?
@louity the latter
Cool, thanks very much @louity for adding this backend!

As far as I know the main speed-up due to skcuda/cupy backend in the 2D version comes from the kernel implementation of the fourier downsampling.

Do you have such an implementation for the 3D code and is this already in the 15% you mention?

Yes @eickenberg  In the 3D code it's 15% essentially because you have more complex multiplications than fourier transforms : in solid harmonic scattering transform, for 1 "convolution" with psi_l you have 2 FFT and 2L+1 complex multiplications, so speeding up complex multiplication really speeds up the whoel thing

I can't write a skcuda backend for 3d myself, but in the interest of time, what I can do is have Kymatio return a `NotImplementedError` if `KYMATIO_BACKEND=skcuda`.
that would be OK if it throws at instantiation, but not OK if it throws that at import. It would render this environment variable useless until that backend is implemented.

So +1 for `NotImplementedError` at instantiation of a 3D object with inexistent backend.
> In the 3D code it's 15% essentially because you have more complex multiplications than fourier transforms : in solid harmonic scattering transform, for 1 "convolution" with psi_l you have 2 FFT and 2L+1 complex multiplications, so speeding up complex multiplication really speeds up the whoel thing

That makes sense, thanks
OK, I'm tagging this issue (skcuda backend by @louity) for 0.2.0 and will raise another one re: NotImplementedError

Resolving patch:
diff --git a/examples/3d/compute_speed.py b/examples/3d/compute_speed.py
--- a/examples/3d/compute_speed.py
+++ b/examples/3d/compute_speed.py
@@ -66,6 +66,8 @@
 
 if backend.NAME == 'torch':
     devices = ['cpu', 'gpu']
+elif backend.NAME == 'skcuda':
+    devices = ['gpu']
 
 ###############################################################################
 # Set up the scattering object and the test data
diff --git a/kymatio/scattering3d/backend/backend_skcuda.py b/kymatio/scattering3d/backend/backend_skcuda.py
--- a/kymatio/scattering3d/backend/backend_skcuda.py
+++ b/kymatio/scattering3d/backend/backend_skcuda.py
@@ -1,19 +1,23 @@
-# The `skcuda` backend for Scattering3D is currently under development.
-# See issue #179.
+import torch
+from skcuda import cublas
 
 NAME = 'skcuda'
 
 
-skcuda_notimplementederror =\
-"""The \"skcuda\" backend for Scattering3D is currently under development.
-To perform 3D scattering, please set KYMATIO_BACKEND_3D=\"torch\"."""
+def iscomplex(input):
+    return input.size(-1) == 2
+
 
 def to_complex(input):
-    raise NotImplementedError(skcuda_notimplementederror)
+    output = input.new(input.size() + (2,)).fill_(0)
+    output[..., 0] = input
+    return output
 
 
 def complex_modulus(input_array):
-    raise NotImplementedError(skcuda_notimplementederror)
+    modulus = torch.zeros_like(input_array)
+    modulus[..., 0] += torch.sqrt((input_array ** 2).sum(-1))
+    return modulus
 
 
 def fft(input, inverse=False):
@@ -30,10 +34,14 @@ def fft(input, inverse=False):
         inverse : bool
             True for computing the inverse FFT.
     """
-    raise NotImplementedError(skcuda_notimplementederror)
+    if not iscomplex(input):
+        raise(TypeError('The input should be complex (e.g. last dimension is 2)'))
+    if inverse:
+        return torch.ifft(input, 3)
+    return torch.fft(input, 3)
 
 
-def cdgmm3d(A, B):
+def cdgmm3d(A, B, inplace=False):
     """
     Pointwise multiplication of complex tensors.
 
@@ -46,4 +54,29 @@ def cdgmm3d(A, B):
     output : tensor of the same size as A containing the result of the
              elementwise complex multiplication of  A with B
     """
-    raise NotImplementedError(skcuda_notimplementederror)
+    A, B = A.contiguous(), B.contiguous()
+    if A.size()[-4:] != B.size():
+        raise RuntimeError('The filters are not compatible for multiplication.')
+
+    if not iscomplex(A) or not iscomplex(B):
+        raise TypeError('The input, filter and output should be complex.')
+
+    if B.ndimension() != 4:
+        raise RuntimeError('The filters must be simply a complex array.')
+
+    if type(A) is not type(B):
+        raise RuntimeError('A and B should be same type.')
+
+    if not A.is_cuda:
+        raise RuntimeError('Use the torch backend for cpu tensors.')
+
+    C = A.new(A.size()) if not inplace else A
+    m, n = B.nelement() // 2, A.nelement() // B.nelement()
+    lda = m
+    ldc = m
+    incx = 1
+    handle = torch.cuda.current_blas_handle()
+    stream = torch.cuda.current_stream()._as_parameter_
+    cublas.cublasSetStream(handle, stream)
+    cublas.cublasCdgmm(handle, 'l', m, n, A.data_ptr(), lda, B.data_ptr(), incx, C.data_ptr(), ldc)
+    return C
diff --git a/kymatio/scattering3d/scattering3d.py b/kymatio/scattering3d/scattering3d.py
--- a/kymatio/scattering3d/scattering3d.py
+++ b/kymatio/scattering3d/scattering3d.py
@@ -5,13 +5,10 @@
 
 import torch
 from .utils import compute_integrals, subsample
- 
+
 from .backend import cdgmm3d, fft, complex_modulus, to_complex
 from .filter_bank import solid_harmonic_filter_bank, gaussian_filter_bank
 
-# TODO remove "import backend" below after implementing skcuda backend
-from kymatio.scattering3d import backend
-
 
 class Scattering3D(object):
     """3D Solid Harmonic scattering .
@@ -32,8 +29,6 @@ class Scattering3D(object):
 
     """
     def __init__(self, J, shape, L, sigma_0):
-        if backend.NAME == "skcuda":
-            raise NotImplementedError(backend.skcuda_notimplementederror)
         super(Scattering3D, self).__init__()
         self.J = J
         self.shape = shape


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
@@ -8,10 +8,13 @@
 from kymatio.scattering3d import backend
 from kymatio.scattering3d.utils import generate_weighted_sum_of_gaussians, compute_integrals, sqrt
 
-if torch.cuda.is_available():
-    devices = ['gpu', 'cpu']
-else:
-    devices = ['cpu']
+devices = []
+if backend.NAME == 'torch':
+    devices.append('cpu')
+if backend.NAME == 'torch' and torch.cuda.is_available():
+    devices.append('gpu')
+if backend.NAME == 'skcuda' and torch.cuda.is_available():
+    devices.append('gpu')
 
 def linfnorm(x,y):
     return torch.max(torch.abs(x-y))
@@ -28,12 +31,6 @@ def relative_difference(a, b):
 
 
 def test_FFT3d_central_freq_batch():
-    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
-
     # Checked the 0 frequency for the 3D FFT
     for device in devices:
         x = torch.zeros(1, 32, 32, 32, 2).float()
@@ -44,14 +41,106 @@ def test_FFT3d_central_freq_batch():
         c = y[:,0,0,0].sum()
         assert (c-a).abs().sum()<1e-6
 
+def test_cdgmm3d():
+    # Not all backends currently implement the inplace variant
+    if backend.NAME == 'torch':
+        has_inplace = False
+    if backend.NAME == 'skcuda':
+        has_inplace = True
 
-def test_against_standard_computations():
-    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
+    if has_inplace:
+        inplace_flags = (False, True)
+    else:
+        inplace_flags = (False,)
 
+    for device in devices:
+        for inplace in inplace_flags:
+            x = torch.zeros(2, 3, 4, 2)
+            x[...,0] = 2
+            x[...,1] = 3
+
+            y = torch.zeros_like(x)
+            y[...,0] = 4
+            y[...,1] = 5
+
+            prod = torch.zeros_like(x)
+            prod[...,0] = x[...,0]*y[...,0] - x[...,1]*y[...,1]
+            prod[...,1] = x[...,0]*y[...,1] + x[...,1]*y[...,0]
+
+            if device == 'gpu':
+                x = x.cuda()
+                y = y.cuda()
+                prod = prod.cuda()
+
+            if has_inplace:
+                z = backend.cdgmm3d(x, y, inplace=inplace)
+            else:
+                z = backend.cdgmm3d(x, y)
+
+            assert (z-prod).norm().cpu().item() < 1e-7
+
+            if inplace:
+                assert (x-z).norm().cpu().item() < 1e-7
+
+    if 'cpu' in devices:
+        tdev = torch.device('cpu')
+    elif 'gpu' in devices:
+        tdev = torch.device('cuda')
+
+    with pytest.raises(RuntimeError) as record:
+        x = torch.randn((3, 3, 3, 2), device=tdev)
+        y = torch.randn((4, 4, 4, 2), device=tdev)
+        backend.cdgmm3d(x, y)
+    assert "not compatible" in record.value.args[0]
+
+    x = torch.randn((2, 3, 3, 3, 2), device=tdev)
+    y = torch.randn((3, 3, 3, 2), device=tdev)
+    backend.cdgmm3d(x, y)
+
+    with pytest.raises(TypeError) as record:
+        x = torch.randn((3, 3, 3, 1), device=tdev)
+        y = torch.randn((3, 3, 3, 1), device=tdev)
+        backend.cdgmm3d(x, y)
+    assert "should be complex" in record.value.args[0]
+
+    # This one is a little tricky. We can't have the number of dimensions be
+    # greater than 4 since that triggers the "not compatible" error.
+    with pytest.raises(RuntimeError) as record:
+        x = torch.randn((3, 3, 2), device=tdev)
+        y = torch.randn((3, 3, 2), device=tdev)
+        backend.cdgmm3d(x, y)
+    assert "must be simply a complex" in record.value.args[0]
+
+    # Create a tensor that behaves like `torch.Tensor` but is technically a
+    # different type.
+    class FakeTensor(torch.Tensor):
+        pass
+
+    with pytest.raises(RuntimeError) as record:
+        x = FakeTensor(3, 3, 3, 2)
+        y = torch.randn(3, 3, 3, 2)
+        backend.cdgmm3d(x, y)
+    assert "should be same type" in record.value.args[0]
+
+
+def test_iscomplex():
+    assert backend.iscomplex(torch.zeros(4, 2))
+    assert not backend.iscomplex(torch.zeros(4, 1))
+
+def test_to_complex():
+    x = torch.randn(4, 3)
+    xc = backend.to_complex(x)
+    assert (xc[...,0] == x).all()
+    assert (xc[...,1] == 0).all()
+
+def test_complex_modulus():
+    x = torch.randn(4, 3, 2)
+    xm = torch.sqrt(x[...,0] ** 2 + x[...,1] ** 2)
+    y = backend.complex_modulus(x)
+    assert (y[...,0] - xm).norm() < 1e-7
+    assert (y[...,1]).norm() < 1e-7
+
+def test_against_standard_computations():
     file_path = os.path.abspath(os.path.dirname(__file__))
     data = torch.load(os.path.join(file_path, 'test_data_3d.pt'))
     x = data['x']
@@ -114,12 +203,6 @@ def test_against_standard_computations():
 
 
 def test_solid_harmonic_scattering():
-    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
-
     # Compare value to analytical formula in the case of a single Gaussian
     centers = torch.FloatTensor(1, 1, 3).fill_(0)
     weights = torch.FloatTensor(1, 1).fill_(1)
@@ -130,22 +213,25 @@ def test_solid_harmonic_scattering():
         np.fft.ifftshift(np.mgrid[-M//2:-M//2+M, -N//2:-N//2+N, -O//2:-O//2+O].astype('float32'), axes=(1,2,3)))
     x = generate_weighted_sum_of_gaussians(grid, centers, weights, sigma_gaussian)
     scattering = Scattering3D(J=J, shape=(M, N, O), L=L, sigma_0=sigma_0_wavelet)
-    s = scattering(x, max_order=1, method='integral', integral_powers=[1])
 
-    for j in range(J+1):
-        sigma_wavelet = sigma_0_wavelet*2**j
-        k = sigma_wavelet / np.sqrt(sigma_wavelet**2 + sigma_gaussian**2)
-        for l in range(1, L+1):
-            err = torch.abs(s[0, 0, j, l] - k ** l).sum()/(1e-6+s[0, 0, j, l].abs().sum())
-            assert err<1e-4
+    for device in devices:
+        if device == 'cpu':
+            x = x.cpu()
+            scattering.cpu()
+        else:
+            x = x.cuda()
+            scattering.cuda()
 
-def test_larger_scales():
-    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
+        s = scattering(x, max_order=1, method='integral', integral_powers=[1])
 
+        for j in range(J+1):
+            sigma_wavelet = sigma_0_wavelet*2**j
+            k = sigma_wavelet / np.sqrt(sigma_wavelet**2 + sigma_gaussian**2)
+            for l in range(1, L+1):
+                err = torch.abs(s[0, 0, j, l] - k ** l).sum()/(1e-6+s[0, 0, j, l].abs().sum())
+                assert err<1e-4
+
+def test_larger_scales():
     shape = (32, 32, 32)
     L = 3
     sigma_0 = 1
@@ -154,15 +240,12 @@ def test_larger_scales():
 
     for J in range(3, 4+1):
         scattering = Scattering3D(J=J, shape=shape, L=L, sigma_0=sigma_0)
+        if not 'cpu' in devices:
+            x = x.cuda()
+            scattering.cuda()
         Sx = scattering(x, method='integral')
 
 def test_scattering_methods():
-    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
@@ -170,6 +253,11 @@ def test_scattering_methods():
     x = torch.randn((1,) + shape)
 
     scattering = Scattering3D(J=J, shape=shape, L=L, sigma_0=sigma_0)
+
+    if not 'cpu' in devices:
+        x = x.cuda()
+        scattering.cuda()
+
     Sx = scattering(x, method='standard')
     Sx = scattering(x, method='standard', rotation_covariant=False)
 
@@ -180,12 +268,6 @@ def test_scattering_methods():
     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
@@ -196,7 +278,8 @@ def test_cpu_cuda():
 
     assert not S.is_cuda
 
-    Sx = S(x)
+    if 'cpu' in devices:
+        Sx = S(x)
 
     if 'gpu' in devices:
         x_gpu = x.cuda()