Add PowerSpectrum CPU operator

dali/kernels/signal/fft/fft_cpu.cc

@@ -54,6 +54,10 @@ void Fft1DCpu<OutputType, InputType, Dims>::Run(
   impl_->Run(context, out, in, args);
 }
 
+// 1 Dim, typically input (time), producing output (frequency)
+template class Fft1DCpu<std::complex<float>, float, 1>;  // complex fft
+template class Fft1DCpu<float, float, 1>;  // magnitude
+
 // 2 Dims, typically input (channels, time), producing output (channels, frequency)
 template class Fft1DCpu<std::complex<float>, float, 2>;  // complex fft
 template class Fft1DCpu<float, float, 2>;  // magnitude

dali/kernels/signal/fft/fft_cpu_impl_ffts.cc

@@ -164,9 +164,14 @@ void Fft1DImplFfts<OutputType, InputType, Dims>::Run(
     });
 }
 
+// 1 Dim, typically input (time), producing output (frequency)
+template class Fft1DImplFfts<std::complex<float>, float, 1>;  // complex fft
+template class Fft1DImplFfts<float, float, 1>;  // magnitude
+
 // 2 Dims, typically input (channels, time), producing output (channels, frequency)
 template class Fft1DImplFfts<std::complex<float>, float, 2>;
 template class Fft1DImplFfts<float, float, 2>;
+
 // 3 Dims, typically input (channels, frames, time), producing output (channels, frames, frequency)
 template class Fft1DImplFfts<std::complex<float>, float, 3>;
 template class Fft1DImplFfts<float, float, 3>;

dali/operators/CMakeLists.txt

@@ -30,6 +30,7 @@ add_subdirectory(paste)
 add_subdirectory(reader)
 add_subdirectory(resize)
 add_subdirectory(sequence)
+add_subdirectory(signal)
 add_subdirectory(support)
 add_subdirectory(transpose)
 add_subdirectory(util)

dali/operators/signal/CMakeLists.txt

@@ -0,0 +1,21 @@
+# Copyright (c) 2019, NVIDIA CORPORATION. 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.
+
+if (BUILD_FFTS)
+  add_subdirectory(fft)
+endif()
+
+collect_headers(DALI_INST_HDRS PARENT_SCOPE)
+collect_sources(DALI_OPERATOR_SRCS PARENT_SCOPE)
+collect_test_sources(DALI_OPERATOR_TEST_SRCS PARENT_SCOPE)

dali/operators/signal/fft/CMakeLists.txt

@@ -0,0 +1,17 @@
+# Copyright (c) 2019, NVIDIA CORPORATION. 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.
+
+collect_headers(DALI_INST_HDRS PARENT_SCOPE)
+collect_sources(DALI_OPERATOR_SRCS PARENT_SCOPE)
+collect_test_sources(DALI_OPERATOR_TEST_SRCS PARENT_SCOPE)

dali/operators/signal/fft/power_spectrum.cc

