# Copyright 2017 The TensorFlow Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== """Tests for FFT via the XLA JIT.""" import itertools import numpy as np import scipy.signal as sps from tensorflow.compiler.tests import xla_test from tensorflow.python.framework import dtypes from tensorflow.python.ops import array_ops from tensorflow.python.ops import gradients_impl from tensorflow.python.ops.signal import signal from tensorflow.python.platform import googletest BATCH_DIMS = (3, 5) RTOL = 0.009 # Eigen/cuFFT differ widely from np, especially for FFT3D ATOL = 1e-4 RTOL_3D = 0.07 ATOL_3D = 4e-4 def pick_10(x): x = list(x) np.random.seed(123) np.random.shuffle(x) return x[:10] def to_32bit(x): if x.dtype == np.complex128: return x.astype(np.complex64) if x.dtype == np.float64: return x.astype(np.float32) return x POWS_OF_2 = 2**np.arange(3, 12) INNER_DIMS_1D = list((x,) for x in POWS_OF_2) POWS_OF_2 = 2**np.arange(3, 8) # To avoid OOM on GPU. INNER_DIMS_2D = pick_10(itertools.product(POWS_OF_2, POWS_OF_2)) INNER_DIMS_3D = pick_10(itertools.product(POWS_OF_2, POWS_OF_2, POWS_OF_2)) class FFTTest(xla_test.XLATestCase): def _VerifyFftMethod(self, inner_dims, complex_to_input, input_to_expected, tf_method, atol=ATOL, rtol=RTOL): for indims in inner_dims: print("nfft =", indims) shape = BATCH_DIMS + indims data = np.arange(np.prod(shape) * 2) / np.prod(indims) np.random.seed(123) np.random.shuffle(data) data = np.reshape(data.astype(np.float32).view(np.complex64), shape) data = to_32bit(complex_to_input(data)) expected = to_32bit(input_to_expected(data)) with self.session() as sess: with self.test_scope(): ph = array_ops.placeholder( dtypes.as_dtype(data.dtype), shape=data.shape) out = tf_method(ph) value = sess.run(out, {ph: data}) self.assertAllClose(expected, value, rtol=rtol, atol=atol) def testContribSignalSTFT(self): ws = 512 hs = 128 dims = (ws * 20,) shape = BATCH_DIMS + dims data = np.arange(np.prod(shape)) / np.prod(dims) np.random.seed(123) np.random.shuffle(data) data = np.reshape(data.astype(np.float32), shape) window = sps.get_window("hann", ws) expected = sps.stft( data, nperseg=ws, noverlap=ws - hs, boundary=None, window=window)[2] expected = np.swapaxes(expected, -1, -2) expected *= window.sum() # scipy divides by window sum with self.session() as sess: with self.test_scope(): ph = array_ops.placeholder( dtypes.as_dtype(data.dtype), shape=data.shape) out = signal.stft(ph, ws, hs) grad = gradients_impl.gradients(out, ph, grad_ys=array_ops.ones_like(out)) # For gradients, we simply verify that they compile & execute. value, _ = sess.run([out, grad], {ph: data}) self.assertAllClose(expected, value, rtol=RTOL, atol=ATOL) def testFFT(self): self._VerifyFftMethod(INNER_DIMS_1D, lambda x: x, np.fft.fft, signal.fft) def testFFT2D(self): self._VerifyFftMethod(INNER_DIMS_2D, lambda x: x, np.fft.fft2, signal.fft2d) def testFFT3D(self): self._VerifyFftMethod(INNER_DIMS_3D, lambda x: x, lambda x: np.fft.fftn(x, axes=(-3, -2, -1)), signal.fft3d, ATOL_3D, RTOL_3D) def testIFFT(self): self._VerifyFftMethod(INNER_DIMS_1D, lambda x: x, np.fft.ifft, signal.ifft) def testIFFT2D(self): self._VerifyFftMethod(INNER_DIMS_2D, lambda x: x, np.fft.ifft2, signal.ifft2d) def testIFFT3D(self): self._VerifyFftMethod(INNER_DIMS_3D, lambda x: x, lambda x: np.fft.ifftn(x, axes=(-3, -2, -1)), signal.ifft3d, ATOL_3D, RTOL_3D) def testRFFT(self): def _to_expected(x): return np.fft.rfft(x, n=x.shape[-1]) def _tf_fn(x): return signal.rfft(x, fft_length=[x.shape[-1]]) self._VerifyFftMethod(INNER_DIMS_1D, np.real, _to_expected, _tf_fn) def testRFFT2D(self): def _tf_fn(x): return signal.rfft2d(x, fft_length=[x.shape[-2], x.shape[-1]]) self._VerifyFftMethod( INNER_DIMS_2D, np.real, lambda x: np.fft.rfft2(x, s=[x.shape[-2], x.shape[-1]]), _tf_fn) def testRFFT3D(self): def _to_expected(x): return np.fft.rfftn( x, axes=(-3, -2, -1), s=[x.shape[-3], x.shape[-2], x.shape[-1]]) def _tf_fn(x): return signal.rfft3d( x, fft_length=[x.shape[-3], x.shape[-2], x.shape[-1]]) self._VerifyFftMethod(INNER_DIMS_3D, np.real, _to_expected, _tf_fn, ATOL_3D, RTOL_3D) def testRFFT3DMismatchedSize(self): def _to_expected(x): return np.fft.rfftn( x, axes=(-3, -2, -1), s=[x.shape[-3] // 2, x.shape[-2], x.shape[-1] * 2]) def _tf_fn(x): return signal.rfft3d( x, fft_length=[x.shape[-3] // 2, x.shape[-2], x.shape[-1] * 2]) self._VerifyFftMethod(INNER_DIMS_3D, np.real, _to_expected, _tf_fn) def testIRFFT(self): def _tf_fn(x): return signal.irfft(x, fft_length=[2 * (x.shape[-1] - 1)]) self._VerifyFftMethod( INNER_DIMS_1D, lambda x: np.fft.rfft(np.real(x), n=x.shape[-1]), lambda x: np.fft.irfft(x, n=2 * (x.shape[-1] - 1)), _tf_fn) def testIRFFT2D(self): def _tf_fn(x): return signal.irfft2d(x, fft_length=[x.shape[-2], 2 * (x.shape[-1] - 1)]) self._VerifyFftMethod( INNER_DIMS_2D, lambda x: np.fft.rfft2(np.real(x), s=[x.shape[-2], x.shape[-1]]), lambda x: np.fft.irfft2(x, s=[x.shape[-2], 2 * (x.shape[-1] - 1)]), _tf_fn) def testIRFFT3D(self): def _to_input(x): return np.fft.rfftn( np.real(x), axes=(-3, -2, -1), s=[x.shape[-3], x.shape[-2], x.shape[-1]]) def _to_expected(x): return np.fft.irfftn( x, axes=(-3, -2, -1), s=[x.shape[-3], x.shape[-2], 2 * (x.shape[-1] - 1)]) def _tf_fn(x): return signal.irfft3d( x, fft_length=[x.shape[-3], x.shape[-2], 2 * (x.shape[-1] - 1)]) self._VerifyFftMethod(INNER_DIMS_3D, _to_input, _to_expected, _tf_fn, ATOL_3D, RTOL_3D) def testIRFFT3DMismatchedSize(self): def _to_input(x): return np.fft.rfftn( np.real(x), axes=(-3, -2, -1), s=[x.shape[-3] // 2, x.shape[-2], x.shape[-1] * 2]) def _to_expected(x): return np.fft.irfftn( x, axes=(-3, -2, -1), s=[x.shape[-3] // 2, x.shape[-2], x.shape[-1] * 2]) def _tf_fn(x): return signal.irfft3d( x, fft_length=[x.shape[-3] // 2, x.shape[-2], x.shape[-1] * 2]) self._VerifyFftMethod(INNER_DIMS_3D, _to_input, _to_expected, _tf_fn, ATOL_3D, RTOL_3D) if __name__ == "__main__": googletest.main()