Spaces:
Sleeping
Sleeping
| import numpy as np | |
| import pytest | |
| import pywt | |
| from pytorch_wavelets import DWT1DForward, DWT1DInverse | |
| import torch | |
| from contextlib import contextmanager | |
| PREC_FLT = 3 | |
| PREC_DBL = 7 | |
| HAVE_GPU = torch.cuda.is_available() | |
| if HAVE_GPU: | |
| dev = torch.device('cuda') | |
| else: | |
| dev = torch.device('cpu') | |
| def set_double_precision(): | |
| old_prec = torch.get_default_dtype() | |
| try: | |
| torch.set_default_dtype(torch.float64) | |
| yield | |
| finally: | |
| torch.set_default_dtype(old_prec) | |
| def test_ok(wave, J, mode): | |
| x = torch.randn(5, 4, 64).to(dev) | |
| dwt = DWT1DForward(J=J, wave=wave, mode=mode).to(dev) | |
| iwt = DWT1DInverse(wave=wave, mode=mode).to(dev) | |
| yl, yh = dwt(x) | |
| x2 = iwt((yl, yh)) | |
| # Can have data errors sometimes | |
| assert yl.is_contiguous() | |
| for j in range(J): | |
| assert yh[j].is_contiguous() | |
| assert x2.is_contiguous() | |
| def test_equal(wave, J, mode): | |
| x = torch.randn(5, 4, 64).to(dev) | |
| dwt = DWT1DForward(J=J, wave=wave, mode=mode).to(dev) | |
| yl, yh = dwt(x) | |
| # Test it is the same as doing the PyWavelets wavedec with reflection padding | |
| coeffs = pywt.wavedec(x.cpu().numpy(), wave, level=J, mode=mode) | |
| np.testing.assert_array_almost_equal(yl.cpu(), coeffs[0], decimal=PREC_FLT) | |
| for j in range(J): | |
| np.testing.assert_array_almost_equal(coeffs[J-j], yh[j].cpu(), decimal=PREC_FLT) | |
| # Test the forward and inverse worked | |
| iwt = DWT1DInverse(wave=wave, mode=mode).to(dev) | |
| x2 = iwt((yl, yh)) | |
| np.testing.assert_array_almost_equal(x.cpu(), x2.detach().cpu(), decimal=PREC_FLT) | |
| def test_equal_oddshape(length, mode): | |
| wave = 'db3' | |
| J = 3 | |
| x = torch.randn(5, 4, length).to(dev) | |
| dwt1 = DWT1DForward(J=J, wave=wave, mode=mode).to(dev) | |
| iwt1 = DWT1DInverse(wave=wave, mode=mode).to(dev) | |
| yl1, yh1 = dwt1(x) | |
| x1 = iwt1((yl1, yh1)) | |
| # Test it is the same as doing the PyWavelets wavedec | |
| coeffs = pywt.wavedec(x.cpu().numpy(), wave, level=J, mode=mode) | |
| X = pywt.waverec(coeffs, wave, mode=mode) | |
| np.testing.assert_array_almost_equal(X, x1.detach().cpu(), decimal=PREC_FLT) | |
| np.testing.assert_array_almost_equal(yl1.cpu(), coeffs[0], decimal=PREC_FLT) | |
| for j in range(J): | |
| np.testing.assert_array_almost_equal(coeffs[J-j], yh1[j].cpu(), decimal=PREC_FLT) | |
| def test_equal_double(wave, J, mode): | |
| with set_double_precision(): | |
| x = torch.randn(5, 4, 100).to(dev) | |
| assert x.dtype == torch.float64 | |
| dwt = DWT1DForward(J=J, wave=wave, mode=mode).to(dev) | |
| iwt = DWT1DInverse(wave=wave, mode=mode).to(dev) | |
| yl, yh = dwt(x) | |
| x2 = iwt((yl, yh)) | |
| # Test the forward and inverse worked | |
| np.testing.assert_array_almost_equal(x.cpu(), x2.detach().cpu(), decimal=PREC_DBL) | |
| coeffs = pywt.wavedec(x.cpu().numpy(), wave, level=J, mode=mode) | |
| np.testing.assert_array_almost_equal(yl.cpu(), coeffs[0], decimal=7) | |
| for j in range(J): | |
| np.testing.assert_array_almost_equal(coeffs[J-j], yh[j].cpu(), decimal=PREC_DBL) | |
| # Test gradients | |
| def test_gradients_fwd(wave, J, mode): | |
| """ Gradient of forward function should be inverse function with filters | |
| swapped """ | |
| im = np.random.randn(5, 6, 128).astype('float32') | |
| imt = torch.tensor(im, dtype=torch.float32, requires_grad=True, device=dev) | |
| wave = pywt.Wavelet(wave) | |
| fwd_filts = (wave.dec_lo, wave.dec_hi) | |
| inv_filts = (wave.dec_lo[::-1], wave.dec_hi[::-1]) | |
| dwt = DWT1DForward(J=J, wave=fwd_filts, mode=mode).to(dev) | |
| iwt = DWT1DInverse(wave=inv_filts, mode=mode).to(dev) | |
| yl, yh = dwt(imt) | |
| # Test the lowpass | |
| ylg = torch.randn(*yl.shape, device=dev) | |
| yl.backward(ylg, retain_graph=True) | |
| zeros = [torch.zeros_like(yh[i]) for i in range(J)] | |
| ref = iwt((ylg, zeros)) | |
| if (imt.grad.detach().cpu() - ref.cpu()).abs().sum() > 1e-3: | |
| import pdb; pdb.set_trace() | |
| np.testing.assert_array_almost_equal(imt.grad.detach().cpu(), ref.cpu(), decimal=PREC_FLT) | |
| # Test the bandpass | |
| for j, y in enumerate(yh): | |
| imt.grad.zero_() | |
| g = torch.randn(*y.shape, device=dev) | |
| y.backward(g, retain_graph=True) | |
| hps = [zeros[i] for i in range(J)] | |
| hps[j] = g | |
| ref = iwt((torch.zeros_like(yl), hps)) | |
| np.testing.assert_array_almost_equal(imt.grad.detach().cpu(), ref.cpu(), decimal=PREC_FLT) | |
| # Test gradients | |
| def test_gradients_inv(wave, J, mode): | |
| """ Gradient of inverse function should be forward function with filters | |
| swapped """ | |
| wave = pywt.Wavelet(wave) | |
| fwd_filts = (wave.dec_lo, wave.dec_hi) | |
| inv_filts = (wave.dec_lo[::-1], wave.dec_hi[::-1]) | |
| dwt = DWT1DForward(J=J, wave=fwd_filts, mode=mode).to(dev) | |
| iwt = DWT1DInverse(wave=inv_filts, mode=mode).to(dev) | |
| # Get the shape of the pyramid | |
| temp = torch.zeros(5,6,128).to(dev) | |
| l, h = dwt(temp) | |
| # Create our inputs | |
| yl = torch.randn(*l.shape, requires_grad=True, device=dev) | |
| yh = [torch.randn(*h[i].shape, requires_grad=True, device=dev) for i in range(J)] | |
| y = iwt((yl, yh)) | |
| # Test the gradients | |
| yg = torch.randn(*y.shape, device=dev) | |
| y.backward(yg, retain_graph=True) | |
| dyl, dyh = dwt(yg) | |
| # test the lowpass | |
| np.testing.assert_array_almost_equal(yl.grad.detach().cpu(), dyl.cpu(), decimal=PREC_FLT) | |
| # Test the bandpass | |
| for j in range(J): | |
| np.testing.assert_array_almost_equal(yh[j].grad.detach().cpu(), dyh[j].cpu(), decimal=PREC_FLT) | |