@@ -0,0 +1,104 @@
+// Copyright (c) 2019, NVIDIA CORPORATION. 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.
+
+#include "dali/operators/signal/fft/power_spectrum.h"
+#include "dali/core/static_switch.h"
+#include "dali/kernels/signal/fft/fft_cpu.h"
+#include "dali/pipeline/data/views.h"
+
+#define FFT_SUPPORTED_NDIMS (1, 2, 3)
+
+static constexpr int kNumInputs = 1;
+static constexpr int kNumOutputs = 1;
+
+namespace dali {
+
+DALI_SCHEMA(PowerSpectrum)
+    .DocStr(R"code(Power spectrum of signal.)code")
+    .NumInput(kNumInputs)
+    .NumOutput(kNumOutputs)
+    .AddOptionalArg("nfft",
+      R"code(Size of the FFT. By default nfft is selected to match the lenght of the data in the
+transformation axis. The number of bins created in the output is `nfft // 2 + 1` (positive part
+of the spectrum only).)code",
+      -1)
+    .AddOptionalArg("axis",
+      R"code(Index of the dimension to be transformed to the frequency domain. By default, the
+last dimension is selected.)code",
+      -1)
+    .AddOptionalArg("power",
+      R"code(Exponent of the fft magnitude: Supported values are `2` for power spectrum
+(`real*real + imag*imag`) and `1` for complex magnitude (`sqrt(real*real + imag*imag)`).)code",
+      2);
+
+template <>
+bool PowerSpectrum<CPUBackend>::SetupImpl(std::vector<OutputDesc> &output_desc,
+                                          const workspace_t<CPUBackend> &ws) {
+  output_desc.resize(kNumOutputs);
+  const auto &input = ws.InputRef<CPUBackend>(0);
+  auto &output = ws.OutputRef<CPUBackend>(0);
+  kernels::KernelContext ctx;
+  auto in_shape = input.shape();
+  int nsamples = input.size();
+  auto nthreads = ws.GetThreadPool().size();
+
+  // Other types not supported for now
+  using InputType = float;
+  using OutputType = float;
+  VALUE_SWITCH(in_shape.sample_dim(), Dims, FFT_SUPPORTED_NDIMS, (
+    using FftKernel = kernels::signal::fft::Fft1DCpu<OutputType, InputType, Dims>;
+    kmgr_.Initialize<FftKernel>();
+    kmgr_.Resize<FftKernel>(nthreads, nsamples);
+    output_desc[0].type = TypeInfo::Create<OutputType>();
+    output_desc[0].shape.resize(nsamples, Dims);
+    for (int i = 0; i < nsamples; i++) {
+      const auto in_view = view<const InputType, Dims>(input[i]);
+      auto &req = kmgr_.Setup<FftKernel>(i, ctx, in_view, fft_args_);
+      output_desc[0].shape.set_tensor_shape(i, req.output_shapes[0][0].shape);
+    }
+  ), DALI_FAIL(make_string("Unsupported number of dimensions ", in_shape.size())));  // NOLINT
+
+  return true;
+}
+
+template <>
+void PowerSpectrum<CPUBackend>::RunImpl(workspace_t<CPUBackend> &ws) {
+  const auto &input = ws.InputRef<CPUBackend>(0);
+  auto &output = ws.OutputRef<CPUBackend>(0);
+  auto in_shape = input.shape();
+  int nsamples = input.size();
+  auto& thread_pool = ws.GetThreadPool();
+  // Other types not supported for now
+  using InputType = float;
+  using OutputType = float;
+  VALUE_SWITCH(in_shape.sample_dim(), Dims, FFT_SUPPORTED_NDIMS, (
+    using FftKernel = kernels::signal::fft::Fft1DCpu<OutputType, InputType, Dims>;
+
+    for (int i = 0; i < input.shape().num_samples(); i++) {
+      thread_pool.DoWorkWithID(
+        [this, &input, &output, i](int thread_id) {
+          kernels::KernelContext ctx;
+          auto in_view = view<const InputType, Dims>(input[i]);
+          auto out_view = view<OutputType, Dims>(output[i]);
+          kmgr_.Run<FftKernel>(thread_id, i, ctx, out_view, in_view, fft_args_);
+        });
+    }
+  ), DALI_FAIL(make_string("Not supported number of dimensions: ", in_shape.size())));  // NOLINT
+
+  thread_pool.WaitForWork();
+}
+
+DALI_REGISTER_OPERATOR(PowerSpectrum, PowerSpectrum<CPUBackend>, CPU);
+
+}  // namespace dali

dali/operators/signal/fft/power_spectrum.h

@@ -0,0 +1,63 @@
+// Copyright (c) 2019, NVIDIA CORPORATION. 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.
+
+#ifndef DALI_OPERATORS_SIGNAL_FFT_POWER_SPECTRUM_H_
+#define DALI_OPERATORS_SIGNAL_FFT_POWER_SPECTRUM_H_
+
+#include <string>
+#include <vector>
+#include "dali/core/common.h"
+#include "dali/kernels/kernel_manager.h"
+#include "dali/kernels/signal/fft/fft_cpu.h"
+#include "dali/pipeline/operator/common.h"
+#include "dali/pipeline/operator/operator.h"
+
+namespace dali {
+
+template <typename Backend>
+class PowerSpectrum : public Operator<Backend> {
+ public:
+  explicit PowerSpectrum(const OpSpec &spec)
+      : Operator<Backend>(spec) {
+    fft_args_.nfft = spec.GetArgument<int>("nfft");
+    fft_args_.transform_axis = spec.GetArgument<int>("axis");
+    int power = spec.GetArgument<int>("power");
+    switch (power) {
+      case 1:
+        fft_args_.spectrum_type = kernels::signal::fft::FFT_SPECTRUM_MAGNITUDE;
+        break;
+      case 2:
+        fft_args_.spectrum_type = kernels::signal::fft::FFT_SPECTRUM_POWER;
+        break;
+      default:
+        DALI_FAIL(make_string("Power argument should be either `2` for power spectrum or `1` "
+          "for complex magnitude. Received: ", power));
+    }
+  }
+
+ protected:
+  bool CanInferOutputs() const override { return true; }
+  bool SetupImpl(std::vector<OutputDesc> &output_desc, const workspace_t<Backend> &ws) override;
+  void RunImpl(workspace_t<CPUBackend> &ws) override;
+
+  USE_OPERATOR_MEMBERS();
+  using Operator<Backend>::RunImpl;
+
+  kernels::KernelManager kmgr_;
+  kernels::signal::fft::FftArgs fft_args_;
+};
+
+}  // namespace dali
+
+#endif  // DALI_OPERATORS_SIGNAL_FFT_POWER_SPECTRUM_H_

dali/test/python/test_operator_power_spectrum.py

