added gradient tests for APRNet layers

Author: joeljonsson

pyapr/__init__.py

@@ -30,4 +30,4 @@
 from .numerics import *
 #from .viewer import *
 
-__all__ = ['data_containers', 'io', 'nn', 'viewer', 'converter', 'numerics']
+__all__ = ['data_containers', 'io', 'nn', 'viewer', 'converter', 'numerics', 'tests']

pyapr/nn/__init__.py

@@ -1 +1,3 @@
 from .APRNet import *
+
+__all__ = ['APRNet', 'testing']

pyapr/nn/testing.py

@@ -0,0 +1,90 @@
+import torch
+from torch.autograd.gradcheck import get_analytical_jacobian, get_numerical_jacobian, zero_gradients, make_jacobian
+
+
+def get_analytical_jacobian_params(output, target):
+    """
+    Computes the analytical jacobian with respect to all tensors in `target`, which can hold some or all of the
+    parameters of a module used to compute `output`.
+
+    output: torch.tensor output from which to backpropagate the gradients
+    target: torch.tensor or iterable containing torch.tensor for which to compute the gradients
+    """
+
+    jacobian = make_jacobian(target, output.numel())
+    grad_output = torch.zeros_like(output)
+    flat_grad_output = grad_output.view(-1)
+
+    for i in range(flat_grad_output.numel()):
+        flat_grad_output.zero_()
+        flat_grad_output[i] = 1
+
+        zero_gradients(target)
+        torch.autograd.backward(output, grad_output, retain_graph=True)
+
+        for j in range(len(jacobian)):
+            jacobian[j][:, i] = target[j].grad.clone().flatten()
+
+    return jacobian
+
+
+def gradcheck(m, input, eps=1e-6, atol=1e-5, rtol=1e-3, raise_exception=True):
+    """
+    Compare analytical gradients of a module to numerical gradients computed via central finite differences.
+
+    Disclaimer: this is a modified version of torch.autograd.gradcheck::gradcheck
+    (https://pytorch.org/docs/stable/_modules/torch/autograd/gradcheck.html, 2019-06-04) distributed under license
+    https://github.com/pytorch/pytorch/blob/master/LICENSE
+    :param m:
+    :param input:
+    :param eps:
+    :param atol:
+    :param rtol:
+    :param raise_exception:
+    :return:
+    """
+    def fail_test(msg):
+        if raise_exception:
+            raise RuntimeError(msg)
+        return False
+
+    def fn(input):
+        return m(*input)
+
+    output = fn(input)
+
+    for i, o in enumerate(output):
+        if not o.requires_grad:
+            continue
+
+        # compare input gradients
+        analytical, reentrant, correct_grad_sizes = get_analytical_jacobian(input, o)
+        numerical = get_numerical_jacobian(fn, input, eps=eps)
+
+        if not correct_grad_sizes:
+            return fail_test('Analytical gradient has incorrect size')
+
+        for j, (a, n) in enumerate(zip(analytical, numerical)):
+            if a.numel() != 0 or n.numel() != 0:
+                if not torch.allclose(a, n, rtol, atol):
+                    return fail_test('Jacobian mismatch for output %d with respect to input %d,\n'
+                                     'numerical:%s\nanalytical:%s\n' % (i, j, n, a))
+
+        if not reentrant:
+            return fail_test('Backward is not reentrant, i.e., running backward with same '
+                             'input and grad_output multiple times gives different values, '
+                             'although analytical gradient matches numerical gradient')
+
+        # compare parameter gradients
+        pars = [t for t in m.parameters()]
+
+        if pars:
+            numerical = get_numerical_jacobian(fn, input, target=pars)
+            analytical = get_analytical_jacobian_params(output, pars)
+
+            for j, (a, n) in enumerate(zip(analytical, numerical)):
+                if a.numel() != 0 or n.numel() != 0:
+                    if not torch.allclose(a, n, rtol, atol):
+                        return fail_test('Jacobian mismatch for output %d with respect to parameter %d,\n'
+                                         'numerical:%s\nanalytical:%s\n' % (i, j, n, a))
+    return True

pyapr/tests/__init__.py

@@ -0,0 +1 @@
+from .test_nn_modules import *

pyapr/tests/test_nn_modules.py

@@ -0,0 +1,76 @@
+import pyapr
+import pyapr.nn as aprnn
+import numpy as np
+from demo.io import read_tiff
+import unittest
+import pyapr.nn.testing as testing
+
+
+class TestAPRNetModules(unittest.TestCase):
+
+    def __init__(self, *args, **kwargs):
+        super(TestAPRNetModules, self).__init__(*args, **kwargs)
+
+        # Load in an image and extract a small tile
+        img = read_tiff('../../LibAPR/test/files/Apr/sphere_120/sphere_original.tif')
+        img = img[20:52, 40:72, 53].astype(np.float32)
+
+        # Ensure that pixel values are distinct to avoid problems with finite difference of max pooling
+        for i in range(32):
+            for j in range(32):
+                img[i, j] += 0.001 * (i*32 + j)
+
+        # Initialize objects for APR conversion
+        apr = pyapr.APR()
+        parts = pyapr.FloatParticles()
+        par = pyapr.APRParameters()
+        converter = pyapr.converter.FloatConverter()
+
+        # Set some parameters
+        par.auto_parameters = False
+        par.rel_error = 0.1
+        par.Ip_th = 0
+        par.gradient_smoothing = 2
+        par.sigma_th = 50
+        par.sigma_th_max = 20
+        converter.set_parameters(par)
+        converter.set_verbose(False)
+
+        # Compute APR and sample particle values
+        converter.get_apr(apr, img)
+        apr.init_tree()
+        parts.sample_image(apr, img)
+
+        apr_arr = np.empty(1, dtype=object)
+        parts_arr = np.empty(1, dtype=object)
+        apr_arr[0] = apr
+        parts_arr[0] = parts
+
+        x, dlvl = aprnn.APRInputLayer()(apr_arr, parts_arr, dtype=np.float64)
+        x.requires_grad = True
+
+        self.aprs = apr_arr
+        self.x = x
+        self.dlvl = dlvl
+
+    def test_gradients_conv1x1(self):
+        m = aprnn.APRConv(1, 4, 1, 2)
+        assert testing.gradcheck(m, (self.x, self.aprs, self.dlvl))
+
+    def test_gradients_conv3x3(self):
+        m = aprnn.APRConv(1, 4, 3, 2)
+        assert testing.gradcheck(m, (self.x, self.aprs, self.dlvl))
+
+    def test_gradients_conv5x5(self):
+        m = aprnn.APRConv(1, 4, 5, 2)
+        assert testing.gradcheck(m, (self.x, self.aprs, self.dlvl))
+
+    def test_gradients_maxpool(self):
+        m = aprnn.APRMaxPool(increment_level_delta=False)
+        assert testing.gradcheck(m, (self.x, self.aprs, self.dlvl))
+
+
+if __name__ == '__main__':
+    unittest.main()
+
+