| from math import pi |
|
|
| import pytest |
| import torch |
| from torch.autograd import gradcheck |
|
|
| import kornia.testing as utils |
| from kornia.feature.mkd import ( |
| COEFFS, |
| EmbedGradients, |
| ExplicitSpacialEncoding, |
| get_grid_dict, |
| get_kron_order, |
| MKDDescriptor, |
| MKDGradients, |
| SimpleKD, |
| spatial_kernel_embedding, |
| VonMisesKernel, |
| Whitening, |
| ) |
| from kornia.testing import assert_close |
|
|
|
|
| @pytest.mark.parametrize("ps", [5, 13, 25]) |
| def test_get_grid_dict(ps): |
| grid_dict = get_grid_dict(ps) |
| param_keys = ['x', 'y', 'phi', 'rho'] |
| assert set(grid_dict.keys()) == set(param_keys) |
| for k in param_keys: |
| assert grid_dict[k].shape == (ps, ps) |
|
|
|
|
| @pytest.mark.parametrize("d1,d2", [(1, 1), (1, 2), (2, 1), (5, 6)]) |
| def test_get_kron_order(d1, d2): |
| out = get_kron_order(d1, d2) |
| assert out.shape == (d1 * d2, 2) |
|
|
|
|
| class TestMKDGradients: |
| @pytest.mark.parametrize("ps", [5, 13, 25]) |
| def test_shape(self, ps, device): |
| inp = torch.ones(1, 1, ps, ps).to(device) |
| gradients = MKDGradients().to(device) |
| out = gradients(inp) |
| assert out.shape == (1, 2, ps, ps) |
|
|
| @pytest.mark.parametrize("bs", [1, 5, 13]) |
| def test_batch_shape(self, bs, device): |
| inp = torch.ones(bs, 1, 15, 15).to(device) |
| gradients = MKDGradients().to(device) |
| out = gradients(inp) |
| assert out.shape == (bs, 2, 15, 15) |
|
|
| def test_print(self, device): |
| gradients = MKDGradients().to(device) |
| gradients.__repr__() |
|
|
| def test_toy(self, device): |
| patch = torch.ones(1, 1, 6, 6).to(device).float() |
| patch[0, 0, :, 3:] = 0 |
| gradients = MKDGradients().to(device) |
| out = gradients(patch) |
| expected_mags_1 = torch.Tensor([0, 0, 1.0, 1.0, 0, 0]).to(device) |
| expected_mags = expected_mags_1.unsqueeze(0).repeat(6, 1) |
| expected_oris_1 = torch.Tensor([-pi, -pi, 0, 0, -pi, -pi]).to(device) |
| expected_oris = expected_oris_1.unsqueeze(0).repeat(6, 1) |
| assert_close(out[0, 0, :, :], expected_mags, atol=1e-3, rtol=1e-3) |
| assert_close(out[0, 1, :, :], expected_oris, atol=1e-3, rtol=1e-3) |
|
|
| def test_gradcheck(self, device): |
| batch_size, channels, height, width = 1, 1, 13, 13 |
| patches = torch.rand(batch_size, channels, height, width).to(device) |
| patches = utils.tensor_to_gradcheck_var(patches) |
|
|
| def grad_describe(patches): |
| mkd_grads = MKDGradients() |
| mkd_grads.to(device) |
| return mkd_grads(patches) |
|
|
| assert gradcheck(grad_describe, (patches), raise_exception=True, nondet_tol=1e-4) |
|
|
| @pytest.mark.jit |
| def test_jit(self, device, dtype): |
| B, C, H, W = 2, 1, 13, 13 |
| patches = torch.rand(B, C, H, W, device=device, dtype=dtype) |
| model = MKDGradients().to(patches.device, patches.dtype).eval() |
| model_jit = torch.jit.script(MKDGradients().to(patches.device, patches.dtype).eval()) |
| assert_close(model(patches), model_jit(patches)) |
|
|
|
|
| class TestVonMisesKernel: |
| @pytest.mark.parametrize("ps", [5, 13, 25]) |
| def test_shape(self, ps, device): |
| inp = torch.ones(1, 1, ps, ps).to(device) |
| vm = VonMisesKernel(patch_size=ps, coeffs=[0.38214156, 0.48090413]).to(device) |
| out = vm(inp) |
| assert out.shape == (1, 3, ps, ps) |
|
|
| @pytest.mark.parametrize("bs", [1, 5, 13]) |
| def test_batch_shape(self, bs, device): |
| inp = torch.ones(bs, 1, 15, 15).to(device) |
| vm = VonMisesKernel(patch_size=15, coeffs=[0.38214156, 0.48090413]).to(device) |
| out = vm(inp) |
| assert out.shape == (bs, 3, 15, 15) |
|
|
| @pytest.mark.parametrize("coeffs", COEFFS.values()) |
| def test_coeffs(self, coeffs, device): |
| inp = torch.ones(1, 1, 15, 15).to(device) |
| vm = VonMisesKernel(patch_size=15, coeffs=coeffs).to(device) |
| out = vm(inp) |
| assert out.shape == (1, 2 * len(coeffs) - 1, 15, 15) |
|
|
| def test_print(self, device): |
| vm = VonMisesKernel(patch_size=32, coeffs=[0.38214156, 0.48090413]).to(device) |
| vm.__repr__() |
|
|
| def test_toy(self, device): |
| patch = torch.ones(1, 1, 6, 6).float().to(device) |
| patch[0, 0, :, 3:] = 0 |
| vm = VonMisesKernel(patch_size=6, coeffs=[0.38214156, 0.48090413]).to(device) |
| out = vm(patch) |
| expected = torch.ones_like(out[0, 0, :, :]).to(device) |
| assert_close(out[0, 0, :, :], expected * 0.6182, atol=1e-3, rtol=1e-3) |
|
|
| expected = torch.Tensor([0.3747, 0.3747, 0.3747, 0.6935, 0.6935, 0.6935]).to(device) |
| expected = expected.unsqueeze(0).repeat(6, 1) |
| assert_close(out[0, 1, :, :], expected, atol=1e-3, rtol=1e-3) |
|
|
| expected = torch.Tensor([0.5835, 0.5835, 0.5835, 0.0000, 0.0000, 0.0000]).to(device) |
| expected = expected.unsqueeze(0).repeat(6, 1) |
| assert_close(out[0, 2, :, :], expected, atol=1e-3, rtol=1e-3) |
|
|
| def test_gradcheck(self, device): |
| batch_size, channels, ps = 1, 1, 13 |
| patches = torch.rand(batch_size, channels, ps, ps).to(device) |
| patches = utils.tensor_to_gradcheck_var(patches) |
|
|
| def vm_describe(patches, ps=13): |
| vmkernel = VonMisesKernel(patch_size=ps, coeffs=[0.38214156, 0.48090413]).double() |
| vmkernel.to(device) |
| return vmkernel(patches.double()) |
|
|
| assert gradcheck(vm_describe, (patches, ps), raise_exception=True, nondet_tol=1e-4) |
|
|
| @pytest.mark.jit |
| def test_jit(self, device, dtype): |
| B, C, H, W = 2, 1, 13, 13 |
| patches = torch.rand(B, C, H, W, device=device, dtype=dtype) |
| model = VonMisesKernel(patch_size=13, coeffs=[0.38214156, 0.48090413]).to(patches.device, patches.dtype).eval() |
| model_jit = torch.jit.script( |
| VonMisesKernel(patch_size=13, coeffs=[0.38214156, 0.48090413]).to(patches.device, patches.dtype).eval() |
| ) |
| assert_close(model(patches), model_jit(patches)) |
|
|
|
|
| class TestEmbedGradients: |
| @pytest.mark.parametrize("ps,relative", [(5, True), (13, True), (25, True), (5, False), (13, False), (25, False)]) |
| def test_shape(self, ps, relative, device): |
| inp = torch.ones(1, 2, ps, ps).to(device) |
| emb_grads = EmbedGradients(patch_size=ps, relative=relative).to(device) |
| out = emb_grads(inp) |
| assert out.shape == (1, 7, ps, ps) |
|
|
| @pytest.mark.parametrize("bs", [1, 5, 13]) |
| def test_batch_shape(self, bs, device): |
| inp = torch.ones(bs, 2, 15, 15).to(device) |
| emb_grads = EmbedGradients(patch_size=15, relative=True).to(device) |
| out = emb_grads(inp) |
| assert out.shape == (bs, 7, 15, 15) |
|
|
| def test_print(self, device): |
| emb_grads = EmbedGradients(patch_size=15, relative=True).to(device) |
| emb_grads.__repr__() |
|
|
| def test_toy(self, device): |
| grads = torch.ones(1, 2, 6, 6).float().to(device) |
| grads[0, 0, :, 3:] = 0 |
| emb_grads = EmbedGradients(patch_size=6, relative=True).to(device) |
| out = emb_grads(grads) |
| expected = torch.ones_like(out[0, 0, :, :3]).to(device) |
| assert_close(out[0, 0, :, :3], expected * 0.3787, atol=1e-3, rtol=1e-3) |
| assert_close(out[0, 0, :, 3:], expected * 0, atol=1e-3, rtol=1e-3) |
|
|
| |
| @pytest.mark.xfail(reason="RuntimeError: Jacobian mismatch for output 0 with respect to input 0,") |
| def test_gradcheck(self, device): |
| batch_size, channels, ps = 1, 2, 13 |
| patches = torch.rand(batch_size, channels, ps, ps).to(device) |
| patches = utils.tensor_to_gradcheck_var(patches) |
|
|
| def emb_grads_describe(patches, ps=13): |
| emb_grads = EmbedGradients(patch_size=ps, relative=True).double() |
| emb_grads.to(device) |
| return emb_grads(patches.double()) |
|
|
| assert gradcheck(emb_grads_describe, (patches, ps), raise_exception=True, nondet_tol=1e-4) |
|
|
| @pytest.mark.jit |
| def test_jit(self, device, dtype): |
| B, C, H, W = 2, 2, 13, 13 |
| patches = torch.rand(B, C, H, W, device=device, dtype=dtype) |
| model = EmbedGradients(patch_size=W, relative=True).to(patches.device, patches.dtype).eval() |
| model_jit = torch.jit.script( |
| EmbedGradients(patch_size=W, relative=True).to(patches.device, patches.dtype).eval() |
| ) |
| assert_close(model(patches), model_jit(patches)) |
|
|
|
|
| @pytest.mark.parametrize("kernel_type,d,ps", [('cart', 9, 9), ('polar', 25, 9), ('cart', 9, 16), ('polar', 25, 16)]) |
| def test_spatial_kernel_embedding(kernel_type, ps, d): |
| grids = get_grid_dict(ps) |
| spatial_kernel = spatial_kernel_embedding(kernel_type, grids) |
| assert spatial_kernel.shape == (d, ps, ps) |
|
|
|
|
| class TestExplicitSpacialEncoding: |
| @pytest.mark.parametrize( |
| "kernel_type,ps,in_dims", [('cart', 9, 3), ('polar', 9, 3), ('cart', 13, 7), ('polar', 13, 7)] |
| ) |
| def test_shape(self, kernel_type, ps, in_dims, device): |
| inp = torch.ones(1, in_dims, ps, ps).to(device) |
| ese = ExplicitSpacialEncoding(kernel_type=kernel_type, fmap_size=ps, in_dims=in_dims).to(device) |
| out = ese(inp) |
| d_ = 9 if kernel_type == 'cart' else 25 |
| assert out.shape == (1, d_ * in_dims) |
|
|
| @pytest.mark.parametrize( |
| "kernel_type,bs", [('cart', 1), ('cart', 5), ('cart', 13), ('polar', 1), ('polar', 5), ('polar', 13)] |
| ) |
| def test_batch_shape(self, kernel_type, bs, device): |
| inp = torch.ones(bs, 7, 15, 15).to(device) |
| ese = ExplicitSpacialEncoding(kernel_type=kernel_type, fmap_size=15, in_dims=7).to(device) |
| out = ese(inp) |
| d_ = 9 if kernel_type == 'cart' else 25 |
| assert out.shape == (bs, d_ * 7) |
|
|
| @pytest.mark.parametrize("kernel_type", ['cart', 'polar']) |
| def test_print(self, kernel_type, device): |
| ese = ExplicitSpacialEncoding(kernel_type=kernel_type, fmap_size=15, in_dims=7).to(device) |
| ese.__repr__() |
|
|
| def test_toy(self, device): |
| inp = torch.ones(1, 2, 6, 6).to(device).float() |
| inp[0, 0, :, :] = 0 |
| cart_ese = ExplicitSpacialEncoding(kernel_type='cart', fmap_size=6, in_dims=2).to(device) |
| out = cart_ese(inp) |
| out_part = out[:, :9] |
| expected = torch.zeros_like(out_part).to(device) |
| assert_close(out_part, expected, atol=1e-3, rtol=1e-3) |
|
|
| polar_ese = ExplicitSpacialEncoding(kernel_type='polar', fmap_size=6, in_dims=2).to(device) |
| out = polar_ese(inp) |
| out_part = out[:, :25] |
| expected = torch.zeros_like(out_part).to(device) |
| assert_close(out_part, expected, atol=1e-3, rtol=1e-3) |
|
|
| @pytest.mark.parametrize("kernel_type", ['cart', 'polar']) |
| def test_gradcheck(self, kernel_type, device): |
| batch_size, channels, ps = 1, 2, 13 |
| patches = torch.rand(batch_size, channels, ps, ps).to(device) |
| patches = utils.tensor_to_gradcheck_var(patches) |
|
|
| def explicit_spatial_describe(patches, ps=13): |
| ese = ExplicitSpacialEncoding(kernel_type=kernel_type, fmap_size=ps, in_dims=2) |
| ese.to(device) |
| return ese(patches) |
|
|
| assert gradcheck(explicit_spatial_describe, (patches, ps), raise_exception=True, nondet_tol=1e-4) |
|
|
| @pytest.mark.jit |
| def test_jit(self, device, dtype): |
| B, C, H, W = 2, 2, 13, 13 |
| patches = torch.rand(B, C, H, W, device=device, dtype=dtype) |
| model = ( |
| ExplicitSpacialEncoding(kernel_type='cart', fmap_size=W, in_dims=2).to(patches.device, patches.dtype).eval() |
| ) |
| model_jit = torch.jit.script( |
| ExplicitSpacialEncoding(kernel_type='cart', fmap_size=W, in_dims=2).to(patches.device, patches.dtype).eval() |
| ) |
| assert_close(model(patches), model_jit(patches)) |
|
|
|
|
| class TestWhitening: |
| @pytest.mark.parametrize( |
| "kernel_type,xform,output_dims", |
| [ |
| ('cart', None, 3), |
| ('polar', None, 3), |
| ('cart', 'lw', 7), |
| ('polar', 'lw', 7), |
| ('cart', 'pca', 9), |
| ('polar', 'pca', 9), |
| ], |
| ) |
| def test_shape(self, kernel_type, xform, output_dims, device): |
| in_dims = 63 if kernel_type == 'cart' else 175 |
| wh = Whitening(xform=xform, whitening_model=None, in_dims=in_dims, output_dims=output_dims).to(device) |
| inp = torch.ones(1, in_dims).to(device) |
| out = wh(inp) |
| assert out.shape == (1, output_dims) |
|
|
| @pytest.mark.parametrize("bs", [1, 3, 7]) |
| def test_batch_shape(self, bs, device): |
| wh = Whitening(xform='lw', whitening_model=None, in_dims=175, output_dims=128).to(device) |
| inp = torch.ones(bs, 175).to(device) |
| out = wh(inp) |
| assert out.shape == (bs, 128) |
|
|
| def test_print(self, device): |
| wh = Whitening(xform='lw', whitening_model=None, in_dims=175, output_dims=128).to(device) |
| wh.__repr__() |
|
|
| def test_toy(self, device): |
| wh = Whitening(xform='lw', whitening_model=None, in_dims=175, output_dims=175).to(device) |
| inp = torch.ones(1, 175).to(device).float() |
| out = wh(inp) |
| expected = torch.ones_like(inp).to(device) * 0.0756 |
| assert_close(out, expected, atol=1e-3, rtol=1e-3) |
|
|
| def test_gradcheck(self, device): |
| batch_size, in_dims = 1, 175 |
| patches = torch.rand(batch_size, in_dims).to(device) |
| patches = utils.tensor_to_gradcheck_var(patches) |
|
|
| def whitening_describe(patches, in_dims=175): |
| wh = Whitening(xform='lw', whitening_model=None, in_dims=in_dims).double() |
| wh.to(device) |
| return wh(patches.double()) |
|
|
| assert gradcheck(whitening_describe, (patches, in_dims), raise_exception=True, nondet_tol=1e-4) |
|
|
| @pytest.mark.jit |
| def test_jit(self, device, dtype): |
| batch_size, in_dims = 1, 175 |
| patches = torch.rand(batch_size, in_dims).to(device) |
| model = Whitening(xform='lw', whitening_model=None, in_dims=in_dims).to(patches.device, patches.dtype).eval() |
| model_jit = torch.jit.script( |
| Whitening(xform='lw', whitening_model=None, in_dims=in_dims).to(patches.device, patches.dtype).eval() |
| ) |
| assert_close(model(patches), model_jit(patches)) |
|
|
|
|
| class TestMKDDescriptor: |
| dims = {'cart': 63, 'polar': 175, 'concat': 238} |
|
|
| @pytest.mark.parametrize( |
| "ps,kernel_type", [(9, 'concat'), (9, 'cart'), (9, 'polar'), (32, 'concat'), (32, 'cart'), (32, 'polar')] |
| ) |
| def test_shape(self, ps, kernel_type, device): |
| mkd = MKDDescriptor(patch_size=ps, kernel_type=kernel_type, whitening=None).to(device) |
| inp = torch.ones(1, 1, ps, ps).to(device) |
| out = mkd(inp) |
| assert out.shape == (1, self.dims[kernel_type]) |
|
|
| @pytest.mark.parametrize( |
| "ps,kernel_type,whitening", |
| [ |
| (9, 'concat', 'lw'), |
| (9, 'cart', 'lw'), |
| (9, 'polar', 'lw'), |
| (9, 'concat', 'pcawt'), |
| (9, 'cart', 'pcawt'), |
| (9, 'polar', 'pcawt'), |
| ], |
| ) |
| def test_whitened_shape(self, ps, kernel_type, whitening, device): |
| mkd = MKDDescriptor(patch_size=ps, kernel_type=kernel_type, whitening=whitening).to(device) |
| inp = torch.ones(1, 1, ps, ps).to(device) |
| out = mkd(inp) |
| output_dims = min(self.dims[kernel_type], 128) |
| assert out.shape == (1, output_dims) |
|
|
| @pytest.mark.parametrize("bs", [1, 3, 7]) |
| def test_batch_shape(self, bs, device): |
| mkd = MKDDescriptor(patch_size=19, kernel_type='concat', whitening=None).to(device) |
| inp = torch.ones(bs, 1, 19, 19).to(device) |
| out = mkd(inp) |
| assert out.shape == (bs, 238) |
|
|
| def test_print(self, device): |
| mkd = MKDDescriptor(patch_size=32, whitening='lw', training_set='liberty', output_dims=128).to(device) |
| mkd.__repr__() |
|
|
| def test_toy(self, device): |
| inp = torch.ones(1, 1, 6, 6).to(device).float() |
| inp[0, 0, :, :] = 0 |
| mkd = MKDDescriptor(patch_size=6, kernel_type='concat', whitening=None).to(device) |
| out = mkd(inp) |
| out_part = out[0, -28:] |
| expected = torch.zeros_like(out_part).to(device) |
| assert_close(out_part, expected, atol=1e-3, rtol=1e-3) |
|
|
| @pytest.mark.skip("Just because") |
| @pytest.mark.parametrize("whitening", [None, 'lw', 'pca']) |
| def test_gradcheck(self, whitening, device): |
| batch_size, channels, ps = 1, 1, 19 |
| patches = torch.rand(batch_size, channels, ps, ps).to(device) |
| patches = utils.tensor_to_gradcheck_var(patches) |
|
|
| def mkd_describe(patches, patch_size=19): |
| mkd = MKDDescriptor(patch_size=patch_size, kernel_type='concat', whitening=whitening).double() |
| mkd.to(device) |
| return mkd(patches.double()) |
|
|
| assert gradcheck(mkd_describe, (patches, ps), raise_exception=True, nondet_tol=1e-4) |
|
|
| @pytest.mark.skip("neither dict, nor nn.ModuleDict works") |
| @pytest.mark.jit |
| def test_jit(self, device, dtype): |
| batch_size, channels, ps = 1, 1, 19 |
| patches = torch.rand(batch_size, channels, ps, ps).to(device) |
| kt = 'concat' |
| wt = 'lw' |
| model = MKDDescriptor(patch_size=ps, kernel_type=kt, whitening=wt).to(patches.device, patches.dtype).eval() |
| model_jit = torch.jit.script( |
| MKDDescriptor(patch_size=ps, kernel_type=kt, whitening=wt).to(patches.device, patches.dtype).eval() |
| ) |
| assert_close(model(patches), model_jit(patches)) |
|
|
|
|
| class TestSimpleKD: |
| dims = {'cart': 63, 'polar': 175} |
|
|
| @pytest.mark.parametrize("ps,kernel_type", [(9, 'cart'), (9, 'polar'), (32, 'cart'), (32, 'polar')]) |
| def test_shape(self, ps, kernel_type, device): |
| skd = SimpleKD(patch_size=ps, kernel_type=kernel_type).to(device) |
| inp = torch.ones(1, 1, ps, ps).to(device) |
| out = skd(inp) |
| assert out.shape == (1, min(128, self.dims[kernel_type])) |
|
|
| @pytest.mark.parametrize("bs", [1, 3, 7]) |
| def test_batch_shape(self, bs, device): |
| skd = SimpleKD(patch_size=19, kernel_type='polar').to(device) |
| inp = torch.ones(bs, 1, 19, 19).to(device) |
| out = skd(inp) |
| assert out.shape == (bs, 128) |
|
|
| def test_print(self, device): |
| skd = SimpleKD(patch_size=19, kernel_type='polar').to(device) |
| skd.__repr__() |
|
|
| def test_gradcheck(self, device): |
| batch_size, channels, ps = 1, 1, 19 |
| patches = torch.rand(batch_size, channels, ps, ps).to(device) |
| patches = utils.tensor_to_gradcheck_var(patches) |
|
|
| def skd_describe(patches, patch_size=19): |
| skd = SimpleKD(patch_size=ps, kernel_type='polar', whitening='lw').double() |
| skd.to(device) |
| return skd(patches.double()) |
|
|
| assert gradcheck(skd_describe, (patches, ps), raise_exception=True, nondet_tol=1e-4) |
|
|
| @pytest.mark.jit |
| def test_jit(self, device, dtype): |
| batch_size, channels, ps = 1, 1, 19 |
| patches = torch.rand(batch_size, channels, ps, ps).to(device) |
| model = SimpleKD(patch_size=ps, kernel_type='polar', whitening='lw').to(patches.device, patches.dtype).eval() |
| model_jit = torch.jit.script( |
| SimpleKD(patch_size=ps, kernel_type='polar', whitening='lw').to(patches.device, patches.dtype).eval() |
| ) |
| assert_close(model(patches), model_jit(patches)) |
|
|