@@ -0,0 +1,106 @@
+# Copyright (c) 2019, NVIDIA CORPORATION. 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.
+
+from __future__ import division
+from nvidia.dali.pipeline import Pipeline
+import nvidia.dali.ops as ops
+import nvidia.dali.types as types
+import nvidia.dali as dali
+import numpy as np
+from numpy.testing import assert_array_equal, assert_allclose
+from functools import partial
+from test_utils import check_batch
+from test_utils import compare_pipelines
+from test_utils import RandomDataIterator
+
+class PowerSpectrumPipeline(Pipeline):
+    def __init__(self, device, batch_size, iterator, axis, nfft, num_threads=1, device_id=0):
+        super(PowerSpectrumPipeline, self).__init__(batch_size, num_threads, device_id)
+        self.device = device
+        self.iterator = iterator
+        self.inputs = ops.ExternalSource()
+        self.fft = ops.PowerSpectrum(device=self.device, axis=axis, nfft=nfft)
+
+    def define_graph(self):
+        self.data = self.inputs()
+        out = self.data.gpu() if self.device == 'gpu' else self.data
+        out = self.fft(out)
+        return out
+
+    def iter_setup(self):
+        data = self.iterator.next()
+        self.feed_input(self.data, data)
+
+def power_spectrum_numpy(nfft, axis, waveform):
+    fft_out = np.fft.fft(waveform, axis=axis, n=nfft)
+    power_spectrum = fft_out.real ** 2 + fft_out.imag ** 2
+    shape = waveform.shape
+
+    out_shape = list(shape)
+    out_shape[axis] = nfft//2+1
+    out_shape = tuple(out_shape)
+
+    if len(out_shape) == 1:
+        out = power_spectrum[0:out_shape[0]]
+    elif len(out_shape) == 2:
+        out = power_spectrum[0:out_shape[0], 0:out_shape[1]]
+    elif len(out_shape) == 3:
+        out = power_spectrum[0:out_shape[0], 0:out_shape[1], 0:out_shape[2]]
+    return out
+
+class PowerSpectrumNumpyPipeline(Pipeline):
+    def __init__(self, device, batch_size, iterator, axis, nfft,
+                 num_threads=1, device_id=0):
+        super(PowerSpectrumNumpyPipeline, self).__init__(
+              batch_size, num_threads, device_id,
+              seed=12345, exec_async=False, exec_pipelined=False)
+        self.device = "cpu"
+        self.iterator = iterator
+        self.inputs = ops.ExternalSource()
+
+        function = partial(power_spectrum_numpy, nfft, axis)
+        self.power_spectrum = ops.PythonFunction(function=function)
+
+    def define_graph(self):
+        self.data = self.inputs()
+        out = self.power_spectrum(self.data)
+        return out
+
+    def iter_setup(self):
+        data = self.iterator.next()
+        self.feed_input(self.data, data)
+
+def check_operator_power_spectrum(device, batch_size, input_shape, nfft, axis):
+    eii1 = RandomDataIterator(batch_size, shape=input_shape, dtype=np.float32)
+    eii2 = RandomDataIterator(batch_size, shape=input_shape, dtype=np.float32)
+    compare_pipelines(
+        PowerSpectrumPipeline(device, batch_size, iter(eii1), axis=axis, nfft=nfft),
+        PowerSpectrumNumpyPipeline(device, batch_size, iter(eii2), axis=axis, nfft=nfft),
+        batch_size=batch_size, N_iterations=5, eps=1e-04)
+
+def test_operator_power_spectrum():
+    for device in ['cpu']:
+        for batch_size in [3]:
+            for nfft, axis, shape in [(16, 1, (2, 16)),
+                                      (1024, 1, (1, 1024)),
+                                      (1024, 0, (1024,)),
+                                      (128, 1, (1, 100)),
+                                      (128, 0, (100,)),
+                                      (16, 0, (16, 2)),
+                                      (8, 1, (2, 8, 2))]:
+                yield check_operator_power_spectrum, device, batch_size, shape, nfft, axis
+
+if __name__ == "__main__":
+    check_operator_power_spectrum(device='cpu', batch_size=3, input_shape=(2, 1024),
+                                      nfft=1024, axis=1)

dali/test/python/test_utils.py

@@ -121,13 +121,14 @@ def compare_pipelines(pipe1, pipe2, batch_size, N_iterations, eps = 1e-07):
 
 class RandomDataIterator(object):
     import_numpy()
-    def __init__(self, batch_size, shape=(10, 600, 800, 3)):
+    def __init__(self, batch_size, shape=(10, 600, 800, 3), dtype=np.uint8):
         self.batch_size = batch_size
         self.test_data = []
         for _ in range(self.batch_size):
             np.random.seed(0)
-            self.test_data.append(np.array(np.random.rand(*shape) * 255,
-                                  dtype = np.uint8 ) )
+            self.test_data.append(
+                np.array(np.random.rand(*shape) * (1.0 if dtype == np.float32 else 255),
+                dtype=dtype ) )
 
     def __iter__(self):
         self.i = 0